Cleaning in ERI

Basis/ERI.ml

@@ -8,84 +8,13 @@ let max_ao = 1 lsl 14
 
 type index_pair = { first : int ; second : int }
 
-type t = 
-  | Dense  of (float, float64_elt, fortran_layout) Bigarray.Genarray.t
-  | Sparse of (int, float) Hashtbl.t
+type t = FourIdxStorage.t
 
+let get_chem = FourIdxStorage.get_chem
+let get_phys = FourIdxStorage.get_phys
 
-let key_of_indices ~r1 ~r2 =
-  let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in
-  let i,k = if i<=k then i,k else k,i
-  and j,l = if j<=l then j,l else l,j in
-  let i,k,j,l = if k<=l then i,k,j,l else j,l,i,k in
-  ((((((i lsl 15) lor k) lsl 15) lor j) lsl 15) lor l)
-
-
-let get ~r1 ~r2 = function
-  | Dense  t -> let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in t.{i,j,k,l}
-  | Sparse t -> let key = key_of_indices ~r1 ~r2 in
-                try Hashtbl.find t key
-                with Not_found -> 0.
-
-
-let set ~r1 ~r2 ~value = function
-  | Dense  t -> let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in
-                t.{i,j,k,l} <- value;
-                t.{k,j,i,l} <- value;
-                t.{i,l,k,j} <- value;
-                t.{k,l,i,j} <- value;
-                t.{j,i,l,k} <- value;
-                t.{j,k,l,i} <- value;
-                t.{l,i,j,k} <- value;
-                t.{l,k,j,i} <- value;
-
-  | Sparse t -> let key = key_of_indices ~r1 ~r2 in
-                Hashtbl.replace t key value
-    
-
-let increment ~r1 ~r2 ~value = function
-  | Dense  t -> let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in
-                t.{i,j,k,l} <- t.{i,j,k,l} +. value;
-                t.{k,j,i,l} <- t.{k,j,i,l} +. value;
-                t.{i,l,k,j} <- t.{i,l,k,j} +. value;
-                t.{k,l,i,j} <- t.{k,l,i,j} +. value;
-                t.{j,i,l,k} <- t.{j,i,l,k} +. value;
-                t.{j,k,l,i} <- t.{j,k,l,i} +. value;
-                t.{l,i,j,k} <- t.{l,i,j,k} +. value;
-                t.{l,k,j,i} <- t.{l,k,j,i} +. value;
-  | Sparse t -> let key = key_of_indices ~r1 ~r2 in
-                let old_value =
-                  try Hashtbl.find t key
-                  with Not_found -> 0.
-                in
-                Hashtbl.replace t key (old_value +. value)
-    
-
-let create = function
-  | `Dense n -> 
-    let eri_array = 
-      Genarray.create Float64 fortran_layout [| n ; n ; n ; n|]
-    in
-    Genarray.fill eri_array 0.;
-    Dense eri_array
-
-  | `Sparse n -> 
-    let eri_array = Hashtbl.create (n*n+13) in
-    Hashtbl.add eri_array (-1) (float_of_int n);
-    Sparse eri_array
-
-let size = function
-  | Dense t -> Genarray.nth_dim t 3
-  | Sparse t -> Hashtbl.find t (-1) |> int_of_float
-
-(** TODO : remove epsilons *)
-
-
-let get_chem t i j k l = get ~r1:{ first=i ; second=j } ~r2:{ first=k ; second=l } t
-let get_phys t i j k l = get ~r1:{ first=i ; second=k } ~r2:{ first=j ; second=l } t
-
-let set_chem t i j k l value = set ~r1:{ first=i ; second=j } ~r2:{ first=k ; second=l } ~value t
-let set_phys t i j k l value = set ~r1:{ first=i ; second=k } ~r2:{ first=j ; second=l } ~value t 
+let set_chem = FourIdxStorage.set_chem
+let set_phys = FourIdxStorage.set_phys
 
 
 module Am  = AngularMomentum
@@ -95,7 +24,6 @@ module Bs  = Basis
 module Cs  = ContractedShell
 module Csp = ContractedShellPair
 
