Parallel integrals with Parmap and Zmq

Basis/ERI.ml

@@ -98,7 +98,7 @@ let filter_contracted_shell_pair_couples ?(cutoff=integrals_cutoff) shell_pair_c
 *)
 
 
-let store_class ?(cutoff=integrals_cutoff) push_socket contracted_shell_pair_couple cls = 
+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
@@ -110,7 +110,6 @@ let store_class ?(cutoff=integrals_cutoff) push_socket contracted_shell_pair_cou
   and shell_q = Cspc.shell_pair_q contracted_shell_pair_couple
   in
 
-  let msg = 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
@@ -125,19 +124,119 @@ let store_class ?(cutoff=integrals_cutoff) push_socket contracted_shell_pair_cou
                   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
-                  msg := (i_c,j_c,k_c,l_c,value) :: !msg;
+                  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));
+    ) (Cs.zkey_array (Csp.shell_a shell_p))
-    Zmq.Socket.send_all push_socket ["0" ; Bytes.to_string (Marshal.to_bytes !msg []) ]
 
 
 
 
 
 
-let of_basis basis = 
+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
+
+
+
+
+let of_basis_parallel basis = 
+
+  let store_class ?(cutoff=integrals_cutoff) push_socket 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 msg = 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
+                    msg := (i_c,j_c,k_c,l_c,value) :: !msg;
+                  ) (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));
+      Zmq.Socket.send_all push_socket ["0" ; Bytes.to_string (Marshal.to_bytes !msg []) ]
+  in
+
 
   let n = Bs.size basis
   and shell = Bs.contracted_shells basis
@@ -172,7 +271,7 @@ let of_basis basis =
     | 0 ->    begin
                 let zmq = ref None in
 
-                Parmap.pariter ~chunksize:1 ~ncores:4
+                Parmap.pariter ~chunksize:1 
                   ~init:(fun _ ->
                     let zmq_context = 
                       Zmq.Context.create ()
@@ -271,3 +370,4 @@ let of_basis basis =
 
 
 
+let of_basis = of_basis_parallel