Merge branch 'master' of github.com:LCPQ/quantum_package

ocaml/Input_determinants.ml

@@ -131,9 +131,11 @@ end = struct
     |> States_number.of_int
   ;;
 
-  let write_n_states_diag n =
-    States_number.to_int n
-    |> Ezfio.set_determinants_n_states_diag
+  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)
   ;;
 
 
@@ -232,20 +234,27 @@ end = struct
 
   let read_psi_coef () =
     if not (Ezfio.has_determinants_psi_coef ()) then
-        Ezfio.ezfio_array_of_list ~rank:1 ~dim:[| 1 |] ~data:[1.]
-        |> Ezfio.set_determinants_psi_coef 
-      ;
+      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 c =
+  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:1 ~dim:[| n_det |] ~data:c
+    Ezfio.ezfio_array_of_list ~rank:2 ~dim:[| n_det ; n_states |] ~data:c
     |> Ezfio.set_determinants_psi_coef 
   ;;
 
@@ -339,14 +348,14 @@ end = struct
      write_mo_label mo_label;
      write_n_det n_det;
      write_n_states n_states;
-     write_n_states_diag n_states_diag;
+     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;
+     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;
   ;;
 

ocaml/qp_create_ezfio_from_xyz.ml

@@ -6,29 +6,85 @@ let spec =
   let open Command.Spec in
   empty 
   +> flag "o"  (optional string)
-     ~doc:"file Name of the created EZFIO file"
+     ~doc:"file Name of the created EZFIO file."
   +> flag "b" (required string)
-     ~doc:"name Basis set."
+     ~doc:"definition of basis set."
   +> flag "c" (optional_with_default 0 int)
      ~doc:"int Total charge of the molecule. Default is 0."
   +> flag "m"  (optional_with_default 1 int)
-     ~doc:"int Spin multiplicity (2S+1) of the molecule. Default is 1"
+     ~doc:"int Spin multiplicity (2S+1) of the molecule. Default is 1."
   +> anon ("xyz_file" %: string)
 ;;
 
 let run ?o b c m xyz_file =
 
-  (* Open basis set channel *)
-  let basis_channel =
-    In_channel.create
-      (Qpackage.root ^ "/data/basis/" ^ (String.lowercase b))
-  in
-
   (* Read molecule *)
   let molecule =
     (Molecule.of_xyz_file xyz_file ~charge:(Charge.of_int c)
       ~multiplicity:(Multiplicity.of_int m) )
   in
+  let nuclei = molecule.Molecule.nuclei in
+
+  let basis_table = Hashtbl.Poly.create () in
+  (* Open basis set channels *)
+  let basis_channel element =
+    let key = Element.to_string element in
+    match Hashtbl.find basis_table key with
+    | Some in_channel -> 
+        in_channel
+    | None ->
+        begin
+         Printf.printf "%s is not defined in basis %s.\nEnter alternate basis : %!"
+          (Element.to_long_string element) b ;
+          let bas =
+             match In_channel.input_line stdin with
+             | Some line -> String.strip line |> String.lowercase
+             | None -> failwith "Aborted"
+          in
+          let new_channel = In_channel.create 
+            (Qpackage.root ^ "/data/basis/" ^ bas)
+          in
+          Hashtbl.add_exn basis_table ~key:key ~data:new_channel;
+          new_channel
+        end
+  in
+  
+  let rec build_basis = function
+  | [] -> ()
+  | elem_and_basis_name :: rest -> 
+    begin
+      match (String.lsplit2 ~on:':' elem_and_basis_name) with
+      | None -> (* Principal basis *)
+        begin
+          let new_channel = In_channel.create
+            (Qpackage.root ^ "/data/basis/" ^ elem_and_basis_name)
+          in
+          List.iter nuclei ~f:(fun x->
+            let key = Element.to_string x.Atom.element
+            in
+            match Hashtbl.add basis_table ~key:key ~data:new_channel with
+            | `Ok -> ()
+            | `Duplicate -> ()
+          )
+        end
+      | Some (key, basis) -> (*Aux basis *)
+        begin
+          let elem  = Element.of_string key
+          and basis = String.lowercase basis
+          in
+          let new_channel = In_channel.create
+            (Qpackage.root ^ "/data/basis/" ^ basis)
+          in
+          match Hashtbl.add basis_table ~key:key ~data:new_channel with
+          | `Ok -> ()
+          | `Duplicate -> failwith ("Duplicate definition of basis for "^(Element.to_long_string elem))
+        end
+    end;
+    build_basis rest
+  in
+  String.split ~on:' ' b
+  |> List.rev_map ~f:String.strip 
+  |> build_basis;
 
   (* Build EZFIO File name *)
   let ezfio_file =
@@ -54,7 +110,6 @@ let run ?o b c m xyz_file =
     molecule.Molecule.elec_beta  ) ;
 
   (* Write Nuclei *)
-  let nuclei = molecule.Molecule.nuclei in
   let labels =
      List.map ~f:(fun x->Element.to_string x.Atom.element) nuclei
   and charges = 
@@ -74,17 +129,31 @@ let run ?o b c m xyz_file =
 
   (* Write Basis set *)
   let basis =
+
     let nmax = Nucl_number.get_max () in
     let rec do_work (accu:(Atom.t*Nucl_number.t) list) (n:int) = function
     | [] -> accu
-    | e::tail -> let new_accu = (e,(Nucl_number.of_int ~max:nmax n))::accu in
+    | e::tail ->
+      let new_accu = (e,(Nucl_number.of_int ~max:nmax n))::accu in
       do_work new_accu (n+1) tail
     in
-    do_work [] 1  nuclei
+    let result = do_work [] 1  nuclei
     |> List.rev
     |> List.map ~f:(fun (x,i) ->
-       Basis.read_element basis_channel i x.Atom.element) 
+       try 
+         Basis.read_element (basis_channel x.Atom.element) i x.Atom.element
+       with
+       | End_of_file -> 
+         begin
+           let alt_channel = basis_channel x.Atom.element in
+           Basis.read_element alt_channel i x.Atom.element 
+         end
+       | _ -> assert false
+       ) 
     |> List.concat
+    in
+    (* close all in_channels *)
+    result
   in
   let long_basis = Long_basis.of_basis basis in
   let ao_num = List.length long_basis in
@@ -150,8 +219,14 @@ let run ?o b c m xyz_file =
 let command = 
     Command.basic 
     ~summary: "Quantum Package command"
-    ~readme:(fun () ->
-      "Creates an EZFIO directory from a standard xyz file
+    ~readme:(fun () -> "
+Creates an EZFIO directory from a standard xyz file.
+The basis set is defined as a single string if all the
+atoms are taken from the same basis set, otherwise specific
+elements can be defined as follows:
+
+ -b \"cc-pcvdz H:cc-pvdz C:6-31g\"
+
       ")
     spec
     (fun o b c m xyz_file () ->