Update omm plugin and save_dmff2tf.py for aux (U_ind especially) compatible.

.github/workflows/test_openmm_dmff_plugin.yml

@@ -41,10 +41,11 @@ jobs:
           export OPENMM_INSTALLED_DIR=$CONDA_PREFIX
           export CPPFLOW_INSTALLED_DIR=$CONDA_PREFIX
           export LIBTENSORFLOW_INSTALLED_DIR=$CONDA_PREFIX
-          cmake .. -DOPENMM_DIR=${OPENMM_INSTALLED_DIR} -DCPPFLOW_DIR=${CPPFLOW_INSTALLED_DIR} -DTENSORFLOW_DIR=${LIBTENSORFLOW_INSTALLED_DIR} -DUSE_HIGH_PRECISION=ON
+          cmake .. -DOPENMM_DIR=${OPENMM_INSTALLED_DIR} -DCPPFLOW_DIR=${CPPFLOW_INSTALLED_DIR} -DTENSORFLOW_DIR=${LIBTENSORFLOW_INSTALLED_DIR} -DUSE_HIGH_PRECISION=OFF
           make && make install 
-          make PythonInstall
+          make PythonInstall          
       - name: Run Tests
         run: |
           source $CONDA/bin/activate dmff_omm
-          python -m OpenMMDMFFPlugin.tests.test_dmff_plugin_nve -n 100
+          cd ${GITHUB_WORKSPACE}/backend/
+          python -m OpenMMDMFFPlugin.tests.test_dmff_plugin_nve -n 100 --pdb ../examples/water_fullpol/water_dimer.pdb --model ./openmm_dmff_plugin/python/OpenMMDMFFPlugin/data/admp_water_dimer_aux --has_aux True

backend/openmm_dmff_plugin/openmmapi/include/DMFFForce.h

@@ -50,7 +50,7 @@ typedef double ENERGYTYPE;
 #else
 typedef float FORCETYPE;
 typedef float COORDTYPE;
