Less allocations in localization

This commit is contained in:
Anthony Scemama 2020-07-09 23:46:38 +02:00
parent 158d149036
commit 5456852b0e
3 changed files with 77 additions and 29 deletions

View File

@ -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)

View File

@ -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

View File

@ -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 b12 g = Mat.init_cols n_mo 2 ( fun i j ->
if j = 1
then (g.{i,i} -. g.{j1,j1}) *. g.{i,j1}
else (g.{i,i} -. g.{j2,j2}) *. g.{i,j2})
let m_b12 =
Mat.init_cols n_mo 2 (fun i j ->
if j = 1 then
(x.{i,i} -. x.{j1,j1}) *. x.{i,j1} +.
(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 ->
if j = 1
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})))
in
Mat.add (a12 phi_x_phi) ( Mat.add (a12 phi_y_phi) (a12 phi_z_phi))
Mat.init_cols n_mo 2 (fun i j ->
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.{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.init n_mo ( fun i ->
(phi_x_phi.{i,i})**2. +. (phi_y_phi.{i,i})**2. +. (phi_z_phi.{i,i})**2.))
Vec.sum @@ Vec.init n_mo ( fun i ->
x.{i,i} *. x.{i,i} +.
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 ->
if j = 1
then (g.{i,i} -. g.{j1,j1})**2. -. 4. *. g.{i,j1}**2.
else (g.{i,i} -. g.{j2,j2})**2. -. 4. *. g.{i,j2}**2.)
in Mat.add (beta phi_x_phi) (Mat.add (beta phi_y_phi) (beta phi_z_phi))
let x = phi_x_phi
and y = phi_y_phi
and z = phi_z_phi in
Mat.init_cols n_mo 2 (fun i j ->
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 =