Skip to content

Ani numeric differentiation #5

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 8 commits into
base: ANI-Integration-Check
Choose a base branch
from
Open
39 changes: 36 additions & 3 deletions Code/ForceField/ANI/AtomicContrib.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -121,8 +121,9 @@ double ANIAtomContrib::forwardProp(ArrayXXd &aev) const {
}

double ANIAtomContrib::getEnergy(double *pos) const {
auto aev = RDKit::Descriptors::ANI::AtomicEnvironmentVector(
pos, this->d_speciesVec, this->d_numAtoms, &(this->d_aevParams));
ArrayXXd aev;
RDKit::Descriptors::ANI::AtomicEnvironmentVector(
aev, pos, this->d_speciesVec, this->d_numAtoms, &(this->d_aevParams));
ArrayXXd row = aev.row(this->d_atomIdx);
return this->ANIAtomContrib::forwardProp(row) + this->d_selfEnergy;
}
Expand All @@ -132,7 +133,39 @@ double ANIAtomContrib::getEnergy(Eigen::ArrayXXd &aev) const {
return this->ANIAtomContrib::forwardProp(row) + this->d_selfEnergy;
}

void ANIAtomContrib::getGrad(double *pos, double *grad) const {}
void ANIAtomContrib::getGrad(double *pos, double *grad) const {
auto initEnergy = this->dp_forceField->calcEnergy(pos);

// + - x movement
pos[3 * this->d_atomIdx] += 1e-5;

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please make the displacement amount a const so that you don't have to keep repeating it.

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How did you pick the value 1e-5?

Copy link
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I came up with 1e-5 with a little bit of experimentation.
I tried smaller values but anything lesser didn't really improve the accuracy from the torchani derivatives.

auto posXEnergy = this->dp_forceField->calcEnergy(pos);

pos[3 * this->d_atomIdx] -= 2e-5;
auto negXEnergy = this->dp_forceField->calcEnergy(pos);

grad[3 * this->d_atomIdx] = (posXEnergy - negXEnergy) / (2e-5);
pos[3 * this->d_atomIdx] += 1e-5;

// + - Y movement
pos[3 * this->d_atomIdx + 1] += 1e-5;
auto posYEnergy = this->dp_forceField->calcEnergy(pos);

pos[3 * this->d_atomIdx + 1] -= 2e-5;
auto negYEnergy = this->dp_forceField->calcEnergy(pos);

grad[3 * this->d_atomIdx + 1] = (posYEnergy - negYEnergy) / (2e-5);
pos[3 * this->d_atomIdx + 1] += 1e-5;

// + - Z movement
pos[3 * this->d_atomIdx + 2] += 1e-5;
auto posZEnergy = this->dp_forceField->calcEnergy(pos);

pos[3 * this->d_atomIdx + 2] -= 2e-5;
auto negZEnergy = this->dp_forceField->calcEnergy(pos);

grad[3 * this->d_atomIdx + 2] = (posZEnergy - negZEnergy) / (2e-5);
pos[3 * this->d_atomIdx + 2] += 1e-5;
}