-let cutoff = integrals_cutoff
 
 (** (00|00)^m : Fundamental electron repulsion integral
     $ \int \int \phi_p(r1) 1/r_{12} \phi_q(r2) dr_1 dr_2 $
@@ -125,49 +53,70 @@ let zero_m ~maxm ~expo_pq_inv ~norm_pq_sq =
   result
 
 
-(** 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
-
-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
 
 
-let contracted_class_atomic_shell_pairs ?schwartz_p ?schwartz_q shell_p shell_q : float Zmap.t =
-  TwoElectronRR.contracted_class_atomic_shell_pairs ~zero_m ?schwartz_p ?schwartz_q shell_p shell_q
 
-(*
-let contracted_class_atomic_shell_pairs_vec ?schwartz_p ?schwartz_q shell_p shell_q : float Zmap.t =
-  TwoElectronRRVectorized.contracted_class_atomic_shell_pairs ~zero_m ?schwartz_p ?schwartz_q shell_p shell_q
-*)
+let class_of_contracted_shell_pairs shell_p shell_q = 
+  if Array.length (Csp.shell_pairs shell_p) + (Array.length (Csp.shell_pairs shell_q))  < 4 then 
+    TwoElectronRR.contracted_class_shell_pairs ~zero_m shell_p shell_q
+  else
+    TwoElectronRRVectorized.contracted_class_shell_pairs ~zero_m shell_p shell_q
 
 
-let cutoff2 = cutoff *. cutoff
-(*
-type n_cls = { n : int ; cls : Zkey.t array }
-*)
-exception NullIntegral
 
-(*
-(** Unique index for integral <ij|kl> *)
-let index i j k l =
-  let f i k = 
-    let (p,r) = 
-      if i <= k then (i,k) else (k,i)
-    in p+ (r*r-r)/2
+
+
+let filter_contracted_shell_pairs ?(cutoff=integrals_cutoff) shell_pairs = 
+  let t0 = Unix.gettimeofday () in
+
+  let schwartz = 
+    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
   in
-  let p = f i k and q = f j l in
-  f p q
-*)
+
+  Printf.printf "%d shell pairs computed in %f seconds\n"
+    (List.length schwartz) (Unix.gettimeofday () -. t0);
+  schwartz
+
+let store_class ?(cutoff=integrals_cutoff) data cls shell_p shell_q = 
+    let to_powers x = 
+      let open Zkey in
+      match to_powers x with
+      | Three x -> x
+      | _ -> assert false
+    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 basis = 
-  let to_powers x = 
-    let open Zkey in
-    match to_powers x with
-    | Three x -> x
-    | _ -> assert false
-  in
 
   let n = Bs.size basis
   and shell = Bs.contracted_shells basis
@@ -181,41 +130,18 @@ let of_basis basis =
   let shell_pairs =
     Csp.of_contracted_shell_array shell
   in
-  (*TODO
-    let atomic_shell_pairs =
-    Asp.of_atomic_shell_array ~cutoff atomic_shells
-    in
-  *)
 
   (* Pre-compute diagonal integrals for Schwartz *)
-  let t0 = Unix.gettimeofday () in
-
   let schwartz = 
-    List.map (fun pair -> 
-        let cls = 
-          contracted_class_shell_pairs pair pair
-          (*TODO
-            contracted_class_atomic_shell_pairs pair pair
-          *)
-        in
-        (pair, cls, 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)
+    filter_contracted_shell_pairs shell_pairs
   in
 
-  Printf.printf "%d shell pairs computed in %f seconds\n"
-    (List.length schwartz) (Unix.gettimeofday () -. t0);
-
-
-  (* Group shell pairs by common pairs of atoms *)
-
-
   (* 4D data initialization *)
   let eri_array = 
+    FourIdxStorage.create ~size:n `Dense
     (*
-    create (`Dense n)
+    FourIdxStorage.create ~size:n `Sparse
     *)
-    create (`Sparse n)
   in
 
   (* Compute ERIs *)
@@ -223,11 +149,8 @@ let of_basis basis =
   let t0 = Unix.gettimeofday () in
   let inn = ref 0 and out = ref 0 in
 
-  (*TODO
-    for i=0 to (Array.length atomic_shells) - 1 do
-  *)
   let ishell = ref 0 in
-  List.iter (fun (shell_p, schwartz_p, schwartz_p_max) ->
+  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 ())
@@ -235,85 +158,32 @@ let of_basis basis =
 
       let sp =
         Csp.shell_pairs shell_p
-        (*TODO
-          Asp.atomic_shell_pairs shell_p
-        *)
       in
 
       try
-        List.iter (fun (shell_q, schwartz_q, schwartz_q_max) ->
+        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
 
-            try
-              if schwartz_p_max *. schwartz_q_max < cutoff2 then
-                raise NullIntegral;
+            let swap = 
+              Array.length sp > Array.length sq 
+            in
 
-              let sq =
-                Csp.shell_pairs shell_q
-                (*TODO
-                  Asp.atomic_shell_pairs shell_q
-                *)
-              in
-
-              let swap = 
-                Array.length sp > Array.length sq 
-              in
-
-              (* Compute all the integrals of the class *)
-              let cls =
-                if swap then
-                  (*TODO
-                    contracted_class_atomic_shell_pairs ~schwartz_p:schwartz_q ~schwartz_q:schwartz_p shell_q shell_p 
-                  *)
-                  if (Array.length sp) + (Array.length sq)  < 4 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)  < 4 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
-
-
-              (* Write the data in the output file *)
-              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 = 
-                                if swap then
-                                  Zkey.of_powers_twelve  xk xl xi xj 
-                                else
-                                  Zkey.of_powers_twelve  xi xj xk xl 
-                              in
-                              let value = 
-                                Zmap.find cls key 
-                              in
-                              set_chem eri_array i_c j_c k_c l_c value;
-                              if (abs_float value > cutoff) then
-                                (inn := !inn + 1;
-                                )
-                              else
-                                out := !out + 1;
-                            ) (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))
-            with NullIntegral ->  ()
+            (* Compute all the integrals of the class *)
+            let f p q = 
+              let cls = class_of_contracted_shell_pairs p q in
+              store_class ~cutoff:integrals_cutoff eri_array cls p q
+            in
+            if swap then
+              f shell_q shell_p
+            else
+              f shell_p shell_q
           ) schwartz
       with Exit -> ()
     ) schwartz;
