QCaml/CI/CI.ml

open Lacaml.D

module Ds = Determinant_space

type t =
  {
    det_space   : Ds.t ;
    m_H         : Matrix.t lazy_t ;
    m_S2        : Matrix.t lazy_t ;
    eigensystem : (Mat.t * Vec.t) lazy_t;
    n_states    : int;
  }

let det_space    t = t.det_space

let n_states     t = t.n_states

let m_H          t = Lazy.force t.m_H

let m_S2         t = Lazy.force t.m_S2

let eigensystem  t = Lazy.force t.eigensystem 

let eigenvectors t = 
  let (x,_) = eigensystem t in x

let eigenvalues t = 
  let (_,x) = eigensystem t in x

let h_integrals mo_basis = 
  let one_e_ints  = MOBasis.one_e_ints mo_basis
  and two_e_ints  = MOBasis.two_e_ints mo_basis
  in
  ( (fun i j _ -> one_e_ints.{i,j}), 
    (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)
      [ h_integrals ]
  in
  CIMatrixElement.make integrals ki kj 
  |> List.hd



let make ?(n_states=1) det_space = 
  let ndet = Ds.size det_space in
  let mo_basis = Ds.mo_basis det_space in

  let m_H = 
    let m_H_arbitrary det = lazy (
      let v = Vec.make0 ndet in
      Array.init ndet (fun i -> let ki = det.(i) in
        Printf.eprintf "%8d / %8d\r%!" i ndet;
        let j = ref 1 in
        Ds.determinant_stream det_space
        |> Stream.iter (fun kj ->
          v.{!j} <- h_ij mo_basis ki kj ; incr j);
        Vector.sparse_of_vec v)
      |> Matrix.sparse_of_vector_array)
    in

    let m_H_spin a b = lazy (
      let n_alfa = Array.length a in
      let n_beta = Array.length b in
      let result = Array.init ndet (fun _ -> []) in

      (** Update function when ki and kj are connected *)
      let update i j ki kj =
        let x = h_ij mo_basis ki kj in
        if x <> 0. then
          result.(i) <- (j, x) :: result.(i) ;
      in

      (** Array of (list of singles, list of doubles) in the beta spin *)
      let degree_bb = 
        Array.map (fun det_i ->
          let deg = Spindeterminant.degree det_i in
          let doubles =
            Array.mapi (fun i det_j -> 
              let d = deg det_j in
              if d < 3 then
                Some (i,d,det_j)
              else
                None
              ) b
            |> Array.to_list
            |> Util.list_some
          in
          let singles = 
            List.filter (fun (i,d,det_j) -> d < 2) doubles
            |> List.map (fun (i,_,det_j) -> (i,det_j))
          in
          let doubles = 
            List.map (fun (i,_,det_j) -> (i,det_j)) doubles
          in
          (singles, doubles)
        ) b
      in
      let a = Array.to_list a
      and b = Array.to_list b
      in
      let i = ref 0 in
      List.iteri (fun ia i_alfa ->
          Printf.eprintf "%8d / %8d\r%!" ia n_alfa;
          let j = ref 1 in
          let deg_a = Spindeterminant.degree i_alfa in
          List.iter (fun j_alfa ->
              let degree_a = deg_a j_alfa in
              begin
                match degree_a with
                  | 2 ->
                      let i' = ref !i in
                      List.iteri (fun ib i_beta ->
                          let ki = Determinant.of_spindeterminants i_alfa i_beta in
                          let kj = Determinant.of_spindeterminants j_alfa i_beta in
                          update !i' (ib + !j) ki kj;
                          incr i';
                        ) b;
                  | 1 -> 
                      let i' = ref !i in
                      List.iteri (fun ib i_beta ->
                          let ki = Determinant.of_spindeterminants i_alfa i_beta in
                          let singles, _ = degree_bb.(ib) in
                          List.iter (fun (j', j_beta) ->
                              let kj = Determinant.of_spindeterminants j_alfa j_beta in
                              update !i' (j' + !j) ki kj
                            ) singles;
                          incr i';
                        ) b;
                  | 0 ->
                      let i' = ref !i in
                      List.iteri (fun ib i_beta ->
                          let ki = Determinant.of_spindeterminants i_alfa i_beta in
                          let _singles, doubles = degree_bb.(ib) in
                          List.iter (fun (j', j_beta) ->
                                let kj = Determinant.of_spindeterminants j_alfa j_beta in
                                update !i' (j' + !j) ki kj
                            ) doubles;
                          incr i';
                        ) b;
                  | _ -> ();
              end;
              j := !j + n_beta
            ) a;
          i := !i + n_beta
        ) a;

      Array.map (fun l ->
        List.sort compare l
        |> Vector.sparse_of_assoc_list ndet ) result
      |> Matrix.sparse_of_vector_array
    )
    in

    match Ds.determinants det_space with
    | Ds.Arbitrary a -> m_H_arbitrary a
    | Ds.Spin  (a,b) -> m_H_spin a b 
  in

  let m_S2 = lazy (
    match Ds.determinants det_space with
    | Ds.Spin  (a,b) -> failwith "Not implemented"
    | Ds.Arbitrary det -> 
      let v = Vec.make0 ndet in
      Array.init ndet (fun i -> let ki = det.(i) in
        Array.iteri (fun j kj ->
          v.{j+1} <- CIMatrixElement.make_s2 ki kj) det;
        Vector.sparse_of_vec v)
      |> Matrix.sparse_of_vector_array
  )
  in

  let eigensystem = lazy (
    let m_H =
      Lazy.force m_H
    in
    let diagonal =
      Vec.init (Matrix.dim1 m_H) (fun i -> Matrix.get m_H i i)
    in
    let matrix_prod psi = 
      Matrix.mm ~transa:`T m_H psi
    in
    Davidson.make ~n_states diagonal matrix_prod
  )
  in
  { det_space ; m_H ; m_S2 ; eigensystem ; n_states }