From d5a71707c884684b7f2749ce79807b1929a91605 Mon Sep 17 00:00:00 2001
From: Anthony Scemama <scemama@irsamc.ups-tlse.fr>
Date: Sat, 17 Jun 2023 02:07:35 +0200
Subject: [PATCH] Improved parallel integrals

---
 gaussian/lib/primitive_shell_pair.ml          |  3 +-
 .../lib/two_electron_integrals.ml             | 70 ++++++++++---------
 2 files changed, 37 insertions(+), 36 deletions(-)

diff --git a/gaussian/lib/primitive_shell_pair.ml b/gaussian/lib/primitive_shell_pair.ml
index d3aa31b..83c08c8 100644
--- a/gaussian/lib/primitive_shell_pair.ml
+++ b/gaussian/lib/primitive_shell_pair.ml
@@ -131,8 +131,7 @@ let make p_a p_b =
 let norm_scales x =
   try
       Lazy.force x.norm_scales
-  with Lazy.Undefined -> (Printf.printf "UNDEFINED\n%!" ; compute_norm_scales x.shell_a x.shell_b
-  )
+  with Lazy.Undefined -> compute_norm_scales x.shell_a x.shell_b
 
 let exponent_inv x = x.exponent_inv
 
diff --git a/gaussian_integrals/lib/two_electron_integrals.ml b/gaussian_integrals/lib/two_electron_integrals.ml
index a4a0bbb..8bebf6e 100644
--- a/gaussian_integrals/lib/two_electron_integrals.ml
+++ b/gaussian_integrals/lib/two_electron_integrals.ml
@@ -45,30 +45,40 @@ module Make(T : Two_ei_structure) = struct
     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 expand i_c powers_i shell = (
+       Cs.index shell + i_c + 1,
+       to_powers powers_i)
+    in
+
+    let l1 =
+        Array.mapi (fun i x -> expand i x (Csp.shell_a shell_p))
+           (Cs.zkey_array (Csp.shell_a shell_p))
+    in
+    let l2 =
+        Array.mapi (fun i x -> expand i x (Csp.shell_b shell_p))
+           (Cs.zkey_array (Csp.shell_b shell_p))
+    in
+    let l3 =
+        Array.mapi (fun i x -> expand i x (Csp.shell_a shell_q))
+           (Cs.zkey_array (Csp.shell_a shell_q))
+    in
+    let l4 =
+        Array.mapi (fun i x -> expand i x (Csp.shell_b shell_q))
+                (Cs.zkey_array (Csp.shell_b shell_q))
+    in
+
+    Array.iter (fun (l_c,xl) ->
+       Array.iter (fun (k_c,xk) ->
+          Array.iter (fun (j_c,xj) ->
+             Array.iter (fun (i_c,xi) ->
+                   let key = Zkey.of_powers_twelve  xi xj xk xl in
                     let value = Zmap.find cls key in
                     if abs_float value > cutoff then
                       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))
-
-
-
-
+                  ) l1
+              ) l2
+          ) l3
+    ) l4
 
 
 
@@ -126,31 +136,23 @@ module Make(T : Two_ei_structure) = struct
                 in
                 store_class ~cutoff eri_array cspc cls
               | None -> ()
-            ) shell_pairs;
-        with Exit -> ()
+          ) shell_pairs;
+      with Exit -> ()
     in
 
     let pool =  Domainslib.Task.setup_pool ~num_domains:Qcaml.num_domains () in
     let _ =
       Domainslib.Task.run pool (fun _ ->
-              (*
-              let n = Array.length shell_pairs in
-              Domainslib.Task.parallel_for pool ~start:1 ~finish:n
-                ~chunk_size:2
-                ~body:(fun i -> f shell_pairs.(n-i))
-      *)
       let n = Array.length shell_pairs in
       let i_prev = ref 0 in
       shell_pairs
-      |> Array.map (fun sp ->
-             Domainslib.Task.async pool (fun _ -> f sp) )
+      |> Array.map (fun sp -> Domainslib.Task.async pool (fun _ -> f sp) )
       |> Array.iteri (fun i task ->
                       let i = ((10 * i+1)/n) in
                       if !i_prev <> i then (
                          i_prev := i; 
-                         Printf.printf "%3d %%\n%!" (i*10)
-                      );
-                      ignore (Domainslib.Task.await pool task) ) ;
+                         Printf.printf "%3d %%\n%!" (i*10) );
+                      ignore (Domainslib.Task.await pool task)) 
       )
     in
     Domainslib.Task.teardown_pool pool;