@@ -323,19 +193,7 @@ let of_basis basis =
 
 
 (** 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 size eri_array 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 = get_phys 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
+let to_file = FourIdxStorage.to_file
+
+
 

Basis/ERI.mli

@@ -1,28 +1,28 @@
+(** Electron repulsion intergals. *)
+
 type t 
+
+(*
+val filter_atomic_shell_pairs : ?cutoff:float -> AtomicShellPair.t list -> AtomicShellPair.t  list
+(** Uses Schwartz screening on atomic shell pairs. *)
+*)
+
+val filter_contracted_shell_pairs : ?cutoff:float -> ContractedShellPair.t list -> ContractedShellPair.t list
+(** Uses Schwartz screening on contracted shell pairs. *)
+
+val class_of_contracted_shell_pairs : ContractedShellPair.t -> ContractedShellPair.t -> float Zmap.t
+(** Computes all the ERI of the class built from th combination of two contracted shell pairs. *)
+
+
 val get_chem : t -> int -> int -> int -> int -> float
+(** Get an integral using the Chemist's convention { \[ (ij|kl) \] }. *)
+
 val get_phys : t -> int -> int -> int -> int -> float
-
-val zero_m : maxm:int -> expo_pq_inv:float -> norm_pq_sq:float -> float array
-
-val contracted_class_shell_pairs :
-  ?schwartz_p:float Zmap.t ->
-  ?schwartz_q:float Zmap.t ->
-  ContractedShellPair.t ->
-  ContractedShellPair.t -> float Zmap.t
-
-val contracted_class_shell_pairs_vec :
-  ?schwartz_p:float Zmap.t ->
-  ?schwartz_q:float Zmap.t ->
-  ContractedShellPair.t ->
-  ContractedShellPair.t -> float Zmap.t
-
-val contracted_class_atomic_shell_pairs :
-  ?schwartz_p:float Zmap.t ->
-  ?schwartz_q:float Zmap.t ->
-  AtomicShellPair.t ->
-  AtomicShellPair.t -> float Zmap.t
+(** Get an integral using the Physicist's convention { \[ \langle ij|kl \rangle \] }. *)
 
 val of_basis : Basis.t -> t
+(** Compute all ERI's for a given {!Basis.t}. *)
 
-val to_file : filename:string -> t -> unit
+val to_file : ?cutoff:float -> filename:string -> t -> unit
+(** Write the ERI's to a file, using the Physicist's convention. *)
 

Utils/FourIdxStorage.ml

@@ -0,0 +1,123 @@
+let max_index = 1 lsl 14
+
+type index_pair = { first : int ; second : int }
+
+type storage_t = 
+  | Dense  of (float, Bigarray.float64_elt, Bigarray.fortran_layout) Bigarray.Genarray.t
+  | Sparse of (int, float) Hashtbl.t
+
+type t =
+  {
+    size         : int ;
+    four_index   : storage_t ;
+  }
+
+let key_of_indices ~r1 ~r2 =
+  let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in
+  let f i k = 
+    let p, r = 
+      if i <= k then i, k else k, i
+    in p+ (r*r-r)/2
+  in
+  let p = f i k and q = f j l in
+  f p q
+
+
+let get_four_index ~r1 ~r2 t = 
+  match t.four_index with
+  | Dense  t -> let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in t.{i,j,k,l}
+  | Sparse t -> let key = key_of_indices ~r1 ~r2 in
+                try Hashtbl.find t key
+                with Not_found -> 0.
+
+
+let set_four_index ~r1 ~r2 ~value t = 
+  match t.four_index with
+  | Dense  t -> let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in
+                t.{i,j,k,l} <- value;
+                t.{k,j,i,l} <- value;
+                t.{i,l,k,j} <- value;
+                t.{k,l,i,j} <- value;
+                t.{j,i,l,k} <- value;
+                t.{j,k,l,i} <- value;
+                t.{l,i,j,k} <- value;
+                t.{l,k,j,i} <- value;
+
+  | Sparse t -> let key = key_of_indices ~r1 ~r2 in
+                Hashtbl.replace t key value
+    
+
+let increment_four_index ~r1 ~r2 ~value x = 
+  match x.four_index with
+  | Dense  t -> let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in
+                t.{i,j,k,l} <- t.{i,j,k,l} +. value;
+                t.{k,j,i,l} <- t.{k,j,i,l} +. value;
+                t.{i,l,k,j} <- t.{i,l,k,j} +. value;
+                t.{k,l,i,j} <- t.{k,l,i,j} +. value;
+                t.{j,i,l,k} <- t.{j,i,l,k} +. value;
+                t.{j,k,l,i} <- t.{j,k,l,i} +. value;
+                t.{l,i,j,k} <- t.{l,i,j,k} +. value;
+                t.{l,k,j,i} <- t.{l,k,j,i} +. value
+
+  | Sparse t -> let key = key_of_indices ~r1 ~r2 in
+                let old_value =
+                  try Hashtbl.find t key
+                  with Not_found -> 0.
+                in
+                Hashtbl.replace t key (old_value +. value)
+
+let get ~r1 ~r2 = 
+  get_four_index ~r1 ~r2 
+
+let set ~r1 ~r2 = 
+  set_four_index ~r1 ~r2 
+
+let increment ~r1 ~r2 = 
+  increment_four_index ~r1 ~r2 
+
+let create ~size sparsity =
+  assert (size < max_index);
+  let four_index = 
+    match sparsity with
+    | `Dense ->
+      let result = 
+        Bigarray.Genarray.create Float64 Bigarray.fortran_layout [| size ; size ; size ; size |]
+      in
+      Bigarray.Genarray.fill result 0.;
+      Dense result
+
+    | `Sparse -> 
+      let result = Hashtbl.create (size*size+13) in
+      Sparse result
+  in
+  { size ; four_index }
+
+
+let size t = t.size
+
+(** TODO : remove epsilons *)
+
+
+let get_chem t i j k l = get ~r1:{ first=i ; second=j } ~r2:{ first=k ; second=l } t
+let get_phys t i j k l = get ~r1:{ first=i ; second=k } ~r2:{ first=j ; second=l } t
+
+let set_chem t i j k l value = set ~r1:{ first=i ; second=j } ~r2:{ first=k ; second=l } ~value t
+let set_phys t i j k l value = set ~r1:{ first=i ; second=k } ~r2:{ first=j ; second=l } ~value t 
+
+
+(** Write all integrals to a file with the <ij|kl> convention *)
+let to_file ?(cutoff=Constants.epsilon) ~filename data =
+  let oc = open_out filename in
+  for l_c=1 to size data 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 = get_phys data 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
+