Fixed bug in CustomIntegrator when atoms get reordered

platforms/opencl/src/OpenCLContext.cpp

@@ -203,6 +203,8 @@ OpenCLContext::OpenCLContext(int numParticles, int platformIndex, int deviceInde
 OpenCLContext::~OpenCLContext() {
     for (int i = 0; i < (int) forces.size(); i++)
         delete forces[i];
+    for (int i = 0; i < (int) reorderListeners.size(); i++)
+        delete reorderListeners[i];
     if (posq != NULL)
         delete posq;
     if (velm != NULL)
@@ -728,6 +730,12 @@ void OpenCLContext::reorderAtoms() {
     posq->upload();
     velm->upload();
     atomIndex->upload();
+    for (int i = 0; i < (int) reorderListeners.size(); i++)
+        reorderListeners[i]->execute();
+}
+
+void OpenCLContext::addReorderListener(ReorderListener* listener) {
+    reorderListeners.push_back(listener);
 }
 
 struct OpenCLContext::WorkThread::ThreadData {

platforms/opencl/src/OpenCLContext.h

@@ -143,6 +143,7 @@ class OPENMM_EXPORT OpenCLContext {
 public:
     class WorkTask;
     class WorkThread;
+    class ReorderListener;
     static const int ThreadBlockSize;
     static const int TileSize;
     OpenCLContext(int numParticles, int platformIndex, int deviceIndex, OpenCLPlatform::PlatformData& platformData);
@@ -467,6 +468,11 @@ public:
      * together in the arrays.
      */
     void reorderAtoms();
+    /**
+     * Add a listener that should be called whenever atoms get reordered.  The OpenCLContext
+     * assumes ownership of the object, and deletes it when the context itself is deleted.
+     */
+    void addReorderListener(ReorderListener* listener);
 private:
     struct Molecule;
     struct MoleculeGroup;
@@ -513,6 +519,7 @@ private:
     OpenCLArray<cl_int>* atomIndex;
     std::vector<cl::Memory*> autoclearBuffers;
     std::vector<int> autoclearBufferSizes;
+    std::vector<ReorderListener*> reorderListeners;
     OpenCLIntegrationUtilities* integration;
     OpenCLBondedUtilities* bonded;
     OpenCLNonbondedUtilities* nonbonded;
@@ -562,6 +569,16 @@ private:
     pthread_t thread;
 };
 
+/**
+ * This abstract class defines a function to be executed whenever atoms get reordered.
+ * Objects that need to know when reordering happens should create a reorderListener
+ * and register it by calling addReorderListener().
+ */
+class OpenCLContext::ReorderListener {
+public:
+    virtual void execute() = 0;
+};
+
 } // namespace OpenMM
 
 #endif /*OPENMM_OPENCLCONTEXT_H_*/

platforms/opencl/src/OpenCLKernels.cpp

@@ -3500,6 +3500,43 @@ void OpenCLIntegrateVariableLangevinStepKernel::execute(ContextImpl& context, co
     cl.setStepCount(cl.getStepCount()+1);
 }
 
+class OpenCLIntegrateCustomStepKernel::ReorderListener : public OpenCLContext::ReorderListener {
+public:
+    ReorderListener(OpenCLContext& cl, OpenCLParameterSet& perDofValues, vector<vector<cl_float> >& localPerDofValues) :
+            cl(cl), perDofValues(perDofValues), localPerDofValues(localPerDofValues) {
+        int numAtoms = cl.getNumAtoms();
+        lastAtomOrder.resize(numAtoms);
+        for (int i = 0; i < numAtoms; i++)
+            lastAtomOrder[i] = cl.getAtomIndex()[i];
+    }
+    void execute() {
+        // Reorder the per-DOF variables to reflect the new atom order.
+
+        int numAtoms = cl.getNumAtoms();
+        perDofValues.getParameterValues(localPerDofValues);
+        vector<vector<cl_float> > swap(3*numAtoms);
+        for (int i = 0; i < numAtoms; i++) {
+            swap[3*lastAtomOrder[i]] = localPerDofValues[3*i];
+            swap[3*lastAtomOrder[i]+1] = localPerDofValues[3*i+1];
+            swap[3*lastAtomOrder[i]+2] = localPerDofValues[3*i+2];
+        }
+        OpenCLArray<cl_int>& order = cl.getAtomIndex();
+        for (int i = 0; i < numAtoms; i++) {
+            localPerDofValues[3*i] = swap[3*order[i]];
+            localPerDofValues[3*i+1] = swap[3*order[i]+1];
+            localPerDofValues[3*i+2] = swap[3*order[i]+2];
+        }
+        perDofValues.setParameterValues(localPerDofValues);
+        for (int i = 0; i < numAtoms; i++)
+            lastAtomOrder[i] = order[i];
+    }
+private:
+    OpenCLContext& cl;
+    OpenCLParameterSet& perDofValues;
+    vector<vector<cl_float> >& localPerDofValues;
+    std::vector<int> lastAtomOrder;
+};
+
 OpenCLIntegrateCustomStepKernel::~OpenCLIntegrateCustomStepKernel() {
     if (globalValues != NULL)
         delete globalValues;
@@ -3619,6 +3656,7 @@ void OpenCLIntegrateCustomStepKernel::execute(ContextImpl& context, CustomIntegr
         invalidatesForces.resize(numSteps, false);
         merged.resize(numSteps, false);
         modifiesParameters = false;
+        cl.addReorderListener(new ReorderListener(cl, *perDofValues, localPerDofValues));
         map<string, string> defines;
         defines["NUM_ATOMS"] = intToString(cl.getNumAtoms());
         defines["WORK_GROUP_SIZE"] = intToString(OpenCLContext::ThreadBlockSize);
@@ -3962,9 +4000,10 @@ void OpenCLIntegrateCustomStepKernel::getPerDofVariable(ContextImpl& context, in
         localValuesAreCurrent = true;
     }
     values.resize(perDofValues->getNumObjects()/3);
+    OpenCLArray<cl_int>& order = cl.getAtomIndex();
     for (int i = 0; i < (int) values.size(); i++)
         for (int j = 0; j < 3; j++)
-            values[i][j] = localPerDofValues[3*i+j][variable];
+            values[order[i]][j] = localPerDofValues[3*i+j][variable];
 }
 
 void OpenCLIntegrateCustomStepKernel::setPerDofVariable(ContextImpl& context, int variable, const vector<Vec3>& values) {
@@ -3972,9 +4011,10 @@ void OpenCLIntegrateCustomStepKernel::setPerDofVariable(ContextImpl& context, in
         perDofValues->getParameterValues(localPerDofValues);
         localValuesAreCurrent = true;
     }
+    OpenCLArray<cl_int>& order = cl.getAtomIndex();
     for (int i = 0; i < (int) values.size(); i++)
         for (int j = 0; j < 3; j++)
-            localPerDofValues[3*i+j][variable] = (float) values[i][j];
+            localPerDofValues[3*i+j][variable] = (float) values[order[i]][j];
     deviceValuesAreCurrent = false;
 }
 

platforms/opencl/src/OpenCLKernels.h

@@ -946,6 +946,7 @@ public:
      */
     void setPerDofVariable(ContextImpl& context, int variable, const std::vector<Vec3>& values);
 private:
+    class ReorderListener;
     std::string createGlobalComputation(const std::string& variable, const Lepton::ParsedExpression& expr, CustomIntegrator& integrator);
     std::string createPerDofComputation(const std::string& variable, const Lepton::ParsedExpression& expr, int component, CustomIntegrator& integrator);
     void recordChangedParameters(ContextImpl& context);

platforms/opencl/tests/TestOpenCLCustomIntegrator.cpp

@@ -467,6 +467,63 @@ void testRandomDistributions() {
     ASSERT_EQUAL_TOL(0.0, c4, 3.0/pow(numValues, 1.0/4.0));
 }
 
+/**
+ * Test getting and setting per-DOF variables.
+ */
+void testPerDofVariables() {
+    const int numParticles = 200;
+    const double boxSize = 10;
+    OpenCLPlatform platform;
+    System system;
+    system.setDefaultPeriodicBoxVectors(Vec3(boxSize, 0, 0), Vec3(0, boxSize, 0), Vec3(0, 0, boxSize));
+    NonbondedForce* nb = new NonbondedForce();
+    system.addForce(nb);
+    nb->setNonbondedMethod(NonbondedForce::CutoffNonPeriodic);
+    vector<Vec3> positions(numParticles);
+    OpenMM_SFMT::SFMT sfmt;
+    init_gen_rand(0, sfmt);
+    for (int i = 0; i < numParticles; i++) {
+        system.addParticle(1.5);
+        nb->addParticle(i%2 == 0 ? 1 : -1, 0.1, 1);
+        bool close = true;
+        while (close) {
+            positions[i] = Vec3(boxSize*genrand_real2(sfmt), boxSize*genrand_real2(sfmt), boxSize*genrand_real2(sfmt));
+            close = false;
+            for (int j = 0; j < i; ++j) {
+                Vec3 delta = positions[i]-positions[j];
+                if (delta.dot(delta) < 0.1)
+                    close = true;
+            }
+        }
+    }
+    CustomIntegrator integrator(0.01);
+    integrator.addPerDofVariable("temp", 0);
+    integrator.addPerDofVariable("pos", 0);
+    integrator.addComputePerDof("v", "v+dt*f/m");
+    integrator.addComputePerDof("x", "x+dt*v");
+    integrator.addComputePerDof("pos", "x");
+    Context context(system, integrator, platform);
+    context.setPositions(positions);
+    vector<Vec3> initialValues(numParticles);
+    for (int i = 0; i < numParticles; i++)
+        initialValues[i] = Vec3(i+0.1, i+0.2, i+0.3);
+    integrator.setPerDofVariable(0, initialValues);
+    
+    // Run a simulation, then query per-DOF values and see if they are correct.
+    
+    vector<Vec3> values;
+    for (int i = 0; i < 100; ++i) {
+        integrator.step(1);
+        State state = context.getState(State::Positions);
+        integrator.getPerDofVariable(0, values);
+        for (int j = 0; j < numParticles; j++)
+            ASSERT_EQUAL_VEC(initialValues[j], values[j], 1e-5);
+        integrator.getPerDofVariable(1, values);
+        for (int j = 0; j < numParticles; j++)
+            ASSERT_EQUAL_VEC(state.getPositions()[j], values[j], 1e-5);
+    }
+}
+
 int main() {
     try {
         testSingleBond();
@@ -477,6 +534,7 @@ int main() {
         testSum();
         testParameter();
         testRandomDistributions();
+        testPerDofVariables();
     }
     catch(const exception& e) {
         cout << "exception: " << e.what() << endl;

platforms/reference/tests/TestReferenceCustomIntegrator.cpp

@@ -453,6 +453,63 @@ void testRandomDistributions() {
     ASSERT_EQUAL_TOL(0.0, c4, 3.0/pow(numValues, 1.0/4.0));
 }
 
+/**
+ * Test getting and setting per-DOF variables.
+ */
+void testPerDofVariables() {
+    const int numParticles = 200;
+    const double boxSize = 10;
+    ReferencePlatform platform;
+    System system;
+    system.setDefaultPeriodicBoxVectors(Vec3(boxSize, 0, 0), Vec3(0, boxSize, 0), Vec3(0, 0, boxSize));
+    NonbondedForce* nb = new NonbondedForce();
+    system.addForce(nb);
+    nb->setNonbondedMethod(NonbondedForce::CutoffNonPeriodic);
+    vector<Vec3> positions(numParticles);
+    OpenMM_SFMT::SFMT sfmt;
+    init_gen_rand(0, sfmt);
+    for (int i = 0; i < numParticles; i++) {
+        system.addParticle(1.5);
+        nb->addParticle(i%2 == 0 ? 1 : -1, 0.1, 1);
+        bool close = true;
+        while (close) {
+            positions[i] = Vec3(boxSize*genrand_real2(sfmt), boxSize*genrand_real2(sfmt), boxSize*genrand_real2(sfmt));
+            close = false;
+            for (int j = 0; j < i; ++j) {
+                Vec3 delta = positions[i]-positions[j];
+                if (delta.dot(delta) < 0.1)
+                    close = true;
+            }
+        }
+    }
+    CustomIntegrator integrator(0.01);
+    integrator.addPerDofVariable("temp", 0);
+    integrator.addPerDofVariable("pos", 0);
+    integrator.addComputePerDof("v", "v+dt*f/m");
+    integrator.addComputePerDof("x", "x+dt*v");
+    integrator.addComputePerDof("pos", "x");
+    Context context(system, integrator, platform);
+    context.setPositions(positions);
+    vector<Vec3> initialValues(numParticles);
+    for (int i = 0; i < numParticles; i++)
+        initialValues[i] = Vec3(i+0.1, i+0.2, i+0.3);
+    integrator.setPerDofVariable(0, initialValues);
+    
+    // Run a simulation, then query per-DOF values and see if they are correct.
+    
+    vector<Vec3> values;
+    for (int i = 0; i < 100; ++i) {
+        integrator.step(1);
+        State state = context.getState(State::Positions);
+        integrator.getPerDofVariable(0, values);
+        for (int j = 0; j < numParticles; j++)
+            ASSERT_EQUAL_VEC(initialValues[j], values[j], 1e-5);
+        integrator.getPerDofVariable(1, values);
+        for (int j = 0; j < numParticles; j++)
+            ASSERT_EQUAL_VEC(state.getPositions()[j], values[j], 1e-5);
+    }
+}
+
 int main() {
     try {
         testSingleBond();
@@ -463,6 +520,7 @@ int main() {
         testSum();
         testParameter();
         testRandomDistributions();
+        testPerDofVariables();
     }
     catch(const exception& e) {
         cout << "exception: " << e.what() << endl;