Skip to content

Commit

Permalink
implemented gradients of pairs in ESD approximation
Browse files Browse the repository at this point in the history
  • Loading branch information
@hochej
hochej committed May 27, 2021
1 parent 3341c2e commit fcede88
Show file tree
Hide file tree
Showing 11 changed files with 256 additions and 89 deletions.
160 changes: 86 additions & 74 deletions src/fmo/gradients/embedding.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,69 +7,62 @@ use nalgebra::Vector3;
use ndarray::prelude::*;
use ndarray::RawData;
use std::iter::FromIterator;
use std::ops::AddAssign;
use crate::utils::array_helper::ToOwnedF;
use std::ops::{AddAssign, SubAssign};

impl SuperSystem {

/// Computes and returns the gradient of the embedding energy.
pub fn embedding_gradient(&mut self) -> Array1<f64> {
// initialize the gradient of the embedding energy array with zeros

// The gradient of the embedding energy is initialized as an array with zeros.
let mut gradient: Array1<f64> = Array1::zeros([3 * self.atoms.len()]);

// A reference to the charge differences and gamma matrix for all atoms is created.
let dq: ArrayView1<f64> = self.properties.dq().unwrap();
let gamma: ArrayView2<f64> = self.properties.gamma().unwrap();

// broadcast the charge differences into the shape of the gradients (n_atoms -> 3 * n_atoms)
// The charge differences are broadcast into the shape the gradients.
let dq_f: Array1<f64> = dq
.broadcast([3, self.atoms.len()])
.unwrap()
.reversed_axes()
.to_owned_f()
.to_owned()
.into_shape([3 * self.atoms.len()])
.unwrap();

// combine all derivatives of the charge differences into one Array that has the shape of
// the total gradient. Only the diagonal elements are used.
let mut grad_dq: Array1<f64> = Array1::zeros([3 * self.atoms.len()]);

// PARALLEL or this could also be done at once instead of getting the diag in the loop. However
// the disadvantage would be that we need to create the 3D Array and store it temporarily.
for mol in self.monomers.iter() {
let mol_grad_dq: ArrayView3<f64> = mol
.properties
.grad_dq()
.unwrap()
.into_shape([3, mol.n_atoms, mol.n_atoms])
.unwrap();
grad_dq
.slice_mut(s![mol.slice.grad])
.assign(&diag_of_last_dimensions(mol_grad_dq, mol.n_atoms));
}
// The derivative of the charge differences is initialized as an array with zeros.
let mut grad_dq: ArrayView1<f64> = self.properties.grad_dq_diag().unwrap();

// TODO: it is not neccessary to calculate the derivative of gamma two times. this should be
// improved! it is already computed in the gradient of the monomer/pair
let grad_gamma_sparse: Array2<f64> =
gamma_gradients_atomwise_2d(&self.gammafunction, &self.atoms, self.atoms.len());
let mut grad_gamma_dot_dq: Array1<f64> = grad_gamma_sparse.dot(&dq);

// compute the gradient of the embedding energy for each pair, that is not treated within
// the ES-dimer approximation PARALLEL
// Begin of the loop to compute the gradient of the embedding energy for each pair.
// PARALLEL
for pair in self.pairs.iter() {
// get references to the corresponding monomers
// References to the corresponding monomers.
let m_i: &Monomer = &self.monomers[pair.i];
let m_j: &Monomer = &self.monomers[pair.j];

// if the derivative is w.r.t to an atom that is within this pair:
// the first part of the equation reads:
// If the derivative is w.r.t to an atom that is within this pair:
// The first part of the equation reads:
// dDeltaE_IJ^V/dR_a x = DDq_a^IJ sum_(K!=I,J)^(N) sum_(C in K) Dq_C^K dgamma_(a C)/dR_(a x)

// Reference to the DDq_a^IJ (difference of charge difference between pair and monomer)
let delta_dq: ArrayView1<f64> = pair.properties.delta_dq().unwrap();
// broadcast Delta dq into the shape of the gradients

// DDq is broadcasted into the shape of the gradients.
let delta_dq_f: Array1<f64> = delta_dq
.broadcast([3, pair.n_atoms])
.unwrap()
.reversed_axes()
.to_owned_f()
.to_owned()
.into_shape([3 * pair.n_atoms])
.unwrap();

// The gradient for a in I is computed and assigned.
gradient.slice_mut(s![m_i.slice.grad]).add_assign(
&(&delta_dq_f.slice(s![..3 * m_i.n_atoms])
* &(&grad_gamma_dot_dq.slice(s![m_i.slice.grad])
Expand All @@ -78,6 +71,7 @@ impl SuperSystem {
.dot(&dq.slice(s![m_i.slice.atom])))),
);

// The gradient for a in J is computed and assigned.
gradient.slice_mut(s![m_j.slice.grad]).add_assign(
&(&delta_dq_f.slice(s![3 * m_i.n_atoms..])
* &(&grad_gamma_dot_dq.slice(s![m_j.slice.grad])
Expand All @@ -86,51 +80,64 @@ impl SuperSystem {
.dot(&dq.slice(s![m_j.slice.atom])))),
);

// the second term...
let grad_delta_dq: Array1<f64> = get_grad_delta_dq(pair, m_i, m_j);
let mut esp_ij: Array1<f64> = Array1::zeros([pair.n_atoms]);
esp_ij
.slice_mut(s![..m_i.n_atoms])
.assign(&m_i.properties.esp_q().unwrap());
esp_ij
.slice_mut(s![m_i.n_atoms..])
.assign(&m_j.properties.esp_q().unwrap());
let esp_ij: Array1<f64> = esp_ij
.broadcast([3, pair.n_atoms])
.unwrap()
.reversed_axes()
.to_owned_f()
.into_shape([3 * pair.n_atoms])
.unwrap();
let gddq_esp: Array1<f64> = &grad_delta_dq * &esp_ij;
gradient
.slice_mut(s![m_i.slice.grad])
.add_assign(&gddq_esp.slice(s![..3 * m_i.n_atoms]));

gradient
.slice_mut(s![m_j.slice.grad])
.add_assign(&gddq_esp.slice(s![3 * m_i.n_atoms..]));

// if the derivative is w.r.t to an atom that is not in this pair, a -> K where K != I,J
// left hand side

// A in monomer I
// Right hand side of Eq. 24, but a is still in I,J
// The difference between the derivative of the charge differences between the monomer
// and the dimer is computed.
// let grad_delta_dq: Array1<f64> = get_grad_delta_dq(pair, m_i, m_j);
//
// // The electrostatic potential (ESP) is collected from the corresponding monomers.
// let mut esp_ij: Array1<f64> = Array1::zeros([pair.n_atoms]);
// esp_ij
// .slice_mut(s![..m_i.n_atoms])
// .assign(&m_i.properties.esp_q().unwrap());
// esp_ij
// .slice_mut(s![m_i.n_atoms..])
// .assign(&m_j.properties.esp_q().unwrap());
//
// // The ESP is transformed into the shape of the gradient.
// let esp_ij: Array1<f64> = esp_ij
// .broadcast([3, pair.n_atoms])
// .unwrap()
// .reversed_axes()
// .to_owned_f()
// .into_shape([3 * pair.n_atoms])
// .unwrap();
//
// // The (elementwise) product of the ESP with the derivative of the pair charge
// // differences is computed.
// let gddq_esp: Array1<f64> = &grad_delta_dq * &esp_ij;
//
// // The gradient of the rhs for a in I is assigned.
// gradient
// .slice_mut(s![m_i.slice.grad])
// .add_assign(&gddq_esp.slice(s![..3 * m_i.n_atoms]));
//
// // The gradient of the rhs for a in J is assigned.
// gradient
// .slice_mut(s![m_j.slice.grad])
// .add_assign(&gddq_esp.slice(s![3 * m_i.n_atoms..]));

// Start of the computation if the derivative is w.r.t to an atom that is not in
// this pair. So that a in K where K != I,J.

// The matrix vector product of the gamma matrix derivative with DDq is computed for a in I
let mut dg_ddq: Array1<f64> = grad_gamma_sparse
.slice(s![0.., m_i.slice.atom])
.dot(&delta_dq.slice(s![..m_i.n_atoms]));

// A in monomer J
// and for a in J.
dg_ddq += &grad_gamma_sparse
.slice(s![0.., m_j.slice.atom])
.dot(&delta_dq.slice(s![m_i.n_atoms..]));

// since K != I,J => the elements were K = I,J are set to zero
// Since K != I,J => the elements where K = I,J are set to zero.
dg_ddq.slice_mut(s![m_i.slice.grad]).assign(&Array1::zeros([3 * m_i.n_atoms]));
dg_ddq.slice_mut(s![m_j.slice.grad]).assign(&Array1::zeros([3 * m_j.n_atoms]));

// The (elementwise) product with the charge differences is computed and assigned.
gradient += &(&dg_ddq * &dq_f);

// right hand side
// Start of the computation of the right hand side of Eq. 25.
// A in monomer I
let mut ddq_gamma: Array1<f64> = delta_dq
.slice(s![..m_i.n_atoms])
Expand All @@ -141,7 +148,7 @@ impl SuperSystem {
.slice(s![m_i.n_atoms..])
.dot(&gamma.slice(s![m_j.slice.atom, 0..]));

// since K != I,J => the elements were K = I,J are set to zero
// Since K != I,J => the elements were K = I,J are set to zero.
ddq_gamma.slice_mut(s![m_i.slice.atom]).assign(&Array1::zeros([m_i.n_atoms]));
ddq_gamma.slice_mut(s![m_j.slice.atom]).assign(&Array1::zeros([m_j.n_atoms]));

Expand All @@ -152,7 +159,7 @@ impl SuperSystem {
.broadcast([3, self.atoms.len()])
.unwrap()
.reversed_axes()
.to_owned_f()
.to_owned()
.into_shape([3 * self.atoms.len()])
.unwrap());
}
Expand All @@ -167,38 +174,43 @@ fn get_grad_delta_dq(pair: &Pair, m_i: &Monomer, m_j: &Monomer) -> Array1<f64> {
let grad_dq_i: ArrayView2<f64> = m_i.properties.grad_dq().unwrap();
let grad_dq_j: ArrayView2<f64> = m_j.properties.grad_dq().unwrap();


// compute the difference between dimers and monomers and take the diagonal values
let mut grad_delta_dq_2d: Array2<f64> = Array2::zeros([3 * pair.n_atoms, pair.n_atoms]);
let mut grad_delta_dq_2d: Array2<f64> = grad_dq.to_owned();

// difference for monomer i
grad_delta_dq_2d
.slice_mut(s![..(3 * m_i.n_atoms), ..m_i.n_atoms])
.assign(&(&grad_dq.slice(s![..(3 * m_i.n_atoms), ..m_i.n_atoms]) - &grad_dq_i));
.sub_assign(&grad_dq_i);

// difference for monomer j
grad_delta_dq_2d
.slice_mut(s![(3 * m_i.n_atoms).., m_i.n_atoms..])
.assign(&(&grad_dq.slice(s![(3 * m_i.n_atoms).., m_i.n_atoms..]) - &grad_dq_j));
.sub_assign(&grad_dq_j);

let grad_delta_dq_3d: Array3<f64> = grad_delta_dq_2d
.into_shape([3, pair.n_atoms, pair.n_atoms])
.unwrap();

diag_of_last_dimensions(grad_delta_dq_3d, pair.n_atoms)
diag_of_last_dimensions(grad_delta_dq_3d)
}

fn diag_of_last_dimensions<S>(data: ArrayBase<S, Ix3>, n_atoms: usize) -> Array1<f64>
pub fn diag_of_last_dimensions<S>(data: ArrayBase<S, Ix3>) -> Array1<f64>
where
S: ndarray::Data<Elem = f64>,
{
// create a temporary array to store the values it will be flattened afterwards
let mut grad_charge: Array2<f64> = Array2::zeros([3, n_atoms]);
let (a, b, c): (usize, usize, usize) = data.dim();
assert_eq!(b, c, "The last two dimension should have the same length");

// A temporary array to store the values is created.
let mut grad_charge: Array2<f64> = Array2::zeros([a, b]);

// take the diagonal of each of the three dimensions
for i in 0..3 {
// The diagonal of each of the three dimensions is saved.
for i in 0..a {
grad_charge
.slice_mut(s![i, ..])
.assign(&data.slice(s![i, .., ..]).diag());
}
// reshape it into a one dimensional array
grad_charge.into_shape([3 * n_atoms]).unwrap()
// The gradient of the charges are reshaped into a one dimensional array.
grad_charge.into_shape([a * b]).unwrap()
}
88 changes: 88 additions & 0 deletions src/fmo/gradients/es_dimer.rs
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,88 @@
use crate::fmo::*;
use crate::initialization::*;
use crate::scc::gamma_approximation::gamma_gradients_atomwise_2d;
use crate::scc::*;
use crate::utils::array_helper::ToOwnedF;
use crate::utils::Timer;
use nalgebra::Vector3;
use ndarray::prelude::*;
use ndarray::RawData;
use std::iter::FromIterator;
use std::ops::{AddAssign, SubAssign};

impl SuperSystem {
/// Computes and returns the gradient of the embedding energy.
pub fn es_dimer_gradient(&mut self) -> Array1<f64> {
let mut gradient: Array1<f64> = Array1::zeros([3 * self.atoms.len()]);

// A reference to the charge differences and gamma matrix for all atoms is created.
let dq: ArrayView1<f64> = self.properties.dq().unwrap();
let grad_dq: ArrayView1<f64> = self.properties.grad_dq_diag().unwrap();

let gamma: ArrayView2<f64> = self.properties.gamma().unwrap();

// The charge differences are broadcast into the shape the gradients.
let dq_f: Array1<f64> = dq
.broadcast([3, self.atoms.len()])
.unwrap()
.reversed_axes()
.to_owned()
.into_shape([3 * self.atoms.len()])
.unwrap();

// TODO: it is not neccessary to calculate the derivative of gamma two times. this should be
// improved! it is already computed in the gradient of the monomer/pair
let grad_gamma_sparse: Array2<f64> =
gamma_gradients_atomwise_2d(&self.gammafunction, &self.atoms, self.atoms.len());

for pair in self.esd_pairs.iter() {
// References to the corresponding monomers
let m_i: &Monomer = &self.monomers[pair.i];
let m_j: &Monomer = &self.monomers[pair.j];

// a in I
let mut gradient_slice: ArrayViewMut1<f64> = gradient.slice_mut(s![m_i.slice.grad]);

// lhs. of Eq. 29 in Ref. [1]
gradient_slice += &(&dq_f.slice(s![m_i.slice.grad])
* &grad_gamma_sparse
.slice(s![m_i.slice.grad, m_j.slice.atom])
.dot(&dq.slice(s![m_j.slice.atom])));

// rhs of Eq. 29 in Ref. [1]
gradient_slice += &(&grad_dq.slice(s![m_i.slice.grad])
* &gamma
.slice(s![m_i.slice.atom, m_j.slice.atom])
.dot(&dq.slice(s![m_j.slice.atom]))
.broadcast([3, m_i.n_atoms])
.unwrap()
.reversed_axes()
.to_owned_f()
.into_shape([3 * m_i.n_atoms])
.unwrap());

// a in J
let mut gradient_slice: ArrayViewMut1<f64> = gradient.slice_mut(s![m_j.slice.grad]);

// lhs of Eq. 30 in Ref. [1]
gradient_slice += &(&dq_f.slice(s![m_j.slice.grad])
* &grad_gamma_sparse
.slice(s![m_j.slice.grad, m_i.slice.atom])
.dot(&dq.slice(s![m_i.slice.atom])));

// rhs of Eq. 30 in Ref [1]
gradient_slice += &(&grad_dq.slice(s![m_j.slice.grad])
* &gamma
.slice(s![m_j.slice.atom, m_i.slice.atom])
.dot(&dq.slice(s![m_i.slice.atom]))
.broadcast([3, m_j.n_atoms])
.unwrap()
.reversed_axes()
.to_owned_f()
.into_shape([3 * m_j.n_atoms])
.unwrap());
}

gradient
}
}
Loading

0 comments on commit fcede88

Please sign in to comment.