quantum_package/ocaml/Input_determinants_by_hand.ml

open Qptypes;;
open Qputils;;
open Core.Std;;

module Determinants_by_hand : sig
  type t = 
    { n_int                  : N_int_number.t;
      bit_kind               : Bit_kind.t;
      n_det                  : Det_number.t;
      expected_s2            : Positive_float.t;
      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;
      n_det                  : Det_number.t;
      expected_s2            : Positive_float.t;
      psi_coef               : Det_coef.t array;
      psi_det                : Determinant.t array;
    } with sexp
  ;;

  let get_default = Qpackage.get_ezfio_default "determinants";;

  let n_det_read_max = 50_000_000 ;;

  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_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_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 () =
    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 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 = 
      read_n_states () |> States_number.to_int
    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 alpha =
        Ezfio.get_electrons_elec_alpha_num ()
        |> Elec_alpha_number.of_int 
      and 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 alpha)
                |> Bitlist.of_mo_number_list n_int
            and det_b =
                build_data [] (Elec_beta_number.to_int 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_i =
          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_i = dim.(0));
      assert (dim.(1) = 2);
      assert (dim.(2) = (Det_number.to_int (read_n_det ())));
      let len = 
        2 * n_int_i
      in
      Array.init dim.(2) ~f:(fun i ->
         Array.sub ~pos:(len * i) ~len data
         |> Determinant.of_int64_array ~n_int ~alpha ~beta 
      )
  ;;

  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
      let n_det = 
         read_n_det ()
      in
      if ( (Det_number.to_int n_det) > n_det_read_max ) then
        None
      else
        Some
        { n_int                  = read_n_int ()                ;
          bit_kind               = read_bit_kind ()             ;
          n_det                  = read_n_det ()                ;
          expected_s2            = read_expected_s2 ()          ;
          psi_coef               = read_psi_coef ()             ;
          psi_det                = read_psi_det ()              ;
        }
    else
      None
  ;;

  let write { n_int                ;
              bit_kind             ;
              n_det                ;
              expected_s2          ;
              psi_coef             ;
              psi_det              ;
            } =
     write_n_int n_int ;
     write_bit_kind bit_kind;
     write_n_det n_det;
     write_expected_s2 expected_s2;
     write_psi_coef ~n_det:n_det psi_coef ;
     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 =
        read_n_states () |> States_number.to_int
      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 "
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

Determinants ::

%s
"
     (b.expected_s2   |> Positive_float.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
n_det                  = %s
expected_s2            = %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.expected_s2   |> Positive_float.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 -> x.(i)) buffer
        in
        let rec build_result accu = function
        | 0 -> accu
        | i -> 
            let new_accu =
               (extract_state i) :: accu
            in
            build_result new_accu (i-1)
        in
        build_result [] nstates 
      in
      List.concat a
      |> Array.of_list
    in

(*
    let dets = match ( dets
      |> String.split ~on:'\n'
      |> List.map ~f:(String.strip)
    ) with 
    | _::lines -> lines 
    | _ -> failwith "Error in determinants"
    in
*)

    (* Handle determinants *)
    let psi_det = 
      let alpha = Ezfio.get_electrons_elec_alpha_num ()
        |> Elec_alpha_number.of_int 
      and beta = Ezfio.get_electrons_elec_beta_num ()
        |> Elec_beta_number.of_int 
      and n_int =
        N_int_number.get_max ()
        |> N_int_number.of_int
      in
      let rec read_dets accu = function
      | [] -> List.rev  accu
      | ""::_::alpha_str::beta_str::tail -> 
          begin
            let newdet =
               (Bitlist.of_string_mp alpha_str, Bitlist.of_string_mp beta_str)
               |> Determinant.of_bitlist_couple ~n_int ~alpha ~beta
            in
            read_dets (newdet::accu) tail
          end
      | _::tail -> read_dets accu tail
      in
      let dets = 
        List.map ~f:String.rev dets
      in
      read_dets [] dets
      |> Array.of_list
    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 =
      [ header ; bitkind ; n_int ; "(psi_coef ())" ; "(psi_det ())"]
      |> String.concat 
    in
    let result = 
      Generic_input_of_rst.evaluate_sexp t_of_sexp s
    in
    match result with
    | Some x -> Some { x with psi_coef ; psi_det }
    | None -> None
  ;;

end