-typedef float ENERGYTYPE;
+typedef double ENERGYTYPE;
 #endif
 
 namespace DMFFPlugin {
@@ -71,6 +71,12 @@ class OPENMM_EXPORT_DMFF DMFFForce : public OpenMM::Force {
      * @param energyCoefficient : the energy transform coefficient.
      */
     void setUnitTransformCoefficients(const double coordCoefficient, const double forceCoefficient, const double energyCoefficient);
+    /**
+     * @brief Set the has_aux flag when model was saved with auxilary input.
+     * 
+     * @param hasAux  : true if model was saved with auxilary input.
+     */
+    void setHasAux(const bool hasAux);
 
     const std::string& getDMFFGraphFile() const;
     /**
@@ -97,6 +103,13 @@ class OPENMM_EXPORT_DMFF DMFFForce : public OpenMM::Force {
      * @return double
      */
     double getCutoff() const;
+    /**
+     * @brief Get the Has Aux object
+     * 
+     * @return true 
+     * @return false 
+     */
+    bool getHasAux() const;
     void updateParametersInContext(OpenMM::Context& context);
     bool usesPeriodicBoundaryConditions() const {
         return use_pbc;
@@ -106,6 +119,7 @@ class OPENMM_EXPORT_DMFF DMFFForce : public OpenMM::Force {
 private:
     string graph_file;
     bool use_pbc = true;
+    bool has_aux = false;
     double cutoff = 1.2;
     double coordCoeff, forceCoeff, energyCoeff;
 

backend/openmm_dmff_plugin/openmmapi/src/DMFFForce.cpp

@@ -54,12 +54,18 @@ void DMFFForce::setUnitTransformCoefficients(const double coordCoefficient, cons
     energyCoeff = energyCoefficient;
 }
 
+void DMFFForce::setHasAux(const bool hasAux){
+    this->has_aux = hasAux;
+}
+
 double DMFFForce::getCoordUnitCoefficient() const {return coordCoeff;}
 double DMFFForce::getForceUnitCoefficient() const {return forceCoeff;}
 double DMFFForce::getEnergyUnitCoefficient() const {return energyCoeff;}
 
 double DMFFForce::getCutoff() const {return cutoff;}
 
+bool DMFFForce::getHasAux() const {return has_aux;}
+
 const string& DMFFForce::getDMFFGraphFile() const{return graph_file;}
 
 

backend/openmm_dmff_plugin/platforms/cuda/include/CudaDMFFKernels.h

@@ -58,19 +58,23 @@ class CudaCalcDMFFForceKernel : public CalcDMFFForceKernel{
     std::string graph_file;
     cppflow::model jax_model;
     vector<int64_t> coord_shape = vector<int64_t>(2);
+    vector<int64_t> U_ind_shape = vector<int64_t>(2);
     vector<int64_t> box_shape{3, 3};
     vector<int64_t> pair_shape = vector<int64_t>(2);
     vector<int32_t> pairs_v;
-    cppflow::tensor coord_tensor, box_tensor, pair_tensor;
+    cppflow::tensor coord_tensor, box_tensor, pair_tensor, U_ind_tensor;
     vector<cppflow::tensor> output_tensors;
+    vector<double> last_U_ind;
     vector<std::string> operations;
     vector<std::string> input_node_names = vector<std::string>(3);
+    vector<std::string> output_node_names = vector<std::string>(2);
 
     OpenMM::NeighborList neighborList;
     vector<std::set<int>> exclusions;
 
     int natoms;
     double cutoff;
+    bool has_aux;
     ENERGYTYPE dener;
     vector<FORCETYPE> dforce;
     vector<COORDTYPE> dcoord;

backend/openmm_dmff_plugin/platforms/cuda/src/CudaDMFFKernels.cpp

@@ -50,25 +50,50 @@ void CudaCalcDMFFForceKernel::initialize(const System& system, const DMFFForce&
     energyUnitCoeff = force.getEnergyUnitCoefficient();
     coordUnitCoeff = force.getCoordUnitCoefficient();
     cutoff = force.getCutoff();
+    this->has_aux = force.getHasAux();
 
     natoms = system.getNumParticles();
     coord_shape[0] = natoms;
     coord_shape[1] = 3;
     exclusions.resize(natoms);
 
+    if (this->has_aux){
+        U_ind_shape[0] = natoms;
+        U_ind_shape[1] = 3;
+        // Initialize the last_U_ind.
+        for(int ii = 0; ii < natoms * 3; ii ++){
+            last_U_ind.push_back(0.0);
+        }
+    }
+
     // Load the ordinary graph firstly.
     jax_model.init(graph_file);
 
     operations = jax_model.get_operations();
     for (int ii = 0; ii < operations.size(); ii++){
-        if (operations[ii].find("serving")!= std::string::npos){
-            if (operations[ii].find("0")!= std::string::npos){
+        if (operations[ii].find("serving") != std::string::npos){
+            if (operations[ii].find("0") != std::string::npos){
                 input_node_names[0] = operations[ii] + ":0";
             } else if (operations[ii].find("1") != std::string::npos){
                 input_node_names[1] = operations[ii] + ":0";
             } else if (operations[ii].find("2") != std::string::npos){
                 input_node_names[2] = operations[ii] + ":0";
             }
+            // Set up the auxilary input node name. For U_ind
+            if(this->has_aux){
+                if (operations[ii].find("3") != std::string::npos){
+                    input_node_names.push_back(operations[ii] + ":0");
+                }
+            }
+        }
+        // Set up the output names.
+        if (operations[ii].find("PartitionedCall") != std::string::npos){
+            output_node_names[0] = operations[ii] + ":0";
+            output_node_names[1] = operations[ii] + ":1";
+            if(this->has_aux){
+                output_node_names.push_back(operations[ii] + ":2");
+            }
+            break;
         }
     }
 
@@ -123,6 +148,9 @@ double CudaCalcDMFFForceKernel::execute(ContextImpl& context, bool includeForces
     }
     coord_tensor = cppflow::tensor(dcoord, coord_shape);
 
+    // Set input U_ind
+    U_ind_tensor = cppflow::tensor(last_U_ind, U_ind_shape);
+
     // Fetch the neighbor list for input pairs tensor.
     computeNeighborListVoxelHash(
         neighborList,
@@ -145,10 +173,27 @@ double CudaCalcDMFFForceKernel::execute(ContextImpl& context, bool includeForces
     pair_tensor = cppflow::tensor(pairs_v, pair_shape);
 
     // Calculate the energy and forces.
-    output_tensors = jax_model({{input_node_names[0], coord_tensor}, {input_node_names[1], box_tensor}, {input_node_names[2], pair_tensor}}, {"PartitionedCall:0", "PartitionedCall:1"});
-
-    dener = output_tensors[0].get_data<ENERGYTYPE>()[0];
-    dforce = output_tensors[1].get_data<FORCETYPE>();    
+    if (!this->has_aux){
+        output_tensors = jax_model({
+            {input_node_names[0], coord_tensor}, 
+            {input_node_names[1], box_tensor}, 
+            {input_node_names[2], pair_tensor}}, 
+            {output_node_names[0], output_node_names[1]});
+        dener = output_tensors[0].get_data<ENERGYTYPE>()[0];
+        dforce = output_tensors[1].get_data<FORCETYPE>();    
+    } else {
+        output_tensors = jax_model({
+            {input_node_names[0], coord_tensor}, 
+            {input_node_names[1], box_tensor}, 
+            {input_node_names[2], U_ind_tensor}, 
+            {input_node_names[3], pair_tensor}}, 
+            {output_node_names[0], output_node_names[1], output_node_names[2]});
+
+        dener = output_tensors[0].get_data<ENERGYTYPE>()[0];
+        dforce = output_tensors[1].get_data<FORCETYPE>();
+        // Save last U_ind for next step usage.
+        last_U_ind = output_tensors[2].get_data<double>();
+    }
 
 
     // Transform the unit from eV/A to KJ/(mol*nm)

backend/openmm_dmff_plugin/platforms/reference/include/ReferenceDMFFKernels.h

@@ -75,21 +75,25 @@ class ReferenceCalcDMFFForceKernel : public CalcDMFFForceKernel {
     cppflow::model jax_model;
 
     std::vector<int64_t> coord_shape = vector<int64_t>(2); 
-    std::vector<int64_t> coord_shape_1 = vector<int64_t>(2); 
-    std::vector<int64_t> coord_shape_2 = vector<int64_t>(2);
+    vector<int64_t> U_ind_shape = vector<int64_t>(2);
+
     std::vector<int64_t> box_shape{3, 3};
     std::vector<int64_t> pair_shape = vector<int64_t>(2);
 
     std::vector<cppflow::tensor> output;
-    cppflow::tensor coord_tensor, box_tensor, pair_tensor;
+    cppflow::tensor coord_tensor, box_tensor, pair_tensor, U_ind_tensor;
+    vector<cppflow::tensor> output_tensors;
     vector<std::string> operations;
+    vector<double> last_U_ind;
     vector<std::string> input_node_names = vector<std::string>(3);
+    vector<std::string> output_node_names = vector<std::string>(2);
 
     OpenMM::NeighborList neighborList;
     std::vector<std::set<int>> exclusions;
 
     int natoms;
     double cutoff;
+    bool has_aux;
     ENERGYTYPE dener;
     vector<FORCETYPE> dforce;
     vector<COORDTYPE> dcoord;

backend/openmm_dmff_plugin/platforms/reference/src/ReferenceDMFFKernels.cpp

@@ -70,12 +70,22 @@ void ReferenceCalcDMFFForceKernel::initialize(const System& system, const DMFFFo
     energyUnitCoeff = force.getEnergyUnitCoefficient();
     coordUnitCoeff = force.getCoordUnitCoefficient();
     cutoff = force.getCutoff();
+    this->has_aux = force.getHasAux();
 
     natoms = system.getNumParticles();
     coord_shape[0] = natoms;
     coord_shape[1] = 3;
     exclusions.resize(natoms);
 
+    if (this->has_aux){
+        U_ind_shape[0] = natoms;
+        U_ind_shape[1] = 3;
+        // Initialize the last_U_ind.
+        for(int ii = 0; ii < natoms * 3; ii ++){
+            last_U_ind.push_back(0.0);
+        }
+    }
+
     // Load the ordinary graph firstly.
     jax_model.init(graph_file);
 
@@ -88,9 +98,22 @@ void ReferenceCalcDMFFForceKernel::initialize(const System& system, const DMFFFo
                 input_node_names[1] = operations[ii]+":0";
             } else if (operations[ii].find("2") != std::string::npos){
                 input_node_names[2] = operations[ii]+":0";
-            } else {
-                std::cout << "Warning: Unknown input node name: " << operations[ii] << std::endl;
+            } 
+            // Set up the auxilary input node name. For U_ind
+            if(this->has_aux){
+                if (operations[ii].find("3") != std::string::npos){
+                    input_node_names.push_back(operations[ii] + ":0");
+                }
+            }
+        }
+        // Set up the output names.
+        if (operations[ii].find("PartitionedCall") != std::string::npos){
+            output_node_names[0] = operations[ii] + ":0";
+            output_node_names[1] = operations[ii] + ":1";
+            if(this->has_aux){
+                output_node_names.push_back(operations[ii] + ":2");
             }
+            break;
         }
     }
 
@@ -134,6 +157,9 @@ double ReferenceCalcDMFFForceKernel::execute(ContextImpl& context, bool includeF
     }
     coord_tensor = cppflow::tensor(dcoord, coord_shape);
 
+    // Set input U_ind.
+    U_ind_tensor = cppflow::tensor(last_U_ind, U_ind_shape);
+
     // Set input pairs.
     computeNeighborListVoxelHash(
         neighborList, 
@@ -156,12 +182,30 @@ double ReferenceCalcDMFFForceKernel::execute(ContextImpl& context, bool includeF
     }
     pair_shape[0] = totpairs;
     pair_shape[1] = 2;
-    cppflow::tensor pair_tensor = cppflow::tensor(dpairs, pair_shape);
-
-    output = jax_model({{input_node_names[0], coord_tensor}, {input_node_names[1], box_tensor}, {input_node_names[2], pair_tensor}}, {"PartitionedCall:0", "PartitionedCall:1"});
+    pair_tensor = cppflow::tensor(dpairs, pair_shape);
+
+    if (!this->has_aux){
+        output_tensors = jax_model({
+            {input_node_names[0], coord_tensor}, 
+            {input_node_names[1], box_tensor}, 
+            {input_node_names[2], pair_tensor}}, 
+            {output_node_names[0], output_node_names[1]});
+        dener = output_tensors[0].get_data<ENERGYTYPE>()[0];
+        dforce = output_tensors[1].get_data<FORCETYPE>();    
+    } else {
+        output_tensors = jax_model({
+            {input_node_names[0], coord_tensor}, 
+            {input_node_names[1], box_tensor}, 
+            {input_node_names[2], U_ind_tensor}, 
+            {input_node_names[3], pair_tensor}}, 
+            {output_node_names[0], output_node_names[1], output_node_names[2]});
+
+        dener = output_tensors[0].get_data<ENERGYTYPE>()[0];
+        dforce = output_tensors[1].get_data<FORCETYPE>();
+        // Save last U_ind for next step usage.
+        last_U_ind = output_tensors[2].get_data<double>();
+    }
 
-    dener = output[0].get_data<ENERGYTYPE>()[0];
-    dforce = output[1].get_data<FORCETYPE>();
 
     // Transform the unit from output units to KJ/(mol*nm)
     for(int ii = 0; ii < natoms; ii ++){

backend/openmm_dmff_plugin/python/OpenMMDMFFPlugin.i

@@ -47,6 +47,7 @@ public:
     DMFFForce(const string& GraphFile);
 
     void setUnitTransformCoefficients(const double coordCoefficient, const double forceCoefficient, const double energyCoefficient);
+    void setHasAux(const bool hasAux);
 
     /*
      * Add methods for casting a Force to a DMFFForce.

backend/openmm_dmff_plugin/python/OpenMMDMFFPlugin/data/admp_water_dimer_aux/fingerprint.pb

@@ -0,0 +1 @@
+‘Þжµ¾“¯Ó
ãÞúؓžÂæ÷
ûÔ¬Ö°Ô÷®
 ³Ž¢”šŽÿ¶-(°—¶úۄÿƒ
2