QCaml/Basis/TwoElectronIntegrals.ml

(** Two electron integral functor, parameterized by the zero_m function.
*)

open Constants
let cutoff = integrals_cutoff

module Am  = AngularMomentum
module As  = AtomicShell
module Asp = AtomicShellPair
module Bs  = Basis
module Cs  = ContractedShell
module Csp = ContractedShellPair
module Cspc = ContractedShellPairCouple
module Fis = FourIdxStorage


module type Zero_mType =
  sig
    val name : string
    val zero_m : maxm:int -> expo_p_inv:float -> expo_q_inv:float ->
        norm_pq_sq:float -> float array
  end


module Make(Zero_m : Zero_mType) = struct

  include FourIdxStorage

  let zero_m = Zero_m.zero_m

  let class_of_contracted_shell_pair_couple shell_pair_couple = 
    let shell_p = Cspc.shell_pair_p shell_pair_couple
    and shell_q = Cspc.shell_pair_q shell_pair_couple
    in
    if Array.length (Csp.shell_pairs shell_p) +
       (Array.length (Csp.shell_pairs shell_q))  < 4 then 
      TwoElectronRR.contracted_class_shell_pair_couple
        ~zero_m shell_pair_couple
    else
      TwoElectronRRVectorized.contracted_class_shell_pairs
        ~zero_m shell_p shell_q



  let filter_contracted_shell_pairs ?(cutoff=integrals_cutoff) shell_pairs = 
    List.map (fun pair -> 
        match Cspc.make ~cutoff pair pair with
        | Some cspc -> 
          let cls = class_of_contracted_shell_pair_couple cspc in
          (pair, Zmap.fold (fun key value accu -> max (abs_float value) accu) cls 0. )
        (* TODO \sum_k |coef_k * integral_k| *)
        | None -> (pair, -1.) 
      ) shell_pairs
    |> List.filter (fun (_, schwartz_p_max) -> schwartz_p_max >= cutoff)
    |> List.map fst


  (*  TODO
      let filter_contracted_shell_pair_couples
      ?(cutoff=integrals_cutoff) shell_pair_couples = 
      List.map (fun pair -> 
          let cls = 
            class_of_contracted_shell_pairs pair pair
          in
          (pair, Zmap.fold (fun key value accu -> max (abs_float value) accu) cls 0. )
        ) shell_pairs
      |> List.filter (fun (_, schwartz_p_max) -> schwartz_p_max >= cutoff)
      |> List.map fst

  *)


  let store_class ?(cutoff=integrals_cutoff) data contracted_shell_pair_couple cls = 
    let to_powers x = 
      let open Zkey in
      match to_powers x with
      | Three x -> x
      | _ -> assert false
    in

    let shell_p = Cspc.shell_pair_p contracted_shell_pair_couple
    and shell_q = Cspc.shell_pair_q contracted_shell_pair_couple
    in

    Array.iteri (fun i_c powers_i ->
        let i_c = Cs.index (Csp.shell_a shell_p) + i_c + 1 in
        let xi = to_powers powers_i in
        Array.iteri (fun j_c powers_j ->
            let j_c = Cs.index (Csp.shell_b shell_p) + j_c + 1 in
            let xj = to_powers powers_j in
            Array.iteri (fun k_c powers_k ->
                let k_c = Cs.index (Csp.shell_a shell_q) + k_c + 1 in 
                let xk = to_powers powers_k in
                Array.iteri (fun l_c powers_l ->
                    let l_c = Cs.index (Csp.shell_b shell_q) + l_c + 1 in
                    let xl = to_powers powers_l in
                    let key = Zkey.of_powers_twelve  xi xj xk xl in
                    let value = Zmap.find cls key in
                    set_chem data i_c j_c k_c l_c value
                  ) (Cs.zkey_array (Csp.shell_b shell_q))
              ) (Cs.zkey_array (Csp.shell_a shell_q))
          ) (Cs.zkey_array (Csp.shell_b shell_p))
      ) (Cs.zkey_array (Csp.shell_a shell_p))






  let of_basis_serial basis = 

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

    let eri_array = 
      Fis.create ~size:n `Dense
(*
    Fis.create ~size:n `Sparse
