Merge pull request #1048 from peastman/flowcontrol

CustomIntegrator supports if and while blocks

openmmapi/include/openmm/CustomIntegrator.h

@@ -178,6 +178,30 @@ namespace OpenMM {
  * integrator.addComputePerDof("v", "v+0.5*dt*f1/m");
  * </pre></tt>
  *
+ * The sequence of computations in a CustomIntegrator can include flow control in
+ * the form of "if" and "while" blocks.  The computations inside an "if" block
+ * are executed either zero or one times, depending on whether a condition is
+ * true.  The computations inside a "while" block are executed repeatedly for as
+ * long as the condition remains true.  Be very careful when writing "while"
+ * blocks; there is nothing to stop you from creating an infinite loop!
+ *
+ * For example, suppose you are writing a Monte Carlo algorithm.  Assume you have
+ * already computed a new set of particle coordinates "xnew" and a step acceptance
+ * probability "acceptanceProbability".  The following lines use an "if" block
+ * to decide whether to accept the step, and if it is accepted, store the new
+ * positions into "x".
+ *
+ * <tt><pre>
+ * integrator.beginIfBlock("uniform < acceptanceProbability");
+ * integrator.computePerDof("x", "xnew");
+ * integrator.endBlock();
+ * </pre></tt>
+ *
+ * The condition in an "if" or "while" block is evaluated globally, so it may
+ * only involve global variables, not per-DOF ones.  It may use any of the
+ * following comparison operators: =, <. >, !=, <=, >=.  Blocks may be nested
+ * inside each other.
+ *
  * An Integrator has one other job in addition to evolving the equations of motion:
  * it defines how to compute the kinetic energy of the system.  Depending on the
  * integration method used, simply summing mv<sup>2</sup>/2 over all degrees of
@@ -238,7 +262,19 @@ public:
         /**
          * Allow Forces to update the context state.
          */
-        UpdateContextState = 5
+        UpdateContextState = 5,
+        /**
+         * Begin an "if" block.
+         */
+        BeginIfBlock = 6,
+        /**
+         * Begin a while" block.
+         */
+        BeginWhileBlock = 7,
+        /**
+         * End an "if" or "while" block.
+         */
+        EndBlock = 8
     };
     /**
      * Create a CustomIntegrator.
@@ -407,6 +443,35 @@ public:
      * @return the index of the step that was added
      */
     int addUpdateContextState();
+    /**
+     * Add a step which begins a new "if" block.
+     *
+     * @param expression  a mathematical expression involving a comparison operator
+     *                    and global variables.  All steps between this one and
+     *                    the end of the block are executed only if the condition
+     *                    is true.
+     *
+     * @return the index of the step that was added
+     */
+    int beginIfBlock(const std::string& condition);
+    /**
+     * Add a step which begins a new "while" block.
+     *
+     * @param expression  a mathematical expression involving a comparison operator
+     *                    and global variables.  All steps between this one and
+     *                    the end of the block are executed repeatedly as long as
+     *                    the condition remains true.
+     *
+     * @return the index of the step that was added
+     */
+    int beginWhileBlock(const std::string& condition);
+    /**
+     * Add a step which marks the end of the most recently begun "if" or "while"
+     * block.
+     *
+     * @return the index of the step that was added
+     */
+    int endBlock();
     /**
      * Get the details of a computation step that has been added to the integration algorithm.
      * 

platforms/cpu/include/CompiledExpressionSet.h → openmmapi/include/openmm/internal/CompiledExpressionSet.h

@@ -9,7 +9,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2014 Stanford University and the Authors.           *
+ * Portions copyright (c) 2014-2015 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -33,7 +33,7 @@
  * -------------------------------------------------------------------------- */
 
 #include "lepton/CompiledExpression.h"
-#include "windowsExportCpu.h"
+#include "windowsExport.h"
 #include <string>
 #include <vector>
 
