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;
    f12_amplitudes : Mat.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 f12_integrals mo_basis =
  let two_e_ints  = MOBasis.two_e_ints mo_basis in
  ( (fun i j _ -> 0.),
    (fun i j k l s s' ->
       if s' = Spin.other s then
         ERI.get_phys two_e_ints i j k l
       else
         (ERI.get_phys two_e_ints i j k l) -.
         (ERI.get_phys two_e_ints i j l k)
    ) )




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 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 dressing_vector f12_amplitudes ci =

  let mo_basis  = DeterminantSpace.mo_basis ci.CI.det_space in

  let i_o1_alfa = h_ij mo_basis in
  let alfa_o2_i = f_ij mo_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_particles1 = MOClass.auxiliary_mos mo_class
  and list_particles2 = 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
        i_o1_alfa  alfa_o2_i  w_alfa  f12_amplitudes
    |> Vec.of_list
  in
  Matrix.sparse_of_vector_array [| Vector.sparse_of_vec result |]






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

  let mo_num = MOBasis.size mo_basis in

  (* 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 ~charge ~multiplicity ~nuclei 
  in

  let aux_basis = 
    MOBasis.of_mo_basis s mo_basis
  in

  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 = Parallel.broadcast ci.eigensystem in

  let f12_amplitudes = 
    (* While in a sequential region, initiate the parallel
       4-idx transformation to avoid nested parallel jobs
    *)
    ignore @@ MOBasis.f12_ints mo_basis;

    let f = fun ki kj ->
      if ki <> kj then
        f_ij mo_basis ki kj
      else
        f_ij mo_basis ki kj +. 1.
    in
    let m_F = 
      CI.create_matrix_spin f det_space
      |> Lazy.force
    in
    Matrix.ax_eq_b m_F (Matrix.dense_of_mat ci_coef)
    |> Matrix.to_mat
  in

  let f_11, e_shift = 
    let det =
      DeterminantSpace.determinant_stream det_space 
      |> Stream.next
    in
    f_ij mo_basis det det,
    h_ij mo_basis det det
  in

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

    let rec iteration ?(state=1) psi = 
      let diagonal =
        Vec.init (Matrix.dim1 m_H) (fun i -> Matrix.get m_H i i +. if i=1 then f_11 /. (Matrix.get psi 1 state) else 0. )
      in
      let matrix_prod psi = 
        Matrix.add
          (Matrix.mm ~transa:`T m_H psi)
          (dressing_vector f12_amplitudes ci)
      in
      let eigenvectors, eigenvalues = 
        Parallel.broadcast (lazy (
            Davidson.make ~threshold:1.e-6 ~guess:(Matrix.to_mat psi) ~n_states diagonal matrix_prod
          ))
      in
      let m =
        Matrix.mm ~transa:`T psi (Matrix.dense_of_mat eigenvectors)
        |> Matrix.to_mat
      in
      let conv = Mat.sum m -. (Vec.sum (Mat.copy_diag m)) in
      Printf.printf "Convergence : %f  %f\n" conv eigenvalues.{1};
      if conv > threshold then
        iteration (Matrix.dense_of_mat eigenvectors)
      else
        let eigenvalues =
          Vec.map (fun x -> x +. e_shift) eigenvalues
        in
        eigenvectors, eigenvalues
    in
    iteration (Matrix.dense_of_mat ci_coef)
  )
  in
  { mo_basis ; aux_basis ; det_space ; ci ; f12_amplitudes ; eigensystem }