Separated 4, 3 and 2-idx integrals

Utils/FourIdxStorage.ml

@@ -6,15 +6,20 @@ let max_index = 1 lsl 14
 
 type index_pair = { first : int ; second : int }
 
+type array2 = (float, Bigarray.float64_elt, Bigarray.fortran_layout) Bigarray.Array2.t
 
 type storage_t = 
-  | Dense  of (float, Bigarray.float64_elt, Bigarray.fortran_layout) Bigarray.Array2.t
+  | Dense  of array2
   | Sparse of (int, float) Hashtbl.t
 
 type t =
   {
-    size         : int ;
+    size             : int ;
-    four_index   : storage_t ;
+    two_index        : array2;
+    two_index_anti   : array2;
+    three_index      : array2;
+    three_index_anti : array2;
+    four_index       : storage_t ;
   }
 
 let key_of_indices ~r1 ~r2 =
@@ -28,85 +33,164 @@ let key_of_indices ~r1 ~r2 =
   f p q
 
 
+let check_bounds r1 r2 t =
+  let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in 
+  let size = t.size in
+  assert ( (i lor j lor k lor l) > 0 );
+  assert ( i <= size && j <= size && k <= size && l <= size )
+
+
 let dense_index i j size =
   (j-1)*size + i
 
 
-let get_four_index ~r1 ~r2 t = 
+let unsafe_get_four_index ~r1 ~r2 t = 
-  match t.four_index with
-  | Dense  a -> (let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in 
-                 let size = t.size in
-                 assert ( (i lor j lor k lor l) > 0 );
-                 assert ( i <= size && j <= size && k <= size && l <= size );
-                 Bigarray.Array2.unsafe_get a (dense_index i j size) (dense_index k l size)
-                )
-  | Sparse a -> let key = key_of_indices ~r1 ~r2 in
-                try Hashtbl.find a key
-                with Not_found -> 0.
 
+  let open Bigarray.Array2 in
+  let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in 
 
-let set_four_index ~r1 ~r2 ~value t = 
+  if i=k then
-  match t.four_index with
+    if j=l then
-  | Dense  a -> (let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in
+      unsafe_get t.two_index i j
-                 let size = t.size in
+    else
-                 assert ( (i lor j lor k lor l) > 0 );
+      unsafe_get t.three_index (dense_index j l t.size) i
-                 assert ( i <= size && j <= size && k <= size && l <= size);
+  else if j=l then
-                 let ij = (dense_index i j size)
+      unsafe_get t.three_index (dense_index i k t.size) j
-                 and kl = (dense_index k l size)
-                 and il = (dense_index i l size)
-                 and kj = (dense_index k j size)
-                 and ji = (dense_index j i size)
-                 and lk = (dense_index l k size)
-                 and li = (dense_index l i size)
-                 and jk = (dense_index j k size)
-                 in
-                 let open Bigarray.Array2 in
-                   unsafe_set a ij kl value;
-                   unsafe_set a kj il value;
-                   unsafe_set a il kj value;
-                   unsafe_set a kl ij value;
-                   unsafe_set a ji lk value;
-                   unsafe_set a li jk value;
-                   unsafe_set a jk li value;
-                   unsafe_set a lk ji value
-                )
 
-  | Sparse a -> let key = key_of_indices ~r1 ~r2 in
+  else if i=l then
-                Hashtbl.replace a key value
+    if k=j then 
-    
+      unsafe_get t.two_index_anti i j
+    else
+      unsafe_get t.three_index_anti (dense_index j k t.size) i 
 
