Reference implementation of flow control in CustomIntegrator

openmmapi/include/openmm/internal/CustomIntegratorUtilities.h

@@ -60,16 +60,17 @@ public:
      * Analyze the sequence of steps in a CustomIntegrator.  For each step:
      *
      * 1. Parse all expressions involved in the step, and identify the comparison operator (for conditional steps).
-     * 2. Determine whether the step causes previously computed forces and energies to become invalid.
-     * 3. Determine whether the step itself needs forces and/or energies.
-     * 4. Decide whether forces and energies should both be computed at that step (because
+     * 2. For each conditional block, identify what step marks the end of the block.
+     * 3. Determine whether the step causes previously computed forces and energies to become invalid.
+     * 4. Determine whether the step itself needs forces and/or energies.
+     * 5. Decide whether forces and energies should both be computed at that step (because
      *    it is more efficient to compute both at once, even if one won't be needed
      *    until a later step).
-     * 5. Identify what force group each step needs forces and/or energies for.
+     * 6. Identify what force group each step needs forces and/or energies for.
      */
     static void analyzeComputations(const ContextImpl& context, const CustomIntegrator& integrator, std::vector<std::vector<Lepton::ParsedExpression> >& expressions,
-            std::vector<Comparison>& comparisons, std::vector<bool>& invalidatesForces, std::vector<bool>& needsForces, std::vector<bool>& needsEnergy,
-            std::vector<bool>& computeBoth, std::vector<int>& forceGroup);
+            std::vector<Comparison>& comparisons, std::vector<int>& blockEnd, std::vector<bool>& invalidatesForces, std::vector<bool>& needsForces,
+            std::vector<bool>& needsEnergy, std::vector<bool>& computeBoth, std::vector<int>& forceGroup);
     /**
      * Determine whether an expression involves a particular variable.
      */

openmmapi/src/CustomIntegratorUtilities.cpp

@@ -69,7 +69,7 @@ bool CustomIntegratorUtilities::usesVariable(const Lepton::ParsedExpression& exp
 }
 
 void CustomIntegratorUtilities::analyzeComputations(const ContextImpl& context, const CustomIntegrator& integrator, vector<vector<Lepton::ParsedExpression> >& expressions,
-            vector<Comparison>& comparisons, vector<bool>& invalidatesForces, vector<bool>& needsForces, vector<bool>& needsEnergy,
+            vector<Comparison>& comparisons, vector<int>& blockEnd, vector<bool>& invalidatesForces, vector<bool>& needsForces, vector<bool>& needsEnergy,
             vector<bool>& computeBoth, vector<int>& forceGroup) {
     int numSteps = integrator.getNumComputations();
     expressions.resize(numSteps);
@@ -138,13 +138,13 @@ void CustomIntegratorUtilities::analyzeComputations(const ContextImpl& context,
                     if (forceGroup[step] != -2)
                         throw OpenMMException("A single computation step cannot depend on multiple force groups");
                     needsForces[step] = true;
-                    forceGroup[step] = 1<<i;
+                    forceGroup[step] = i;
                 }
                 if (usesVariable(expressions[step][expr], 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;
+                    forceGroup[step] = i;
                 }
             }
         }
