open Qptypes;; open Qputils;; open Core.Std;; module Determinants : sig type t = { n_int : N_int_number.t; bit_kind : Bit_kind.t; mo_label : MO_label.t; n_det : Det_number.t; n_states : States_number.t; n_states_diag : States_number.t; n_det_max_jacobi : Strictly_positive_int.t; threshold_generators : Threshold.t; threshold_selectors : Threshold.t; read_wf : bool; expected_s2 : Positive_float.t; s2_eig : bool; psi_coef : Det_coef.t array; psi_det : Determinant.t array; } with sexp val read : unit -> t option val write : t -> unit val to_string : t -> string val to_rst : t -> Rst_string.t val of_rst : Rst_string.t -> t option end = struct type t = { n_int : N_int_number.t; bit_kind : Bit_kind.t; mo_label : MO_label.t; n_det : Det_number.t; n_states : States_number.t; n_states_diag : States_number.t; n_det_max_jacobi : Strictly_positive_int.t; threshold_generators : Threshold.t; threshold_selectors : Threshold.t; read_wf : bool; expected_s2 : Positive_float.t; s2_eig : bool; psi_coef : Det_coef.t array; psi_det : Determinant.t array; } with sexp ;; let get_default = Qpackage.get_ezfio_default "determinants";; let read_n_int () = if not (Ezfio.has_determinants_n_int()) then Ezfio.get_mo_basis_mo_tot_num () |> Bitlist.n_int_of_mo_tot_num |> N_int_number.to_int |> Ezfio.set_determinants_n_int ; Ezfio.get_determinants_n_int () |> N_int_number.of_int ;; let write_n_int n = N_int_number.to_int n |> Ezfio.set_determinants_n_int ;; let read_bit_kind () = if not (Ezfio.has_determinants_bit_kind ()) then Lazy.force Qpackage.bit_kind |> Bit_kind.to_int |> Ezfio.set_determinants_bit_kind ; Ezfio.get_determinants_bit_kind () |> Bit_kind.of_int ;; let write_bit_kind b = Bit_kind.to_int b |> Ezfio.set_determinants_bit_kind ;; let read_mo_label () = if (not (Ezfio.has_determinants_mo_label ())) then Ezfio.get_mo_basis_mo_label () |> Ezfio.set_determinants_mo_label ; Ezfio.get_determinants_mo_label () |> MO_label.of_string ;; let write_mo_label l = MO_label.to_string l |> Ezfio.set_determinants_mo_label ;; let read_n_det () = if not (Ezfio.has_determinants_n_det ()) then Ezfio.set_determinants_n_det 1 ; Ezfio.get_determinants_n_det () |> Det_number.of_int ;; let write_n_det n = Det_number.to_int n |> Ezfio.set_determinants_n_det ;; let read_n_states () = if not (Ezfio.has_determinants_n_states ()) then Ezfio.set_determinants_n_states 1 ; Ezfio.get_determinants_n_states () |> States_number.of_int ;; let write_n_states n = States_number.to_int n |> Ezfio.set_determinants_n_states ;; let read_n_states_diag () = if not (Ezfio.has_determinants_n_states_diag ()) then read_n_states () |> States_number.to_int |> Ezfio.set_determinants_n_states_diag ; Ezfio.get_determinants_n_states_diag () |> States_number.of_int ;; let write_n_states_diag ~n_states n = let n_states = States_number.to_int n_states and n = States_number.to_int n in Ezfio.set_determinants_n_states_diag (max n_states n) ;; let read_n_det_max_jacobi () = if not (Ezfio.has_determinants_n_det_max_jacobi ()) then get_default "n_det_max_jacobi" |> Int.of_string |> Ezfio.set_determinants_n_det_max_jacobi ; Ezfio.get_determinants_n_det_max_jacobi () |> Strictly_positive_int.of_int ;; let write_n_det_max_jacobi n = Strictly_positive_int.to_int n |> Ezfio.set_determinants_n_det_max_jacobi ;; let read_threshold_generators () = if not (Ezfio.has_determinants_threshold_generators ()) then get_default "threshold_generators" |> Float.of_string |> Ezfio.set_determinants_threshold_generators ; Ezfio.get_determinants_threshold_generators () |> Threshold.of_float ;; let write_threshold_generators t = Threshold.to_float t |> Ezfio.set_determinants_threshold_generators ;; let read_threshold_selectors () = if not (Ezfio.has_determinants_threshold_selectors ()) then get_default "threshold_selectors" |> Float.of_string |> Ezfio.set_determinants_threshold_selectors ; Ezfio.get_determinants_threshold_selectors () |> Threshold.of_float ;; let write_threshold_selectors t = Threshold.to_float t |> Ezfio.set_determinants_threshold_selectors ;; let read_read_wf () = if not (Ezfio.has_determinants_read_wf ()) then get_default "read_wf" |> Bool.of_string |> Ezfio.set_determinants_read_wf ; Ezfio.get_determinants_read_wf () ;; let write_read_wf = Ezfio.set_determinants_read_wf ;; let read_expected_s2 () = if not (Ezfio.has_determinants_expected_s2 ()) then begin let na = Ezfio.get_electrons_elec_alpha_num () and nb = Ezfio.get_electrons_elec_beta_num () in let s = 0.5 *. (Float.of_int (na - nb)) in Ezfio.set_determinants_expected_s2 ( s *. (s +. 1.) ) end ; Ezfio.get_determinants_expected_s2 () |> Positive_float.of_float ;; let write_expected_s2 s2 = Positive_float.to_float s2 |> Ezfio.set_determinants_expected_s2 ;; let read_s2_eig () = if not (Ezfio.has_determinants_s2_eig ()) then get_default "s2_eig" |> Bool.of_string |> Ezfio.set_determinants_s2_eig ; Ezfio.get_determinants_s2_eig () ;; let write_s2_eig = Ezfio.set_determinants_s2_eig ;; let read_psi_coef () = if not (Ezfio.has_determinants_psi_coef ()) then begin let n_states = read_n_states () |> States_number.to_int in Ezfio.ezfio_array_of_list ~rank:2 ~dim:[| 1 ; n_states |] ~data:(List.init n_states ~f:(fun i -> if (i=0) then 1. else 0. )) |> Ezfio.set_determinants_psi_coef end; Ezfio.get_determinants_psi_coef () |> Ezfio.flattened_ezfio |> Array.map ~f:Det_coef.of_float ;; let write_psi_coef ~n_det ~n_states c = let n_det = Det_number.to_int n_det and c = Array.to_list c |> List.map ~f:Det_coef.to_float and n_states = States_number.to_int n_states in Ezfio.ezfio_array_of_list ~rank:2 ~dim:[| n_det ; n_states |] ~data:c |> Ezfio.set_determinants_psi_coef ;; let read_psi_det () = let n_int = read_n_int () and n_alpha = Ezfio.get_electrons_elec_alpha_num () |> Elec_alpha_number.of_int and n_beta = Ezfio.get_electrons_elec_beta_num () |> Elec_beta_number.of_int in if not (Ezfio.has_determinants_psi_det ()) then begin let mo_tot_num = MO_number.get_max () in let rec build_data accu = function | 0 -> accu | n -> build_data ((MO_number.of_int ~max:mo_tot_num n)::accu) (n-1) in let det_a = build_data [] (Elec_alpha_number.to_int n_alpha) |> Bitlist.of_mo_number_list n_int and det_b = build_data [] (Elec_beta_number.to_int n_beta) |> Bitlist.of_mo_number_list n_int in let data = ( (Bitlist.to_int64_list det_a) @ (Bitlist.to_int64_list det_b) ) in Ezfio.ezfio_array_of_list ~rank:3 ~dim:[| N_int_number.to_int n_int ; 2 ; 1 |] ~data:data |> Ezfio.set_determinants_psi_det ; end ; let n_int = N_int_number.to_int n_int in let psi_det_array = Ezfio.get_determinants_psi_det () in let dim = psi_det_array.Ezfio.dim and data = Ezfio.flattened_ezfio psi_det_array in assert (n_int = dim.(0)); assert (dim.(1) = 2); assert (dim.(2) = (Det_number.to_int (read_n_det ()))); List.init dim.(2) ~f:(fun i -> Array.sub ~pos:(2*n_int*i) ~len:(2*n_int) data) |> List.map ~f:(Determinant.of_int64_array ~n_int:(N_int_number.of_int n_int) ~alpha:n_alpha ~beta:n_beta ) |> Array.of_list ;; let write_psi_det ~n_int ~n_det d = let data = Array.to_list d |> Array.concat |> Array.to_list in Ezfio.ezfio_array_of_list ~rank:3 ~dim:[| N_int_number.to_int n_int ; 2 ; Det_number.to_int n_det |] ~data:data |> Ezfio.set_determinants_psi_det ;; let read () = if (Ezfio.has_mo_basis_mo_tot_num ()) then Some { n_int = read_n_int () ; bit_kind = read_bit_kind () ; mo_label = read_mo_label () ; n_det = read_n_det () ; n_states = read_n_states () ; n_states_diag = read_n_states_diag () ; n_det_max_jacobi = read_n_det_max_jacobi () ; threshold_generators = read_threshold_generators () ; threshold_selectors = read_threshold_selectors () ; read_wf = read_read_wf () ; expected_s2 = read_expected_s2 () ; s2_eig = read_s2_eig () ; psi_coef = read_psi_coef () ; psi_det = read_psi_det () ; } else None ;; let write { n_int ; bit_kind ; mo_label ; n_det ; n_states ; n_states_diag ; n_det_max_jacobi ; threshold_generators ; threshold_selectors ; read_wf ; expected_s2 ; s2_eig ; psi_coef ; psi_det ; } = write_n_int n_int ; write_bit_kind bit_kind; write_mo_label mo_label; write_n_det n_det; write_n_states n_states; write_n_states_diag ~n_states:n_states n_states_diag; write_n_det_max_jacobi n_det_max_jacobi; write_threshold_generators threshold_generators; write_threshold_selectors threshold_selectors; write_read_wf read_wf; write_expected_s2 expected_s2; write_s2_eig s2_eig; write_psi_coef ~n_det:n_det psi_coef ~n_states:n_states; write_psi_det ~n_int:n_int ~n_det:n_det psi_det; ;; let to_rst b = let mo_tot_num = Ezfio.get_mo_basis_mo_tot_num () in let mo_tot_num = MO_number.of_int mo_tot_num ~max:mo_tot_num in let det_text = let nstates = States_number.to_int b.n_states and ndet = Det_number.to_int b.n_det in let coefs_string i = Array.init nstates (fun j -> let ishift = j*ndet in if (ishift < Array.length b.psi_coef) then b.psi_coef.(i+ishift) |> Det_coef.to_float |> Float.to_string else "0." ) |> String.concat_array ~sep:"\t" in Array.init ndet ~f:(fun i -> Printf.sprintf " %s\n%s\n" (coefs_string i) (Determinant.to_string ~mo_tot_num:mo_tot_num b.psi_det.(i) |> String.split ~on:'\n' |> List.map ~f:(fun x -> " "^x) |> String.concat ~sep:"\n" ) ) |> String.concat_array ~sep:"\n" in Printf.sprintf " Read the current wave function :: read_wf = %s Label of the MOs on which the determinants were computed :: mo_label = %s Force the selected wave function to be an eigenfunction of S^2. If true, input the expected value of S^2 :: s2_eig = %s expected_s2 = %s Thresholds on generators and selectors (fraction of the norm) :: threshold_generators = %s threshold_selectors = %s Number of requested states, and number of states used for the Davidson diagonalization :: n_states = %s n_states_diag = %s Maximum size of the Hamiltonian matrix that will be fully diagonalized :: n_det_max_jacobi = %s Number of determinants :: n_det = %s Determinants :: %s " (b.read_wf |> Bool.to_string) (b.mo_label |> MO_label.to_string) (b.s2_eig |> Bool.to_string) (b.expected_s2 |> Positive_float.to_string) (b.threshold_generators |> Threshold.to_string) (b.threshold_selectors |> Threshold.to_string) (b.n_states |> States_number.to_string) (b.n_states_diag |> States_number.to_string) (b.n_det_max_jacobi |> Strictly_positive_int.to_string) (b.n_det |> Det_number.to_string) det_text |> Rst_string.of_string ;; let to_string b = let mo_tot_num = Ezfio.get_mo_basis_mo_tot_num () in let mo_tot_num = MO_number.of_int mo_tot_num ~max:mo_tot_num in Printf.sprintf " n_int = %s bit_kind = %s mo_label = \"%s\" n_det = %s n_states = %s n_states_diag = %s n_det_max_jacobi = %s threshold_generators = %s threshold_selectors = %s read_wf = %s expected_s2 = %s s2_eig = %s psi_coef = %s psi_det = %s " (b.n_int |> N_int_number.to_string) (b.bit_kind |> Bit_kind.to_string) (b.mo_label |> MO_label.to_string) (b.n_det |> Det_number.to_string) (b.n_states |> States_number.to_string) (b.n_states_diag |> States_number.to_string) (b.n_det_max_jacobi |> Strictly_positive_int.to_string) (b.threshold_generators |> Threshold.to_string) (b.threshold_selectors |> Threshold.to_string) (b.read_wf |> Bool.to_string) (b.expected_s2 |> Positive_float.to_string) (b.s2_eig |> Bool.to_string) (b.psi_coef |> Array.to_list |> List.map ~f:Det_coef.to_string |> String.concat ~sep:", ") (b.psi_det |> Array.to_list |> List.map ~f:(Determinant.to_string ~mo_tot_num:mo_tot_num) |> String.concat ~sep:"\n\n") ;; let of_rst r = let r = Rst_string.to_string r in (* Split into header and determinants data *) let idx = String.substr_index_exn r ~pos:0 ~pattern:"\nDeterminants" in let (header, dets) = (String.prefix r idx, String.suffix r ((String.length r)-idx) ) in (* Handle header *) let header = r |> String.split ~on:'\n' |> List.filter ~f:(fun line -> if (line = "") then false else ( (String.contains line '=') && (line.[0] = ' ') ) ) |> List.map ~f:(fun line -> "("^( String.tr line ~target:'=' ~replacement:' ' |> String.strip )^")" ) |> String.concat in (* Handle determinant coefs *) let dets = match ( dets |> String.split ~on:'\n' |> List.map ~f:(String.strip) ) with | _::lines -> lines | _ -> failwith "Error in determinants" in let psi_coef = let rec read_coefs accu = function | [] -> List.rev accu | ""::""::tail -> read_coefs accu tail | ""::c::tail -> let c = String.split ~on:'\t' c |> List.map ~f:(fun x -> Det_coef.of_float (Float.of_string x)) |> Array.of_list in read_coefs (c::accu) tail | _::tail -> read_coefs accu tail in let a = let buffer = read_coefs [] dets in let nstates = List.hd_exn buffer |> Array.length in let extract_state i = let i = i-1 in List.map ~f:(fun x -> Det_coef.to_string x.(i)) buffer |> String.concat ~sep:" " in let rec build_result = function | 1 -> extract_state 1 | i -> (build_result (i-1))^" "^(extract_state i) in build_result nstates in "(psi_coef ("^a^"))" in (* Handle determinants *) let psi_det = let n_alpha = Ezfio.get_electrons_elec_alpha_num () |> Elec_alpha_number.of_int and n_beta = Ezfio.get_electrons_elec_beta_num () |> Elec_beta_number.of_int in let rec read_dets accu = function | [] -> List.rev accu | ""::c::alpha::beta::tail -> begin let alpha = String.rev alpha |> Bitlist.of_string ~zero:'-' ~one:'+' and beta = String.rev beta |> Bitlist.of_string ~zero:'-' ~one:'+' in let newdet = Determinant.of_bitlist_couple ~alpha:n_alpha ~beta:n_beta (alpha,beta) |> Determinant.sexp_of_t |> Sexplib.Sexp.to_string in read_dets (newdet::accu) tail end | _::tail -> read_dets accu tail in let a = read_dets [] dets |> String.concat in "(psi_det ("^a^"))" in let bitkind = Printf.sprintf "(bit_kind %d)" (Lazy.force Qpackage.bit_kind |> Bit_kind.to_int) and n_int = Printf.sprintf "(n_int %d)" (N_int_number.get_max ()) in let s = String.concat [ header ; bitkind ; n_int ; psi_coef ; psi_det] in Generic_input_of_rst.evaluate_sexp t_of_sexp s ;; end