*)
    in

    let t0 = Unix.gettimeofday () in

    let shell_pairs =
      Csp.of_contracted_shell_array shell
      |> filter_contracted_shell_pairs ~cutoff
    in

    Printf.printf "%d significant shell pairs computed in %f seconds\n"
      (List.length shell_pairs) (Unix.gettimeofday () -. t0);


    let t0 = Unix.gettimeofday () in
    let ishell = ref 0 in

    List.iter (fun shell_p ->
        let () = 
          if (Cs.index (Csp.shell_a shell_p) > !ishell) then
            (ishell := Cs.index (Csp.shell_a shell_p) ; print_int !ishell ; print_newline ())
        in

        let sp =
          Csp.shell_pairs shell_p
        in

        try
          List.iter (fun shell_q ->
              let () = 
                if  Cs.index (Csp.shell_a shell_q) >
                    Cs.index (Csp.shell_a shell_p) then
                  raise Exit
              in
              let sq = Csp.shell_pairs shell_q in
              let cspc = 
                if Array.length sp < Array.length sq then
                  Cspc.make ~cutoff shell_p shell_q
                else
                  Cspc.make ~cutoff shell_q shell_p
              in

              match cspc with
              | Some cspc ->
                let cls =
                  class_of_contracted_shell_pair_couple cspc
                in
                store_class ~cutoff eri_array cspc cls
              | None -> ()
            ) shell_pairs
        with Exit -> ()
      ) shell_pairs ;
    Printf.printf "Computed ERIs in %f seconds\n%!" (Unix.gettimeofday () -. t0);
    eri_array







  (* Parallel functions *)



  let of_basis_parallel basis = 

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

    let store_class_parallel
        ?(cutoff=integrals_cutoff) contracted_shell_pair_couple cls = 
      let to_powers x = 
        let open Zkey in
        match to_powers x with
        | Three x -> x
        | _ -> assert false
      in

      let shell_p = Cspc.shell_pair_p contracted_shell_pair_couple
      and shell_q = Cspc.shell_pair_q contracted_shell_pair_couple
      in

      let result = ref [] in
      Array.iteri (fun i_c powers_i ->
          let i_c = Cs.index (Csp.shell_a shell_p) + i_c + 1 in
          let xi = to_powers powers_i in
          Array.iteri (fun j_c powers_j ->
              let j_c = Cs.index (Csp.shell_b shell_p) + j_c + 1 in
              let xj = to_powers powers_j in
              Array.iteri (fun k_c powers_k ->
                  let k_c = Cs.index (Csp.shell_a shell_q) + k_c + 1 in 
                  let xk = to_powers powers_k in
                  Array.iteri (fun l_c powers_l ->
                      let l_c = Cs.index (Csp.shell_b shell_q) + l_c + 1 in
                      let xl = to_powers powers_l in
                      let key = Zkey.of_powers_twelve  xi xj xk xl in
                      let value = Zmap.find cls key in
                      result := (i_c, j_c, k_c, l_c, value) :: !result
                    ) (Cs.zkey_array (Csp.shell_b shell_q))
                ) (Cs.zkey_array (Csp.shell_a shell_q))
            ) (Cs.zkey_array (Csp.shell_b shell_p))
        ) (Cs.zkey_array (Csp.shell_a shell_p));
      !result
    in



    let t0 = Unix.gettimeofday () in

    let shell_pairs =
      Csp.of_contracted_shell_array shell
      |> filter_contracted_shell_pairs ~cutoff
    in

    if Parallel.master then
      Printf.printf "%d significant shell pairs computed in %f seconds\n"
        (List.length shell_pairs) (Unix.gettimeofday () -. t0);


    let t0 = Unix.gettimeofday () in
    let ishell = ref max_int in

    let input_stream = Stream.of_list (List.rev shell_pairs) in

    let f shell_p =
      let () = 
        if Parallel.rank < 2 && Cs.index (Csp.shell_a shell_p) < !ishell then
          (ishell := Cs.index (Csp.shell_a shell_p) ; print_int !ishell ; print_newline ())
      in

      let sp =
        Csp.shell_pairs shell_p
      in

      let result = ref [] in
      try
        List.iter (fun shell_q ->
            let () = 
              if  Cs.index (Csp.shell_a shell_q) >
                  Cs.index (Csp.shell_a shell_p) then
                raise Exit
            in
            let sq = Csp.shell_pairs shell_q in
            let cspc = 
              if Array.length sp < Array.length sq then
                Cspc.make ~cutoff shell_p shell_q
              else
                Cspc.make ~cutoff shell_q shell_p
            in

            match cspc with
            | Some cspc ->
              let cls =
                class_of_contracted_shell_pair_couple cspc
              in
              result := (store_class_parallel ~cutoff cspc cls) :: !result;
            | None -> ()
          ) shell_pairs;
        raise Exit
      with Exit -> List.concat !result |> Array.of_list
    in

    let eri_array = 
      if Parallel.master then
        Fis.create ~size:n `Dense
      else
        Fis.create ~size:0 `Dense
    in
    Farm.run ~ordered:false ~f input_stream
    |> Stream.iter (fun l ->
        Array.iter (fun (i_c,j_c,k_c,l_c,value) ->
            set_chem eri_array i_c j_c k_c l_c value) l);

    if Parallel.master then
      Printf.printf
        "Computed %s Integrals in parallel in %f seconds\n%!" Zero_m.name (Unix.gettimeofday () -. t0);
    Parallel.broadcast (lazy eri_array)



  let of_basis = 
    match Parallel.size with
    | 1 -> of_basis_serial
    | _ -> of_basis_parallel


end