CustomIntegrator can depend on the energy of a force group

openmmapi/include/openmm/CustomIntegrator.h

@@ -91,6 +91,10 @@ namespace OpenMM {
  * <li>dt: (global) This is the step size being used by the integrator.</li>
  * <li>energy: (global, read-only) This is the current potential energy of the
  * system.</li>
+ * <li>energy0, energy1, energy2, ...: (global, read-only) This is similar to
+ * energy, but includes only the contribution from forces in one force group.
+ * A single computation step may only depend on a single energy variable
+ * (energy, energy0, energy1, etc.).</li>
  * <li>x: (per-DOF) This is the current value of the degree of freedom (the x,
  * y, or z coordinate of a particle).</li>
  * <li>v: (per-DOF) This is the current velocity associated with the degree of

platforms/opencl/src/OpenCLKernels.cpp

@@ -3606,7 +3606,7 @@ void OpenCLIntegrateCustomStepKernel::initialize(const System& system, const Cus
     SimTKOpenMMUtilities::setRandomNumberSeed(integrator.getRandomNumberSeed());
 }
 
-string OpenCLIntegrateCustomStepKernel::createGlobalComputation(const string& variable, const Lepton::ParsedExpression& expr, CustomIntegrator& integrator) {
+string OpenCLIntegrateCustomStepKernel::createGlobalComputation(const string& variable, const Lepton::ParsedExpression& expr, CustomIntegrator& integrator, const string& energyName) {
     map<string, Lepton::ParsedExpression> expressions;
     if (variable == "dt")
         expressions["dt[0].y = "] = expr;
@@ -3626,7 +3626,7 @@ string OpenCLIntegrateCustomStepKernel::createGlobalComputation(const string& va
     variables["dt"] = "dt[0].y";
     variables["uniform"] = "uniform";
     variables["gaussian"] = "gaussian";
-    variables["energy"] = "energy[0]";
+    variables[energyName] = "energy[0]";
     for (int i = 0; i < integrator.getNumGlobalVariables(); i++)
         variables[integrator.getGlobalVariableName(i)] = "globals["+intToString(i)+"]";
     for (int i = 0; i < (int) parameterNames.size(); i++)
@@ -3635,7 +3635,7 @@ string OpenCLIntegrateCustomStepKernel::createGlobalComputation(const string& va
     return OpenCLExpressionUtilities::createExpressions(expressions, variables, functions, "temp", "");
 }
 
-string OpenCLIntegrateCustomStepKernel::createPerDofComputation(const string& variable, const Lepton::ParsedExpression& expr, int component, CustomIntegrator& integrator, const string& forceName) {
+string OpenCLIntegrateCustomStepKernel::createPerDofComputation(const string& variable, const Lepton::ParsedExpression& expr, int component, CustomIntegrator& integrator, const string& forceName, const string& energyName) {
     const string suffixes[] = {".x", ".y", ".z"};
     string suffix = suffixes[component];
     map<string, Lepton::ParsedExpression> expressions;
@@ -3660,7 +3660,7 @@ string OpenCLIntegrateCustomStepKernel::createPerDofComputation(const string& va
     variables["uniform"] = "uniform"+suffix;
     variables["m"] = "mass";
     variables["dt"] = "stepSize";
-    variables["energy"] = "energy[0]";
+    variables[energyName] = "energy[0]";
     for (int i = 0; i < integrator.getNumGlobalVariables(); i++)
         variables[integrator.getGlobalVariableName(i)] = "globals["+intToString(i)+"]";
     for (int i = 0; i < integrator.getNumPerDofVariables(); i++)
@@ -3736,12 +3736,17 @@ void OpenCLIntegrateCustomStepKernel::execute(ContextImpl& context, CustomIntegr
         vector<string> variable(numSteps);
         vector<Lepton::ParsedExpression> expression(numSteps);
         vector<string> forceGroupName;
+        vector<string> energyGroupName;
         for (int i = 0; i < 32; i++) {
-            stringstream str;
-            str << "f" << i;
-            forceGroupName.push_back(str.str());
+            stringstream fname;
+            fname << "f" << i;
+            forceGroupName.push_back(fname.str());
+            stringstream ename;
+            ename << "energy" << i;
+            energyGroupName.push_back(ename.str());
         }
         vector<string> forceName(numSteps, "f");
+        vector<string> energyName(numSteps, "energy");
         for (int step = 0; step < numSteps; step++) {
             string expr;
             integrator.getComputationStep(step, stepType[step], variable[step], expr);
@@ -3763,6 +3768,13 @@ void OpenCLIntegrateCustomStepKernel::execute(ContextImpl& context, CustomIntegr
                         forceGroup[step] = 1<<i;
                         forceName[step] = forceGroupName[i];
                     }
+                    if (usesVariable(expression[step], energyGroupName[i])) {
+                        if (forceGroup[step] != -2)
+                            throw OpenMMException("A single computation step cannot depend on multiple force groups");
+                        needsEnergy[step] = true;
+                        forceGroup[step] = 1<<i;
+                        energyName[step] = energyGroupName[i];
+                    }
                 }
             }
             invalidatesForces[step] = (stepType[step] == CustomIntegrator::ConstrainPositions || affectsForce.find(variable[step]) != affectsForce.end());
@@ -3820,7 +3832,7 @@ void OpenCLIntegrateCustomStepKernel::execute(ContextImpl& context, CustomIntegr
                 for (int j = step; j < numSteps && (j == step || merged[j]); j++) {
                     compute << "{\n";
                     for (int i = 0; i < 3; i++)
-                        compute << createPerDofComputation(stepType[j] == CustomIntegrator::ComputePerDof ? variable[j] : "", expression[j], i, integrator, forceName[j]);
+                        compute << createPerDofComputation(stepType[j] == CustomIntegrator::ComputePerDof ? variable[j] : "", expression[j], i, integrator, forceName[j], energyName[j]);
                     if (variable[j] == "x") {
                         if (storePosAsDelta[j]) {
                             if (cl.getSupportsDoublePrecision())
@@ -3910,7 +3922,7 @@ void OpenCLIntegrateCustomStepKernel::execute(ContextImpl& context, CustomIntegr
 
                 stringstream compute;
                 for (int i = step; i < numSteps && (i == step || merged[i]); i++)
-                    compute << "{\n" << createGlobalComputation(variable[i], expression[i], integrator) << "}\n";
+                    compute << "{\n" << createGlobalComputation(variable[i], expression[i], integrator, energyName[i]) << "}\n";
                 map<string, string> replacements;
                 replacements["COMPUTE_STEP"] = compute.str();
                 cl::Program program = cl.createProgram(cl.replaceStrings(OpenCLKernelSources::customIntegratorGlobal, replacements), defines);

platforms/opencl/src/OpenCLKernels.h

@@ -979,8 +979,8 @@ 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, const std::string& forceName);
+    std::string createGlobalComputation(const std::string& variable, const Lepton::ParsedExpression& expr, CustomIntegrator& integrator, const std::string& energyName);
+    std::string createPerDofComputation(const std::string& variable, const Lepton::ParsedExpression& expr, int component, CustomIntegrator& integrator, const std::string& forceName, const std::string& energyName);
     void recordChangedParameters(ContextImpl& context);
     OpenCLContext& cl;
     double prevStepSize;

platforms/opencl/tests/TestOpenCLCustomIntegrator.cpp

@@ -539,9 +539,15 @@ void testForceGroups() {
     integrator.addPerDofVariable("outf", 0);
     integrator.addPerDofVariable("outf1", 0);
     integrator.addPerDofVariable("outf2", 0);
+    integrator.addGlobalVariable("oute", 0);
+    integrator.addGlobalVariable("oute1", 0);
+    integrator.addGlobalVariable("oute2", 0);
     integrator.addComputePerDof("outf", "f");
     integrator.addComputePerDof("outf1", "f1");
     integrator.addComputePerDof("outf2", "f2");
+    integrator.addComputeGlobal("oute", "energy");
+    integrator.addComputeGlobal("oute1", "energy1");
+    integrator.addComputeGlobal("oute2", "energy2");
     HarmonicBondForce* bonds = new HarmonicBondForce();
     bonds->addBond(0, 1, 1.5, 1.1);
     bonds->setForceGroup(1);
@@ -561,6 +567,8 @@ void testForceGroups() {
     
     integrator.step(1);
     vector<Vec3> f, f1, f2;
+    double e1 = 0.5*1.1*0.5*0.5;
+    double e2 = 138.935456*0.2*0.2/2.0;
     integrator.getPerDofVariable(0, f);
     integrator.getPerDofVariable(1, f1);
     integrator.getPerDofVariable(2, f2);
@@ -570,9 +578,15 @@ void testForceGroups() {
     ASSERT_EQUAL_VEC(Vec3(138.935456*0.2*0.2/4.0, 0, 0), f2[1], 1e-5);
     ASSERT_EQUAL_VEC(f1[0]+f2[0], f[0], 1e-5);
     ASSERT_EQUAL_VEC(f1[1]+f2[1], f[1], 1e-5);
+    ASSERT_EQUAL_TOL(e1, integrator.getGlobalVariable(1), 1e-5);
+    ASSERT_EQUAL_TOL(e2, integrator.getGlobalVariable(2), 1e-5);
+    ASSERT_EQUAL_TOL(e1+e2, integrator.getGlobalVariable(0), 1e-5);
     
     // Make sure they also match the values returned by the Context.
     
+    State s = context.getState(State::Forces | State::Energy, false);
+    State s1 = context.getState(State::Forces | State::Energy, false, 2);
+    State s2 = context.getState(State::Forces | State::Energy, false, 4);
     vector<Vec3> c, c1, c2;
     c = context.getState(State::Forces, false).getForces();
     c1 = context.getState(State::Forces, false, 2).getForces();
@@ -583,6 +597,9 @@ void testForceGroups() {
     ASSERT_EQUAL_VEC(f1[1], c1[1], 1e-5);
     ASSERT_EQUAL_VEC(f2[0], c2[0], 1e-5);
     ASSERT_EQUAL_VEC(f2[1], c2[1], 1e-5);
+    ASSERT_EQUAL_TOL(s.getPotentialEnergy(), integrator.getGlobalVariable(0), 1e-5);
+    ASSERT_EQUAL_TOL(s1.getPotentialEnergy(), integrator.getGlobalVariable(1), 1e-5);
+    ASSERT_EQUAL_TOL(s2.getPotentialEnergy(), integrator.getGlobalVariable(2), 1e-5);
 }
 
 /**

platforms/reference/src/SimTKReference/ReferenceCustomDynamics.cpp

@@ -103,6 +103,7 @@ void ReferenceCustomDynamics::update(ContextImpl& context, int numberOfAtoms, ve
         needsEnergy.resize(numSteps, false);
         forceGroup.resize(numSteps, -2);
         forceName.resize(numSteps, "f");
+        energyName.resize(numSteps, "energy");
         set<string> affectsForce;
         affectsForce.insert("x");
         for (vector<ForceImpl*>::const_iterator iter = context.getForceImpls().begin(); iter != context.getForceImpls().end(); ++iter) {
@@ -116,10 +117,14 @@ void ReferenceCustomDynamics::update(ContextImpl& context, int numberOfAtoms, ve
         // Make a list of which steps require valid forces or energy to be known.
         
         vector<string> forceGroupName;
+        vector<string> energyGroupName;
         for (int i = 0; i < 32; i++) {
-            stringstream str;
-            str << "f" << i;
-            forceGroupName.push_back(str.str());
+            stringstream fname;
+            fname << "f" << i;
+            forceGroupName.push_back(fname.str());
+            stringstream ename;
+            ename << "energy" << i;
+            energyGroupName.push_back(ename.str());
         }
         for (int i = 0; i < numSteps; i++) {
             if (stepType[i] == CustomIntegrator::ComputeGlobal || stepType[i] == CustomIntegrator::ComputePerDof || stepType[i] == CustomIntegrator::ComputeSum) {
@@ -132,6 +137,17 @@ void ReferenceCustomDynamics::update(ContextImpl& context, int numberOfAtoms, ve
                             needsEnergy[i] = true;
                             forceGroup[i] = -1;
                         }
+                        else if (op.getName().substr(0, 6) == "energy") {
+                            for (int k = 0; k < (int) energyGroupName.size(); k++)
+                                if (op.getName() == energyGroupName[k]) {
+                                    if (forceGroup[i] != -2)
+                                        throw OpenMMException("A single computation step cannot depend on multiple force groups");
+                                    needsForces[i] = true;
+                                    forceGroup[i] = 1<<k;
+                                    energyName[i] = energyGroupName[k];
+                                    break;
+                                }
+                        }
                         else if (op.getName() == "f") {
                             if (forceGroup[i] != -2)
                                 throw OpenMMException("A single computation step cannot depend on multiple force groups");
@@ -191,7 +207,7 @@ void ReferenceCustomDynamics::update(ContextImpl& context, int numberOfAtoms, ve
                 energy = e;
             forcesAreValid = true;
         }
-        globals["energy"] = energy;
+        globals[energyName[i]] = energy;
         
         // Execute the step.
         

platforms/reference/src/SimTKReference/ReferenceCustomDynamics.h

@@ -43,7 +43,7 @@ private:
     std::vector<RealOpenMM> inverseMasses;
     std::vector<OpenMM::RealVec> sumBuffer, oldPos;
     std::vector<OpenMM::CustomIntegrator::ComputationType> stepType;
-    std::vector<std::string> stepVariable, forceName;
+    std::vector<std::string> stepVariable, forceName, energyName;
     std::vector<Lepton::ExpressionProgram> stepExpression;
     std::vector<bool> invalidatesForces, needsForces, needsEnergy;
     std::vector<int> forceGroup;

platforms/reference/tests/TestReferenceCustomIntegrator.cpp

@@ -524,9 +524,15 @@ void testForceGroups() {
     integrator.addPerDofVariable("outf", 0);
     integrator.addPerDofVariable("outf1", 0);
     integrator.addPerDofVariable("outf2", 0);
+    integrator.addGlobalVariable("oute", 0);
+    integrator.addGlobalVariable("oute1", 0);
+    integrator.addGlobalVariable("oute2", 0);
     integrator.addComputePerDof("outf", "f");
     integrator.addComputePerDof("outf1", "f1");
     integrator.addComputePerDof("outf2", "f2");
+    integrator.addComputeGlobal("oute", "energy");
+    integrator.addComputeGlobal("oute1", "energy1");
+    integrator.addComputeGlobal("oute2", "energy2");
     HarmonicBondForce* bonds = new HarmonicBondForce();
     bonds->addBond(0, 1, 1.5, 1.1);
     bonds->setForceGroup(1);
@@ -546,6 +552,8 @@ void testForceGroups() {
     
     integrator.step(1);
     vector<Vec3> f, f1, f2;
+    double e1 = 0.5*1.1*0.5*0.5;
+    double e2 = 138.935456*0.2*0.2/2.0;
     integrator.getPerDofVariable(0, f);
     integrator.getPerDofVariable(1, f1);
     integrator.getPerDofVariable(2, f2);
@@ -555,9 +563,15 @@ void testForceGroups() {
     ASSERT_EQUAL_VEC(Vec3(138.935456*0.2*0.2/4.0, 0, 0), f2[1], 1e-5);
     ASSERT_EQUAL_VEC(f1[0]+f2[0], f[0], 1e-5);
     ASSERT_EQUAL_VEC(f1[1]+f2[1], f[1], 1e-5);
+    ASSERT_EQUAL_TOL(e1, integrator.getGlobalVariable(1), 1e-5);
+    ASSERT_EQUAL_TOL(e2, integrator.getGlobalVariable(2), 1e-5);
+    ASSERT_EQUAL_TOL(e1+e2, integrator.getGlobalVariable(0), 1e-5);
     
     // Make sure they also match the values returned by the Context.
     
+    State s = context.getState(State::Forces | State::Energy, false);
+    State s1 = context.getState(State::Forces | State::Energy, false, 2);
+    State s2 = context.getState(State::Forces | State::Energy, false, 4);
     vector<Vec3> c, c1, c2;
     c = context.getState(State::Forces, false).getForces();
     c1 = context.getState(State::Forces, false, 2).getForces();
@@ -568,6 +582,9 @@ void testForceGroups() {
     ASSERT_EQUAL_VEC(f1[1], c1[1], 1e-5);
     ASSERT_EQUAL_VEC(f2[0], c2[0], 1e-5);
     ASSERT_EQUAL_VEC(f2[1], c2[1], 1e-5);
+    ASSERT_EQUAL_TOL(s.getPotentialEnergy(), integrator.getGlobalVariable(0), 1e-5);
+    ASSERT_EQUAL_TOL(s1.getPotentialEnergy(), integrator.getGlobalVariable(1), 1e-5);
+    ASSERT_EQUAL_TOL(s2.getPotentialEnergy(), integrator.getGlobalVariable(2), 1e-5);
 }
 
 /**