Less allocations in localization

Basis/Multipole.ml

@@ -21,7 +21,6 @@ let matrix_z  t =  t.(2)
 let matrix_x2 t =  t.(3)
 let matrix_y2 t =  t.(4)
 let matrix_z2 t =  t.(5)
-let matrix_r2 t =  Mat.add t.(3) @@ Mat.add t.(4) t.(5)
 let matrix_xy t =  t.(6)
 let matrix_yz t =  t.(7)
 let matrix_zx t =  t.(8)
@@ -32,6 +31,13 @@ let matrix_x4 t =  t.(12)
 let matrix_y4 t =  t.(13)
 let matrix_z4 t =  t.(14)
 
+let add3 a1 a2 a3 = 
+  let result = Mat.add a1 a2 in
+  Mat.add a3 result ~c:result
+
+let matrix_r2 t =
+  add3 t.(3) t.(4) t.(5)
+
 
 
 

Basis/Multipole.mli

@@ -31,9 +31,6 @@ val matrix_y2 : t -> Mat.t
 val matrix_z2 : t -> Mat.t
 (** {% $$ \langle \chi_i | z^2 | \chi_j \rangle $$ %} *)
 
-val matrix_r2 : t -> Mat.t
-(** {% $$ \langle \chi_i | r^2 | \chi_j \rangle $$ %} *)
-
 val matrix_x3 : t -> Mat.t
 (** {% $$ \langle \chi_i | x^3 | \chi_j \rangle $$ %} *)
 
@@ -52,4 +49,9 @@ val matrix_y4 : t -> Mat.t
 val matrix_z4 : t -> Mat.t
 (** {% $$ \langle \chi_i | z^4 | \chi_j \rangle $$ %} *)
 
+
+val matrix_r2 : t -> Mat.t
+(** {% $$ \langle \chi_i | r^2 | \chi_j \rangle $$ %} *)
+
+
 val of_basis : Basis.t -> t

MOBasis/Localisation.ml

@@ -326,36 +326,56 @@ let localize mo_basis =
     let multipoles = AOBasis.multipole ao_basis in
     let n_mo = Mat.dim2 m_C
     in
-    let phi_x_phi =
+    let x =
           Multipole.matrix_x multipoles
           |> MOBasis.mo_matrix_of_ao_matrix ~mo_coef:m_C
     in
-    let phi_y_phi =
+    let y =
           Multipole.matrix_y multipoles
           |> MOBasis.mo_matrix_of_ao_matrix ~mo_coef:m_C
     in
-    let phi_z_phi =
+    let z =
           Multipole.matrix_z multipoles
           |> MOBasis.mo_matrix_of_ao_matrix ~mo_coef:m_C
     in
 
-    let m_b12=
+    let m_b12 = 
-        let b12 g =  Mat.init_cols n_mo 2 ( fun i j ->
+      Mat.init_cols n_mo 2 (fun i j ->
-            if j = 1
+        if j = 1 then
-                then  (g.{i,i} -. g.{j1,j1}) *. g.{i,j1}
+          (x.{i,i} -. x.{j1,j1}) *. x.{i,j1} +. 
-                else  (g.{i,i} -. g.{j2,j2}) *. g.{i,j2})
+          (y.{i,i} -. y.{j1,j1}) *. y.{i,j1} +. 
+          (z.{i,i} -. z.{j1,j1}) *. z.{i,j1}
+        else
+          (x.{i,i} -. x.{j2,j2}) *. x.{i,j2} +. 
+          (y.{i,i} -. y.{j2,j2}) *. y.{i,j2} +. 
+          (z.{i,i} -. z.{j2,j2}) *. z.{i,j2}
+        )
+    in
 
-        in
-        Mat.add (b12 phi_x_phi) ( Mat.add (b12 phi_y_phi) (b12 phi_z_phi))
-    in
     let m_a12 =
