QCaml/CI/F12CI.ml

open Lacaml.D

type t = 
  {
    mo_basis  : MOBasis.t ;
    aux_basis : MOBasis.t ;
    det_space : DeterminantSpace.t ;
    ci        : CI.t ;
    eigensystem : (Mat.t * Vec.t) lazy_t;
  }

let ci t        =  t.ci
let mo_basis t  =  t.mo_basis
let det_space t =  t.det_space
let mo_class t  =  DeterminantSpace.mo_class @@ det_space t
let eigensystem t =  Lazy.force t.eigensystem


let f12_integrals mo_basis =
  let two_e_ints  = MOBasis.f12_ints mo_basis in
  ( (fun _ _ _ -> 0.),
    (fun i j k l s s' ->
       if (i=k && j<>l) || (j=l && i<>k) then
          0.
       else
        begin
          if s' = Spin.other s then
            0.5 *. (F12.get_phys two_e_ints i j k l)
          else
            0.25 *. ((F12.get_phys two_e_ints i j k l) -.
                      (F12.get_phys two_e_ints i j l k))
        end
    ) )




let h_ij mo_basis ki kj =
  let integrals =
    List.map (fun f -> f mo_basis)
      [ CI.h_integrals ]
  in
  CIMatrixElement.make integrals ki kj 
  |> List.hd


let hf_ij mo_basis ki kj =
  let integrals =
    List.map (fun f -> f mo_basis)
      [ CI.h_integrals ; f12_integrals ]
  in
  CIMatrixElement.make integrals ki kj 


let f_ij mo_basis ki kj =
  let integrals =
    List.map (fun f -> f mo_basis)
      [ f12_integrals ]
  in
  CIMatrixElement.make integrals ki kj 
  |> List.hd


let is_internal det_space =
    let mo_class      = DeterminantSpace.mo_class det_space in
    let mo_num = Array.length @@ MOClass.mo_class_array mo_class in
    let m l =
      List.fold_left (fun accu i ->
          let j = i-1 in Bitstring.logor accu (Bitstring.shift_left_one mo_num j)
        ) (Bitstring.zero mo_num) l
    in
    let aux_mask   = m (MOClass.auxiliary_mos mo_class) in
    fun a ->
      let alfa =
        Determinant.alfa a
        |> Spindeterminant.bitstring
      in
      let beta =
        Determinant.beta a
        |> Spindeterminant.bitstring
      in
      let a = Bitstring.logand aux_mask alfa
      and b = Bitstring.logand aux_mask beta
      in
      match Bitstring.popcount a + Bitstring.popcount b with
      | 1 | 2 -> false
      | _ -> true

      (*
      if not (Bitstring.logand aux_mask alfa |> Bitstring.is_zero ) then 
        false
      else
        let beta =
          Determinant.beta a
          |> Spindeterminant.bitstring
        in
        Bitstring.logand aux_mask beta |> Bitstring.is_zero 
      *)


let dressing_vector ~frozen_core aux_basis f12_amplitudes ci =

