From 8ded68a447f453925b457869207c8086c62d2a3a Mon Sep 17 00:00:00 2001
From: Anthony Scemama <scemama@irsamc.ups-tlse.fr>
Date: Thu, 28 Jun 2018 14:43:24 +0200
Subject: [PATCH] Added stream in Four_idx

---
 SCF/Fock.ml              | 15 +++++++---
 Utils/FourIdxStorage.ml  | 60 +++++++++++++++++++++++++++++++++-------
 Utils/FourIdxStorage.mli | 11 ++++++++
 Utils/Util.ml            | 13 +++++++++
 Utils/Util.mli           |  2 ++
 5 files changed, 87 insertions(+), 14 deletions(-)

diff --git a/SCF/Fock.ml b/SCF/Fock.ml
index 35ca589..e801dca 100644
--- a/SCF/Fock.ml
+++ b/SCF/Fock.ml
@@ -30,8 +30,13 @@ let make ~density ao_basis =
   for sigma = 1 to nBas do
     for nu = 1 to nBas do
       let m_Jnu = m_J.(nu-1) in
-      for lambda = 1 to nBas do
-        let p = m_P.{lambda,sigma} in
+      for lambda = 1 to sigma do
+        let p = 
+          if lambda < sigma then
+            2. *. m_P.{lambda,sigma}
+          else
+            m_P.{lambda,sigma}
+        in
         if abs_float p > epsilon then
           for mu = 1 to nu do
             m_Jnu.(mu-1) <- m_Jnu.(mu-1) +.  p *. 
@@ -44,10 +49,12 @@ let make ~density ao_basis =
     let m_Knu = m_K.(nu-1) in
     for sigma = 1 to nBas do
       for lambda = 1 to nBas do
-        let p = 0.5 *. m_P.{lambda,sigma} in
+        let p = 
+            0.5 *. m_P.{lambda,sigma}
+        in
         if abs_float p > epsilon then
           for mu = 1 to nu do
-            m_Knu.(mu-1) <- m_Knu.(mu-1) -.  p *. 
+            m_Knu.(mu-1) <- m_Knu.(mu-1) -. p *.
                           ERI.get_phys m_G mu lambda sigma nu
           done
       done
diff --git a/Utils/FourIdxStorage.ml b/Utils/FourIdxStorage.ml
index 7a4f4e3..2b7d6ad 100644
--- a/Utils/FourIdxStorage.ml
+++ b/Utils/FourIdxStorage.ml
@@ -1,3 +1,5 @@
+open Util
+
 let max_index = 1 lsl 14
 
 type index_pair = { first : int ; second : int }
@@ -153,19 +155,57 @@ let set_chem t i j k l value = set ~r1:{ first=i ; second=j } ~r2:{ first=k ; se
 let set_phys t i j k l value = set ~r1:{ first=i ; second=k } ~r2:{ first=j ; second=l } ~value t 
 
 
+
+type element = (** Element for the stream *)
+{
+  i_r1: int ;
+  j_r2: int ;
+  k_r1: int ;
+  l_r2: int ;
+  value: float
+}
+
+
+let to_stream d =
+
+  let i = ref 0
+  and j = ref 1
+  and k = ref 1
+  and l = ref 1
+  in
+  let f _ =
+    i := !i+1;
+    if !i > !k then begin
+      i := 1;
+      j := !j + 1;
+      if !j > !l then begin
+        j := 1;
+        k := !k + 1;
+        if !k > !l then begin
+          k := 1;
+          l := !l + 1;
+        end;
+      end;
+     end;
+     if !l <= d.size then 
+       Some { i_r1 = !i ; j_r2 = !j ;
+             k_r1 = !k ; l_r2 = !l ;
+             value = get_phys d !i !j !k !l
+           }
+     else
+       None
+  in
+  Stream.from f
+    
+
 (** 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
+  to_stream data 
+  |> Stream.iter (fun {i_r1 ; j_r2 ; k_r1 ; l_r2 ; value} ->
            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;
+            Printf.fprintf oc " %5d %5d %5d %5d%20.15f\n" i_r1 j_r2 k_r1 l_r2 value);
   close_out oc
 
+
+
diff --git a/Utils/FourIdxStorage.mli b/Utils/FourIdxStorage.mli
index 0b5fa3d..065c31d 100644
--- a/Utils/FourIdxStorage.mli
+++ b/Utils/FourIdxStorage.mli
@@ -9,6 +9,15 @@ There are two kinds of ordering of indices:
 
 type t
 
+type element = (** Element for the stream *)
+{
+  i_r1: int ;
+  j_r2: int ;
+  k_r1: int ;
+  l_r2: int ;
+  value: float
+}
+
 val create : size:int -> [< `Dense | `Sparse ] -> t
 (** If [`Dense] is chosen, internally the data is stored as a 4-dimensional
     [Bigarray]. Else, it is stored as a hash table.
@@ -20,6 +29,8 @@ val get_phys : t -> int -> int -> int -> int -> float
 val set_chem : t -> int -> int -> int -> int -> float -> unit
 val set_phys : t -> int -> int -> int -> int -> float -> unit
 
+val to_stream : t -> element Stream.t
+
 (** {2 I/O} *)
 val to_file : ?cutoff:float -> filename:string -> t -> unit
 (** Write the data to file, using the physicist's ordering. *)
diff --git a/Utils/Util.ml b/Utils/Util.ml
index bcbf32f..789b3ec 100644
--- a/Utils/Util.ml
+++ b/Utils/Util.ml
@@ -167,6 +167,19 @@ let list_some l =
   |> List.map (function Some x -> x | _ -> assert false)
 
 
+(** {2 Stream functions} *)
+
+let range ?(start=0) n = 
+  Stream.from (fun i ->
+    let result = i+start in
+    if result < n then
+      Some result
+    else None
+  )
+
+
+
+
 (** {2 Linear algebra} *)
 
 
diff --git a/Utils/Util.mli b/Utils/Util.mli
index 21c3cfd..7d9d818 100644
--- a/Utils/Util.mli
+++ b/Utils/Util.mli
@@ -63,6 +63,8 @@ val list_some : 'a option list -> 'a list
 (** Filters out all [None] elements of the list, and returns the elements without
     the [Some]. *)
 
+(** {2 Useful streams} *)
+val range : ?start:int -> int -> int Stream.t
 
 (** {2 Linear algebra } *)