-        let a12 g  = Mat.init_cols n_mo 2 ( fun i j ->
+      Mat.init_cols n_mo 2 (fun i j ->
-            if j = 1
+        if j = 1 then
-                then  g.{i,j1} *. g.{i,j1} -. 0.25 *. ((g.{i,i} -. g.{j1,j1}) *. (g.{i,i} -. g.{j1,j1}))
+
-                else  g.{i,j2} *. g.{i,j2} -. 0.25 *. ((g.{i,i} -. g.{j2,j2}) *. (g.{i,i} -. g.{j2,j2})))
+          let x' = (x.{i,i} -. x.{j1,j1})
-        in
+          and y' = (y.{i,i} -. y.{j1,j1})
-        Mat.add (a12 phi_x_phi) ( Mat.add (a12 phi_y_phi) (a12 phi_z_phi))
+          and z' = (z.{i,i} -. z.{j1,j1}) in
+          x.{i,j1} *. x.{i,j1} +.
+          y.{i,j1} *. y.{i,j1} +. 
+          z.{i,j1} *. z.{i,j1} -.
+          0.25 *. (x' *. x' +. y' *. y' +. z' *. z' )
+
+        else
+
+          let x' = (x.{i,i} -. x.{j2,j2})
+          and y' = (y.{i,i} -. y.{j2,j2})
+          and z' = (z.{i,i} -. z.{j2,j2}) in
+          x.{i,j2} *. x.{i,j2} +.
+          y.{i,j2} *. y.{i,j2} +. 
+          z.{i,j2} *. z.{i,j2} -.
+          0.25 *. (x' *. x' +. y' *. y' +. z' *. z' )
+       )
     in
+
     (Mat.init_cols n_mo 2 ( fun i j ->
         if i=j1 && j=1
             then 0.
@@ -366,9 +386,11 @@ let localize mo_basis =
           if x >= 0.  then 0.25 *. (pi -. x)
           else -. 0.25 *. (pi +. x)
         ),
-        Vec.sum @@
+        Vec.sum @@ Vec.init n_mo ( fun i ->
-        Vec.init n_mo ( fun i ->
+           x.{i,i} *. x.{i,i} +.
-           (phi_x_phi.{i,i})**2. +. (phi_y_phi.{i,i})**2. +. (phi_z_phi.{i,i})**2.))
+           y.{i,i} *. y.{i,i} +.
+           z.{i,i} *. z.{i,i} )
+        )
   in
 
   (* Function to compute the new angle after rotation between the 2 orbitals j1 and j2, by the Boys localization *)
@@ -396,11 +418,29 @@ let localize mo_basis =
         in
 
         let m_beta =
-            let beta g = Mat.init_cols n_mo 2 (fun i j ->
+            let x = phi_x_phi
-                if j = 1
+            and y = phi_y_phi
-                    then (g.{i,i} -. g.{j1,j1})**2. -. 4. *. g.{i,j1}**2.
+            and z = phi_z_phi in
-                    else (g.{i,i} -. g.{j2,j2})**2. -. 4. *. g.{i,j2}**2.)
+            Mat.init_cols n_mo 2 (fun i j ->
-            in Mat.add (beta phi_x_phi) (Mat.add (beta phi_y_phi) (beta phi_z_phi))
+              if j = 1 then
+           
+                let x' = (x.{i,i} -. x.{j1,j1})
+                and y' = (y.{i,i} -. y.{j1,j1})
+                and z' = (z.{i,i} -. z.{j1,j1}) in
+                x' *. x' +. y' *. y' +. z' *. z' 
+                -. 4. *. ( x.{i,j1} *. x.{i,j1} +.
+                           y.{i,j1} *. y.{i,j1} +. 
+                           z.{i,j1} *. z.{i,j1} )
+              else
+           
+                let x' = (x.{i,i} -. x.{j2,j2})
+                and y' = (y.{i,i} -. y.{j2,j2})
+                and z' = (z.{i,i} -. z.{j2,j2}) in
+                x' *. x' +. y' *. y' +. z' *. z' 
+                -. 4. *. ( x.{i,j2} *. x.{i,j2} +.
+                           y.{i,j2} *. y.{i,j2} +. 
+                           z.{i,j2} *. z.{i,j2} )
+            )
         in
 
         let m_gamma =