(*
  let i_o1_alfa = h_ij aux_basis in
  let alfa_o2_i = f_ij aux_basis in

  let w_alfa _ _ = 1. in

  let mo_class = CI.mo_class ci in

  let list_holes = List.concat
    [ MOClass.inactive_mos mo_class ; MOClass.active_mos mo_class ]
  and list_particles2 = MOClass.auxiliary_mos mo_class
  and list_particles1 = List.concat
    [ MOClass.active_mos mo_class ; MOClass.virtual_mos mo_class ; MOClass.auxiliary_mos mo_class ]
  in
  (* Single state here *)
  let result = 
    CI.second_order_sum  ci  list_holes  list_particles1  list_holes  list_particles2
        (is_internal ci.det_space) i_o1_alfa  alfa_o2_i  w_alfa  f12_amplitudes ~unique:false
    |> Vec.of_list
  in
  Matrix.sparse_of_vector_array [| Vector.sparse_of_vec result |]

*)

  Printf.printf "Building matrix\n%!";

  (* Determinants of the FCI space as a list *)
  let in_dets =
    DeterminantSpace.determinants_array ci.CI.det_space 
    |> Array.to_list
  in

  (* Stream that generates only singly and doubly excited determinants
      wrt FCI space *)
  let out_dets_stream =

    (* Stream that generates all determinants of FCI space *)
    let s = 
      DeterminantSpace.fci_of_mo_basis ~frozen_core aux_basis
      |> DeterminantSpace.determinant_stream

    in
    (* Select only singly and doubly excited determinants
        wrt FCI space *)
    Stream.from (fun _ ->
      try
        let rec result () = 
          let ki = Stream.next s in
          if is_internal ci.CI.det_space ki then
              result ()
          else
              Some ki
        in
        result ()
      with Stream.Failure -> None
    )
  in

  let make_h_and_f n = 

    let rec col_vecs_list accu_H accu_F = function
    | 0 -> 
        List.rev accu_H, 
        List.rev accu_F 
    | n ->
      try
        let ki = Stream.next out_dets_stream in
        let h, f = 
          List.map (fun kj ->
          match hf_ij aux_basis ki kj with
          | [ a ; b ] -> a, b
          | _ -> assert false ) in_dets
          |> List.split
        in
        let h = 
          Vec.of_list h
        and f = 
          Vec.of_list f
        in
        col_vecs_list (h::accu_H) (f::accu_F) (n-1)
      with
     | Stream.Failure -> col_vecs_list accu_H accu_F 0
    in
    let h, f = 
      col_vecs_list [] [] n
    in
    Mat.of_col_vecs_list h,
    Mat.of_col_vecs_list f
  in


  Printf.printf "Matrix product\n%!";
  let m_HF = 
    let batch_size =  10_000_000 / (Mat.dim1 f12_amplitudes) in 
    let result = 
      let m_H_aux, m_F_aux = make_h_and_f batch_size in
      gemm m_H_aux m_F_aux ~transb:`T 
    in
    while (Stream.peek out_dets_stream <> None)
    do
      Printf.printf "gemm\n%!";
      let m_H_aux, m_F_aux = make_h_and_f batch_size in
      let hf = 
        gemm m_H_aux m_F_aux ~transb:`T 
      in
      ignore @@ Mat.add result hf ~c:result     
    done;
    result
  in


(*
  let m_HF =

    let in_dets =
      DeterminantSpace.determinants_array ci.CI.det_space 
    in

    let fci_space = DeterminantSpace.fci_of_mo_basis ~frozen_core aux_basis in
    Array.mapi (fun i ki ->
      Printf.printf "%d / %d\r%!"  i (Array.length in_dets);
      Array.map (fun kj ->
        match Determinant.degrees ki kj with
        | (0,0) | (0,1) | (0,2) | (0,3) | (0,4) 
        |         (1,0) | (2,0) | (3,0) | (4,0) 
        |                 (1,1) | (2,1) | (3,1) 
        |                         (1,2) | (1,3) 
        |                                 (2,2) -> 
          ( let x = ref 0. in
            DeterminantSpace.determinant_stream fci_space
            |> Stream.iter (fun k_alfa ->
                if not (is_internal ci.CI.det_space ki) then 
                    let f = f_ij aux_basis k_alfa kj in
                    if f <> 0. then
                      let h = h_ij aux_basis ki k_alfa in
                      x := !x +. f *. h
                );
            !x )
        | _ -> 0.
        ) in_dets
      ) in_dets
    |> Mat.of_array
  in
  *)

  Printf.printf "Done\n%!";
  gemm m_HF f12_amplitudes
  |> Matrix.dense_of_mat





let make ~simulation ?(threshold=1.e-12) ~frozen_core ~mo_basis ~aux_basis_filename () =

  let f12 = Util.of_some @@ Simulation.f12 simulation in
  let mo_num = MOBasis.size mo_basis in

Printf.printf "Add aux basis\n%!";
  (* Add auxiliary basis set *)
  let s =
    let charge        = Charge.to_int @@ Simulation.charge simulation 
    and multiplicity  = Electrons.multiplicity @@ Simulation.electrons simulation
    and nuclei        = Simulation.nuclei simulation
    in
    let general_basis = 
      Basis.general_basis @@ Simulation.basis simulation
    in
    GeneralBasis.combine [
      general_basis ; GeneralBasis.read aux_basis_filename
    ]
    |> Basis.of_nuclei_and_general_basis nuclei
    |> Simulation.make ~f12 ~charge ~multiplicity ~nuclei 
  in

  let aux_basis = 
    MOBasis.of_mo_basis s mo_basis
  in
  let () = 
