ANI Backward Pass

Code/ForceField/ANI/AtomicContrib.h

@@ -0,0 +1,148 @@
+//
+//  Copyright (C) 2020 Manan Goel
+//
+//   @@ All Rights Reserved @@
+//  This file is part of the RDKit.
+//  The contents are covered by the terms of the BSD license
+//  which is included in the file license.txt, found at the root
+//  of the RDKit source tree.
+//
+
+#include <RDGeneral/export.h>
+#ifndef __RD_ANI_H__
+#define __RD_ANI_H__
+
+#include <ForceField/ForceField.h>
+#include <ForceField/Contrib.h>
+
+#ifdef RDK_HAS_EIGEN3
+#include <Eigen/Dense>
+#include <boost/tokenizer.hpp>
+using namespace Eigen;
+
+namespace ForceFields {
+namespace ANI {
+
+//! Atomic contribution term from Atomic Environment Vectors using ANI models
+class RDKIT_FORCEFIELD_EXPORT ANIAtomContrib : public ForceFieldContrib {
+ public:
+  ANIAtomContrib(){};
+  //! Constructor
+  /*!
+    The contribution is found according to the AEV, atom type along with
+    hyperparameters and weights, biases of the nueral network
+
+    \param owner        pointer to the owning ForceField
+    \param atomType     ANI currently supports H.C, N and O
+    \param atomIdx      index of atom in ForceField's positions
+    \param speciesVec   vector with atom wise encoding
+    \param numAtoms     number of atoms in the molecule
+    \param numLayers    number of layers in the neural network
+    \param ensembleSize number of models in the ensemble
+    \param modelType    model types like ANI-1x and ANI-1ccx
+  */
+  ANIAtomContrib(ForceField *owner, int atomType, unsigned int atomIdx,
+                 VectorXi &speciesVec, unsigned int numAtoms,
+                 unsigned int numLayers, unsigned int ensembleSize,
+                 std::string modelType);
+  double getEnergy(double *pos) const;
+  double getEnergy(Eigen::ArrayXXd &aev) const;
+  void getGrad(double *pos, double *grad) const;
+
+  /*!
+    Find atomic contribution according to atom's interactions with other atoms
+    by forward prpogation through the neural network
+
+    \param aev      Atomic Environment Vector of the atom
+
+    \return         Contribtution to the total energy of molecule according to
+                    interactions
+  */
+  double forwardProp(ArrayXXd &aev) const;
+
+  virtual ANIAtomContrib *copy() const { return new ANIAtomContrib(*this); };
+
+ private:
+  int d_atomType;
+  int d_atomIdx;
+  int d_numAtoms;
+  VectorXi d_speciesVec;
+  std::vector<std::vector<MatrixXd>> d_weights;
+  std::vector<std::vector<MatrixXd>> d_biases;
+  double d_selfEnergy;
+  unsigned int d_ensembleSize;
+  std::string d_modelType;
+  std::map<int, std::string> d_atomEncoding = {
+      {0, "H"}, {1, "C"}, {2, "N"}, {3, "O"}};
+  std::map<std::string, ArrayXXd> d_aevParams;
+};
+
+namespace Utils {
+void AngularTerms_d(double cutoff, std::vector<ArrayXXd> &derivatives, ArrayXXd coordinates,
+                    ArrayXXd &vectors12, unsigned int atomIdx, ArrayXXi &atomIndex12Angular, ArrayXXi centralAtomIndex,
+                    const std::map<std::string, ArrayXXd> *params);
+void RadialTerms_d(double cutoff, std::vector<ArrayXXd> &derivatives,
+                   std::map<int, std::vector<int>> &addedMapping,
+                   ArrayXXd &selectedCoordinates,
+                   const std::map<std::string, Eigen::ArrayXXd> *params,
+                   ArrayXXd &distances, ArrayXi &atomIndex12,
+                   unsigned int atomIdx);
+//! CELU activation function
+/*!
+  Continuously Differentiable Exponential Linear Unit Activation function
+  \param input      Vector from hidden layer
+  \param alpha      hyperparameter for CELU
+*/
+void CELU(MatrixXd &input, double alpha);
+
+void CELUGrad(MatrixXd &input, double alpha);
+
+/*!
+  Load model weights from CSV file
+  \param weights    Pointer to array of weights of neural network
+  \param model      Index of model in the ensemble
+  \param weightType Type of weights "weight" or "bias"
+  \param layer      Index of layer of NN
+  \param atomType   Atomic Symbol of atom
+  \param modelType  Architecture being used
+*/
+void loadFromCSV(std::vector<MatrixXd> *weights, unsigned int model,
+                 std::string weightType, unsigned int layer,
+                 std::string atomType, std::string modelType);
+
+/*!
+  Load model weights boost serialized files
+  \param weights    Pointer to array of weights of neural network
+  \param model      Index of model in the ensemble
+  \param weightType Type of weights "weight" or "bias"
+  \param layer      Index of layer of NN
+  \param atomType   Atomic Symbol of atom
+  \param modelType  Architecture being used
+*/
+void loadFromBin(std::vector<MatrixXd> *weights, unsigned int model,
+                 std::string weightType, unsigned int layer,
+                 std::string atomType, std::string modelType);
+
+/*!
+  Load all model weights from a single boost serialized file
+  \param weights    Pointer to array of weights of neural network
+  \param biases     Pointer to array of biases of neural network
+  \param model      Index of model in the ensemble
+  \param atomType   Atomic Symbol of atom
+  \param modelType  Architecture being used
+*/
+void loadFromBin(std::vector<MatrixXd> *weights, std::vector<MatrixXd> *biases,
+                 unsigned int model, std::string atomType,
+                 std::string modelType);
+
+//! Load self energy of atom from modelType/selfEnergies file
+void loadSelfEnergy(double *energy, std::string atomType,
+                    std::string modelType);
+
+std::vector<std::string> tokenize(const std::string &s);
+
+}  // namespace Utils
+}  // namespace ANI
+}  // namespace ForceFields
+#endif
+#endif

