bugfix in gradient of charge differences

the problem was the use of reversed_axes and into_shape as mentioned in rust-ndarray/ndarray#390
mitric-lab · May 28, 2021 · 6dfd3a7 · 6dfd3a7
1 parent fcede88
commit 6dfd3a7
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 34 deletions.
diff --git a/src/fmo/gradients/monomer.rs b/src/fmo/gradients/monomer.rs
@@ -29,6 +29,7 @@ use rayon::iter::ParallelIterator;
 use rayon::prelude::IntoParallelRefMutIterator;
 use std::iter::FromIterator;
 use std::ops::{AddAssign, SubAssign};
+use crate::utils::ToOwnedF;
 
 impl GroundStateGradient for Monomer {
     fn scc_gradient(&mut self, atoms: &[Atom]) -> Array1<f64> {
@@ -137,28 +138,27 @@ impl GroundStateGradient for Monomer {
         // X1.T[nu, f * rho]       --reshape--> X2[nu * f, rho]
         // X2[nu * f, rho]    . P[mu, rho]   -> X3[nu * f, mu];  since P is symmetric -> P = P.T
         // X3[nu * f, mu]          --reshape--> X3[nu, f * mu]
-        // X3.T[f * mu, nu]        --reshape--> W[f, mu, nu]
-        let w: Array3<f64> = -0.5
+        // W.T[f * mu, nu]    . S[mu, nu|    -> WS[f * mu, mu] since S is symmetric -> S = S.T
+        let w_s: Array2<f64> = -0.5
             * grad_s_2d
-                .dot(&p)
-                .reversed_axes()
-                .into_shape([n_orb * f, n_orb])
-                .unwrap()
-                .dot(&p)
-                .into_shape([n_orb, f * n_orb])
-                .unwrap()
-                .reversed_axes()
-                .into_shape([f, n_orb, n_orb])
-                .unwrap();
-
-        // Compute W . S, and contract their last dimension
-        let w_s: Array2<f64> = w.into_shape([f * n_orb, n_orb]).unwrap().dot(&s);
-
-        // compute P . S'; it is necessary to broadcast P into the shape of S'
+            .dot(&p)
+            .reversed_axes()
+            .as_standard_layout()
+            .to_owned()
+            .into_shape([n_orb * f, n_orb])
+            .unwrap()
+            .dot(&p)
+            .into_shape([n_orb, f * n_orb])
+            .unwrap()
+            .reversed_axes()
+            .as_standard_layout()
+            .to_owned()
+            .dot(&s);
+
+        // compute P . S' and contract their last dimension
         let d_grad_s: Array2<f64> = grad_s_2d.dot(&p);
-        //println!("pair d_grad_s {}", w);
 
-        // do the sum of both terms and sum over nu
+        // do the sum of both terms
         let w_plus_ps: Array3<f64> = (&w_s + &d_grad_s).into_shape([f, n_orb, n_orb]).unwrap();
 
         // sum over mu where mu is on atom a
@@ -173,7 +173,6 @@ impl GroundStateGradient for Monomer {
             }
         }
 
-        println!("grad dq {}", grad_dq);
         // Shape of returned Array: [f, n_atoms], f = 3 * n_atoms
         return grad_dq;
     }

diff --git a/src/fmo/gradients/pair.rs b/src/fmo/gradients/pair.rs
@@ -138,22 +138,23 @@ impl GroundStateGradient for Pair {
         // X1.T[nu, f * rho]       --reshape--> X2[nu * f, rho]
         // X2[nu * f, rho]    . P[mu, rho]   -> X3[nu * f, mu];  since P is symmetric -> P = P.T
         // X3[nu * f, mu]          --reshape--> X3[nu, f * mu]
-        // X3.T[f * mu, nu]        --reshape--> W[f, mu, nu]
-        let w: Array3<f64> = -0.5
+        // W.T[f * mu, nu]    . S[mu, nu|    -> WS[f * mu, mu] since S is symmetric -> S = S.T
+        let w_s: Array2<f64> = -0.5
             * grad_s_2d
-                .dot(&p)
-                .reversed_axes()
-                .into_shape([n_orb * f, n_orb])
-                .unwrap()
-                .dot(&p)
-                .into_shape([n_orb, f * n_orb])
-                .unwrap()
-                .reversed_axes()
-                .into_shape([f, n_orb, n_orb])
-                .unwrap();
+            .dot(&p)
+            .reversed_axes()
+            .as_standard_layout()
+            .to_owned()
+            .into_shape([n_orb * f, n_orb])
+            .unwrap()
+            .dot(&p)
+            .into_shape([n_orb, f * n_orb])
+            .unwrap()
+            .reversed_axes()
+            .as_standard_layout()
+            .to_owned()
+            .dot(&s);
 
-        // Compute W . S, and contract their last dimension
-        let w_s: Array2<f64> = w.into_shape([f * n_orb, n_orb]).unwrap().dot(&s);
 
         // compute P . S'; it is necessary to broadcast P into the shape of S'
         let d_grad_s: Array2<f64> = grad_s_2d.dot(&p);
@@ -173,6 +174,7 @@ impl GroundStateGradient for Pair {
                 mu += 1;
             }
         }
+        //println!("grad dq pair {}", grad_dq);
 
         // Shape of returned Array: [f, n_atoms], f = 3 * n_atoms
         return grad_dq;