@@ -153,7 +153,7 @@ void CustomIntegratorUtilities::analyzeComputations(const ContextImpl& context,
     // Find the end point of each block.
 
     vector<int> blockStart;
-    vector<int> blockEnd(numSteps);
+    blockEnd.resize(numSteps, -1);
     for (int step = 0; step < numSteps; step++) {
         if (stepType[step] == CustomIntegrator::BeginIfBlock || stepType[step] == CustomIntegrator::BeginWhileBlock)
             blockStart.push_back(step);

platforms/reference/include/ReferenceCustomDynamics.h

 
-/* Portions copyright (c) 2011 Stanford University and Simbios.
+/* Portions copyright (c) 2011-2015 Stanford University and Simbios.
  * Contributors: Peter Eastman
  *
  * Permission is hereby granted, free of charge, to any person obtaining
@@ -28,6 +28,7 @@
 #include "ReferenceDynamics.h"
 #include "openmm/CustomIntegrator.h"
 #include "openmm/internal/ContextImpl.h"
+#include "openmm/internal/CustomIntegratorUtilities.h"
 #include "lepton/ExpressionProgram.h"
 
 #include <map>
@@ -44,9 +45,10 @@ private:
     std::vector<OpenMM::RealVec> sumBuffer, oldPos;
     std::vector<OpenMM::CustomIntegrator::ComputationType> stepType;
     std::vector<std::string> stepVariable, forceName, energyName;
-    std::vector<Lepton::ExpressionProgram> stepExpression;
-    std::vector<bool> invalidatesForces, needsForces, needsEnergy;
-    std::vector<int> forceGroup;
+    std::vector<std::vector<Lepton::ExpressionProgram> > stepExpressions;
+    std::vector<CustomIntegratorUtilities::Comparison> comparisons;
+    std::vector<bool> invalidatesForces, needsForces, needsEnergy, computeBothForceAndEnergy;
+    std::vector<int> forceGroupFlags, blockEnd;
     RealOpenMM energy;
     Lepton::ExpressionProgram kineticEnergyExpression;
     bool kineticEnergyNeedsForce;
@@ -58,6 +60,8 @@ private:
     
     void recordChangedParameters(OpenMM::ContextImpl& context, std::map<std::string, RealOpenMM>& globals);
 
+    bool evaluateCondition(int step, std::map<std::string, RealOpenMM>& globals);
+      
 public:
 
       /**---------------------------------------------------------------------------------------

platforms/reference/src/SimTKReference/ReferenceCustomDynamics.cpp

 
-/* Portions copyright (c) 2011-2013 Stanford University and Simbios.
+/* Portions copyright (c) 2011-2015 Stanford University and Simbios.
  * Contributors: Peter Eastman
  *
  * Permission is hereby granted, free of charge, to any person obtaining
@@ -52,12 +52,9 @@ ReferenceCustomDynamics::ReferenceCustomDynamics(int numberOfAtoms, const Custom
     oldPos.resize(numberOfAtoms);
     stepType.resize(integrator.getNumComputations());
     stepVariable.resize(integrator.getNumComputations());
-    stepExpression.resize(integrator.getNumComputations());
     for (int i = 0; i < integrator.getNumComputations(); i++) {
         string expression;
         integrator.getComputationStep(i, stepType[i], stepVariable[i], expression);
-        if (expression.length() > 0)
-            stepExpression[i] = Lepton::Parser::parse(expression).optimize().createProgram();
     }
     kineticEnergyExpression = Lepton::Parser::parse(integrator.getKineticEnergyExpression()).optimize().createProgram();
     kineticEnergyNeedsForce = false;
@@ -100,27 +97,25 @@ void ReferenceCustomDynamics::update(ContextImpl& context, int numberOfAtoms, ve
     globals.insert(context.getParameters().begin(), context.getParameters().end());
     oldPos = atomCoordinates;
     if (invalidatesForces.size() == 0) {
-        // The first time this is called, work out when to recompute forces and energy.  First build a
-        // list of every step that invalidates the forces.
+        // Some initialization can't be done in the constructor, since we need a ContextImpl from which to get the list of
+        // Context parameters.  Instead, we do it the first time this method is called.
 
-        invalidatesForces.resize(numSteps, false);
-        needsForces.resize(numSteps, false);
-        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) {
-            const map<string, double> params = (*iter)->getDefaultParameters();
-            for (map<string, double>::const_iterator param = params.begin(); param != params.end(); ++param)
-                affectsForce.insert(param->first);
+        vector<int> forceGroup;
+        vector<vector<Lepton::ParsedExpression> > expressions;
+        CustomIntegratorUtilities::analyzeComputations(context, integrator, expressions, comparisons, blockEnd, invalidatesForces, needsForces, needsEnergy, computeBothForceAndEnergy, forceGroup);
+        stepExpressions.resize(expressions.size());
+        for (int i = 0; i < numSteps; i++) {
+            for (int j = 0; j < (int) expressions[i].size(); j++)
+                stepExpressions[i].push_back(expressions[i][j].createProgram());
+            if (stepType[i] == CustomIntegrator::BeginWhileBlock)
+                blockEnd[blockEnd[i]] = i; // Record where to branch back to.
         }
-        for (int i = 0; i < numSteps; i++)
-            invalidatesForces[i] = (stepType[i] == CustomIntegrator::ConstrainPositions || affectsForce.find(stepVariable[i]) != affectsForce.end());
 
-        // Make a list of which steps require valid forces or energy to be known.
+        // Record the variable names and flags for the force and energy in each step.
 
+        forceGroupFlags.resize(numSteps, -1);
+        forceName.resize(numSteps, "f");
+        energyName.resize(numSteps, "energy");
         vector<string> forceGroupName;
         vector<string> energyGroupName;
         for (int i = 0; i < 32; i++) {
@@ -132,47 +127,12 @@ void ReferenceCustomDynamics::update(ContextImpl& context, int numberOfAtoms, ve
             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) {
-                for (int j = 0; j < stepExpression[i].getNumOperations(); j++) {
-                    const Lepton::Operation& op = stepExpression[i].getOperation(j);
-                    if (op.getId() == Lepton::Operation::VARIABLE) {
-                        if (op.getName() == "energy") {
-                            if (forceGroup[i] != -2)
-                                throw OpenMMException("A single computation step cannot depend on multiple force groups");
-                            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");
-                            needsForces[i] = true;
-                            forceGroup[i] = -1;
-                        }
-                        else if (op.getName()[0] == 'f') {
-                            for (int k = 0; k < (int) forceGroupName.size(); k++)
-                                if (op.getName() == forceGroupName[k]) {
-                                    if (forceGroup[i] != -2)
-                                        throw OpenMMException("A single computation step cannot depend on multiple force groups");
-                                    needsForces[i] = true;
-                                    forceGroup[i] = 1<<k;
-                                    forceName[i] = forceGroupName[k];
-                                    break;
-                                }
-                        }
-                    }
-                }
-            }
+            if (needsForces[i] && forceGroup[i] > -1)
+                forceName[i] = forceGroupName[forceGroup[i]];
+            if (needsEnergy[i] && forceGroup[i] > -1)
+                energyName[i] = energyGroupName[forceGroup[i]];
+            if (forceGroup[i] > -1)
+                forceGroupFlags[i] = 1<<forceGroup[i];
         }
         
         // Build the list of inverse masses.
@@ -188,65 +148,54 @@ void ReferenceCustomDynamics::update(ContextImpl& context, int numberOfAtoms, ve
     
     // Loop over steps and execute them.
     
-    for (int i = 0; i < numSteps; i++) {
-        if ((needsForces[i] || needsEnergy[i]) && (!forcesAreValid || context.getLastForceGroups() != forceGroup[i])) {
-            // Recompute forces and or energy.  Figure out what is actually needed
-            // between now and the next time they get invalidated again.
-            
-            bool computeForce = false, computeEnergy = false;
-            for (int j = i; ; j++) {
-                if (needsForces[j])
-                    computeForce = true;
-                if (needsEnergy[j])
-                    computeEnergy = true;
-                if (invalidatesForces[j])
-                    break;
-                if (j == numSteps-1)
-                    j = -1;
-                if (j == i-1)
-                    break;
-            }
+    for (int step = 0; step < numSteps; ) {
+        if ((needsForces[step] || needsEnergy[step]) && (!forcesAreValid || context.getLastForceGroups() != forceGroupFlags[step])) {
+            // Recompute forces and/or energy.
+            
+            bool computeForce = needsForces[step] || computeBothForceAndEnergy[step];
+            bool computeEnergy = needsEnergy[step] || computeBothForceAndEnergy[step];
             recordChangedParameters(context, globals);
-            RealOpenMM e = context.calcForcesAndEnergy(computeForce, computeEnergy, forceGroup[i]);
+            RealOpenMM e = context.calcForcesAndEnergy(computeForce, computeEnergy, forceGroupFlags[step]);
             if (computeEnergy)
                 energy = e;
             forcesAreValid = true;
         }
-        globals[energyName[i]] = energy;
+        globals[energyName[step]] = energy;
         
         // Execute the step.
 
-        switch (stepType[i]) {
+        int nextStep = step+1;
+        switch (stepType[step]) {
             case CustomIntegrator::ComputeGlobal: {
                 map<string, RealOpenMM> variables = globals;
                 variables["uniform"] = SimTKOpenMMUtilities::getUniformlyDistributedRandomNumber();
                 variables["gaussian"] = SimTKOpenMMUtilities::getNormallyDistributedRandomNumber();
-                globals[stepVariable[i]] = stepExpression[i].evaluate(variables);
+                globals[stepVariable[step]] = stepExpressions[step][0].evaluate(variables);
                 break;
             }
             case CustomIntegrator::ComputePerDof: {
                 vector<RealVec>* results = NULL;
-                if (stepVariable[i] == "x")
+                if (stepVariable[step] == "x")
                     results = &atomCoordinates;
-                else if (stepVariable[i] == "v")
+                else if (stepVariable[step] == "v")
                     results = &velocities;
                 else {
                     for (int j = 0; j < integrator.getNumPerDofVariables(); j++)
-                        if (stepVariable[i] == integrator.getPerDofVariableName(j))
+                        if (stepVariable[step] == integrator.getPerDofVariableName(j))
                             results = &perDof[j];
                 }
                 if (results == NULL)
-                    throw OpenMMException("Illegal per-DOF output variable: "+stepVariable[i]);
-                computePerDof(numberOfAtoms, *results, atomCoordinates, velocities, forces, masses, globals, perDof, stepExpression[i], forceName[i]);
+                    throw OpenMMException("Illegal per-DOF output variable: "+stepVariable[step]);
+                computePerDof(numberOfAtoms, *results, atomCoordinates, velocities, forces, masses, globals, perDof, stepExpressions[step][0], forceName[step]);
                 break;
             }
             case CustomIntegrator::ComputeSum: {
-                computePerDof(numberOfAtoms, sumBuffer, atomCoordinates, velocities, forces, masses, globals, perDof, stepExpression[i], forceName[i]);
+                computePerDof(numberOfAtoms, sumBuffer, atomCoordinates, velocities, forces, masses, globals, perDof, stepExpressions[step][0], forceName[step]);
                 RealOpenMM sum = 0.0;
                 for (int j = 0; j < numberOfAtoms; j++)
                     if (masses[j] != 0.0)
                         sum += sumBuffer[j][0]+sumBuffer[j][1]+sumBuffer[j][2];
-                globals[stepVariable[i]] = sum;
+                globals[stepVariable[step]] = sum;
                 break;
             }
             case CustomIntegrator::ConstrainPositions: {
@@ -262,10 +211,27 @@ void ReferenceCustomDynamics::update(ContextImpl& context, int numberOfAtoms, ve
                 recordChangedParameters(context, globals);
                 context.updateContextState();
                 globals.insert(context.getParameters().begin(), context.getParameters().end());
+                break;
+            }
+            case CustomIntegrator::BeginIfBlock: {
+                if (!evaluateCondition(step, globals))
+                    nextStep = blockEnd[step]+1;
+                break;
+            }
+            case CustomIntegrator::BeginWhileBlock: {
+                if (!evaluateCondition(step, globals))
+                    nextStep = blockEnd[step]+1;
+                break;
+            }
+            case CustomIntegrator::EndBlock: {
+                if (blockEnd[step] != -1)
+                    nextStep = blockEnd[step]; // Return to the start of a while block.
+                break;
             }
         }
-        if (invalidatesForces[i])
+        if (invalidatesForces[step])
             forcesAreValid = false;
+        step = nextStep;
     }
     ReferenceVirtualSites::computePositions(context.getSystem(), atomCoordinates);
     incrementTimeStep();
@@ -298,6 +264,29 @@ void ReferenceCustomDynamics::computePerDof(int numberOfAtoms, vector<RealVec>&
     }
 }
 
+bool ReferenceCustomDynamics::evaluateCondition(int step, map<string, RealOpenMM>& globals) {
+    map<string, RealOpenMM> variables = globals;
+    variables["uniform"] = SimTKOpenMMUtilities::getUniformlyDistributedRandomNumber();
+    variables["gaussian"] = SimTKOpenMMUtilities::getNormallyDistributedRandomNumber();
+    double lhs = stepExpressions[step][0].evaluate(variables);
+    double rhs = stepExpressions[step][1].evaluate(variables);
+    switch (comparisons[step]) {
+        case CustomIntegratorUtilities::EQUAL:
+            return (lhs == rhs);
+        case CustomIntegratorUtilities::LESS_THAN:
+            return (lhs < rhs);
+        case CustomIntegratorUtilities::GREATER_THAN:
+            return (lhs > rhs);
+        case CustomIntegratorUtilities::NOT_EQUAL:
+            return (lhs != rhs);
+        case CustomIntegratorUtilities::LESS_THAN_OR_EQUAL:
+            return (lhs <= rhs);
+        case CustomIntegratorUtilities::GREATER_THAN_OR_EQUAL:
+            return (lhs >= rhs);
+    }
+    throw OpenMMException("ReferenceCustomDynamics: Invalid comparison operator");
+}
+
 /**
  * Check which context parameters have changed and register them with the context.
  */

platforms/reference/tests/TestReferenceCustomIntegrator.cpp

@@ -6,7 +6,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2008-2013 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2015 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -688,6 +688,67 @@ void testRespa() {
     }
 }
 
+void testIfBlock() {
+    System system;
+    system.addParticle(2.0);
+    system.addParticle(2.0);
+    const double dt = 0.01;
+    CustomIntegrator integrator(dt);
+    integrator.addGlobalVariable("a", 0);
+    integrator.addGlobalVariable("b", 0);
+    integrator.addComputeGlobal("b", "1");
+    integrator.beginIfBlock("a < 3.5");
+    integrator.addComputeGlobal("b", "a+1");
+    integrator.endBlock();
+    Context context(system, integrator, platform);
+
+    // Set "a" to 1.7 and verify that "b" gets set to a+1.
+
+    integrator.setGlobalVariable(0, 1.7);
+    integrator.step(1);
+    ASSERT_EQUAL_TOL(2.7, integrator.getGlobalVariable(1), 1e-6);
+
+    // Now set it to a value that should cause the block to be skipped.
+
+    integrator.setGlobalVariable(0, 4.1);
+    integrator.step(1);
+    ASSERT_EQUAL_TOL(1.0, integrator.getGlobalVariable(1), 1e-6);
+}
+
+void testWhileBlock() {
+    System system;
+    system.addParticle(2.0);
+    system.addParticle(2.0);
+    const double dt = 0.01;
+    CustomIntegrator integrator(dt);
+    integrator.addGlobalVariable("a", 0);
+    integrator.addGlobalVariable("b", 0);
+    integrator.addComputeGlobal("b", "1");
+    integrator.beginWhileBlock("b <= a");
+    integrator.addComputeGlobal("b", "b+1");
+    integrator.endBlock();
+    Context context(system, integrator, platform);
+
+    // Try a case where the loop should be skipped.
+
+    integrator.setGlobalVariable(0, -3.3);
+    integrator.step(1);
+    ASSERT_EQUAL_TOL(1.0, integrator.getGlobalVariable(1), 1e-6);
+
+    // In this case it should be executed exactly once.
+
+    integrator.setGlobalVariable(0, 1.2);
+    integrator.step(1);
+    ASSERT_EQUAL_TOL(2.0, integrator.getGlobalVariable(1), 1e-6);
+
+    // In this case, it should be executed several times.
+
+    integrator.setGlobalVariable(0, 5.3);
+    integrator.step(1);
+    ASSERT_EQUAL_TOL(6.0, integrator.getGlobalVariable(1), 1e-6);
+}
+
+
 int main() {
     try {
         testSingleBond();
@@ -702,6 +763,8 @@ int main() {
         testPerDofVariables();
         testForceGroups();
         testRespa();
+        testIfBlock();
+        testWhileBlock();
     }
     catch(const exception& e) {
         cout << "exception: " << e.what() << endl;