@@ -42,7 +42,7 @@ namespace OpenMM {
 /**
  * This class simplifies the management of a set of related CompiledExpressions that share variables.
  */
-class OPENMM_EXPORT_CPU CompiledExpressionSet {
+class OPENMM_EXPORT CompiledExpressionSet {
 public:
     CompiledExpressionSet();
     /**
@@ -60,6 +60,10 @@ public:
      * @param value    the value to set it to
      */
     void setVariable(int index, double value);
+    /**
+     * Get the total number of variables for which indices have been allocated.
+     */
+    int getNumVariables() const;
 private:
     std::vector<Lepton::CompiledExpression*> expressions;
     std::vector<std::string> variables;

openmmapi/include/openmm/internal/CustomIntegratorUtilities.h

+#ifndef OPENMM_CUSTOMINTEGRATORUTILITIES_H_
+#define OPENMM_CUSTOMINTEGRATORUTILITIES_H_
+
+/* -------------------------------------------------------------------------- *
+ *                                   OpenMM                                   *
+ * -------------------------------------------------------------------------- *
+ * This is part of the OpenMM molecular simulation toolkit originating from   *
+ * Simbios, the NIH National Center for Physics-Based Simulation of           *
+ * Biological Structures at Stanford, funded under the NIH Roadmap for        *
+ * Medical Research, grant U54 GM072970. See https://simtk.org.               *
+ *                                                                            *
+ * Portions copyright (c) 2015 Stanford University and the Authors.           *
+ * Authors: Peter Eastman                                                     *
+ * Contributors:                                                              *
+ *                                                                            *
+ * Permission is hereby granted, free of charge, to any person obtaining a    *
+ * copy of this software and associated documentation files (the "Software"), *
+ * to deal in the Software without restriction, including without limitation  *
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,   *
+ * and/or sell copies of the Software, and to permit persons to whom the      *
+ * Software is furnished to do so, subject to the following conditions:       *
+ *                                                                            *
+ * The above copyright notice and this permission notice shall be included in *
+ * all copies or substantial portions of the Software.                        *
+ *                                                                            *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR *
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,   *
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL    *
+ * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,    *
+ * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR      *
+ * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE  *
+ * USE OR OTHER DEALINGS IN THE SOFTWARE.                                     *
+ * -------------------------------------------------------------------------- */
+
+#include "openmm/CustomIntegrator.h"
+#include "openmm/internal/ContextImpl.h"
+#include "lepton/ParsedExpression.h"
+#include <map>
+#include <vector>
+
+namespace OpenMM {
+
+class System;
+
+/**
+ * This class defines a set of utility functions that are useful in implementing CustomIntegrator.
+ */
+
+class OPENMM_EXPORT CustomIntegratorUtilities {
+public:
+    enum Comparison {
+        EQUAL = 0, LESS_THAN = 1, GREATER_THAN = 2, NOT_EQUAL = 3, LESS_THAN_OR_EQUAL = 4, GREATER_THAN_OR_EQUAL = 5
+    };
+    /**
+     * Parse the expression for the condition of an "if" or "while" block, and split it into
+     * a left hand side, right hand side, and comparison operator.
+     */
+    static void parseCondition(const std::string& expression, std::string& lhs, std::string& rhs, Comparison& comparison);
+    /**
+     * 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. 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).
+     * 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<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.
+     */
+    static bool usesVariable(const Lepton::ParsedExpression& expression, const std::string& variable);
+private:
+    static bool usesVariable(const Lepton::ExpressionTreeNode& node, const std::string& variable);
+    static void enumeratePaths(int firstStep, std::vector<int> steps, std::vector<int> jumps, const std::vector<int>& blockEnd,
+            const std::vector<CustomIntegrator::ComputationType>& stepType, const std::vector<bool>& needsForces, const std::vector<bool>& needsEnergy,
+            const std::vector<bool>& invalidatesForces, const std::vector<int>& forceGroup, std::vector<bool>& computeBoth);
+    static void analyzeForceComputationsForPath(std::vector<int>& steps, const std::vector<bool>& needsForces, const std::vector<bool>& needsEnergy,
+            const std::vector<bool>& invalidatesForces, const std::vector<int>& forceGroup, std::vector<bool>& computeBoth);
+};
+
+} // namespace OpenMM
+
+#endif /*OPENMM_CUSTOMINTEGRATORUTILITIES_H_*/

platforms/cpu/src/CompiledExpressionSet.cpp → openmmapi/src/CompiledExpressionSet.cpp

-/* Portions copyright (c) 2014 Stanford University and Simbios.
+/* Portions copyright (c) 2014-2015 Stanford University and Simbios.
  * Contributors: Peter Eastman
  *
  * Permission is hereby granted, free of charge, to any person obtaining
@@ -21,7 +21,7 @@
  * WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
  */
 
-#include "CompiledExpressionSet.h"
+#include "openmm/internal/CompiledExpressionSet.h"
 
 using namespace OpenMM;
 using namespace Lepton;
@@ -54,3 +54,7 @@ void CompiledExpressionSet::setVariable(int index, double value) {
     for (int i = 0; i < (int) variableReferences[index].size(); i++)
         *variableReferences[index][i] = value;
 }
+
+int CompiledExpressionSet::getNumVariables() const {
+    return variables.size();
+}

openmmapi/src/CustomIntegrator.cpp

@@ -245,6 +245,27 @@ int CustomIntegrator::addUpdateContextState() {
     return computations.size()-1;
 }
 
+int CustomIntegrator::beginIfBlock(const string& expression) {
+    if (owner != NULL)
+        throw OpenMMException("The integrator cannot be modified after it is bound to a context");
+    computations.push_back(ComputationInfo(BeginIfBlock, "", expression));
+    return computations.size()-1;
+}
+
+int CustomIntegrator::beginWhileBlock(const string& expression) {
+    if (owner != NULL)
+        throw OpenMMException("The integrator cannot be modified after it is bound to a context");
+    computations.push_back(ComputationInfo(BeginWhileBlock, "", expression));
+    return computations.size()-1;
+}
+
+int CustomIntegrator::endBlock() {
+    if (owner != NULL)
+        throw OpenMMException("The integrator cannot be modified after it is bound to a context");
+    computations.push_back(ComputationInfo(EndBlock, "", ""));
+    return computations.size()-1;
+}
+
 void CustomIntegrator::getComputationStep(int index, ComputationType& type, string& variable, string& expression) const {
     ASSERT_VALID_INDEX(index, computations);
     type = computations[index].type;

openmmapi/src/CustomIntegratorUtilities.cpp

+/* -------------------------------------------------------------------------- *
+ *                                   OpenMM                                   *
+ * -------------------------------------------------------------------------- *
+ * This is part of the OpenMM molecular simulation toolkit originating from   *
+ * Simbios, the NIH National Center for Physics-Based Simulation of           *
+ * Biological Structures at Stanford, funded under the NIH Roadmap for        *
+ * Medical Research, grant U54 GM072970. See https://simtk.org.               *
+ *                                                                            *
+ * Portions copyright (c) 2015 Stanford University and the Authors.           *
+ * Authors: Peter Eastman                                                     *
+ * Contributors:                                                              *
+ *                                                                            *
+ * Permission is hereby granted, free of charge, to any person obtaining a    *
+ * copy of this software and associated documentation files (the "Software"), *
+ * to deal in the Software without restriction, including without limitation  *
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,   *
+ * and/or sell copies of the Software, and to permit persons to whom the      *
+ * Software is furnished to do so, subject to the following conditions:       *
+ *                                                                            *
+ * The above copyright notice and this permission notice shall be included in *
+ * all copies or substantial portions of the Software.                        *
+ *                                                                            *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR *
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,   *
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL    *
+ * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,    *
+ * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR      *
+ * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE  *
+ * USE OR OTHER DEALINGS IN THE SOFTWARE.                                     *
+ * -------------------------------------------------------------------------- */
+
+#include "openmm/internal/CustomIntegratorUtilities.h"
+#include "openmm/OpenMMException.h"
+#include "openmm/internal/ForceImpl.h"
+#include "lepton/Operation.h"
+#include "lepton/Parser.h"
+#include <set>
+#include <sstream>
+
+using namespace OpenMM;
+using namespace std;
+
+void CustomIntegratorUtilities::parseCondition(const string& expression, string& lhs, string& rhs, Comparison& comparison) {
+    string operators[] = {"=", "<", ">", "!=", "<=", ">="};
+    for (int i = 5; i >= 0; i--) {
+        int index = expression.find(operators[i]);
+        if (index != string::npos) {
+            lhs = expression.substr(0, index);
+            rhs = expression.substr(index+operators[i].size());
+            comparison = Comparison(i);
+            return;
+        }
+    }
+    throw OpenMMException("No comparison operator found in condition: "+expression);
+}
+
+bool CustomIntegratorUtilities::usesVariable(const Lepton::ExpressionTreeNode& node, const string& variable) {
+    const Lepton::Operation& op = node.getOperation();
+    if (op.getId() == Lepton::Operation::VARIABLE && op.getName() == variable)
+        return true;
+    for (int i = 0; i < (int) node.getChildren().size(); i++)
+        if (usesVariable(node.getChildren()[i], variable))
+            return true;
+    return false;
+}
+
+bool CustomIntegratorUtilities::usesVariable(const Lepton::ParsedExpression& expression, const string& variable) {
+    return usesVariable(expression.getRootNode(), variable);
+}
+
+void CustomIntegratorUtilities::analyzeComputations(const ContextImpl& context, const CustomIntegrator& integrator, vector<vector<Lepton::ParsedExpression> >& expressions,
+            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);
+    comparisons.resize(numSteps);
+    invalidatesForces.resize(numSteps, false);
+    needsForces.resize(numSteps, false);
+    needsEnergy.resize(numSteps, false);
+    computeBoth.resize(numSteps, false);
+    forceGroup.resize(numSteps, -2);
+    vector<CustomIntegrator::ComputationType> stepType(numSteps);
+    vector<string> stepVariable(numSteps);
+
+    // Parse the expressions.
+
+    for (int step = 0; step < numSteps; step++) {
+        string expression;
+        integrator.getComputationStep(step, stepType[step], stepVariable[step], expression);
+        if (stepType[step] == CustomIntegrator::BeginIfBlock || stepType[step] == CustomIntegrator::BeginWhileBlock) {
+            // This step involves a condition.
+
+            string lhs, rhs;
+            parseCondition(expression, lhs, rhs, comparisons[step]);
+            expressions[step].push_back(Lepton::Parser::parse(lhs).optimize());
+            expressions[step].push_back(Lepton::Parser::parse(rhs).optimize());
+        }
+        else if (expression.size() > 0)
+            expressions[step].push_back(Lepton::Parser::parse(expression).optimize());
+    }
+
+    // Identify which steps invalidate the forces.
+
+    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);
+    }
+    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.
+
+    vector<string> forceGroupName;
+    vector<string> energyGroupName;
+    for (int i = 0; i < 32; i++) {
+        stringstream fname;
+        fname << "f" << i;
+        forceGroupName.push_back(fname.str());
+        stringstream ename;
+        ename << "energy" << i;
+        energyGroupName.push_back(ename.str());
+    }
+    for (int step = 0; step < numSteps; step++) {
+        for (int expr = 0; expr < expressions[step].size(); expr++) {
+            if (usesVariable(expressions[step][expr], "f")) {
+                needsForces[step] = true;
+                forceGroup[step] = -1;
+            }
+            if (usesVariable(expressions[step][expr], "energy")) {
+                needsEnergy[step] = true;
+                forceGroup[step] = -1;
+            }
+            for (int i = 0; i < 32; i++) {
+                if (usesVariable(expressions[step][expr], forceGroupName[i])) {
+                    if (forceGroup[step] != -2)
+                        throw OpenMMException("A single computation step cannot depend on multiple force groups");
+                    needsForces[step] = true;
+                    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] = i;
+                }
+            }
+        }
+    }
+
+    // Find the end point of each block.
+
+    vector<int> blockStart;
+    blockEnd.resize(numSteps, -1);
+    for (int step = 0; step < numSteps; step++) {
+        if (stepType[step] == CustomIntegrator::BeginIfBlock || stepType[step] == CustomIntegrator::BeginWhileBlock)
+            blockStart.push_back(step);
+        else if (stepType[step] == CustomIntegrator::EndBlock) {
+            if (blockStart.size() == 0) {
+                stringstream error("CustomIntegrator: Unexpected end of block at computation ");
+                error << step;
+                throw OpenMMException(error.str());
+            }
+            blockEnd[blockStart.back()] = step;
+            blockStart.pop_back();
+        }
+    }
+    if (blockStart.size() > 0)
+        throw OpenMMException("CustomIntegrator: Missing EndBlock");
+
+    // If a step requires either forces or energy, and a later step will require the other one, it's most efficient
+    // to compute both at the same time.  Figure out whether we should do that.  In principle it's easy: step through
+    // the sequence of computations and see if the other one is used before the next time they get invalidated.
+    // Unfortunately, flow control makes this much more complicated, because there are many possible paths to
+    // consider.
+    //
+    // The cost of computing both when we really only needed one is much less than the cost of computing only one,
+    // then later finding we need to compute the other separately.  So we always err on the side of computing both.
+    // If there is any possible path that would lead to us needing it, go ahead and compute it.
+    //
+    // So we need to enumerate all possible paths.  For each "if" block, there are two possibilities: execute it
+    // or don't.  For each "while" block there are three possibilities: don't execute it; execute it and then
+    // continue on; or execute it and then jump back to the beginning.  I'm assuming the number of blocks will
+    // always remain small.  Otherwise, this could become very expensive!
+
+    vector<int> jumps(numSteps, -1);
+    vector<int> stepsInPath;
+    enumeratePaths(0, stepsInPath, jumps, blockEnd, stepType, needsForces, needsEnergy, invalidatesForces, forceGroup, computeBoth);
+}
+
+void CustomIntegratorUtilities::enumeratePaths(int firstStep, vector<int> steps, vector<int> jumps, const vector<int>& blockEnd,
+            const vector<CustomIntegrator::ComputationType>& stepType, const vector<bool>& needsForces, const vector<bool>& needsEnergy,
+            const vector<bool>& invalidatesForces, const vector<int>& forceGroup, vector<bool>& computeBoth) {
+    int step = firstStep;
+    int numSteps = stepType.size();
+    while (step < numSteps) {
+        steps.push_back(step);
+        if (jumps[step] > 0) {
+            // Follow the jump and remove it from the list.
+
+            int nextStep = jumps[step];
+            jumps[step] = -1;
+            step = nextStep;
+        }
+        else if (stepType[step] == CustomIntegrator::BeginIfBlock) {
+            // Consider skipping the block.
+
+            enumeratePaths(blockEnd[step]+1, steps, jumps, blockEnd, stepType, needsForces, needsEnergy, invalidatesForces, forceGroup, computeBoth);
+
+            // Continue on to execute the block.
+
+            step++;
+        }
+        else if (stepType[step] == CustomIntegrator::BeginWhileBlock && jumps[step] != -2) {
+            // Consider skipping the block.
+
+            enumeratePaths(blockEnd[step]+1, steps, jumps, blockEnd, stepType, needsForces, needsEnergy, invalidatesForces, forceGroup, computeBoth);
+
+            // Consider executing the block once.
+
+            enumeratePaths(step+1, steps, jumps, blockEnd, stepType, needsForces, needsEnergy, invalidatesForces, forceGroup, computeBoth);
+
+            // Continue on to execute the block twice.
+
+            jumps[step] = -2; // Mark this "while" block as already processed.
+            jumps[blockEnd[step]] = step;
+            step++;
+        }
+        else
+            step++;
+    }
+    analyzeForceComputationsForPath(steps, needsForces, needsEnergy, invalidatesForces, forceGroup, computeBoth);
+}
+
+void CustomIntegratorUtilities::analyzeForceComputationsForPath(vector<int>& steps, const vector<bool>& needsForces, const vector<bool>& needsEnergy,
+            const vector<bool>& invalidatesForces, const vector<int>& forceGroup, vector<bool>& computeBoth) {
+    vector<int> candidatePoints;
+    int currentGroup = -1;
+    for (int i = 0; i < (int) steps.size(); i++) {
+        int step = steps[i];
+        if (invalidatesForces[step] || ((needsForces[step] || needsEnergy[step]) && forceGroup[step] != currentGroup)) {
+            // Forces and energies are invalidated at this step, or it changes to a different force group,
+            // so anything from this point on won't affect what we do at earlier steps.
+
+            candidatePoints.clear();
+        }
+        if (needsForces[step] || needsEnergy[step]) {
+            // See if this step affects what we do at earlier points.
+
+            for (int j = 0; j < (int) candidatePoints.size(); j++) {
+                int candidate = candidatePoints[j];
+                if ((needsForces[candidate] && needsEnergy[step]) || (needsEnergy[candidate] && needsForces[step]))
+                    computeBoth[candidate] = true;
+            }
+
+            // Add this to the list of candidates that might be affected by later steps.
+
+            candidatePoints.push_back(step);
+            currentGroup = forceGroup[step];
+        }
+    }
+}
\ No newline at end of file

