QCaml/Basis/KinInt.ml

open Util
open Constants
open Lacaml.D

module Am  = AngularMomentum
module Bs  = Basis
module Co  = Coordinate
module Cs  = ContractedShell
module Csp = ContractedShellPair
module Po  = Powers
module Sp  = ShellPair

type t = Mat.t

let cutoff = integrals_cutoff

let to_powers x = 
  let open Zkey in
  match to_powers x with
  | Six x -> x
  | _ -> assert false

(** Computes all the kinetic integrals of the contracted shell pair *)
let contracted_class shell_a shell_b : float Zmap.t = 

  let shell_p =
    Csp.create shell_a shell_b
  in

  (* Pre-computation of integral class indices *)
  let class_indices = 
    Am.zkey_array (Am.Doublet Cs.(totAngMom shell_a, totAngMom shell_b))
  in

  let contracted_class = 
    Array.make (Array.length class_indices) 0.
  in

  (* Compute all integrals in the shell for each pair of significant shell pairs *)

  let sp = Csp.shell_pairs shell_p in
  let center_ab = 
    Csp.center_ab shell_p
  in
  let norm_coef_scale =
    Csp.norm_coef_scale shell_p
  in

  for ab=0 to (Array.length sp - 1)
  do
    let coef_prod =
      (Csp.coef shell_p).(ab) 
    in
    (** Screening on thr product of coefficients *)
    if (abs_float coef_prod) > 1.e-4*.cutoff then
      begin
        let center_pa = 
          sp.(ab).Sp.center_a
        in
        let expo_inv = 
          (Csp.expo_inv shell_p).(ab)
        in
        let i, j = 
          sp.(ab).Sp.i, sp.(ab).Sp.j 
        in
        let expo_a = 
          (Cs.expo sp.(ab).Sp.shell_a).(i)
        and expo_b = 
          (Cs.expo sp.(ab).Sp.shell_b).(j)
        in

        let xyz_of_int k =
          match k with
          | 0 -> Co.X
          | 1 -> Co.Y
          | _ -> Co.Z
        in
        Array.iteri (fun i key ->
            let (angMomA,angMomB) = to_powers key in
            let ov a b k = 
              let xyz = xyz_of_int k in
              Overlap_primitives.hvrr (a, b)
                expo_inv
                (Co.get xyz center_ab,
                 Co.get xyz center_pa) 
            in
            let f k = 
               let xyz = xyz_of_int k in 
               ov (Po.get xyz angMomA) (Po.get xyz angMomB) k
            and g k =
               let xyz = xyz_of_int k in 
               let s1 = ov (Po.get xyz angMomA - 1) (Po.get xyz angMomB - 1) k
               and s2 = ov (Po.get xyz angMomA + 1) (Po.get xyz angMomB - 1) k
               and s3 = ov (Po.get xyz angMomA - 1) (Po.get xyz angMomB + 1) k
               and s4 = ov (Po.get xyz angMomA + 1) (Po.get xyz angMomB + 1) k
               and a = float_of_int (Po.get xyz angMomA)
               and b = float_of_int (Po.get xyz angMomB) 
               in
               0.5 *. a *. b *. s1 -. expo_a *. b *. s2 -. expo_b *. a *. s3 +. 
                 2.0 *. expo_a *. expo_b *. s4
            in
            let s = Array.init 3 f
            and k = Array.init 3 g
            in
            let norm =  norm_coef_scale.(i) in
            let integral = chop norm (fun () ->
                k.(0)*.s.(1)*.s.(2) +.
                s.(0)*.k.(1)*.s.(2) +.
                s.(0)*.s.(1)*.k.(2) 
              ) in 
            contracted_class.(i) <- contracted_class.(i) +. coef_prod *. integral
          ) class_indices
      end
  done;
  let result = 
    Zmap.create (Array.length contracted_class)
  in
  Array.iteri (fun i key -> Zmap.add result key contracted_class.(i)) class_indices;
  result


(** Create kinetic energy matrix *)
let of_basis basis =
  let to_powers x =
    let open Zkey in
    match to_powers x with
    | Three x -> x
    | _ -> assert false
  in

  let n     = basis.Bs.size 
  and shell = basis.Bs.contracted_shells
  in

  let result = Mat.create n n in
  for j=0 to (Array.length shell) - 1 do
    for i=0 to j do
      (* Compute all the integrals of the class *)
      let cls =
        contracted_class shell.(i) shell.(j)
      in

      Array.iteri (fun j_c powers_j ->
          let j_c = Cs.index shell.(j) + j_c + 1 in
          let xj = to_powers powers_j in
          Array.iteri (fun i_c powers_i ->
              let i_c = Cs.index shell.(i) + i_c + 1 in
              let xi = to_powers powers_i in
              let key =
                Zkey.of_powers_six xi xj
              in
              let value =
                try Zmap.find cls key
                with Not_found -> failwith "Bug in kinetic integrals"
              in
              result.{i_c,j_c} <- value;
              result.{j_c,i_c} <- value;
            ) (Am.zkey_array (Singlet (Cs.totAngMom shell.(i))))
        ) (Am.zkey_array (Singlet (Cs.totAngMom shell.(j)))) 
    done;
  done;
  Mat.detri result;
  result


(** Write all kinetic integrals to a file *)
let to_file ~filename kinetic =

  let oc = open_out filename in
  let n =
    Mat.dim1 kinetic
  in

  for j=1 to n  do
    for i=1 to j do
      if (abs_float kinetic.{i,j} > cutoff) then
        Printf.fprintf oc "%4d %4d %20.12e\n" i j kinetic.{i,j}
    done;
  done;
  close_out oc