Printf.printf "F12 ints\n%!";
    ignore @@ MOBasis.f12_ints aux_basis
  in
  let () =
Printf.printf "2e ints\n%!";
    ignore @@ MOBasis.two_e_ints aux_basis
  in

Printf.printf "det space\n%!";
  let det_space = 
    DeterminantSpace.fci_f12_of_mo_basis  aux_basis  ~frozen_core  mo_num
  in

  let ci = CI.make det_space in

  let ci_coef, ci_energy =
    let x = Lazy.force ci.eigensystem in
    Parallel.broadcast (lazy x)
  in


  (*
  let f12_amplitudes = 
    fun c -> 
      let result = lacpy c in
      Mat.scal (0.5) result ; 
      result
    (* While in a sequential region, initiate the parallel
       4-idx transformation to avoid nested parallel jobs
    *)
    ignore @@ MOBasis.f12_ints aux_basis;

    let two_gamma_inv = -2. *. f12.F12.expo_s  in

    let f = fun ki kj ->
      if ki <> kj then
        (f_ij aux_basis ki kj)
      else
        two_gamma_inv +. (f_ij aux_basis ki kj)
    in
    let m_F = 
      CI.create_matrix_spin f det_space
      |> Lazy.force
    in
    fun ci_coef -> 
      let result = 
        Matrix.ax_eq_b m_F (Matrix.dense_of_mat ci_coef)
        |> Matrix.to_mat
      in
      let norm = gemm ~transa:`T result result in
      Printf.printf "Norm of |F> : %f\n%!" norm.{1,1};
      if abs_float norm.{1,1}  > 1. then
        failwith "Norm of |F> > 1";
      result
  in
      *)

  let e_shift = 
    let det =
      DeterminantSpace.determinant_stream det_space 
      |> Stream.next
    in
    h_ij aux_basis det det
  in

  let eigensystem = lazy (
    let m_H =
      Lazy.force ci.CI.m_H
    in

    let rec iteration ?(state=1) psi = 
      let delta = 
         dressing_vector ~frozen_core aux_basis psi ci
      in

      let f = 1.0 /. psi.{1,1} in
      let delta_00 =
        Util.list_range 2 (Mat.dim1 psi)
        |> List.fold_left (fun accu i -> accu +.
            (Matrix.get delta i 1) *. psi.{i,1} *. f) 0. 
      in
      let delta = Matrix.to_mat delta in
      delta.{1,1} <- delta.{1,1} -. delta_00;

(*------
TODO SINGLE STATE HERE
*)
      let n_states = ci.CI.n_states in
      let diagonal =
        Vec.init (Matrix.dim1 m_H) (fun i ->
          if i = 1 then
            Matrix.get m_H i i +. delta.{1,1} *. f 
          else
            Matrix.get m_H i i
          )
      in

      let matrix_prod c = 
        let w = 
          Matrix.mm ~transa:`T m_H c
          |> Matrix.to_mat
        in
        let c11 = Matrix.get c 1 1 in
        Util.list_range 1 (Mat.dim1 w)
        |> List.iter (fun i ->
          let dci = 
            delta.{i,1} *. f ;
          in
          w.{i,1} <- w.{i,1} +. dci *.  c11;
          if (i <> 1) then
            w.{1,1} <- w.{1,1} +. dci *. (Matrix.get c i 1);
          );
        Matrix.dense_of_mat w
      in

      let eigenvectors, eigenvalues = 
        Parallel.broadcast (lazy (
          Davidson.make ~threshold:1.e-6 ~guess:psi ~n_states diagonal matrix_prod
          ))
      in

      let conv =
        1.0 -. abs_float ( dot
          (Mat.to_col_vecs psi).(0)
          (Mat.to_col_vecs eigenvectors).(0) )
      in
      if Parallel.master then
        Printf.printf "F12 Convergence : %e  %f\n" conv (eigenvalues.{1} +. e_shift 
          +. Simulation.nuclear_repulsion simulation);

      if conv > threshold then
        iteration eigenvectors
      else
        let eigenvalues =
          Vec.map (fun x -> x +. e_shift) eigenvalues
        in
        eigenvectors, eigenvalues
    in
    iteration ci_coef

  )
  in
  { mo_basis ; aux_basis ; det_space ; ci ; eigensystem }