Add option to restrict the particles the force applies to

Author: peastman

README.md

@@ -108,6 +108,20 @@ torch.onnx.export(model=PeriodicForce(),
                   dynamic_axes={"positions":[0], "forces":[0]})
 ```
 
+## Applying to a Subset of Particles
+
+In some cases one wants to model part of a system with a machine learning potential and the rest with a
+conventional force field.  You can restrict which particles the `OnnxForce` acts on by calling `setParticleIndices()`.
+For example, the following applies it only to the first 50 particles in the system.
+
+```python
+particles = list(range(50))
+force.setParticleIndices(particles)
+```
+
+The `positions` tensor passed to the model will contain only the positions of the specified particles.
+Likewise, the `forces` tensor returned by the model should contain only the forces on those particles.
+
 ## Global Parameters
 
 An `OnnxForce` can define global parameters that the model depends on.  The model should have an additional

openmmapi/include/OnnxForce.h

@@ -101,6 +101,16 @@ class OPENMM_EXPORT_ONNX OnnxForce : public OpenMM::Force {
      * Set the execution provider to be used for computing the model.
      */
     void setExecutionProvider(ExecutionProvider provider);
+    /**
+     * Get the indices of the particles this force is applied to.  If this is empty, the
+     * force is applied to all particles in the system.
+     */
+    const std::vector<int>& getParticleIndices() const;
+    /**
+     * Set the indices of the particles this force is applied to.  If this is empty, the
+     * force is applied to all particles in the system.
+     */
+    void setParticleIndices(const std::vector<int>& indices);
     /**
      * Get whether this force uses periodic boundary conditions.
      */
@@ -169,6 +179,7 @@ class OPENMM_EXPORT_ONNX OnnxForce : public OpenMM::Force {
     class GlobalParameterInfo;
     void initProperties(const std::map<std::string, std::string>& properties);
     std::vector<uint8_t> model;
+    std::vector<int> particleIndices;
     ExecutionProvider provider;
     bool periodic;
     std::vector<GlobalParameterInfo> globalParameters;

openmmapi/include/internal/OnnxForceImpl.h

@@ -60,6 +60,7 @@ class OPENMM_EXPORT_ONNX OnnxForceImpl : public OpenMM::CustomCPPForceImpl {
     Ort::Session session;
     std::vector<Ort::Value> inputTensors, outputTensors;
     std::vector<const char*> inputNames;
+    std::vector<int> particleIndices;
     std::vector<float> positionVec, paramVec;
     float boxVectors[9];
 };

openmmapi/src/OnnxForce.cpp

@@ -74,6 +74,14 @@ void OnnxForce::setExecutionProvider(OnnxForce::ExecutionProvider provider) {
     this->provider = provider;
 }
 
+const vector<int>& OnnxForce::getParticleIndices() const {
+    return particleIndices;
+}
+
+void OnnxForce::setParticleIndices(const vector<int>& indices) {
+    particleIndices = indices;
+}
+
 bool OnnxForce::usesPeriodicBoundaryConditions() const {
     return periodic;
 }

openmmapi/src/OnnxForceImpl.cpp

@@ -47,6 +47,15 @@ OnnxForceImpl::~OnnxForceImpl() {
 void OnnxForceImpl::initialize(ContextImpl& context) {
     CustomCPPForceImpl::initialize(context);
 
+    // Record which particles the force is applied to.
+
+    particleIndices = owner.getParticleIndices();
+    if (particleIndices.size() == 0) {
+        int numParticles = context.getSystem().getNumParticles();
+        for (int i = 0; i < numParticles; i++)
+            particleIndices.push_back(i);
+    }
+
     // Select the execution provider and set options.
 
     OnnxForce::ExecutionProvider provider = owner.getExecutionProvider();
@@ -97,11 +106,10 @@ void OnnxForceImpl::initialize(ContextImpl& context) {
 
     const vector<uint8_t>& model = owner.getModel();
     session = Session(env, model.data(), model.size(), options);
-    int numParticles = context.getSystem().getNumParticles();
-    positionVec.resize(3*numParticles);
+    positionVec.resize(3*particleIndices.size());
     paramVec.resize(owner.getNumGlobalParameters());
     auto memoryInfo = MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
-    int64_t positionsShape[] = {numParticles, 3};
+    int64_t positionsShape[] = {static_cast<int64_t>(particleIndices.size()), 3};
     int64_t boxShape[] = {3, 3};
     int64_t paramShape[] = {1};
     int numInputs = 1+owner.getNumGlobalParameters();
@@ -129,11 +137,12 @@ map<string, double> OnnxForceImpl::getDefaultParameters() {
 double OnnxForceImpl::computeForce(ContextImpl& context, const vector<Vec3>& positions, vector<Vec3>& forces) {
     // Pass the current state to ONNX Runtime.
 
-    int numParticles = context.getSystem().getNumParticles();
+    int numParticles = particleIndices.size();
     for (int i = 0; i < numParticles; i++) {
-        positionVec[3*i] = (float) positions[i][0];
-        positionVec[3*i+1] = (float) positions[i][1];
-        positionVec[3*i+2] = (float) positions[i][2];
+        int index = particleIndices[i];
+        positionVec[3*i] = (float) positions[index][0];
+        positionVec[3*i+1] = (float) positions[index][1];
+        positionVec[3*i+2] = (float) positions[index][2];
     }
     if (owner.usesPeriodicBoundaryConditions()) {
         Vec3 box[3];
@@ -152,6 +161,6 @@ double OnnxForceImpl::computeForce(ContextImpl& context, const vector<Vec3>& pos
     const float* energy = outputTensors[0].GetTensorData<float>();
     const float* forceData = outputTensors[1].GetTensorData<float>();
     for (int i = 0; i < numParticles; i++)
-        forces[i] = Vec3(forceData[3*i], forceData[3*i+1], forceData[3*i+2]);
+        forces[particleIndices[i]] = Vec3(forceData[3*i], forceData[3*i+1], forceData[3*i+2]);
     return *energy;
 }

python/onnxplugin.i

@@ -34,6 +34,7 @@
 
 namespace std {
     %template(vectorbyte) vector<unsigned char>;
+    %template(vectorint) vector<int>;
     %template(property_map) map<std::string, std::string>;
 };
 
@@ -53,6 +54,8 @@ public:
     const std::vector<uint8_t>& getModel() const;
     ExecutionProvider getExecutionProvider() const;
     void setExecutionProvider(ExecutionProvider provider);
+    const std::vector<int>& getParticleIndices() const;
+    void setParticleIndices(const std::vector<int>& indices);
     bool usesPeriodicBoundaryConditions() const;
     void setUsesPeriodicBoundaryConditions(bool periodic);
     int getNumGlobalParameters() const;

python/tests/TestOnnxForce.py

@@ -15,7 +15,8 @@ def testConstructors():
 
 @pytest.mark.parametrize('use_cv_force', [True, False])
 @pytest.mark.parametrize('platform', [mm.Platform.getPlatform(i).getName() for i in range(mm.Platform.getNumPlatforms())])
-def testForce(use_cv_force, platform):
+@pytest.mark.parametrize('particles', [[], [5,3,0]])
+def testForce(use_cv_force, platform, particles):
 
     # Create a random cloud of particles.
     numParticles = 10
@@ -26,6 +27,7 @@ def testForce(use_cv_force, platform):
 
     # Create a force
     force = openmmonnx.OnnxForce('../../tests/central.onnx', {'UseGraphs': 'false'})
+    force.setParticleIndices(particles)
     assert force.getProperties()['UseGraphs'] == 'false'
     if use_cv_force:
         # Wrap OnnxForce into CustomCVForce
@@ -45,9 +47,14 @@ def testForce(use_cv_force, platform):
     state = context.getState(getEnergy=True, getForces=True)
 
     # See if the energy and forces are correct.  The network defines a potential of the form E(r) = |r|^2
+    if len(particles) > 0:
+        positions = positions[particles]
     expectedEnergy = np.sum(positions*positions)
     assert np.allclose(expectedEnergy, state.getPotentialEnergy().value_in_unit(unit.kilojoules_per_mole))
-    assert np.allclose(-2*positions, state.getForces(asNumpy=True))
+    forces = state.getForces(asNumpy=True)
+    if len(particles) > 0:
+        forces = forces[particles]
+    assert np.allclose(-2*positions, forces)
 
 def testProperties():
     """ Test that the properties are correctly set and retrieved """

serialization/src/OnnxForceProxy.cpp

@@ -70,6 +70,10 @@ void OnnxForceProxy::serialize(const void* object, SerializationNode& node) cons
     node.setStringProperty("model", hexEncode(force.getModel()));
     node.setIntProperty("forceGroup", force.getForceGroup());
     node.setBoolProperty("usesPeriodic", force.usesPeriodicBoundaryConditions());
+    const vector<int>& indices = force.getParticleIndices();
+    auto& indicesNode = node.createChildNode("ParticleIndices");
+    for (int i = 0; i < indices.size(); i++)
+        indicesNode.createChildNode("Particle").setIntProperty("index", indices[i]);
     SerializationNode& globalParams = node.createChildNode("GlobalParameters");
     for (int i = 0; i < force.getNumGlobalParameters(); i++)
         globalParams.createChildNode("Parameter").setStringProperty("name", force.getGlobalParameterName(i)).setDoubleProperty("default", force.getGlobalParameterDefaultValue(i));
@@ -86,6 +90,12 @@ void* OnnxForceProxy::deserialize(const SerializationNode& node) const {
     force->setForceGroup(node.getIntProperty("forceGroup"));
     force->setUsesPeriodicBoundaryConditions(node.getBoolProperty("usesPeriodic"));
     for (const SerializationNode& child : node.getChildren()) {
+        if (child.getName() == "ParticleIndices") {
+            vector<int> indices;
+            for (auto& particle : child.getChildren())
+                indices.push_back(particle.getIntProperty("index"));
+            force->setParticleIndices(indices);
+        }
         if (child.getName() == "GlobalParameters")
             for (auto& parameter : child.getChildren())
                 force->addGlobalParameter(parameter.getStringProperty("name"), parameter.getDoubleProperty("default"));

serialization/tests/TestSerializeOnnxForce.cpp

@@ -49,6 +49,7 @@ void testSerialization() {
     force.addGlobalParameter("y", 2.221);
     force.setUsesPeriodicBoundaryConditions(true);
     force.setProperty("UseGraphs", "true");
+    force.setParticleIndices({0, 2, 4});
 
     // Serialize and then deserialize it.
 
@@ -61,6 +62,7 @@ void testSerialization() {
     OnnxForce& force2 = *copy;
     ASSERT_EQUAL_CONTAINERS(force.getModel(), force2.getModel());
     ASSERT_EQUAL(force.getForceGroup(), force2.getForceGroup());
+    ASSERT_EQUAL_CONTAINERS(force.getParticleIndices(), force2.getParticleIndices());
     ASSERT_EQUAL(force.getNumGlobalParameters(), force2.getNumGlobalParameters());
     for (int i = 0; i < force.getNumGlobalParameters(); i++) {
         ASSERT_EQUAL(force.getGlobalParameterName(i), force2.getGlobalParameterName(i));

tests/TestOnnxForce.cpp

@@ -36,6 +36,7 @@
 #include "openmm/System.h"
 #include "openmm/VerletIntegrator.h"
 #include "sfmt/SFMT.h"
+#include <algorithm>
 #include <cmath>
 #include <iostream>
 #include <vector>
@@ -44,7 +45,7 @@ using namespace OnnxPlugin;
 using namespace OpenMM;
 using namespace std;
 
-void testForce(Platform& platform) {
+void testForce(Platform& platform, vector<int> particleIndices) {
     // Create a random cloud of particles.
 
     const int numParticles = 10;
@@ -57,6 +58,10 @@ void testForce(Platform& platform) {
         positions[i] = Vec3(genrand_real2(sfmt), genrand_real2(sfmt), genrand_real2(sfmt))*10;
     }
     OnnxForce* force = new OnnxForce("tests/central.onnx");
+    force->setParticleIndices(particleIndices);
+    if (particleIndices.size() == 0)
+        for (int i = 0; i < numParticles; i++)
+            particleIndices.push_back(i);
     system.addForce(force);
 
     // Compute the forces and energy.
@@ -70,10 +75,16 @@ void testForce(Platform& platform) {
 
     double expectedEnergy = 0;
     for (int i = 0; i < numParticles; i++) {
-        Vec3 pos = positions[i];
-        double r = sqrt(pos.dot(pos));
-        expectedEnergy += r*r;
-        ASSERT_EQUAL_VEC(pos*(-2.0), state.getForces()[i], 1e-5);
+        if (find(particleIndices.begin(), particleIndices.end(), i) != particleIndices.end()) {
+            Vec3 pos = positions[i];
+            double r = sqrt(pos.dot(pos));
+            expectedEnergy += r*r;
+            ASSERT_EQUAL_VEC(pos*(-2.0), state.getForces()[i], 1e-5);
+        }
+        else {
+            Vec3 zero;
+            ASSERT_EQUAL_VEC(zero, state.getForces()[i], 1e-5);
+        }
     }
     ASSERT_EQUAL_TOL(expectedEnergy, state.getPotentialEnergy(), 1e-5);
 }
@@ -164,7 +175,8 @@ void testGlobal(Platform& platform) {
 }
 
 void testPlatform(Platform& platform) {
-    testForce(platform);
+    testForce(platform, {});
+    testForce(platform, {0, 1, 2, 9, 5});
     testPeriodicForce(platform);
     testGlobal(platform);
 }