platforms/cpu/include/CpuCustomGBForce.h

@@ -25,10 +25,10 @@
 #ifndef OPENMM_CPU_CUSTOM_GB_FORCE_H__
 #define OPENMM_CPU_CUSTOM_GB_FORCE_H__
 
-#include "CompiledExpressionSet.h"
 #include "CpuNeighborList.h"
 #include "lepton/CompiledExpression.h"
 #include "openmm/CustomGBForce.h"
+#include "openmm/internal/CompiledExpressionSet.h"
 #include "openmm/internal/ThreadPool.h"
 #include "openmm/internal/vectorize.h"
 #include <map>

platforms/cpu/include/CpuCustomManyParticleForce.h

@@ -27,9 +27,9 @@
 
 #include "ReferenceForce.h"
 #include "ReferenceBondIxn.h"
-#include "CompiledExpressionSet.h"
 #include "CpuNeighborList.h"
 #include "openmm/CustomManyParticleForce.h"
+#include "openmm/internal/CompiledExpressionSet.h"
 #include "openmm/internal/ThreadPool.h"
 #include "openmm/internal/vectorize.h"
 #include "lepton/CompiledExpression.h"

platforms/cuda/include/CudaKernels.h

@@ -35,6 +35,9 @@
 #include "CudaSort.h"
 #include "openmm/kernels.h"
 #include "openmm/System.h"