-let increment_four_index ~r1 ~r2 ~value t = 
+  else if j=k then
-  match t.four_index with
+      unsafe_get t.three_index_anti (dense_index i l t.size) j
-  | Dense  a -> (let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in
-                 let size = t.size in
-                 assert ( (i lor j lor k lor l) > 0 );
-                 assert ( i <= size && j <= size && k <= size && l <= size);
-                 let ij = (dense_index i j size)
-                 and kl = (dense_index k l size)
-                 and il = (dense_index i l size)
-                 and kj = (dense_index k j size)
-                 and ji = (dense_index j i size)
-                 and lk = (dense_index l k size)
-                 and li = (dense_index l i size)
-                 and jk = (dense_index j k size)
-                 in
-                 let open Bigarray.Array2 in
-                   unsafe_set a ij kl (value +. unsafe_get a ij kl) ;
-                   unsafe_set a kj il (value +. unsafe_get a kj il) ;
-                   unsafe_set a il kj (value +. unsafe_get a il kj) ;
-                   unsafe_set a kl ij (value +. unsafe_get a kl ij) ;
-                   unsafe_set a ji lk (value +. unsafe_get a ji lk) ;
-                   unsafe_set a li jk (value +. unsafe_get a li jk) ;
-                   unsafe_set a jk li (value +. unsafe_get a jk li) ;
-                   unsafe_set a lk ji (value +. unsafe_get a lk ji) 
-                )
 
-  | Sparse a -> let key = key_of_indices ~r1 ~r2 in
+  else if i=j then
-                let old_value =
+    if k=l then 
+      unsafe_get t.two_index_anti i k
+    else
+      unsafe_get t.three_index_anti (dense_index k l t.size) i
+
+  else if k=l then
+      (* <ij|kk> *)
+      unsafe_get t.three_index_anti (dense_index i j t.size) k
+
+  else
+    match t.four_index with
+    | Dense  a -> unsafe_get a (dense_index i j t.size) (dense_index k l t.size)
+    | Sparse a -> let key = key_of_indices ~r1 ~r2 in
                   try Hashtbl.find a key
                   with Not_found -> 0.
-                in
+
-                Hashtbl.replace a key (old_value +. value)
+
+let get_four_index ~r1 ~r2 t = 
+  check_bounds r1 r2 t;
+  unsafe_get_four_index ~r1 ~r2 t  
+
+
+let unsafe_set_four_index ~r1 ~r2 ~value t = 
+
+
+  let open Bigarray.Array2 in
+  let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in 
+
+  if i=k then
+    if j=l then
+      begin
+        unsafe_set t.two_index i j value;
+        unsafe_set t.two_index j i value;
+      end
+    else
+      begin
+        unsafe_set t.three_index (dense_index j l t.size) i value;
+        unsafe_set t.three_index (dense_index l j t.size) i value;
+      end
+  else if j=l then
+      begin
+        unsafe_set t.three_index (dense_index i k t.size) j value;
+        unsafe_set t.three_index (dense_index k i t.size) j value;
+      end
+
+  else if i=l then
+    if j=k then
+      begin
+        unsafe_set t.two_index_anti i j value;
+        unsafe_set t.two_index_anti j i value;
+      end
+    else
+      begin
+        unsafe_set t.three_index_anti (dense_index j k t.size) i value;
+        unsafe_set t.three_index_anti (dense_index k j t.size) i value;
+      end
+  else if j=k then
+      begin
+        unsafe_set t.three_index_anti (dense_index i l t.size) j value;
+        unsafe_set t.three_index_anti (dense_index l i t.size) j value;
+      end
+
+  else if i=j then 
+    if k=l then
+      begin
+        unsafe_set t.two_index_anti i k value;
+        unsafe_set t.two_index_anti k i value;
+      end
+    else 
+      (* <ii|kl> *)
+      begin
+        unsafe_set t.three_index_anti (dense_index k l t.size) i value;
+        unsafe_set t.three_index_anti (dense_index l k t.size) i value;
+      end
+
+  else if k=l then
+      (* <ij|kk> *)
+      begin
+        unsafe_set t.three_index_anti (dense_index i j t.size) k value;
+        unsafe_set t.three_index_anti (dense_index j i t.size) k value;
+      end
+
+  else
+
+    match t.four_index with
+    | Dense  a -> let ij = (dense_index i j t.size)
+                  and kl = (dense_index k l t.size)
+                  and il = (dense_index i l t.size)
+                  and kj = (dense_index k j t.size)
+                  and ji = (dense_index j i t.size)
+                  and lk = (dense_index l k t.size)
+                  and li = (dense_index l i t.size)
+                  and jk = (dense_index j k t.size)
+                  in
+                  begin
+                    unsafe_set a ij kl value;
+                    unsafe_set a kj il value;
+                    unsafe_set a il kj value;
+                    unsafe_set a kl ij value;
+                    unsafe_set a ji lk value;
+                    unsafe_set a li jk value;
+                    unsafe_set a jk li value;
+                    unsafe_set a lk ji value
+                  end
+    | Sparse a -> let key = key_of_indices ~r1 ~r2 in
+                  Hashtbl.replace a key value
+    
+
+let set_four_index ~r1 ~r2 ~value t = 
+  check_bounds r1 r2 t;
+  unsafe_set_four_index ~r1 ~r2 ~value t
+
+
+let unsafe_increment_four_index ~r1 ~r2 ~value t = 
+  let updated_value =
+    value +. unsafe_get_four_index ~r1 ~r2 t
+  in
+  unsafe_set_four_index ~r1 ~r2 ~value:updated_value t
+
+
+let increment_four_index ~r1 ~r2 ~value t = 
+  check_bounds r1 r2 t;
+  unsafe_increment_four_index ~r1 ~r2 ~value t
+
 
 let get ~r1 ~r2 a = 
   get_four_index ~r1 ~r2 a
