QCaml/Basis/AOBasis.ml

open Lacaml.D
open Util

type t =
{
  basis              :  Basis.t                     ;
  overlap            :  Overlap.t             lazy_t;
  multipole          :  Multipole.t           lazy_t;
  ortho              :  Orthonormalization.t  lazy_t;
  eN_ints            :  NucInt.t              lazy_t;
  kin_ints           :  KinInt.t              lazy_t;
  ee_ints            :  ERI.t                 lazy_t;
  ee_lr_ints         :  ERI_lr.t              lazy_t;
  f12_ints           :  F12.t                 lazy_t;
  f12_over_r12_ints  :  ScreenedERI.t         lazy_t;
  cartesian          :  bool                        ;
}

let basis        t = t.basis
let overlap      t = Lazy.force t.overlap
let multipole    t = Lazy.force t.multipole
let ortho        t = Lazy.force t.ortho
let eN_ints      t = Lazy.force t.eN_ints
let kin_ints     t = Lazy.force t.kin_ints
let ee_ints      t = Lazy.force t.ee_ints
let ee_lr_ints   t = Lazy.force t.ee_lr_ints
let f12_ints     t = Lazy.force t.f12_ints
let f12_over_r12_ints t = Lazy.force t.f12_over_r12_ints
let cartesian    t = t.cartesian

module Cs = ContractedShell

let values t point =
  let result = Vec.create (Basis.size t.basis) in
  Array.iter (fun shell ->
      Cs.values shell point
      |> Array.iteri
           (fun i_c value ->
             let i = Cs.index shell + i_c + 1 in
             result.{i} <- value)
    ) (Basis.contracted_shells t.basis);
  result

  
let make ~cartesian ~basis ?f12 nuclei =
  
  let overlap =
    lazy (
        Overlap.of_basis basis
      ) in
  
  let ortho =
    lazy (
        Orthonormalization.make ~cartesian ~basis (Lazy.force overlap)
      ) in

  let eN_ints  =
    lazy (
        NucInt.of_basis_nuclei ~basis nuclei
      ) in

  let kin_ints =
    lazy (
        KinInt.of_basis basis
      ) in

  let ee_ints  =
    lazy (
        ERI.of_basis basis
      ) in

  let ee_lr_ints  =
    lazy (
        ERI_lr.of_basis basis
      ) in

  let f12_ints =
    lazy (
        F12.of_basis basis
      ) in

  let f12_over_r12_ints =
    lazy (
        ScreenedERI.of_basis basis
      ) in

  let multipole =
    lazy (
        Multipole.of_basis basis
      ) in

  { basis ; overlap ; multipole ; ortho ; eN_ints ; kin_ints ; ee_ints ;
    ee_lr_ints ; f12_ints ; f12_over_r12_ints ; cartesian ;
  }



let test_case name t =

  let check_matrix title a r =
    let a = Mat.to_array a in
    Mat.to_array r
    |> Array.iteri (fun i x ->
          let message =
            Printf.sprintf "%s line %d" title i
          in
          Alcotest.(check (array (float 1.e-10))) message a.(i) x
        )
  in

  let check_eri a r =
    let f { ERI.i_r1 ; j_r2 ; k_r1 ; l_r2 ; value } =
      (i_r1, (j_r2, (k_r1, (l_r2, value))))
    in
    let a = ERI.to_list a |> List.rev_map f |> List.rev in
    let r = ERI.to_list r |> List.rev_map f |> List.rev in
    Alcotest.(check (list (pair int (pair int (pair int (pair int (float 1.e-10))))))) "ERI" a r
  in

  let check_eri_lr a r =
    let f { ERI_lr.i_r1 ; j_r2 ; k_r1 ; l_r2 ; value } =
      (i_r1, (j_r2, (k_r1, (l_r2, value))))
    in
    let a = ERI_lr.to_list a |> List.rev_map f |> List.rev in
    let r = ERI_lr.to_list r |> List.rev_map f |> List.rev in
    Alcotest.(check (list (pair int (pair int (pair int (pair int (float 1.e-10))))))) "ERI_lr" a r
  in

  let test_overlap () =
    let reference =
      sym_matrix_of_file ("test_files/"^name^"_overlap.ref")
    in
    let overlap =
      Lazy.force t.overlap |> Overlap.matrix
    in
    check_matrix "Overlap" overlap reference
  in

  let test_eN_ints () =
    let reference =
      sym_matrix_of_file ("test_files/"^name^"_nuc.ref")
    in
    let eN_ints =
      Lazy.force t.eN_ints |> NucInt.matrix
    in
    check_matrix "eN_ints" eN_ints reference
  in

  let test_kin_ints () =
    let reference =
      sym_matrix_of_file ("test_files/"^name^"_kin.ref")
    in
    let kin_ints =
      Lazy.force t.kin_ints |> KinInt.matrix
    in
    check_matrix "kin_ints" kin_ints reference
  in

  let test_ee_ints () =
    let reference =
      ERI.of_file ("test_files/"^name^"_eri.ref") ~sparsity:`Dense ~size:(Basis.size t.basis)
    in
    let ee_ints =
      Lazy.force t.ee_ints
    in
    check_eri ee_ints reference
    ;
  in

  let test_ee_lr_ints () =
    let reference =
      ERI_lr.of_file ("test_files/"^name^"_eri_lr.ref") ~sparsity:`Dense
          ~size:(Basis.size t.basis)
    in
    let ee_lr_ints =
      Lazy.force t.ee_lr_ints
    in
    check_eri_lr ee_lr_ints reference
  in

  [
    "Overlap",     `Quick,  test_overlap;
    "eN_ints",     `Quick,  test_eN_ints;
    "kin_ints",    `Quick,  test_kin_ints;
    "ee_ints",     `Quick,  test_ee_ints;
    "ee_lr_ints",  `Quick,  test_ee_lr_ints;
  ]