9
1
mirror of https://github.com/QuantumPackage/qp2.git synced 2024-06-13 16:15:18 +02:00
qp2/ocaml/Input_determinants_by_hand.ml
Anthony Scemama b25489e14c
Some checks failed
continuous-integration/drone/push Build is failing
Fix modifying determinants before mo_label exists
2023-11-27 14:25:05 +01:00

702 lines
19 KiB
OCaml

open Qptypes;;
open Qputils;;
open Sexplib.Std;;
module Determinants_by_hand : sig
type t =
{ n_int : N_int_number.t;
bit_kind : Bit_kind.t;
n_det : Det_number.t;
n_det_qp_edit : Det_number.t;
n_states : States_number.t;
expected_s2 : Positive_float.t;
psi_coef : Det_coef.t array;
psi_det : Determinant.t array;
state_average_weight : Positive_float.t array;
mo_label : MO_label.t;
} [@@deriving sexp]
val read : ?full:bool -> unit -> t option
val write : ?force:bool -> t -> unit
val to_string : t -> string
val to_rst : t -> Rst_string.t
val of_rst : Rst_string.t -> t option
val read_n_int : unit -> N_int_number.t
val update_ndet : Det_number.t -> unit
val extract_state : States_number.t -> unit
val extract_states : Range.t -> unit
end = struct
type t =
{ n_int : N_int_number.t;
bit_kind : Bit_kind.t;
n_det : Det_number.t;
n_det_qp_edit : Det_number.t;
n_states : States_number.t;
expected_s2 : Positive_float.t;
psi_coef : Det_coef.t array;
psi_det : Determinant.t array;
state_average_weight : Positive_float.t array;
mo_label : MO_label.t;
} [@@deriving sexp]
;;
let get_default = Qpackage.get_ezfio_default "determinants";;
let read_mo_label () =
if not (Ezfio.has_determinants_mo_label ()) then
if Ezfio.has_mo_basis_mo_label () then (
let label = Ezfio.get_mo_basis_mo_label () in
Ezfio.set_determinants_mo_label label) ;
Ezfio.get_determinants_mo_label ()
|> MO_label.of_string
;;
let read_n_int () =
if not (Ezfio.has_determinants_n_int()) then
Ezfio.get_mo_basis_mo_num ()
|> Bitlist.n_int_of_mo_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_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 read_n_det_qp_edit () =
if not (Ezfio.has_determinants_n_det_qp_edit ()) then
begin
let n_det = read_n_det () |> Det_number.to_int in
Ezfio.set_determinants_n_det_qp_edit n_det
end;
Ezfio.get_determinants_n_det_qp_edit ()
|> Det_number.of_int
;;
let write_n_det n =
Det_number.to_int n
|> Ezfio.set_determinants_n_det
;;
let write_n_det_qp_edit n =
let n_det = read_n_det () |> Det_number.to_int in
min n_det (Det_number.to_int n)
|> Ezfio.set_determinants_n_det_qp_edit
;;
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 =
let n_states =
States_number.to_int n
in
(*
let old_nstates, read_wf =
Ezfio.get_determinants_n_states (),
Ezfio.get_determinants_read_wf ()
in
if read_wf && old_nstates <> n_states then
Printf.eprintf "Warning : n_states could not be changed because read_wf is true\n%!"
else
*)
begin
Ezfio.set_determinants_n_states n_states;
let data =
Array.make n_states 1.
|> Array.to_list
in
Ezfio.ezfio_array_of_list ~rank:1 ~dim:[| n_states |] ~data
|> Ezfio.set_determinants_state_average_weight
end
;;
let write_state_average_weight data =
let n_states =
read_n_states ()
|> States_number.to_int
in
let data =
Array.map Positive_float.to_float data
|> Array.to_list
in
Ezfio.ezfio_array_of_list ~rank:1 ~dim:[| n_states |] ~data
|> Ezfio.set_determinants_state_average_weight
;;
let read_state_average_weight () =
let n_states =
read_n_states ()
|> States_number.to_int
in
if not (Ezfio.has_determinants_state_average_weight ()) then
begin
let data =
Array.init n_states (fun _ -> 1./.(float_of_int n_states))
|> Array.map Positive_float.of_float
in
write_state_average_weight data
end;
let result =
Ezfio.get_determinants_state_average_weight ()
|> Ezfio.flattened_ezfio
|> Array.map Positive_float.of_float
in
if Array.length result = n_states then
result
else
let data =
Array.init n_states (fun _ -> 1./.(float_of_int n_states))
|> Array.map Positive_float.of_float
in
(write_state_average_weight data; data)
;;
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_psi_coef ~read_only () =
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 (fun i -> if (i=0) then 1. else 0. ))
|> Ezfio.set_determinants_psi_coef
end;
begin
if read_only then
Ezfio.get_determinants_psi_coef_qp_edit ()
else
Ezfio.get_determinants_psi_coef ()
end
|> Ezfio.flattened_ezfio
|> Array.map 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.map Det_coef.to_float c
|> Array.to_list
and n_states =
States_number.to_int n_states
in
let r =
Ezfio.ezfio_array_of_list ~rank:2 ~dim:[| n_det ; n_states |] ~data:c
in
Ezfio.set_determinants_psi_coef r;
Ezfio.set_determinants_psi_coef_qp_edit r
;;
let read_psi_det ~read_only () =
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_num = MO_number.get_max () in
let rec build_data accu = function
| 0 -> accu
| n -> build_data ((MO_number.of_int ~max:mo_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 =
if read_only then
Ezfio.get_determinants_psi_det_qp_edit ()
else
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);
if read_only then
assert (dim.(2) = (Det_number.to_int (read_n_det_qp_edit ())))
else
assert (dim.(2) = (Det_number.to_int (read_n_det ())));
Array.init dim.(2) (fun i ->
Array.sub data (2*n_int*i) (2*n_int) )
|> Array.map (Determinant.of_int64_array
~n_int:(N_int_number.of_int n_int)
~alpha:n_alpha ~beta:n_beta )
;;
let write_psi_det ~n_int ~n_det d =
let data = Array.to_list d
|> Array.concat
|> Array.to_list
in
let r =
Ezfio.ezfio_array_of_list ~rank:3 ~dim:[| N_int_number.to_int n_int ; 2 ; Det_number.to_int n_det |] ~data:data
in
Ezfio.set_determinants_psi_det r;
Ezfio.set_determinants_psi_det_qp_edit r
;;
let write_mo_label a =
MO_label.to_string a
|> Ezfio.set_determinants_mo_label
let read ?(full=true) () =
let n_det_qp_edit = read_n_det_qp_edit () in
let n_det = read_n_det () in
let read_only =
if full then false else n_det_qp_edit <> n_det
in
if (Ezfio.has_mo_basis_mo_num ()) then
try
Some
{ n_int = read_n_int () ;
bit_kind = read_bit_kind () ;
n_det = read_n_det () ;
n_det_qp_edit = read_n_det_qp_edit () ;
expected_s2 = read_expected_s2 () ;
psi_coef = read_psi_coef ~read_only () ;
psi_det = read_psi_det ~read_only () ;
n_states = read_n_states () ;
state_average_weight = read_state_average_weight () ;
mo_label = read_mo_label () ;
}
with _ -> None
else
(* No molecular orbitals, so no determinants *)
None
;;
let write ?(force=false)
{ n_int ;
bit_kind ;
n_det ;
n_det_qp_edit ;
expected_s2 ;
psi_coef ;
psi_det ;
n_states ;
state_average_weight ;
mo_label ;
} =
write_n_int n_int ;
write_bit_kind bit_kind;
write_n_det n_det;
write_n_states n_states;
write_expected_s2 expected_s2;
if force || (n_det <= n_det_qp_edit) then
begin
write_n_det_qp_edit n_det;
write_psi_coef ~n_det:n_det ~n_states:n_states psi_coef ;
write_psi_det ~n_int:n_int ~n_det:n_det psi_det
end;
write_state_average_weight state_average_weight ;
write_mo_label mo_label ;
()
;;
let to_rst b =
let max =
Ezfio.get_mo_basis_mo_num ()
in
let mo_num =
MO_number.of_int ~max max
in
let det_text =
let nstates =
read_n_states ()
|> States_number.to_int
and ndet_qp_edit =
Det_number.to_int b.n_det_qp_edit
in
let coefs_string i =
Array.init nstates (fun j ->
let ishift =
j*ndet_qp_edit
in
if (ishift < Array.length b.psi_coef) then
b.psi_coef.(i+ishift)
|> Det_coef.to_float
|> Float.to_string
else
"0."
)
|> Array.to_list |> String.concat "\t"
in
Array.init ndet_qp_edit (fun i ->
Printf.sprintf " %s\n%s\n"
(coefs_string i)
(Determinant.to_string ~mo_num:mo_num b.psi_det.(i)
|> String_ext.split ~on:'\n'
|> list_map (fun x -> " "^x)
|> String.concat "\n"
)
)
|> Array.to_list
|> String.concat "\n"
in
Printf.sprintf "
Force the selected wave function to be an eigenfunction of S^2.
If true, input the expected value of S^2 ::
expected_s2 = %s
Number of determinants ::
n_det = %s
State average weights ::
state_average_weight = (%s)
Determinants ::
%s
"
(b.expected_s2 |> Positive_float.to_string)
(b.n_det |> Det_number.to_string)
(b.state_average_weight |> Array.map Positive_float.to_string |> Array.to_list |> String.concat "\t")
det_text
|> Rst_string.of_string
;;
let to_string b =
let mo_num = Ezfio.get_mo_basis_mo_num () in
let mo_num = MO_number.of_int mo_num ~max:mo_num in
Printf.sprintf "
n_int = %s
bit_kind = %s
n_det = %s
n_states = %s
expected_s2 = %s
state_average_weight = %s
psi_coef = %s
psi_det = %s
"
(b.n_int |> N_int_number.to_string)
(b.bit_kind |> Bit_kind.to_string)
(b.n_det |> Det_number.to_string)
(b.n_states |> States_number.to_string)
(b.expected_s2 |> Positive_float.to_string)
(b.state_average_weight |> Array.to_list |> list_map Positive_float.to_string |> String.concat ",")
(b.psi_coef |> Array.map Det_coef.to_string |> Array.to_list
|> String.concat ", ")
(b.psi_det |> Array.map (Determinant.to_string ~mo_num) |> Array.to_list
|> String.concat "\n\n")
;;
let of_rst r =
let r = Rst_string.to_string r
in
(* Split into header and determinants data *)
let idx =
match String_ext.substr_index r ~pos:0 ~pattern:"\nDeterminants" with
| Some x -> x
| None -> assert false
in
let (header, dets) =
(String.sub r 0 idx, String.sub r idx (String.length r - idx) )
in
(* Handle header *)
let header = r
|> String_ext.split ~on:'\n'
|> List.filter (fun line ->
if (line = "") then
false
else
( (String.contains line '=') && (line.[0] = ' ') )
)
|> list_map (fun line ->
"("^(
String_ext.tr line ~target:'=' ~replacement:' '
|> String.trim
)^")" )
|> String.concat ""
in
(* Handle determinant coefs *)
let dets = match ( dets
|> String_ext.split ~on:'\n'
|> list_map String.trim
) 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_ext.split ~on:'\t' c
|> list_map (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 buffer
|> Array.length
in
let extract_state i =
let i =
i-1
in
list_map (fun x -> Det_coef.to_string x.(i)) buffer
|> String.concat " "
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_int = N_int_number.of_int @@ (MO_number.get_max () - 1) / 64 + 1 in
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
| ""::_::alpha::beta::tail ->
begin
let newdet =
(Bitlist.of_string ~zero:'-' ~one:'+' alpha ,
Bitlist.of_string ~zero:'-' ~one:'+' beta)
|> Determinant.of_bitlist_couple ~n_int ~alpha:n_alpha ~beta:n_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 mo_label =
Printf.sprintf "(mo_label %s)" (MO_label.to_string @@ read_mo_label ())
and n_int =
Printf.sprintf "(n_int %d)" (N_int_number.get_max ())
and n_states =
Printf.sprintf "(n_states %d)" (States_number.to_int @@ read_n_states ())
and n_det_qp_edit =
Printf.sprintf "(n_det_qp_edit %d)" (Det_number.to_int @@ read_n_det_qp_edit ())
in
let s =
String.concat "" [ header ; mo_label ; bitkind ; n_int ; n_states ; psi_coef ; psi_det ; n_det_qp_edit ]
in
Generic_input_of_rst.evaluate_sexp t_of_sexp s
;;
let update_ndet n_det_new =
Printf.printf "Reducing n_det to %d\n" (Det_number.to_int n_det_new);
let n_det_new =
Det_number.to_int n_det_new
in
let det =
match read () with
| Some x -> x
| None -> failwith "No determinants in file"
in
let n_det_old, n_states =
Det_number.to_int det.n_det,
States_number.to_int det.n_states
in
if n_det_new = n_det_old then
()
;
if n_det_new > n_det_new then
failwith @@ Printf.sprintf "Requested n_det should be less than %d" n_det_old
;
for j=0 to (n_states-1) do
let ishift_old, ishift_new =
j*n_det_old,
j*n_det_new
in
for i=0 to (n_det_new-1) do
det.psi_coef.(i+ishift_new) <- det.psi_coef.(i+ishift_old)
done
done
;
let new_det =
{ det with n_det = (Det_number.of_int n_det_new) }
in
write ~force:true new_det
;;
let extract_state istate =
Printf.printf "Extracting state %d\n" (States_number.to_int istate);
let det =
match read () with
| Some x -> x
| None -> failwith "No determinants in file"
in
let n_det, n_states =
Det_number.to_int det.n_det,
States_number.to_int det.n_states
in
if (States_number.to_int istate) > n_states then
failwith "State to extract should not be greater than n_states"
;
let j =
(States_number.to_int istate) - 1
in
begin
if (j>0) then
let ishift =
j*n_det
in
for i=0 to (n_det-1) do
det.psi_coef.(i) <- det.psi_coef.(i+ishift)
done
end;
let new_det =
{ det with n_states = (States_number.of_int 1) }
in
write ~force:true new_det
;;
let extract_states range =
Printf.printf "Extracting states %s\n" (Range.to_string range);
let det =
match read () with
| Some x -> x
| None -> failwith "No determinants in file"
in
let n_det, n_states =
Det_number.to_int det.n_det,
States_number.to_int det.n_states
in
Range.to_int_list range
|> List.iter (fun istate ->
if istate > n_states then
failwith "State to extract should not be greater than n_states")
;
let sorted_list =
Range.to_int_list range
|> List.sort compare
in
let state_shift = ref 0 in
List.iter (fun istate ->
let j =
istate - 1
in
begin
if (j>0) then
let ishift =
j*n_det
in
for i=0 to (n_det-1) do
det.psi_coef.(!state_shift+i) <-
det.psi_coef.(i+ishift)
done
; Printf.printf "OK\n%!" ;
end;
state_shift := !state_shift + n_det
) sorted_list
;
let new_det =
{ det with n_states = (States_number.of_int @@ List.length sorted_list) }
in
write ~force:true new_det
;;
end