Can extract multiple states

TODO

@@ -48,10 +48,10 @@ Refaire les benchmarks
 # Documentation de /etc
 
 # Toto
-Selection d'etats avec qp_edit --state
 singles_alpha_csc_idx, singles_alpha_size  | git diff 
 Environment variables dans qp_run (prefix,etc)
 
 Si un provider est un programme, generer une page a lui tout seul avec le man
 
 Options obligatoires dans Command_line.ml
+Documentation des commandes a updater

ocaml/Input_determinants_by_hand.ml

@@ -22,6 +22,7 @@ module Determinants_by_hand : sig
   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;
@@ -592,6 +593,47 @@ psi_det                = %s
     write new_det
   ;;
 
+  let extract_states range =
+    Printf.printf "Extracting states %s\n" (Range.to_string range);
+    let det =
+      read ()
+    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 ~f:(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 ~f:(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
+      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 new_det
+  ;;
+
 end
 
 

ocaml/Range.ml

@@ -14,6 +14,8 @@ open Sexplib.Std
 
 type t = int list [@@deriving sexp] 
 
+let to_int_list r = r
+
 let expand_range r =
   match String_ext.lsplit2 ~on:'-' r with
   | Some (s, f) ->
@@ -50,6 +52,10 @@ let of_string s =
 
 
 let to_string l =
+  "[" ^ 
+  (List.map string_of_int l
+   |> String.concat ",") ^ "]"
+(*
   let rec do_work buf symbol = function
     | [] -> buf
     | a::([] as t) -> 
@@ -58,15 +64,15 @@ let to_string l =
         if (b-a = 1) then
           do_work buf "-" t
         else
-          do_work (buf^symbol^(string_of_int a)^","^(string_of_int b)) "" t
+          do_work (buf^symbol^","^(string_of_int b)) "" t
   in
   let result = 
     match l with
-    | [] ->
-        "[]"
+    | [] -> "[]"
     | h::t  ->
         do_work ("["^(string_of_int h)) "" l in
   (String.sub result 0 ((String.length result)))^"]"
+  *)
 
 
 let test_module () =

ocaml/Range.mli

@@ -8,3 +8,4 @@ type t = int list [@@deriving sexp]
  *)
 val of_string : string -> t
 val to_string : t -> string
+val to_int_list : t -> int list

scripts/ezfio_interface/qp_edit_template

@@ -169,9 +169,9 @@ let run check_only ?ndet ?state ezfio_filename =
   begin
     match state with
     | None -> ()
-    | Some n -> 
+    | Some range -> 
        begin
-         Input.Determinants_by_hand.extract_state (States_number.of_int n)
+         Input.Determinants_by_hand.extract_states range
        end
   end;
 
@@ -273,7 +273,7 @@ let () =
   [ ( 'c', "check", "Checks the input data", Command_line.Without_arg);
     ( 'n', "ndet", "<int> Truncate the wavefunction to the target number of determinants",
       Command_line.With_arg);
-    ( 's', "state", "<int> Pick the state as a new wavefunction", Command_line.With_arg);
+    ( 's', "state", "<range> Extract selected states, for example \"[1,3-5]\"", Command_line.With_arg);
     Command_line.anonymous "<EZFIO_FILE>" "EZFIO directory";
   ]
   |> Command_line.set_specs ;
@@ -292,8 +292,8 @@ let () =
   let state =
     match Command_line.get "state" with
     | None -> None
-    | Some s -> (try Some (int_of_string s)
-            with _ -> failwith "[-s|--state] expects an integer")
+    | Some s -> (try Some (Range.of_string s)
+            with _ -> failwith "[-s|--state] expects a range")
   in
 
   let c =