10
1
mirror of https://gitlab.com/scemama/QCaml.git synced 2024-11-07 14:43:41 +01:00
QCaml/Basis/ERI_parallel.ml
2018-02-14 18:08:43 +01:00

263 lines
8.4 KiB
OCaml

(** Electron-electron repulsion integrals *)
open Util
open Constants
open Bigarray
type t = (float, float32_elt, fortran_layout) Genarray.t
(* Input type *)
type input_data =
{
i : int;
j : int;
shell_pairs : ContractedShellPair.t array array ;
schwartz : (float Zmap.t * float) array array;
cutoff : float
}
(* Output type *)
type output_integral = (* <ij|kl> *)
{
i1 : int ; (* Function i for electron 1 *)
j2 : int ; (* Function j for electron 2 *)
k1 : int ; (* Function k for electron 1 *)
l2 : int ; (* Function l for electron 2 *)
swap : bool ; (* If true, Compute (kl|ij) instead of (ij|kl) *)
cls : float Zmap.t;
}
type output_data = output_integral list
(** (00|00)^m : Fundamental electron repulsion integral
$ \int \int \phi_p(r1) 1/r_{12} \phi_q(r2) dr_1 dr_2 $
maxm : Maximum total angular momentum
expo_pq_inv : $1./p + 1/q$ where $p$ and $q$ are the exponents of
$\phi_p$ and $\phi_q$
norm_pq_sq : square of the distance between the centers of $\phi_p$
and $\phi_q$
*)
let zero_m ~maxm ~expo_pq_inv ~norm_pq_sq =
let exp_pq = 1. /. expo_pq_inv in
let t = norm_pq_sq *. exp_pq in
let f = two_over_sq_pi *. (sqrt exp_pq) in
let result = boys_function ~maxm t in
let rec aux accu k = function
| 0 -> result.(k) <- result.(k) *. accu
| l ->
begin
result.(k) <- result.(k) *. accu;
let new_accu = -. accu *. exp_pq in
aux new_accu (k+1) (l-1)
end
in
aux f 0 maxm;
result
(** Compute all the integrals of a contracted class when shell pairs are not yet available *)
let contracted_class shell_a shell_b shell_c shell_d : float Zmap.t =
TwoElectronRR.contracted_class ~zero_m shell_a shell_b shell_c shell_d
(** Compute all the integrals of a contracted class *)
let contracted_class_shell_pairs_vec ?schwartz_p ?schwartz_q shell_p shell_q : float Zmap.t =
TwoElectronRRVectorized.contracted_class_shell_pairs ~zero_m ?schwartz_p ?schwartz_q shell_p shell_q
(** Compute all the integrals of a contracted class *)
let contracted_class_shell_pairs ?schwartz_p ?schwartz_q shell_p shell_q : float Zmap.t =
TwoElectronRR.contracted_class_shell_pairs ~zero_m ?schwartz_p ?schwartz_q shell_p shell_q
(** Creates the input data structure from the basis set. *)
let make_input basis =
let shell = Basis.contracted_shells basis in
(* Pre-compute all shell pairs *)
let shell_pairs =
ContractedShellPair.shell_pairs shell
in
(* Pre-compute diagonal integrals for Schwartz screening *)
let t0 = Unix.gettimeofday () in
let schwartz =
Array.map (fun pair_array ->
Array.map (fun pair ->
let cls = contracted_class_shell_pairs pair pair in
let m = Zmap.fold (fun _ value accu ->
max (abs_float value) accu) cls 0.
in (cls, m)
) pair_array
) shell_pairs
in
(* Count number of significant shell pairs. *)
let icount = ref 0 in
for i=0 to (Array.length shell) - 1 do
print_int (Contracted_shell.index shell.(i)) ; print_newline ();
for j=0 to i do
let schwartz_p, schwartz_p_max = schwartz.(i).(j) in
if (schwartz_p_max >= cutoff) then
icount := !icount + 1;
done;
done;
Printf.printf "%d shell pairs computed in %f seconds\n" !icount (Unix.gettimeofday () -. t0);
List.init (Array.length shell) (fun i ->
List.init (i+1) (fun j ->
{ i ; j ; shell_pairs ; schwartz ; cutoff } ) )
|> List.fold_left List.rev_append []
exception NullIntegral
let slave_job { i ; j ; shell_pairs ; schwartz ; cutoff} =
let shell_p = shell_pairs.(i).(j) in
let schwartz_p, schwartz_p_max = schwartz.(i).(j) in
if schwartz_p_max < cutoff then [] else
let schwartz_cutoff = cutoff *. cutoff in
let sp = shell_p.ContractedShellPair.shell_pairs in
let f k l =
let schwartz_q, schwartz_q_max = schwartz.(k).(l) in
if schwartz_p_max *. schwartz_q_max < schwartz_cutoff then
raise NullIntegral;
let shell_q = shell_pairs.(k).(l) in
let sq = shell_q.ContractedShellPair.shell_pairs in
let swap = Array.length sp > Array.length sq in
(* Compute all the integrals of the class *)
let cls =
if swap then
if Array.length sp + Array.length sq = 2 then
contracted_class_shell_pairs ~schwartz_p:schwartz_q ~schwartz_q:schwartz_p shell_q shell_p
else
contracted_class_shell_pairs_vec ~schwartz_p:schwartz_q ~schwartz_q:schwartz_p shell_q shell_p
else
if Array.length sp + Array.length sq = 2 then
contracted_class_shell_pairs ~schwartz_p ~schwartz_q shell_p shell_q
else
contracted_class_shell_pairs_vec ~schwartz_p ~schwartz_q shell_p shell_q
in
{ i1=i ; j2=k ; k1=j ; l2=l ; swap ; cls }
in
let rec loop accu k l =
match k, l with
| -1, -1 -> accu
| k, -1 -> loop accu (k-1) (k-1)
| k, l ->
let new_accu =
let _, schwartz_q_max = schwartz.(k).(l) in
if schwartz_p_max *. schwartz_q_max > schwartz_cutoff then
f k l :: accu
else accu
in
loop new_accu k (l-1)
in
loop [] i i
let of_basis basis =
let shell = Basis.contracted_shells basis in
(* 4D data initialization *)
let eri_array =
let n = Basis.size basis in
Genarray.create Float32 fortran_layout [| n ; n ; n ; n|]
in
Genarray.fill eri_array 0.;
(* Compute ERIs *)
let to_int_tuple x =
let open Zkey in
match to_int_tuple Kind_3 x with
| Three x -> x
| _ -> assert false
in
let t0 = Unix.gettimeofday () in
let inn = ref 0 in
let out = ref 0 in
make_input basis
|> List.map slave_job
|> List.iter (fun output_ij ->
List.iter (fun { i1=i ; j2=k ; k1=j ; l2=l ; swap ; cls } ->
Array.iteri (fun i_c powers_i ->
let i_c = (Contracted_shell.index shell.(i)) + i_c + 1 in
let xi = to_int_tuple powers_i in
Array.iteri (fun j_c powers_j ->
let j_c = (Contracted_shell.index shell.(j)) + j_c + 1 in
let xj = to_int_tuple powers_j in
Array.iteri (fun k_c powers_k ->
let k_c = (Contracted_shell.index shell.(k)) + k_c + 1 in
let xk = to_int_tuple powers_k in
Array.iteri (fun l_c powers_l ->
let l_c = (Contracted_shell.index shell.(l)) + l_c + 1 in
let xl = to_int_tuple powers_l in
let key =
if swap then
Zkey.of_int_tuple (Zkey.Twelve (xk,xl,xi,xj))
else
Zkey.of_int_tuple (Zkey.Twelve (xi,xj,xk,xl))
in
let value =
Zmap.find cls key
in
eri_array.{i_c,k_c,j_c,l_c} <- value;
eri_array.{j_c,k_c,i_c,l_c} <- value;
eri_array.{i_c,l_c,j_c,k_c} <- value;
eri_array.{j_c,l_c,i_c,k_c} <- value;
eri_array.{k_c,i_c,l_c,j_c} <- value;
eri_array.{k_c,j_c,l_c,i_c} <- value;
eri_array.{l_c,i_c,k_c,j_c} <- value;
eri_array.{l_c,j_c,k_c,i_c} <- value;
if (abs_float value > cutoff) then
(inn := !inn + 1;
)
else
out := !out + 1;
) (Contracted_shell.powers shell.(l))
) (Contracted_shell.powers shell.(k))
) (Contracted_shell.powers shell.(j))
) (Contracted_shell.powers shell.(i));
) output_ij
);
Printf.printf "In: %d Out:%d\n" !inn !out ;
Printf.printf "Computed ERIs in %f seconds\n%!" (Unix.gettimeofday () -. t0);
eri_array
(** Write all integrals to a file with the <ij|kl> convention *)
let to_file ~filename eri_array =
let oc = open_out filename in
(* Print ERIs *)
for l_c=1 to (Genarray.nth_dim eri_array 3) do
for k_c=1 to l_c do
for j_c=1 to l_c do
for i_c=1 to k_c do
let value = eri_array.{i_c,j_c,k_c,l_c} in
if (abs_float value > cutoff) then
Printf.fprintf oc " %5d %5d %5d %5d%20.15f\n" i_c j_c k_c l_c value;
done;
done;
done;
done;
close_out oc