Code/ForceField/ANI/CMakeLists.txt

@@ -0,0 +1,7 @@
+if(RDK_HAS_EIGEN)
+rdkit_catch_test(ANIForceFieldCatchTest ANIForceField_catch.cpp  catch_main.cpp
+    LINK_LIBRARIES 
+    DistGeomHelpers ForceFieldHelpers
+    FileParsers MolTransforms SmilesParse 
+    SubstructMatch MolAlign ANIDescriptors ANIForceField)
+endif(RDK_HAS_EIGEN)

Code/ForceField/ANI/Params/ANI-1ccx/selfEnergies

@@ -0,0 +1,4 @@
+H,0=-0.5991501324919538
+C,1=-38.03750806057356
+N,2=-54.67448347695333
+O,3=-75.16043537275567

Code/ForceField/ANI/Params/ANI-1x/selfEnergies

@@ -0,0 +1,4 @@
+H,0=-0.600952980000
+C,1=-38.08316124000
+N,2=-54.70775770000
+O,3=-75.19446356000

Code/ForceField/ANI/catch_main.cpp

@@ -0,0 +1,12 @@
+//
+//  Copyright (C) 2020 Manan Goel
+//   @@ All Rights Reserved @@
+//  This file is part of the RDKit.
+//  The contents are covered by the terms of the BSD license
+//  which is included in the file license.txt, found at the root
+//  of the RDKit source tree.
+//
+
+#define CATCH_CONFIG_MAIN  // This tells Catch to provide a main() - only do
+                           // this in one cpp file
+#include "catch.hpp"

Code/ForceField/CMakeLists.txt

@@ -14,6 +14,12 @@ rdkit_library(ForceField
               MMFF/Params.cpp
               LINK_LIBRARIES Optimizer Trajectory)
 
+if(RDK_HAS_EIGEN)
+rdkit_library(ANIForceField
+             ANI/AtomicContrib.cpp
+             LINK_LIBRARIES Optimizer Trajectory ANIDescriptors EigenSerializer)
+endif(RDK_HAS_EIGEN)
+
 rdkit_headers(Contrib.h
               ForceField.h DEST ForceField)
 
@@ -41,9 +47,15 @@ rdkit_headers(MMFF/AngleBend.h
               MMFF/TorsionConstraint.h
               MMFF/PositionConstraint.h
               MMFF/Params.h DEST ForceField/MMFF)
+if(RDK_HAS_EIGEN)
+rdkit_headers(ANI/AtomicContrib.h DEST ForceField/ANI)
+endif(RDK_HAS_EIGEN)
 
 add_subdirectory(UFF)
 add_subdirectory(MMFF)
+if(RDK_HAS_EIGEN)
+add_subdirectory(ANI)
+endif(RDK_HAS_EIGEN)
 if(RDK_BUILD_PYTHON_WRAPPERS)
 add_subdirectory(Wrap)
 endif()

Code/ForceField/Wrap/CMakeLists.txt

@@ -1,5 +1,6 @@
 remove_definitions(-DRDKIT_FORCEFIELD_BUILD)
 rdkit_python_extension(rdForceField ForceField.cpp
                        DEST ForceField
-                       LINK_LIBRARIES ForceFieldHelpers )
+                       LINK_LIBRARIES ForceFieldHelpers ANIForceField)
 add_pytest(pyForceFieldConstraints ${CMAKE_CURRENT_SOURCE_DIR}/testConstraints.py)
+add_pytest(pyForceFieldANI ${CMAKE_CURRENT_SOURCE_DIR}/testANI.py)