From 29d568a47b7215d005c9cd806a1c8980a4fadd4f Mon Sep 17 00:00:00 2001
From: Anthony Scemama <scemama@irsamc.ups-tlse.fr>
Date: Wed, 21 Mar 2018 10:59:22 +0100
Subject: [PATCH] Optimization

---
 Basis/ContractedShellPair.mli    |  2 +-
 Basis/TwoElectronRRVectorized.ml | 52 ++++++++++++--------------------
 2 files changed, 21 insertions(+), 33 deletions(-)

diff --git a/Basis/ContractedShellPair.mli b/Basis/ContractedShellPair.mli
index 087352a..2d2e108 100644
--- a/Basis/ContractedShellPair.mli
+++ b/Basis/ContractedShellPair.mli
@@ -46,7 +46,7 @@ val shell_b : t -> ContractedShell.t
 *)
 
 val coefs_and_shell_pairs : t -> (float * PrimitiveShellPair.t) list
-(** Returns an list of coefficients and of {!PrimitiveShellPair.t}, containing all
+(** Returns an arra of coefficients and of {!PrimitiveShellPair.t}, containing all
     the pairs of primitive functions used to build the contracted shell pair.
 *)
 
diff --git a/Basis/TwoElectronRRVectorized.ml b/Basis/TwoElectronRRVectorized.ml
index 7f8ad47..3ba4031 100644
--- a/Basis/TwoElectronRRVectorized.ml
+++ b/Basis/TwoElectronRRVectorized.ml
@@ -559,8 +559,6 @@ let contracted_class_shell_pairs ~zero_m ?schwartz_p ?schwartz_q shell_p shell_q
   and shell_b = Csp.shell_b     shell_p
   and shell_c = Csp.shell_a     shell_q
   and shell_d = Csp.shell_b     shell_q
-  and sp      = Csp.shell_pairs shell_p
-  and sq      = Csp.shell_pairs shell_q
   in
   let maxm = 
     Am.(Csp.totAngMom shell_p + Csp.totAngMom shell_q |> to_int)
@@ -585,38 +583,34 @@ let contracted_class_shell_pairs ~zero_m ?schwartz_p ?schwartz_q shell_p shell_q
 
   (** Screening on the product of coefficients *)
   let coef_max_p =
-    Array.fold_left (fun accu x ->
+    List.fold_left (fun accu (x,_) ->
       if (abs_float x) > accu then (abs_float x) else accu)
-      0. (Csp.coefficients shell_p)
+      0. (Csp.coefs_and_shell_pairs shell_p)
   and coef_max_q =
-    Array.fold_left (fun accu x ->
+    List.fold_left (fun accu (x,_) ->
       if (abs_float x) > accu then (abs_float x) else accu)
-      0. (Csp.coefficients shell_q)
+      0. (Csp.coefs_and_shell_pairs shell_q)
   in
 
-  let rec build_list cutoff vec accu = function
-  | -1 -> Array.of_list accu
-  | k  -> build_list cutoff vec ( 
-            if (abs_float vec.(k) > cutoff) then (k::accu)
-            else accu ) (k-1)
+  let filter_p =
+    let cutoff_p = cutoff /. coef_max_p in
+    Csp.coefs_and_shell_pairs shell_p
+    |> List.filter (fun (ck,f) -> abs_float ck > cutoff_p) 
+  and filter_q =
+    let cutoff_q = cutoff /. coef_max_q in
+    Csp.coefs_and_shell_pairs shell_q
+    |> List.filter (fun (ck,f) -> abs_float ck > cutoff_q)
   in
-  let p_list =
-    let vec = (Csp.coefficients shell_p) in
-    build_list (cutoff /. coef_max_q) vec [] (Array.length vec - 1)
-  and q_list =
-    let vec = (Csp.coefficients shell_q) in
-    build_list (cutoff /. coef_max_p) vec [] (Array.length vec - 1)
+
+  let sp = List.map snd filter_p |> Array.of_list
+  and sq = List.map snd filter_q |> Array.of_list
+  and cp = List.map fst filter_p |> Array.of_list
+  and cq = List.map fst filter_q |> Array.of_list
   in
 
   let np, nq = 
-    Array.length p_list,
-    Array.length q_list
-  in
-  let filter_p vec = Array.init np (fun k -> vec.(p_list.(k)))
-  and filter_q vec = Array.init nq (fun k -> vec.(q_list.(k)))
-  in
-  let sp = filter_p sp
-  and sq = filter_q sq
+    Array.length sp,
+    Array.length sq
   in
 
 
@@ -637,10 +631,7 @@ let contracted_class_shell_pairs ~zero_m ?schwartz_p ?schwartz_q shell_p shell_q
 
             let coef =
               let result = Mat.make0 nq np in
-              Lacaml.D.ger
-                (Vec.of_array @@ filter_q (Csp.coefficients shell_q))
-                (Vec.of_array @@ filter_p (Csp.coefficients shell_p))
-                result;
+              Lacaml.D.ger (Vec.of_array @@ cq) (Vec.of_array @@ cp) result;
               result
             in
             let zm_array = Mat.init_cols np nq (fun i j ->
@@ -671,9 +662,6 @@ let contracted_class_shell_pairs ~zero_m ?schwartz_p ?schwartz_q shell_p shell_q
     | _ -> 
 
       let coef =
-        let cp = filter_p (Csp.coefficients shell_p)
-        and cq = filter_q (Csp.coefficients shell_q)
-        in
         Array.init np (fun l -> Array.init nq (fun k -> cq.(k) *. cp.(l)) )
       in