QCaml/Basis/Multipole.ml

open Util
open Constants
open Lacaml.D

type t = Mat.t array
(*
[| "x"; "y"; "z"; "x2"; "y2"; "z2" |]
*)

module Am  = AngularMomentum
module Bs  = Basis
module Co  = Coordinate
module Cs  = ContractedShell
module Csp = ContractedShellPair
module Po  = Powers
module Psp = PrimitiveShellPair

let multiply a b =
  let n = Mat.dim1 a in
  let c = Mat.create n n in
  Mat.cpab c a b;
  c

let x0 t = t.(0)
let y0 t = t.(1)
let z0 t = t.(2)
let x1 t = t.(3)
let y1 t = t.(4)
let z1 t = t.(5)
let x2 t = t.(6)
let y2 t = t.(7)
let z2 t = t.(8)

let matrix_x  t = multiply (x1 t) @@ multiply (y0 t) (z0 t)
let matrix_y  t = multiply (x0 t) @@ multiply (y1 t) (z0 t)
let matrix_z  t = multiply (x0 t) @@ multiply (y0 t) (z1 t)
let matrix_x2 t = multiply (x2 t) @@ multiply (y0 t) (z0 t)
let matrix_y2 t = multiply (x0 t) @@ multiply (y2 t) (z0 t)
let matrix_z2 t = multiply (x0 t) @@ multiply (y0 t) (z2 t)
let matrix_xy t = multiply (x1 t) @@ multiply (y1 t) (z0 t)
let matrix_yz t = multiply (x0 t) @@ multiply (y1 t) (z1 t)
let matrix_zx t = multiply (x1 t) @@ multiply (y0 t) (z1 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 integrals of the contracted shell pair *)
let contracted_class shell_a shell_b : float Zmap.t array =

  match Csp.make shell_a shell_b with
  | None -> Array.init 9 (fun _ -> Zmap.create 0)
  | Some shell_p ->
    begin

      (* Pre-computation of integral class indices *)
      let class_indices = Csp.zkey_array shell_p in

      let contracted_class =
        Array.init 9 (fun _ -> Array.make (Array.length class_indices) 0.)
      in

      let a_minus_b =
        Csp.a_minus_b shell_p
      in
      let norm_coef_scales =
        Csp.norm_scales shell_p
      in

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

      let xyz_of_int k =
        match k with
        | 0 -> Co.X
        | 1 -> Co.Y
        | _ -> Co.Z
      in

      List.iter (fun (coef_prod, psp) ->
          (** Screening on the product of coefficients *)
          if (abs_float coef_prod) > 1.e-6*.cutoff then
            begin
              let expo_inv = Psp.exponent_inv psp
              and center_pa = Psp.center_minus_a psp
              and xa = Co.get Co.X @@ Cs.center shell_a
              and ya = Co.get Co.Y @@ Cs.center shell_a
              and za = Co.get Co.Z @@ Cs.center shell_a
              in

              Array.iteri (fun i key ->
                  let (angMomA, angMomB) = to_powers key in
                  (* 1D Overlap <i|j> *)
                  let f k =
                    let xyz = xyz_of_int k in
                    Overlap_primitives.hvrr (Po.get xyz angMomA, Po.get xyz angMomB)
                      expo_inv
                      (Co.get xyz a_minus_b, Co.get xyz center_pa)
                  in
                  (* 1D <i|x-Xa|j> *)
                  let g k =
                    let xyz = xyz_of_int k in
                    Overlap_primitives.hvrr (Po.get xyz angMomA + 1, Po.get xyz angMomB)
                      expo_inv
                      (Co.get xyz a_minus_b, Co.get xyz center_pa)
                  in
                  (* 1D <i|(x-Xa)^2|j> *)
                  let h k =
                    let xyz = xyz_of_int k in
                    Overlap_primitives.hvrr (Po.get xyz angMomA + 2, Po.get xyz angMomB)
                      expo_inv
                      (Co.get xyz a_minus_b, Co.get xyz center_pa)
                  in
                  let norm =  norm_coef_scales.(i) in
                  let f0, f1, f2, g0, g1, g2, h0, h1, h2 =
                    f 0, f 1, f 2, g 0, g 1, g 2, h 0, h 1, h 2
                  in
                  let  x = g0 +. f0 *. xa in
                  let  y = g1 +. f1 *. ya in
                  let  z = g2 +. f2 *. za in
                  let x2 = h0 +. xa *. (2. *. x -. xa *. f0) in
                  let y2 = h1 +. ya *. (2. *. y -. ya *. f1) in
                  let z2 = h2 +. za *. (2. *. z -. za *. f2) in
                  let c = contracted_class in
                  let d = coef_prod *. norm in
                  c.(0).(i) <- c.(0).(i) +. d *. f0 ;
                  c.(1).(i) <- c.(1).(i) +. d *. f1 ;
                  c.(2).(i) <- c.(2).(i) +. d *. f2 ;
                  c.(3).(i) <- c.(3).(i) +. d *.  x ;
                  c.(4).(i) <- c.(4).(i) +. d *.  y ;
                  c.(5).(i) <- c.(5).(i) +. d *.  z ;
                  c.(6).(i) <- c.(6).(i) +. d *. x2 ;
                  c.(7).(i) <- c.(7).(i) +. d *. y2 ;
                  c.(8).(i) <- c.(8).(i) +. d *. z2 ;
                ) class_indices
            end
        ) (Csp.coefs_and_shell_pairs shell_p);
      let result =
        Array.map (fun c -> Zmap.create (Array.length c) ) contracted_class
      in
      for j=0 to Array.length result do
        let rj = result.(j) in
        let cj = contracted_class.(j) in
        Array.iteri (fun i key -> Zmap.add rj key cj.(i)) class_indices
      done;
      result
    end


(** Create multipole matrices *)
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     = Bs.size basis
  and shell = Bs.contracted_shells basis
  in

  let result = Array.init 9 (fun _ -> 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

      for k=0 to 8 do
        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.(k) key
                  with Not_found -> 0.
                in
                result.(k).{i_c,j_c} <- value;
                result.(k).{j_c,i_c} <- value;
              ) (Am.zkey_array (Singlet (Cs.ang_mom shell.(i))))
          ) (Am.zkey_array (Singlet (Cs.ang_mom shell.(j))))
      done;
    done;
  done;
  for k=0 to 8 do
    Mat.detri result.(k);
  done;
  result