From 6dfd3a7ea9a8cc2c78164b8f51935ada11590838 Mon Sep 17 00:00:00 2001
From: hochej <37776098+hochej@users.noreply.github.com>
Date: Fri, 28 May 2021 19:14:06 +0200
Subject: [PATCH] bugfix in gradient of charge differences

the problem was the use of reversed_axes and into_shape as mentioned in https://github.com/rust-ndarray/ndarray/issues/390
---
 src/fmo/gradients/monomer.rs | 39 ++++++++++++++++++------------------
 src/fmo/gradients/pair.rs    | 30 ++++++++++++++-------------
 2 files changed, 35 insertions(+), 34 deletions(-)

diff --git a/src/fmo/gradients/monomer.rs b/src/fmo/gradients/monomer.rs
index 1ef4a21..df3dff9 100644
--- 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
index 68ad219..f5a16a1 100644
--- 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;