+#include "openmm/internal/CompiledExpressionSet.h"
+#include "openmm/internal/CustomIntegratorUtilities.h"
+#include "lepton/CompiledExpression.h"
 #include <cufft.h>
 
 namespace OpenMM {
@@ -1214,9 +1217,10 @@ private:
  */
 class CudaIntegrateCustomStepKernel : public IntegrateCustomStepKernel {
 public:
+    enum GlobalTargetType {DT, VARIABLE, PARAMETER};
     CudaIntegrateCustomStepKernel(std::string name, const Platform& platform, CudaContext& cu) : IntegrateCustomStepKernel(name, platform), cu(cu),
-            hasInitializedKernels(false), localValuesAreCurrent(false), globalValues(NULL), contextParameterValues(NULL), sumBuffer(NULL), potentialEnergy(NULL),
-            kineticEnergy(NULL), uniformRandoms(NULL), randomSeed(NULL), perDofValues(NULL) {
+            hasInitializedKernels(false), localValuesAreCurrent(false), globalValues(NULL), sumBuffer(NULL), summedValue(NULL), uniformRandoms(NULL),
+            randomSeed(NULL), perDofValues(NULL) {
     }
     ~CudaIntegrateCustomStepKernel();
     /**
@@ -1280,21 +1284,21 @@ 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, const std::string& energyName);
+    class GlobalTarget;
     std::string createPerDofComputation(const std::string& variable, const Lepton::ParsedExpression& expr, int component, CustomIntegrator& integrator, const std::string& forceName, const std::string& energyName);
     void prepareForComputation(ContextImpl& context, CustomIntegrator& integrator, bool& forcesAreValid);
+    void recordGlobalValue(double value, GlobalTarget target);
     void recordChangedParameters(ContextImpl& context);
+    bool evaluateCondition(int step);
     CudaContext& cu;
     double prevStepSize, energy;
     float energyFloat;
     int numGlobalVariables;
-    bool hasInitializedKernels, deviceValuesAreCurrent, modifiesParameters, keNeedsForce;
+    bool hasInitializedKernels, deviceValuesAreCurrent, deviceGlobalsAreCurrent, modifiesParameters, keNeedsForce;
     mutable bool localValuesAreCurrent;
     CudaArray* globalValues;
-    CudaArray* contextParameterValues;
     CudaArray* sumBuffer;
-    CudaArray* potentialEnergy;
-    CudaArray* kineticEnergy;
+    CudaArray* summedValue;
     CudaArray* uniformRandoms;
     CudaArray* randomSeed;
     std::map<int, CudaArray*> savedForces;
@@ -1302,21 +1306,43 @@ private:
     CudaParameterSet* perDofValues;
     mutable std::vector<std::vector<float> > localPerDofValuesFloat;
     mutable std::vector<std::vector<double> > localPerDofValuesDouble;
-    std::vector<float> contextValuesFloat;
-    std::vector<double> contextValuesDouble;
+    std::vector<float> globalValuesFloat;
+    std::vector<double> globalValuesDouble;
+    std::vector<double> initialGlobalVariables;
     std::vector<std::vector<CUfunction> > kernels;
     std::vector<std::vector<std::vector<void*> > > kernelArgs;
     std::vector<void*> kineticEnergyArgs;
     CUfunction randomKernel, kineticEnergyKernel, sumKineticEnergyKernel;
     std::vector<CustomIntegrator::ComputationType> stepType;
+    std::vector<CustomIntegratorUtilities::Comparison> comparisons;
+    std::vector<std::vector<Lepton::CompiledExpression> > globalExpressions;
+    CompiledExpressionSet expressionSet;
+    std::vector<bool> needsGlobals;
     std::vector<bool> needsForces;
     std::vector<bool> needsEnergy;
+    std::vector<bool> computeBothForceAndEnergy;
     std::vector<bool> invalidatesForces;
     std::vector<bool> merged;
-    std::vector<int> forceGroup;
+    std::vector<int> forceGroupFlags;
+    std::vector<int> blockEnd;
     std::vector<int> requiredGaussian;
     std::vector<int> requiredUniform;
+    std::vector<int> stepEnergyVariableIndex;
+    std::vector<int> globalVariableIndex;
+    std::vector<int> parameterVariableIndex;
+    int gaussianVariableIndex, uniformVariableIndex, dtVariableIndex;
     std::vector<std::string> parameterNames;
+    std::vector<GlobalTarget> stepTarget;
+};
+
+class CudaIntegrateCustomStepKernel::GlobalTarget {
+public:
+    CudaIntegrateCustomStepKernel::GlobalTargetType type;
+    int variableIndex;
+    GlobalTarget() {
+    }
+    GlobalTarget(CudaIntegrateCustomStepKernel::GlobalTargetType type, int variableIndex) : type(type), variableIndex(variableIndex) {
+    }
 };
 
 /**

platforms/cuda/src/kernels/customIntegrator.cu

@@ -11,7 +11,7 @@ extern "C" __global__ void computeFloatSum(const float* __restrict__ sumBuffer,
             tempBuffer[thread] += tempBuffer[thread+i];
     }
     if (thread == 0)
-        result[SUM_OUTPUT_INDEX] = tempBuffer[0];
+        *result = tempBuffer[0];
 }
 
 extern "C" __global__ void computeDoubleSum(const double* __restrict__ sumBuffer, double* result) {
@@ -27,7 +27,7 @@ extern "C" __global__ void computeDoubleSum(const double* __restrict__ sumBuffer
             tempBuffer[thread] += tempBuffer[thread+i];
     }
     if (thread == 0)
-        result[SUM_OUTPUT_INDEX] = tempBuffer[0];
+        *result = tempBuffer[0];
 }
 
 extern "C" __global__ void applyPositionDeltas(real4* __restrict__ posq, real4* __restrict__ posqCorrection, mixed4* __restrict__ posDelta) {

platforms/cuda/src/kernels/customIntegratorGlobal.cu

-extern "C" __global__ void computeGlobal(mixed2* __restrict__ dt, mixed* __restrict__ globals, mixed* __restrict__ params,
-        float uniform, float gaussian, const real energy) {
-    COMPUTE_STEP
-}

platforms/cuda/src/kernels/customIntegratorPerDof.cu

@@ -33,8 +33,7 @@ inline __device__ mixed4 convertFromDouble4(double4 a) {
 
 extern "C" __global__ void computePerDof(real4* __restrict__ posq, real4* __restrict__ posqCorrection, mixed4* __restrict__ posDelta,
         mixed4* __restrict__ velm, const long long* __restrict__ force, const mixed2* __restrict__ dt, const mixed* __restrict__ globals,
-        const mixed* __restrict__ params, mixed* __restrict__ sum, const float4* __restrict__ gaussianValues,
-        unsigned int gaussianBaseIndex, const float4* __restrict__ uniformValues, const real energy
+        mixed* __restrict__ sum, const float4* __restrict__ gaussianValues, unsigned int gaussianBaseIndex, const float4* __restrict__ uniformValues, const real energy
         PARAMETER_ARGUMENTS) {
     mixed stepSize = dt[0].y;
     int index = blockIdx.x*blockDim.x+threadIdx.x;

platforms/cuda/src/kernels/integrationUtilities.cu

@@ -224,7 +224,6 @@ extern "C" __global__ void applyShakeToVelocities(int numClusters, mixed tol, co
         mixed4 xpj2 = make_mixed4(0);
         float invMassCentral = params.x;
         float avgMass = params.y;
-        float d2 = params.z;
         float invMassPeripheral = params.w;
         if (atoms.z != -1) {
             pos2 = loadPos(oldPos, posCorrection, atoms.z);
@@ -245,9 +244,6 @@ extern "C" __global__ void applyShakeToVelocities(int numClusters, mixed tol, co
         mixed rij1sq = rij1.x*rij1.x + rij1.y*rij1.y + rij1.z*rij1.z;
         mixed rij2sq = rij2.x*rij2.x + rij2.y*rij2.y + rij2.z*rij2.z;
         mixed rij3sq = rij3.x*rij3.x + rij3.y*rij3.y + rij3.z*rij3.z;
-        mixed ld1 = d2-rij1sq;
-        mixed ld2 = d2-rij2sq;
-        mixed ld3 = d2-rij3sq;
 
         // Iterate until convergence.
 
@@ -605,8 +601,6 @@ extern "C" __global__ void computeCCMAVelocityConstraintForce(const int2* __rest
     if (threadIdx.x == 0)
         groupConverged = 1;
     __syncthreads();
-    mixed lowerTol = 1-2*tol+tol*tol;
-    mixed upperTol = 1+2*tol+tol*tol;
     for (int index = blockIdx.x*blockDim.x+threadIdx.x; index < NUM_CCMA_CONSTRAINTS; index += blockDim.x*gridDim.x) {
         // Compute the force due to this constraint.
 

platforms/cuda/tests/TestCudaCustomIntegrator.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:                                                              *
  *                                                                            *
@@ -756,6 +756,92 @@ void testMergedRandoms() {
     }
 }
 
+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);
+}
+
+/**
+ * Test modifying a global variable, then using it in a per-DOF computation.
+ */
+void testChangingGlobal() {
+    System system;
+    system.addParticle(1.0);
+    CustomIntegrator integrator(0.1);
+    integrator.addGlobalVariable("g", 0);
+    integrator.addPerDofVariable("a", 0);
+    integrator.addPerDofVariable("b", 0);
+    integrator.addComputeGlobal("g", "g+1");
+    integrator.addComputePerDof("a", "0.5");
+    integrator.addComputePerDof("b", "a+g");
+    Context context(system, integrator, platform);
+    
+    // See if everything is being calculated correctly..
+    
+    for (int i = 0; i < 10; i++) {
+        integrator.step(1);
+        ASSERT_EQUAL_TOL(i+1, integrator.getGlobalVariable(0), 1e-5);
+        vector<Vec3> values;
+        integrator.getPerDofVariable(1, values);
+        ASSERT_EQUAL_VEC(Vec3(i+1.5, i+1.5, i+1.5), values[0], 1e-5);
+    }
+}
+
 int main(int argc, char* argv[]) {
     try {
         if (argc > 1)
@@ -773,6 +859,9 @@ int main(int argc, char* argv[]) {
         testForceGroups();
         testRespa();
         testMergedRandoms();
+        testIfBlock();
+        testWhileBlock();
+        testChangingGlobal();
     }
     catch(const exception& e) {
         cout << "exception: " << e.what() << endl;

platforms/opencl/include/OpenCLKernels.h

@@ -34,6 +34,9 @@
 #include "OpenCLParameterSet.h"
 #include "OpenCLSort.h"
 #include "openmm/kernels.h"
+#include "openmm/internal/CompiledExpressionSet.h"
+#include "openmm/internal/CustomIntegratorUtilities.h"
+#include "lepton/CompiledExpression.h"
 #include "openmm/System.h"
 
 namespace OpenMM {
@@ -1203,9 +1206,10 @@ private:
  */
 class OpenCLIntegrateCustomStepKernel : public IntegrateCustomStepKernel {
 public:
+    enum GlobalTargetType {DT, VARIABLE, PARAMETER};
     OpenCLIntegrateCustomStepKernel(std::string name, const Platform& platform, OpenCLContext& cl) : IntegrateCustomStepKernel(name, platform), cl(cl),
-            hasInitializedKernels(false), localValuesAreCurrent(false), globalValues(NULL), contextParameterValues(NULL), sumBuffer(NULL), potentialEnergy(NULL),
-            kineticEnergy(NULL), uniformRandoms(NULL), randomSeed(NULL), perDofValues(NULL) {
+            hasInitializedKernels(false), localValuesAreCurrent(false), globalValues(NULL), sumBuffer(NULL), summedValue(NULL), uniformRandoms(NULL),
+            randomSeed(NULL), perDofValues(NULL) {
     }
     ~OpenCLIntegrateCustomStepKernel();
     /**
@@ -1269,20 +1273,21 @@ 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, const std::string& energyName);
+    class GlobalTarget;
     std::string createPerDofComputation(const std::string& variable, const Lepton::ParsedExpression& expr, int component, CustomIntegrator& integrator, const std::string& forceName, const std::string& energyName);
     void prepareForComputation(ContextImpl& context, CustomIntegrator& integrator, bool& forcesAreValid);
+    void recordGlobalValue(double value, GlobalTarget target);
     void recordChangedParameters(ContextImpl& context);
+    bool evaluateCondition(int step);
     OpenCLContext& cl;
     double prevStepSize, energy;
+    float energyFloat;
     int numGlobalVariables;
-    bool hasInitializedKernels, deviceValuesAreCurrent, modifiesParameters, keNeedsForce;
+    bool hasInitializedKernels, deviceValuesAreCurrent, deviceGlobalsAreCurrent, modifiesParameters, keNeedsForce;
     mutable bool localValuesAreCurrent;
     OpenCLArray* globalValues;
-    OpenCLArray* contextParameterValues;
     OpenCLArray* sumBuffer;
-    OpenCLArray* potentialEnergy;
-    OpenCLArray* kineticEnergy;
+    OpenCLArray* summedValue;
     OpenCLArray* uniformRandoms;
     OpenCLArray* randomSeed;
     std::map<int, OpenCLArray*> savedForces;
@@ -1290,20 +1295,41 @@ private:
     OpenCLParameterSet* perDofValues;
     mutable std::vector<std::vector<cl_float> > localPerDofValuesFloat;
     mutable std::vector<std::vector<cl_double> > localPerDofValuesDouble;
-    std::vector<float> contextValuesFloat;
-    std::vector<double> contextValuesDouble;
-    std::vector<float> contextValues;
+    std::vector<float> globalValuesFloat;
+    std::vector<double> globalValuesDouble;
+    std::vector<double> initialGlobalVariables;
     std::vector<std::vector<cl::Kernel> > kernels;
     cl::Kernel randomKernel, kineticEnergyKernel, sumKineticEnergyKernel;
     std::vector<CustomIntegrator::ComputationType> stepType;
+    std::vector<CustomIntegratorUtilities::Comparison> comparisons;
+    std::vector<std::vector<Lepton::CompiledExpression> > globalExpressions;
+    CompiledExpressionSet expressionSet;
+    std::vector<bool> needsGlobals;
     std::vector<bool> needsForces;
     std::vector<bool> needsEnergy;
+    std::vector<bool> computeBothForceAndEnergy;
     std::vector<bool> invalidatesForces;
     std::vector<bool> merged;
-    std::vector<int> forceGroup;
+    std::vector<int> forceGroupFlags;
+    std::vector<int> blockEnd;
     std::vector<int> requiredGaussian;
     std::vector<int> requiredUniform;
+    std::vector<int> stepEnergyVariableIndex;
+    std::vector<int> globalVariableIndex;
+    std::vector<int> parameterVariableIndex;
+    int gaussianVariableIndex, uniformVariableIndex, dtVariableIndex;
     std::vector<std::string> parameterNames;
+    std::vector<GlobalTarget> stepTarget;
+};
+
+class OpenCLIntegrateCustomStepKernel::GlobalTarget {
+public:
+    OpenCLIntegrateCustomStepKernel::GlobalTargetType type;
+    int variableIndex;
+    GlobalTarget() {
+    }
+    GlobalTarget(OpenCLIntegrateCustomStepKernel::GlobalTargetType type, int variableIndex) : type(type), variableIndex(variableIndex) {
+    }
 };
 
 /**

platforms/opencl/src/kernels/customIntegrator.cl

-__kernel void computeFloatSum(__global const float* restrict sumBuffer, __global float* result, unsigned int outputIndex, int bufferSize) {
+__kernel void computeFloatSum(__global const float* restrict sumBuffer, __global float* result, int bufferSize) {
     __local float tempBuffer[WORK_GROUP_SIZE];
     const unsigned int thread = get_local_id(0);
     float sum = 0;
@@ -11,11 +11,11 @@ __kernel void computeFloatSum(__global const float* restrict sumBuffer, __global
             tempBuffer[thread] += tempBuffer[thread+i];
     }
     if (thread == 0)
-        result[outputIndex] = tempBuffer[0];
+        *result = tempBuffer[0];
 }
 
 #ifdef SUPPORTS_DOUBLE_PRECISION
-__kernel void computeDoubleSum(__global const double* restrict sumBuffer, __global double* result, unsigned int outputIndex, int bufferSize) {
+__kernel void computeDoubleSum(__global const double* restrict sumBuffer, __global double* result, int bufferSize) {
     __local double tempBuffer[WORK_GROUP_SIZE];
     const unsigned int thread = get_local_id(0);
     double sum = 0;
@@ -28,7 +28,7 @@ __kernel void computeDoubleSum(__global const double* restrict sumBuffer, __glob
             tempBuffer[thread] += tempBuffer[thread+i];
     }
     if (thread == 0)
-        result[outputIndex] = tempBuffer[0];
+        *result = tempBuffer[0];
 }
 #endif
 

platforms/opencl/src/kernels/customIntegratorGlobal.cl

-__kernel void computeGlobal(__global mixed2* restrict dt, __global mixed* restrict globals, __global mixed* restrict params,
-        float uniform, float gaussian, const real energy) {
-    COMPUTE_STEP
-}

platforms/opencl/src/kernels/customIntegratorPerDof.cl

@@ -25,8 +25,7 @@ void storePos(__global real4* restrict posq, __global real4* restrict posqCorrec
 
 __kernel void computePerDof(__global real4* restrict posq, __global real4* restrict posqCorrection, __global mixed4* restrict posDelta,
         __global mixed4* restrict velm, __global const real4* restrict force, __global const mixed2* restrict dt, __global const mixed* restrict globals,
-        __global const mixed* restrict params, __global mixed* restrict sum, __global const float4* restrict gaussianValues,
-        unsigned int gaussianBaseIndex, __global const float4* restrict uniformValues, const real energy
+        __global mixed* restrict sum, __global const float4* restrict gaussianValues, unsigned int gaussianBaseIndex, __global const float4* restrict uniformValues, const real energy
         PARAMETER_ARGUMENTS) {
     mixed stepSize = dt[0].y;
     int index = get_global_id(0);