Ani numeric differentiation

Code/ForceField/ANI/AtomicContrib.cpp

@@ -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;
 }
@@ -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;
+  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 {
 

Code/GraphMol/Descriptors/AtomicEnvironmentVector.cpp

@@ -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);
 
@@ -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;
@@ -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,
@@ -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,
@@ -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);
       }
@@ -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++) {
@@ -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) {
@@ -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());
@@ -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");
@@ -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);
 
@@ -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) {
@@ -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;
@@ -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);
@@ -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();
@@ -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

Code/GraphMol/Descriptors/AtomicEnvironmentVector.h

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

Code/GraphMol/Descriptors/AtomicEnvironmentVector_catch.cpp

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

Code/GraphMol/ForceFieldHelpers/ANI/ANI.h

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