@@ -124,21 +208,39 @@ let set ~r1 ~r2 ~value =
     in
     raise (Invalid_argument msg)
 
+
 let increment ~r1 ~r2 = 
   increment_four_index ~r1 ~r2 
 
 let create ~size ?(temp_dir="/dev/shm") sparsity =
   assert (size < max_index);
+  let two_index = 
+      SharedMemory.create ~temp_dir Float64 [| size ; size |]
+      |> Bigarray.array2_of_genarray
+  in
+  let two_index_anti = 
+      SharedMemory.create ~temp_dir Float64 [| size ; size |]
+      |> Bigarray.array2_of_genarray
+  in
+  let three_index = 
+      SharedMemory.create ~temp_dir Float64 [| size * size ; size |]
+      |> Bigarray.array2_of_genarray
+  in
+  let three_index_anti = 
+      SharedMemory.create ~temp_dir Float64 [| size * size ; size |]
+      |> Bigarray.array2_of_genarray
+  in
   let four_index = 
     match sparsity with
     | `Dense  -> let result =
-                     SharedMemory.create Float64 [| size*size ; size*size |]
+                     SharedMemory.create ~temp_dir Float64 [| size*size ; size*size |]
                      |> Bigarray.array2_of_genarray
-                 in Dense result
+                 in
+                 Dense result
     | `Sparse -> let result = Hashtbl.create (size*size+13) in
       Sparse result
   in
-  { size ; four_index }
+  { size ; two_index ; two_index_anti ; three_index ; three_index_anti ; four_index }
 
 
 let size t = t.size

run_integrals.ml

@@ -6,7 +6,7 @@ let nuclei_file : string option ref = ref None
 let speclist = [
   ( "-b" , Arg.String (fun x -> basis_file := Some x),
     "File containing the atomic basis set") ;
-  ( "-c" , Arg.String (fun x -> nuclei_file := Some x),
+  ( "-x" , Arg.String (fun x -> nuclei_file := Some x),
     "File containing the nuclear coordinates") ;
   ( "-o" , Arg.String (fun x -> out_file := Some x) ,
     "Output file") ;
@@ -23,7 +23,7 @@ let run ~out =
     | Some x -> x
   and nuclei_file = 
     match !nuclei_file with
-    | None -> raise (Invalid_argument "Coordinate file should be specified with -c")
+    | None -> raise (Invalid_argument "Coordinate file should be specified with -x")
     | Some x -> x
   in
 
@@ -41,12 +41,14 @@ let run ~out =
   let eN_ints  = AOBasis.eN_ints  ao_basis in
   let kin_ints = AOBasis.kin_ints ao_basis in
   let ee_ints  = AOBasis.ee_ints  ao_basis in
-  let f12_ints  = AOBasis.f12_ints  ao_basis in
   Overlap.to_file ~filename:(out_file^".overlap") overlap;
   NucInt.to_file ~filename:(out_file^".nuc") eN_ints;
   KinInt.to_file ~filename:(out_file^".kin") kin_ints;
   ERI.to_file ~filename:(out_file^".eri") ee_ints;
+  (*
+  let f12_ints  = AOBasis.f12_ints  ao_basis in
   F12.to_file ~filename:(out_file^".f12") f12_ints;
+  *)
   ()