Optimizations in localization

Basis/Multipole.ml

@@ -21,6 +21,7 @@ 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)
@@ -215,8 +216,8 @@ let of_basis basis =
       done;
     done;
   done;
-  for k=0 to 14 do
-    Mat.detri result.(k);
+  for k=0 to Array.length result - 1 do
+    Mat.detri result.(k)
   done;
   result
 

Basis/Multipole.mli

@@ -31,6 +31,9 @@ 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 $$ %} *)
 

MOBasis/Localisation.ml

@@ -361,10 +361,14 @@ let localize mo_basis =
             then 0.
         else if i = j2 && j=2
             then 0.
-        else if  +. 0.25 *. atan2 m_b12.{i,j} m_a12.{i,j} >= 0.
-            then pi /. 4. -. 0.25 *. atan2 m_b12.{i,j} m_a12.{i,j}
-            else -. pi /. 4. -. 0.25 *. atan2 m_b12.{i,j} m_a12.{i,j}),
-            Vec.sum(Vec.init n_mo ( fun i -> (phi_x_phi.{i,i})**2. +. (phi_y_phi.{i,i})**2. +. (phi_z_phi.{i,i})**2.)))          
+        else
+          let x = atan2 m_b12.{i,j} m_a12.{i,j} in
+          if x >= 0.  then 0.25 *. (pi -. x)
+          else -. 0.25 *. (pi +. x)
+        ),
+        Vec.sum @@
+        Vec.init n_mo ( fun i ->
+           (phi_x_phi.{i,i})**2. +. (phi_y_phi.{i,i})**2. +. (phi_z_phi.{i,i})**2.))
   in
 
   (* Function to compute the new angle after rotation between the 2 orbitals j1 and j2, by the Boys localization *)
@@ -386,16 +390,8 @@ let localize mo_basis =
               Multipole.matrix_z multipoles
               |> MOBasis.mo_matrix_of_ao_matrix ~mo_coef:m_C
         in
-        let phi_x2_phi =
-              Multipole.matrix_x2 multipoles 
-              |> MOBasis.mo_matrix_of_ao_matrix ~mo_coef:m_C
-        in    
-        let phi_y2_phi =
-              Multipole.matrix_y2 multipoles 
-              |> MOBasis.mo_matrix_of_ao_matrix ~mo_coef:m_C
-        in
-        let phi_z2_phi =
-              Multipole.matrix_z2 multipoles 
+        let phi_r2_phi =
+              Multipole.matrix_r2 multipoles
               |> MOBasis.mo_matrix_of_ao_matrix ~mo_coef:m_C
         in
 
@@ -423,7 +419,7 @@ let localize mo_basis =
                 else +. 0.25 *. atan2 m_gamma.{i,j} m_beta.{i,j})
         in
         let m_critere_B = Vec.init n_mo (fun i ->
-            phi_x2_phi.{i,i} -. phi_x_phi.{i,i}**2. +. phi_y2_phi.{i,i} -. phi_y_phi.{i,i}**2. +. phi_z2_phi.{i,i} -. phi_z_phi.{i,i}**2.) 
+            phi_r2_phi.{i,i} -. phi_x_phi.{i,i}**2. -. phi_y_phi.{i,i}**2. -. phi_z_phi.{i,i}**2.)
 
         in m_theta, Vec.sum( m_critere_B)
    in