namespace Utils {

Expand Down
81 changes: 35 additions & 46 deletions Code/GraphMol/Descriptors/AtomicEnvironmentVector.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,6 @@ VectorXi GenerateSpeciesVector(const ROMol &mol) {
// C : 1
// N : 2
// O : 3
// All other atoms : -1
auto numAtoms = mol.getNumAtoms();
VectorXi species(numAtoms);

Expand Down Expand Up @@ -211,7 +210,7 @@ void NeighborPairs(ArrayXXd *coordinates, const VectorXi *species,
}
}

auto x = dist.array() <= cutoff;
auto x = (dist <= cutoff);

// Store atom pairs distance between which is less than cutoff
std::vector<int> indices;
Expand Down Expand Up @@ -239,13 +238,13 @@ void NeighborPairs(ArrayXXd *coordinates, const VectorXi *species,

//! Computes radial terms of the torchANI style atom features
/*!
\param cutoff Maximum distance between 2 atoms to show if they contribute
\param cutoff Maximum distance between 2 atoms to show if they contribute
to each other's environments
\param distances Distances between pairs of
\param distances Distances between pairs of
atoms which are in each other's neighbourhoods
\return Radial terms according
to each pair of distances in the molecule calculated using hard coded
parameters
\param RadialTerms_ terms according to each pair of distances in the molecule
calculated using hard coded parameters
\param params Symmetry Function parameters
*/
template <typename Derived>
void RadialTerms(double cutoff, ArrayBase<Derived> &distances,
Expand Down Expand Up @@ -279,13 +278,13 @@ void RadialTerms(double cutoff, ArrayBase<Derived> &distances,

//! Computes angular terms of the torchANI style atom features
/*!
\param cutoff Maximum distance between 2 atoms to show if they contribute
to each other's environments
\param vector12 Pairs of vectors generated by a
triplet of 3 atoms which are in each other's neighbourhoods
\return Angular
terms according to each pair of distances in the molecule calculated using
hard coded parameters
\param cutoff Maximum distance between 2 atoms to show if they
contribute to each other's environments
\param vector12 Pairs of vectors generated by a triplet of 3 atoms
which are in each other's neighbourhoods
\param AngularTerms_ terms according to each pair of distances in the
molecule calculated using hard coded parameters
\param params Symmetry Function parameters
*/
template <typename Derived>
void AngularTerms(double cutoff, ArrayBase<Derived> &vectors12,
Expand Down Expand Up @@ -324,14 +323,8 @@ void AngularTerms(double cutoff, ArrayBase<Derived> &vectors12,
unsigned int idx = 0;
for (auto zetaIdx = 0; zetaIdx < zeta.size(); zetaIdx++) {
ArrayXXd intermediate(1, ShfZ.size());
intermediate
<< (((-1 * (ShfZ.rowwise() - angles.row(i)).array()).cos().array() +
1)
.array() /
2)
.array()
.pow(zeta(zetaIdx))
.array();
intermediate << (((-1 * (ShfZ.rowwise() - angles.row(i))).cos() + 1) / 2)
.pow(zeta(zetaIdx));
for (unsigned int j = 0; j < intermediate.size(); j++) {
calculatedRowVector(0, j + idx) = intermediate(j);
}
Expand All @@ -355,9 +348,7 @@ void AngularTerms(double cutoff, ArrayBase<Derived> &vectors12,
ArrayXXd calculatedRowVector(1, ShfA.size() * etaA.size());
for (auto etaAidx = 0; etaAidx < etaA.size(); etaAidx++) {
ArrayXXd intermediate(1, ShfA.size());
intermediate << ((ShfA.rowwise() - (distance12sum.array() / 2).row(i))
.pow(2)
.array() *
intermediate << ((ShfA.rowwise() - (distance12sum / 2).row(i)).pow(2) *
-etaA(etaAidx))
.exp();
for (auto j = 0; j < intermediate.size(); j++) {
Expand Down Expand Up @@ -531,7 +522,7 @@ void TripleByMolecules(ArrayXXi *atomIndex12Angular,
index += (m * (m - 1) / 2);
}
}
ArrayXXi sortedLocalIndex12(2, (mask.array() == 1).count());
ArrayXXi sortedLocalIndex12(2, (mask == 1).count());
index = 0;
for (auto i = 0; i < mask.size(); i++) {
if (mask(0, i) == 1) {
Expand All @@ -550,8 +541,7 @@ void TripleByMolecules(ArrayXXi *atomIndex12Angular,
}

sortedLocalIndex12 =
(sortedLocalIndex12.matrix().rowwise() + extraLocalIndex12.transpose())
.array();
(sortedLocalIndex12.matrix().rowwise() + extraLocalIndex12.transpose());

// unsort from last part
ArrayXXi localIndex12(2, sortedLocalIndex12.cols());
Expand Down Expand Up @@ -614,11 +604,10 @@ void IndexAdd(ArrayXXd &vector1, ArrayXXd &vector2, ArrayBase<Derived> &index,
}
}
}
// return vector1;
}

ArrayXXd AtomicEnvironmentVector(
double *pos, const VectorXi &species, unsigned int numAtoms,
void AtomicEnvironmentVector(
ArrayXXd &AEV, double *pos, const VectorXi &species, unsigned int numAtoms,
const std::map<std::string, Eigen::ArrayXXd> *params) {
PRECONDITION(species.size() == numAtoms,
"Species encoding for each atom is required");
Expand All @@ -631,6 +620,11 @@ ArrayXXd AtomicEnvironmentVector(
// distance 5.2 Angstroms. The constant was obtained by authors of torchANI
ArrayXi atomIndex12;
NeighborPairs(&coordinates, &species, 5.2, numAtoms, &atomIndex12);
if (atomIndex12.size() == 0) {
AEV.resize(numAtoms, 384);
AEV = ArrayXXd::Zero(numAtoms, 384);
return;
}
ArrayXXd selectedCoordinates(atomIndex12.rows(), 3);
IndexSelect(&selectedCoordinates, &coordinates, atomIndex12, 0);

Expand Down Expand Up @@ -680,7 +674,7 @@ ArrayXXd AtomicEnvironmentVector(
// Distance cutoff for angular terms is smaller than for radial terms, so we
// construct a smaller neighbor list. The authors of torchANI found that the
// cutoff 3.5 gave best results
ArrayXi evenCloserIndices((distances.array() <= 3.5).count());
ArrayXi evenCloserIndices((distances <= 3.5).count());
unsigned int idx = 0;
for (auto i = 0; i < distances.size(); i++) {
if (distances(i) <= 3.5) {
Expand Down Expand Up @@ -726,11 +720,8 @@ ArrayXXd AtomicEnvironmentVector(

n = atomIndex12Angular.cols();

// pairIndex12 = pairindex12 % n
auto localIndex = pairIndex12.cast<int>();

pairIndex12 =
(localIndex.array() - (localIndex.array() / n).array() * n).array();
pairIndex12.unaryExpr([&](int val) { return val % n; }).cast<int>();

ArrayXi pairIndex12Flattened(2 * pairIndex12.cols());
idx = 0;
Expand Down Expand Up @@ -796,7 +787,7 @@ ArrayXXd AtomicEnvironmentVector(
index.row(i) = triuIndices(species12_(0, i), species12_(1, i));
}
// The constant 10 comes from 10 pairs that can be formed
index = index + (centralAtomIndexArr.array() * 10).array();
index = index + (centralAtomIndexArr * 10);

ArrayXXd angularAEV = ArrayXXd::Zero(10 * numAtoms, 32);
IndexAdd(angularAEV, AngularTerms_, index, 10, numAtoms);
Expand All @@ -811,16 +802,14 @@ ArrayXXd AtomicEnvironmentVector(
atomIdx++;
}

ArrayXXd finalAEV(finalRadialAEV.rows(),
finalRadialAEV.cols() + finalAngularAEV.cols());
finalAEV << finalRadialAEV, finalAngularAEV;

return finalAEV;
AEV.resize(finalRadialAEV.rows(),
finalRadialAEV.cols() + finalAngularAEV.cols());
AEV << finalRadialAEV, finalAngularAEV;
}

ArrayXXd AtomicEnvironmentVector(
const ROMol &mol, const std::map<std::string, Eigen::ArrayXXd> *params,
int confId) {
void AtomicEnvironmentVector(
ArrayXXd &AEV, const ROMol &mol,
const std::map<std::string, Eigen::ArrayXXd> *params, int confId) {
PRECONDITION(mol.getNumConformers() >= 1, "molecule has no conformers");

auto numAtoms = mol.getNumAtoms();
Expand All @@ -838,7 +827,7 @@ ArrayXXd AtomicEnvironmentVector(
pos[3 * i + 2] = atom.z;
}

return AtomicEnvironmentVector(pos, species, numAtoms, params);
AtomicEnvironmentVector(AEV, pos, species, numAtoms, params);
}
} // namespace ANI
} // namespace Descriptors
Expand Down
20 changes: 10 additions & 10 deletions Code/GraphMol/Descriptors/AtomicEnvironmentVector.h
Original file line number Diff line number Diff line change
Expand Up @@ -36,29 +36,29 @@ RDKIT_DESCRIPTORS_EXPORT Eigen::VectorXi GenerateSpeciesVector(
//! Calculates torchANI style symmetry functions combining both radial and
//! angular terms
/*!
\param AEV Variable in which AEVs are to be stored
\param mol Mol object for which symmetry functions are to be found
\param params Symmetry Function parameters
\param confId Conformer ID for the conformer for which symmetry
functions are to be found

\return numAtoms * 384 shaped matrix containing 384 features for every
atom in the input mol consisting of both radial and angular terms
*/
RDKIT_DESCRIPTORS_EXPORT Eigen::ArrayXXd AtomicEnvironmentVector(
const ROMol &mol, const std::map<std::string, Eigen::ArrayXXd> *params,
int confId = -1);
RDKIT_DESCRIPTORS_EXPORT void AtomicEnvironmentVector(
Eigen::ArrayXXd &AEV, const ROMol &mol,
const std::map<std::string, Eigen::ArrayXXd> *params, int confId = -1);

//! Calculates torchANI style symmetry functions combining both radial and
//! angular terms
/*!
\param AEV Variable in which AEVs are to be stored
\param pos Array of positions of atoms
\param species Encoding of atom types with index
\param numAtoms Number of Atoms

\return numAtoms * 384 shaped matrix containing 384 features for every atom in
the input mol consisting of both radial and angular terms
\param params Symmetry Function parameters
*/
RDKIT_DESCRIPTORS_EXPORT Eigen::ArrayXXd AtomicEnvironmentVector(
double *pos, const Eigen::VectorXi &species, unsigned int numAtoms,
RDKIT_DESCRIPTORS_EXPORT void AtomicEnvironmentVector(
Eigen::ArrayXXd &AEV, double *pos, const Eigen::VectorXi &species,
unsigned int numAtoms,
const std::map<std::string, Eigen::ArrayXXd> *params);

} // namespace ANI
Expand Down
12 changes: 9 additions & 3 deletions Code/GraphMol/Descriptors/AtomicEnvironmentVector_catch.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -119,7 +119,9 @@ TEST_CASE("Symmetry Function Accuracy", "[Symmetry Function]") {
RDKit::ROMOL_SPTR mol(RDKit::MolFileToMol(molFile, true, false));
int confId = -1;

auto aev = RDKit::Descriptors::ANI::AtomicEnvironmentVector(*mol, &params, confId);
ArrayXXd aev;
RDKit::Descriptors::ANI::AtomicEnvironmentVector(aev, *mol, &params,
confId);

CHECK(aev.rows() == mol->getNumAtoms());
CHECK(aev.cols() == 384);
Expand Down Expand Up @@ -166,7 +168,9 @@ TEST_CASE("Symmetry Function Accuracy", "[Symmetry Function]") {
RDKit::ROMOL_SPTR mol(RDKit::MolFileToMol(molFile, true, false));
int confId = -1;

auto aev = RDKit::Descriptors::ANI::AtomicEnvironmentVector(*mol, &params, confId);
ArrayXXd aev;
RDKit::Descriptors::ANI::AtomicEnvironmentVector(aev, *mol, &params,
confId);

CHECK(aev.rows() == mol->getNumAtoms());
CHECK(aev.cols() == 384);
Expand Down Expand Up @@ -213,7 +217,9 @@ TEST_CASE("Symmetry Function Accuracy", "[Symmetry Function]") {
RDKit::ROMOL_SPTR mol(RDKit::MolFileToMol(molFile, false, false));
int confId = -1;

auto aev = RDKit::Descriptors::ANI::AtomicEnvironmentVector(*mol, &params, confId);
ArrayXXd aev;
RDKit::Descriptors::ANI::AtomicEnvironmentVector(aev, *mol, &params,
confId);

CHECK(aev.rows() == mol->getNumAtoms());
CHECK(aev.cols() == 384);
Expand Down
61 changes: 61 additions & 0 deletions Code/GraphMol/ForceFieldHelpers/ANI/ANI.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,61 @@
#include <RDGeneral/export.h>
#ifndef RD_ANICONVENIENCE_H_
#define RD_ANICONVENIENCE_H_

#include <ForceField/ForceField.h>
#include <GraphMol/ForceFieldHelpers/FFConvenience.h>
#include "Builder.h"

namespace RDKit {
class ROMol;
namespace ANI {
//! Convenience function for optimizing a molecule using ANI
/*
\param mol the molecule to use
\param maxIters the maximum number of force-field iterations
\param modelType Name of ANI style model (ANI-1x or ANI-1ccx)
\param ensembleSize Number of Neural Networks inside the model
\param confId the optional conformer id, if this isn't provided, the
molecule's default confId will be used.

\return a pair with:
first: 0 if the optimization converged, 1 if more iterations are required.
second: the energy
*/
std::pair<int, double> ANIOptimizeMolecule(ROMol &mol, std::string modelType,
unsigned int ensembleSize,
int confId = -1,
int maxIters = 1000) {
ForceFields::ForceField *ff =
ANI::constructForceField(mol, modelType, ensembleSize, confId);
std::pair<int, double> res =
ForceFieldsHelper::OptimizeMolecule(*ff, maxIters);
delete ff;
return res;
}

//! Convenience function for optimizing all of a molecule's conformations using
// ANI
/*
\param mol the molecule to use
\param res vector of (needsMore,energy)
\param numThreads the number of simultaneous threads to use (only has an
effect if the RDKit is compiled with thread support).
If set to zero, the max supported by the system will be
used.
\param modelType Name of ANI style model (ANI-1x or ANI-1ccx)
\param ensembleSize Number of Neural Networks inside the model
*/
void ANIOptimizeMoleculeConfs(ROMol &mol,
std::vector<std::pair<int, double>> &res,
std::string modelType, unsigned int ensembleSize,
int numThreads = 1, int maxIters = 1000) {
ForceFields::ForceField *ff =
ANI::constructForceField(mol, modelType, ensembleSize, -1);
ForceFieldsHelper::OptimizeMoleculeConfs(mol, *ff, res, numThreads, maxIters);
delete ff;
}
} // namespace ANI
} // namespace RDKit

#endif
Loading