Merge pull request #2057 from peastman/vector

CustomIntegrator can use vector functions

libraries/lepton/include/lepton/CustomFunction.h

@@ -72,6 +72,38 @@ public:
     virtual CustomFunction* clone() const = 0;
 };
 
+/**
+ * This class is an implementation of CustomFunction that does no computation.  It just returns
+ * 0 for the value and derivatives.  This is useful when using the parser to analyze expressions
+ * rather than to evaluate them.  You can just create PlaceholderFunctions to represent any custom
+ * functions that may appear in expressions.
+ */
+
+class LEPTON_EXPORT PlaceholderFunction : public CustomFunction {
+public:
+    /**
+     * Create a Placeholder function.
+     * 
+     * @param numArgs    the number of arguments the function expects
+     */
+    PlaceholderFunction(int numArgs) : numArgs(numArgs) {
+    }
+    int getNumArguments() const {
+        return numArgs;
+    }
+    double evaluate(const double* arguments) const {
+        return 0.0;
+    }
+    double evaluateDerivative(const double* arguments, const int* derivOrder) const {
+        return 0.0;
+    }
+    CustomFunction* clone() const {
+        return new PlaceholderFunction(numArgs);
+    };
+private:
+    int numArgs;
+};
+
 } // namespace Lepton
 
 #endif /*LEPTON_CUSTOM_FUNCTION_H_*/

openmmapi/include/openmm/CustomIntegrator.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) 2011-2017 Stanford University and the Authors.      *
+ * Portions copyright (c) 2011-2018 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -194,7 +194,7 @@ namespace OpenMM {
  *
  * <tt><pre>
  * integrator.beginIfBlock("uniform < acceptanceProbability");
- * integrator.computePerDof("x", "xnew");
+ * integrator.addComputePerDof("x", "xnew");
  * integrator.endBlock();
  * </pre></tt>
  *
@@ -203,6 +203,21 @@ namespace OpenMM {
  * following comparison operators: =, <. >, !=, <=, >=.  Blocks may be nested
  * inside each other.
  * 
+ * "Per-DOF" computations can also be thought of as per-particle computations
+ * that operate on three component vectors.  For example, "x+dt*v" means to take
+ * the particle's velocity (a vector), multiply it by the step size, and add the
+ * position (also a vector).  The result is a new vector that can be stored into
+ * a per-DOF variable with addComputePerDof(), or it can be summed over all
+ * components of all particles with addComputeSum().  Because the calculation is
+ * done on vectors, you can use functions that operate explicitly on vectors
+ * rather than just computing each component independently.  For example, the
+ * following line uses a cross product to compute the angular momentum of each
+ * particle and stores it into a per-DOF variable.
+ * 
+ * <tt><pre>
+ * integrator.addComputePerDof("angularMomentum", "m*cross(x, v)");
+ * </pre></tt>
+ * 
  * Another feature of CustomIntegrator is that it can use derivatives of the
  * potential energy with respect to context parameters.  These derivatives are
  * typically computed by custom forces, and are only computed if a Force object
@@ -243,6 +258,18 @@ namespace OpenMM {
  * are defined in radians, and log is the natural logarithm.  step(x) = 0 if x is less than 0, 1 otherwise.  delta(x) = 1 if x is 0, 0 otherwise.
  * select(x,y,z) = z if x = 0, y otherwise.  An expression may also involve intermediate quantities that are defined following the main expression, using ";" as a separator.
  * 
+ * Expressions used in ComputePerDof and ComputeSum steps can also use the following
+ * functions that operate on vectors: cross(a, b) is the cross product of two
+ * vectors; dot(a, b) is the dot product of two vectors; _x(a), _y(a), and _z(a)
+ * extract a single component from a vector; and vector(a, b, c) creates a new
+ * vector with the x component of the first argument, the y component of the
+ * second argument, and the z component of the third argument.  Remember that every
+ * quantity appearing in a vector expression is a vector.  Functions that appear
+ * to return a scalar really return a vector whose components are all the same.
+ * For example, _z(a) returns the vector (a.z, a.z, a.z).  Likewise, wherever a
+ * constant appears in the expression, it really means a vector whose components
+ * all have the same value.
+ *
  * In addition, you can call addTabulatedFunction() to define a new function based on tabulated values.  You specify the function by
  * creating a TabulatedFunction object.  That function can then appear in expressions.
  */

openmmapi/include/openmm/internal/CustomIntegratorUtilities.h

@@ -50,7 +50,6 @@ class System;
 
 class OPENMM_EXPORT CustomIntegratorUtilities {
 public:
-    class DerivFunction;
     enum Comparison {
         EQUAL = 0, LESS_THAN = 1, GREATER_THAN = 2, NOT_EQUAL = 3, LESS_THAN_OR_EQUAL = 4, GREATER_THAN_OR_EQUAL = 5
     };
@@ -88,27 +87,6 @@ private:
     static void validateDerivatives(const Lepton::ExpressionTreeNode& node, const std::vector<std::string>& derivNames);
 };
 
-/**
- * This class is used to implement the deriv() function when it appears in expressions.
- */
-class CustomIntegratorUtilities::DerivFunction : public Lepton::CustomFunction {
-public:
-    DerivFunction() {
-    }
-    int getNumArguments() const {
-        return 2;
-    }
-    double evaluate(const double* arguments) const {
-        return 0.0;
-    }
-    double evaluateDerivative(const double* arguments, const int* derivOrder) const {
-        return 0.0;
-    }
-    CustomFunction* clone() const {
-        return new DerivFunction();
-    }
-};
-
 } // namespace OpenMM
 
 #endif /*OPENMM_CUSTOMINTEGRATORUTILITIES_H_*/

openmmapi/include/openmm/internal/VectorExpression.h

+#ifndef OPENMM_VECTOR_EXPRESSION_H_
+#define OPENMM_VECTOR_EXPRESSION_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) 2018 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/Vec3.h"
+#include "lepton/CustomFunction.h"
+#include "lepton/ExpressionProgram.h"
+#include "lepton/ParsedExpression.h"
+#include "windowsExport.h"
+#include <map>
+#include <string>
+#include <vector>
+
+namespace OpenMM {
+
+/**
+ * This class wraps a Lepton expression to support parsing and evaluating expressions
+ * that operate on Vec3s instead of scalars.  In addition to the standard functions
+ * supported by Lepton, vector expressions can use the functions dot(), cross(),
+ * vector(), _x(), _y(), and _z().
+ */
+class OPENMM_EXPORT VectorExpression {
+public:
+    /**
+     * Create a VectorExpression by parsing a mathematical expression.
+     *
+     * @param expression        the mathematical expression to parse
+     * @param customFunctions   a map specifying user defined functions that may appear in the expression.
+     *                          The keys are function names, and the values are corresponding CustomFunction objects.
+     */
+    VectorExpression(const std::string& expression, const std::map<std::string, Lepton::CustomFunction*>& customFunctions);
+    /**
+     * Create a VectorExpression for evaluating a ParsedExpression.
+     *
+     * @param expression        the mathematical expression to evaluate
+     */
+    VectorExpression(const Lepton::ParsedExpression& expression);
+    VectorExpression(const VectorExpression& expression);
+    ~VectorExpression();
+    /**
+     * Evaluate the expression.
+     *
+     * @param variables    a map specifying the values of all variables that appear in the expression.  If any
+     *                     variable appears in the expression but is not included in this map, an exception
+     *                     will be thrown.
+     */
+    Vec3 evaluate(const std::map<std::string, Vec3>& variables) const;
+    VectorExpression& operator=(const VectorExpression& exp);
+private:
+    class VectorOperation;
+    class OpWrapper1;
+    class OpWrapper2;
+    class OpWrapper3;
+    class Variable;
+    class Constant;
+    class Dot;
+    class Cross;
+    class Component;
+    class Vector;
+    Lepton::ParsedExpression parsed;
+    Lepton::ExpressionProgram program;
+    mutable std::vector<Vec3> stack;
+    std::vector<VectorOperation*> operations;
+    void analyzeExpression(const Lepton::ParsedExpression& expression);
+};
+
+} // namespace OpenMM
+
+#endif /*OPENMM_VECTOR_EXPRESSION_H_*/

openmmapi/src/CustomIntegratorUtilities.cpp

@@ -84,8 +84,15 @@ void CustomIntegratorUtilities::analyzeComputations(const ContextImpl& context,
     vector<CustomIntegrator::ComputationType> stepType(numSteps);
     vector<string> stepVariable(numSteps);
     map<string, Lepton::CustomFunction*> customFunctions = functions;
-    DerivFunction derivFunction;
-    customFunctions["deriv"] = &derivFunction;
+    Lepton::PlaceholderFunction fn1(1), fn2(2), fn3(3);
+    customFunctions["deriv"] = &fn2;
+    map<string, Lepton::CustomFunction*> vectorFunctions = customFunctions;
+    vectorFunctions["dot"] = &fn2;
+    vectorFunctions["cross"] = &fn2;
+    vectorFunctions["_x"] = &fn1;
+    vectorFunctions["_y"] = &fn1;
+    vectorFunctions["_z"] = &fn1;
+    vectorFunctions["vector"] = &fn3;
 
     // Parse the expressions.
 
@@ -100,6 +107,8 @@ void CustomIntegratorUtilities::analyzeComputations(const ContextImpl& context,
             expressions[step].push_back(Lepton::Parser::parse(lhs, customFunctions).optimize());
             expressions[step].push_back(Lepton::Parser::parse(rhs, customFunctions).optimize());
         }
+        else if (stepType[step] == CustomIntegrator::ComputePerDof || stepType[step] == CustomIntegrator::ComputeSum)
+            expressions[step].push_back(Lepton::Parser::parse(expression, vectorFunctions).optimize());
         else if (expression.size() > 0)
             expressions[step].push_back(Lepton::Parser::parse(expression, customFunctions).optimize());
     }

openmmapi/src/VectorExpression.cpp

+/* Portions copyright (c) 2018 Stanford University and Simbios.
+ * Contributors: Peter Eastman
+ *
+ * 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/VectorExpression.h"
+#include "lepton/CustomFunction.h"
+#include "lepton/Operation.h"
+#include "lepton/Parser.h"
+#include "lepton/ParsedExpression.h"
+#include "openmm/OpenMMException.h"
+
+using namespace OpenMM;
+using namespace Lepton;
+using namespace std;
+
+static map<string, double> noVariables;
+
+/**
+ * Similar to Lepton::Operation, but it operates on Vec3s instead of doubles.
+ */
+class VectorExpression::VectorOperation {
+public:
+    virtual ~VectorOperation() {
+    }
+    virtual Vec3 evaluate(const Vec3* args, const map<string, Vec3>& variables) const = 0;
+};
+
+class VectorExpression::OpWrapper1 : public VectorExpression::VectorOperation {
+public:
+    OpWrapper1(const Operation& op) : op(op) {
+    }
+    Vec3 evaluate(const Vec3* args, const map<string, Vec3>& variables) const {
+        Vec3 result;
+        double arg = args[0][0];
+        result[0] = op.evaluate(&arg, noVariables);
+        arg = args[0][1];
+        result[1] = op.evaluate(&arg, noVariables);
+        arg = args[0][2];
+        result[2] = op.evaluate(&arg, noVariables);
+        return result;
+    }
+private:
+    const Operation& op;
+};
+
+class VectorExpression::OpWrapper2 : public VectorExpression::VectorOperation {
+public:
+    OpWrapper2(const Operation& op) : op(op) {
+    }
+    Vec3 evaluate(const Vec3* args, const map<string, Vec3>& variables) const {
+        Vec3 result;
+        double args0[] = {args[0][0], args[1][0]};
+        result[0] = op.evaluate(args0, noVariables);
+        double args1[] = {args[0][1], args[1][1]};
+        result[1] = op.evaluate(args1, noVariables);
+        double args2[] = {args[0][2], args[1][2]};
+        result[2] = op.evaluate(args2, noVariables);
+        return result;
+    }
+private:
+    const Operation& op;
+};
+
+class VectorExpression::OpWrapper3 : public VectorExpression::VectorOperation {
+public:
+    OpWrapper3(const Operation& op) : op(op) {
+    }
+    Vec3 evaluate(const Vec3* args, const map<string, Vec3>& variables) const {
+        Vec3 result;
+        double args0[] = {args[0][0], args[1][0], args[2][0]};
+        result[0] = op.evaluate(args0, noVariables);
+        double args1[] = {args[0][1], args[1][1], args[2][1]};
+        result[1] = op.evaluate(args1, noVariables);
+        double args2[] = {args[0][2], args[1][2], args[2][2]};
+        result[2] = op.evaluate(args2, noVariables);
+        return result;
+    }
+private:
+    const Operation& op;
+};
+
+class VectorExpression::Variable : public VectorExpression::VectorOperation {
+public:
+    Variable(const string& name) : name(name) {
+    }
+    Vec3 evaluate(const Vec3* args, const map<string, Vec3>& variables) const {
+        map<string, Vec3>::const_iterator iter = variables.find(name);
+        if (iter == variables.end())
+            throw Exception("No value specified for variable "+name);
+        return iter->second;
+    }
+private:
+    string name;
+};
+
+class VectorExpression::Constant : public VectorExpression::VectorOperation {
+public:
+    Constant(double value) : value(Vec3(value, value, value)) {
+    }
+    Vec3 evaluate(const Vec3* args, const map<string, Vec3>& variables) const {
+        return value;
+    }
+private:
+    Vec3 value;
+};
+
+class VectorExpression::Dot : public VectorExpression::VectorOperation {
+public:
+    Dot() {
+    }
+    Vec3 evaluate(const Vec3* args, const map<string, Vec3>& variables) const {
+        double result = args[0].dot(args[1]);
+        return Vec3(result, result, result);
+    }
+};
+
+class VectorExpression::Cross : public VectorExpression::VectorOperation {
+public:
+    Cross() {
+    }
+    Vec3 evaluate(const Vec3* args, const map<string, Vec3>& variables) const {
+        return args[0].cross(args[1]);
+    }
+};
+
+class VectorExpression::Component : public VectorExpression::VectorOperation {
+public:
+    Component(int index) : index(index) {
+    }
+    Vec3 evaluate(const Vec3* args, const map<string, Vec3>& variables) const {
+        double value = args[0][index];
+        return Vec3(value, value, value);
+    }
+private:
+    int index;
+};
+
+class VectorExpression::Vector : public VectorExpression::VectorOperation {
+public:
+    Vector() {
+    }
+    Vec3 evaluate(const Vec3* args, const map<string, Vec3>& variables) const {
+        return Vec3(args[0][0], args[1][1], args[2][2]);
+    }
+};
+
+VectorExpression::VectorExpression(const string& expression, const map<string, CustomFunction*>& customFunctions) {
+    PlaceholderFunction fn1(1), fn2(2), fn3(3);
+    PlaceholderFunction crossFunction(2);
+    PlaceholderFunction xFunction(1);
+    PlaceholderFunction yFunction(1);
+    PlaceholderFunction zFunction(1);
+    PlaceholderFunction vectorFunction(3);
+    map<string, CustomFunction*> functions = customFunctions;
+    functions["dot"] = &fn2;
+    functions["cross"] = &fn2;
+    functions["_x"] = &fn1;
+    functions["_y"] = &fn1;
+    functions["_z"] = &fn1;
+    functions["vector"] = &fn3;
+    analyzeExpression(Parser::parse(expression, functions));
+}
+
+VectorExpression::VectorExpression(const ParsedExpression& expression) {
+    analyzeExpression(expression);
+}
+
+VectorExpression::VectorExpression(const VectorExpression& expression) {
+    *this = expression;
+}
+
+void VectorExpression::analyzeExpression(const ParsedExpression& expression) {
+    parsed = expression.optimize();
+    program = parsed.createProgram();
+    stack.resize(program.getStackSize()+1);
+    for (int i = 0; i < program.getNumOperations(); i++) {
+        const Operation& op = program.getOperation(i);
+        if (op.getId() == Operation::VARIABLE)
+            operations.push_back(new Variable(op.getName()));
+        else if (op.getId() == Operation::CONSTANT)
+            operations.push_back(new Constant(dynamic_cast<const Operation::Constant&>(op).getValue()));
+        else if (op.getName() == "dot")
+            operations.push_back(new Dot());
+        else if (op.getName() == "cross")
+            operations.push_back(new Cross());
+        else if (op.getName() == "_x")
+            operations.push_back(new Component(0));
+        else if (op.getName() == "_y")
+            operations.push_back(new Component(1));
+        else if (op.getName() == "_z")
+            operations.push_back(new Component(2));
+        else if (op.getName() == "vector")
+            operations.push_back(new Vector());
+        else if (op.getNumArguments() == 1)
+            operations.push_back(new OpWrapper1(op));
+        else if (op.getNumArguments() == 2)
+            operations.push_back(new OpWrapper2(op));
+        else if (op.getNumArguments() == 3)
+            operations.push_back(new OpWrapper3(op));
+        else
+            throw OpenMMException("Unsupported operator in vector expression: "+op.getName());
+    }
+}
+
+VectorExpression::~VectorExpression() {
+    for (int i = 0; i < operations.size(); i++)
+        delete operations[i];
+}
+
+Vec3 VectorExpression::evaluate(const map<string, Vec3>& variables) const {
+    int stackSize = program.getStackSize();
+    int stackPointer = stackSize;
+    for (int i = 0; i < (int) operations.size(); i++) {
+        int numArgs = program.getOperation(i).getNumArguments();
+        Vec3 result = operations[i]->evaluate(&stack[stackPointer], variables);
+        stackPointer += numArgs-1;
+        stack[stackPointer] = result;
+    }
+    return stack[stackSize-1];
+}
+
+VectorExpression& VectorExpression::operator=(const VectorExpression& exp) {
+    analyzeExpression(exp.parsed);
+    return *this;
+}

platforms/cuda/include/CudaExpressionUtilities.h

@@ -122,6 +122,7 @@ private:
             std::vector<const Lepton::ExpressionTreeNode*>& nodes);
     void findRelatedPowers(const Lepton::ExpressionTreeNode& node, const Lepton::ExpressionTreeNode& searchNode,
             std::map<int, const Lepton::ExpressionTreeNode*>& powers);
+    void callFunction(std::stringstream& out, std::string singleFn, std::string doubleFn, const std::string& arg, const std::string& tempType);
     std::vector<std::vector<double> > computeFunctionParameters(const std::vector<const TabulatedFunction*>& functions);
     CudaContext& context;
     FunctionPlaceholder fp1, fp2, fp3, periodicDistance;

platforms/cuda/include/CudaKernels.h

@@ -1495,9 +1495,8 @@ 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), perDofValues(NULL), needsEnergyParamDerivs(false) {
+            hasInitializedKernels(false), needsEnergyParamDerivs(false) {
     }
-    ~CudaIntegrateCustomStepKernel();
     /**
      * Initialize the kernel.
      * 
@@ -1561,7 +1560,7 @@ private:
     class ReorderListener;
     class GlobalTarget;
     class DerivFunction;
-    std::string createPerDofComputation(const std::string& variable, const Lepton::ParsedExpression& expr, int component, CustomIntegrator& integrator,
+    std::string createPerDofComputation(const std::string& variable, const Lepton::ParsedExpression& expr, CustomIntegrator& integrator,
         const std::string& forceName, const std::string& energyName, std::vector<const TabulatedFunction*>& functions,
         std::vector<std::pair<std::string, std::string> >& functionNames);
     void prepareForComputation(ContextImpl& context, CustomIntegrator& integrator, bool& forcesAreValid);
@@ -1574,21 +1573,21 @@ private:
     double energy;
     float energyFloat;
     int numGlobalVariables, sumWorkGroupSize;
-    bool hasInitializedKernels, deviceValuesAreCurrent, deviceGlobalsAreCurrent, modifiesParameters, keNeedsForce, hasAnyConstraints, needsEnergyParamDerivs;
-    mutable bool localValuesAreCurrent;
+    bool hasInitializedKernels, deviceGlobalsAreCurrent, modifiesParameters, keNeedsForce, hasAnyConstraints, needsEnergyParamDerivs;
+    std::vector<bool> deviceValuesAreCurrent;
+    mutable std::vector<bool> localValuesAreCurrent; 
     CudaArray globalValues;
     CudaArray sumBuffer;
     CudaArray summedValue;
     CudaArray uniformRandoms;
     CudaArray randomSeed;
     CudaArray perDofEnergyParamDerivs;
-    std::vector<CudaArray> tabulatedFunctions;
+    std::vector<CudaArray> tabulatedFunctions, perDofValues;
     std::map<int, double> savedEnergy;
     std::map<int, CudaArray> savedForces;
     std::set<int> validSavedForces;
-    CudaParameterSet* perDofValues;
-    mutable std::vector<std::vector<float> > localPerDofValuesFloat;
-    mutable std::vector<std::vector<double> > localPerDofValuesDouble;
+    mutable std::vector<std::vector<float4> > localPerDofValuesFloat;
+    mutable std::vector<std::vector<double4> > localPerDofValuesDouble;
     std::map<std::string, double> energyParamDerivs;
     std::vector<std::string> perDofEnergyParamDerivNames;
     std::vector<float> localPerDofEnergyParamDerivsFloat;

platforms/cuda/src/CudaContext.cpp

@@ -203,12 +203,13 @@ CudaContext::CudaContext(const System& system, int deviceIndex, bool useBlocking
                 break;
             }
         }
-        if (this->deviceIndex == -1)
+        if (this->deviceIndex == -1) {
             if (deviceIndex != -1)
                 throw OpenMMException("The requested CUDA device could not be loaded");
             else
                 throw OpenMMException("No compatible CUDA device is available");
         }
+    }
     else {
         isLinkedContext = true;
         context = originalContext->context;

platforms/cuda/src/kernels/customIntegratorPerDof.cu

@@ -23,10 +23,14 @@ inline __device__ void storePos(real4* __restrict__ posq, real4* __restrict__ po
 #endif
 }
 
-inline __device__ double4 convertToDouble4(mixed4 a) {
+inline __device__ double4 convertToDouble4(float4 a) {
     return make_double4(a.x, a.y, a.z, a.w);
 }
 
+inline __device__ double4 convertToDouble4(double4 a) {
+    return a;
+}
+
 inline __device__ mixed4 convertFromDouble4(double4 a) {
     return make_mixed4(a.x, a.y, a.z, a.w);
 }
@@ -36,7 +40,7 @@ extern "C" __global__ void computePerDof(real4* __restrict__ posq, real4* __rest
         mixed* __restrict__ sum, const float4* __restrict__ gaussianValues, unsigned int gaussianBaseIndex, const float4* __restrict__ uniformValues,
         const mixed energy, mixed* __restrict__ energyParamDerivs
         PARAMETER_ARGUMENTS) {
-    mixed stepSize = dt[0].y;
+    double3 stepSize = make_double3(dt[0].y);
     int index = blockIdx.x*blockDim.x+threadIdx.x;
     const double forceScale = 1.0/0xFFFFFFFF;
     while (index < NUM_ATOMS) {
@@ -47,7 +51,7 @@ extern "C" __global__ void computePerDof(real4* __restrict__ posq, real4* __rest
 #endif
         double4 velocity = convertToDouble4(velm[index]);
         double4 f = make_double4(forceScale*force[index], forceScale*force[index+PADDED_NUM_ATOMS], forceScale*force[index+PADDED_NUM_ATOMS*2], 0.0);
-        double mass = 1.0/velocity.w;
+        double3 mass = make_double3(1.0/velocity.w);
         if (velocity.w != 0.0) {
             int gaussianIndex = gaussianBaseIndex;
             int uniformIndex = 0;

platforms/opencl/include/OpenCLKernels.h

@@ -1484,9 +1484,8 @@ 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), perDofValues(NULL), needsEnergyParamDerivs(false) {
+            hasInitializedKernels(false), needsEnergyParamDerivs(false) {
     }
-    ~OpenCLIntegrateCustomStepKernel();
     /**
      * Initialize the kernel.
      * 
@@ -1550,7 +1549,7 @@ private:
     class ReorderListener;
     class GlobalTarget;
     class DerivFunction;
-    std::string createPerDofComputation(const std::string& variable, const Lepton::ParsedExpression& expr, int component, CustomIntegrator& integrator,
+    std::string createPerDofComputation(const std::string& variable, const Lepton::ParsedExpression& expr, CustomIntegrator& integrator,
         const std::string& forceName, const std::string& energyName, std::vector<const TabulatedFunction*>& functions,
         std::vector<std::pair<std::string, std::string> >& functionNames);
     void prepareForComputation(ContextImpl& context, CustomIntegrator& integrator, bool& forcesAreValid);
@@ -1563,21 +1562,21 @@ private:
     double energy;
     float energyFloat;
     int numGlobalVariables, sumWorkGroupSize;
-    bool hasInitializedKernels, deviceValuesAreCurrent, deviceGlobalsAreCurrent, modifiesParameters, keNeedsForce, hasAnyConstraints, needsEnergyParamDerivs;
-    mutable bool localValuesAreCurrent;
+    bool hasInitializedKernels, deviceGlobalsAreCurrent, modifiesParameters, keNeedsForce, hasAnyConstraints, needsEnergyParamDerivs;
+    std::vector<bool> deviceValuesAreCurrent;
+    mutable std::vector<bool> localValuesAreCurrent;
     OpenCLArray globalValues;
     OpenCLArray sumBuffer;
     OpenCLArray summedValue;
     OpenCLArray uniformRandoms;
     OpenCLArray randomSeed;
     OpenCLArray perDofEnergyParamDerivs;
-    std::vector<OpenCLArray> tabulatedFunctions;
+    std::vector<OpenCLArray> tabulatedFunctions, perDofValues;
     std::map<int, double> savedEnergy;
     std::map<int, OpenCLArray> savedForces;
     std::set<int> validSavedForces;
-    OpenCLParameterSet* perDofValues;
-    mutable std::vector<std::vector<cl_float> > localPerDofValuesFloat;
-    mutable std::vector<std::vector<cl_double> > localPerDofValuesDouble;
+    mutable std::vector<std::vector<mm_float4> > localPerDofValuesFloat;
+    mutable std::vector<std::vector<mm_double4> > localPerDofValuesDouble;
     std::map<std::string, double> energyParamDerivs;
     std::vector<std::string> perDofEnergyParamDerivNames;
     std::vector<cl_float> localPerDofEnergyParamDerivsFloat;

platforms/opencl/src/kernels/customIntegratorPerDof.cl

@@ -37,7 +37,7 @@ __kernel void computePerDof(__global real4* restrict posq, __global real4* restr
         mixed4 position = loadPos(posq, posqCorrection, index);
 #endif
         mixed4 velocity = velm[index];
-        real4 f = force[index];
+        mixed4 f = convert_mixed4(force[index]);
         mixed mass = 1/velocity.w;
         if (velocity.w != 0.0) {
             int gaussianIndex = gaussianBaseIndex;

platforms/reference/include/ReferenceCustomDynamics.h

@@ -30,6 +30,7 @@
 #include "openmm/internal/ContextImpl.h"
 #include "openmm/internal/CustomIntegratorUtilities.h"
 #include "openmm/internal/CompiledExpressionSet.h"
+#include "openmm/internal/VectorExpression.h"
 #include "lepton/CompiledExpression.h"
 
 #include <map>
@@ -48,6 +49,7 @@ private:
     std::vector<OpenMM::CustomIntegrator::ComputationType> stepType;
     std::vector<std::string> stepVariable;
     std::vector<std::vector<Lepton::CompiledExpression> > stepExpressions;
+    std::vector<std::vector<VectorExpression> > stepVectorExpressions;
     std::vector<CustomIntegratorUtilities::Comparison> comparisons;
     std::vector<bool> invalidatesForces, needsForces, needsEnergy, computeBothForceAndEnergy;
     std::vector<int> forceGroupFlags, blockEnd;
@@ -68,6 +70,10 @@ private:
                   const std::vector<OpenMM::Vec3>& velocities, const std::vector<OpenMM::Vec3>& forces, const std::vector<double>& masses,
                   const std::vector<std::vector<OpenMM::Vec3> >& perDof, const Lepton::CompiledExpression& expression);
     
+    void computePerParticle(int numberOfAtoms, std::vector<OpenMM::Vec3>& results, const std::vector<OpenMM::Vec3>& atomCoordinates,
+                  const std::vector<OpenMM::Vec3>& velocities, const std::vector<OpenMM::Vec3>& forces, const std::vector<double>& masses,
+                  const std::vector<std::vector<OpenMM::Vec3> >& perDof, const std::map<std::string, double>& globals, const VectorExpression& expression);
+    
     void recordChangedParameters(OpenMM::ContextImpl& context, std::map<std::string, double>& globals);
 
     bool evaluateCondition(int step);

platforms/reference/src/SimTKReference/ReferenceCustomDynamics.cpp

@@ -60,6 +60,20 @@ private:
     string param;
 };
 
+/**
+ * Determine whether a parsed expression involves any vector functions.
+ */
+static bool isVectorExpression(const ExpressionTreeNode& node) {
+    const Lepton::Operation& op = node.getOperation();
+    if (op.getId() == Lepton::Operation::CUSTOM)
+        if (op.getName() == "dot" || op.getName() == "cross" || op.getName() == "vector" || op.getName() == "_x" || op.getName() == "_y" || op.getName() == "_z")
+            return true;
+    for (auto& child : node.getChildren())
+        if (isVectorExpression(child))
+            return true;
+    return false;
+}
+
 /**---------------------------------------------------------------------------------------
 
    ReferenceCustomDynamics constructor
@@ -127,13 +141,19 @@ void ReferenceCustomDynamics::initialize(ContextImpl& context, vector<double>& m
     vector<vector<ParsedExpression> > expressions;
     CustomIntegratorUtilities::analyzeComputations(context, integrator, expressions, comparisons, blockEnd, invalidatesForces, needsForces, needsEnergy, computeBothForceAndEnergy, forceGroup, functions);
     stepExpressions.resize(expressions.size());
+    stepVectorExpressions.resize(expressions.size());
     for (int i = 0; i < numSteps; i++) {
         stepExpressions[i].resize(expressions[i].size());
         for (int j = 0; j < (int) expressions[i].size(); j++) {
-            stepExpressions[i][j] = ParsedExpression(replaceDerivFunctions(expressions[i][j].getRootNode(), context)).createCompiledExpression();
+            ParsedExpression parsed(replaceDerivFunctions(expressions[i][j].getRootNode(), context));
+            if (isVectorExpression(parsed.getRootNode()))
+                stepVectorExpressions[i].push_back(VectorExpression(parsed));
+            else {
+                stepExpressions[i][j] = parsed.createCompiledExpression();
                 stepExpressions[i][j].setVariableLocations(variableLocations);
                 expressionSet.registerExpression(stepExpressions[i][j]);
             }
+        }
         if (stepType[i] == CustomIntegrator::WhileBlockStart)
             blockEnd[blockEnd[i]] = i; // Record where to branch back to.
     }
@@ -272,10 +292,16 @@ void ReferenceCustomDynamics::update(ContextImpl& context, int numberOfAtoms, ve
                 }
                 if (results == NULL)
                     throw OpenMMException("Illegal per-DOF output variable: "+stepVariable[step]);
+                if (stepVectorExpressions[step].size() > 0)
+                    computePerParticle(numberOfAtoms, *results, atomCoordinates, velocities, stepForces, masses, perDof, globals, stepVectorExpressions[step][0]);
+                else
                     computePerDof(numberOfAtoms, *results, atomCoordinates, velocities, stepForces, masses, perDof, stepExpressions[step][0]);
                 break;
             }
             case CustomIntegrator::ComputeSum: {
+                if (stepVectorExpressions[step].size() > 0)
+                    computePerParticle(numberOfAtoms, sumBuffer, atomCoordinates, velocities, stepForces, masses, perDof, globals, stepVectorExpressions[step][0]);
+                else
                     computePerDof(numberOfAtoms, sumBuffer, atomCoordinates, velocities, stepForces, masses, perDof, stepExpressions[step][0]);
                 double sum = 0.0;
                 for (int j = 0; j < numberOfAtoms; j++)
@@ -354,6 +380,31 @@ void ReferenceCustomDynamics::computePerDof(int numberOfAtoms, vector<Vec3>& res
     }
 }
 
+void ReferenceCustomDynamics::computePerParticle(int numberOfAtoms, vector<Vec3>& results, const vector<Vec3>& atomCoordinates,
+              const vector<Vec3>& velocities, const vector<Vec3>& forces, const vector<double>& masses,
+              const vector<vector<Vec3> >& perDof, const map<string, double>& globals, const VectorExpression& expression) {
+    // Loop over all degrees of freedom.
+
+    map<string, Vec3> variables;
+    for (auto& entry : globals)
+        variables[entry.first] = Vec3(entry.second, entry.second, entry.second);
+    for (int i = 0; i < numberOfAtoms; i++) {
+        if (masses[i] != 0.0) {
+            variables["m"] = Vec3(masses[i], masses[i], masses[i]);
+            variables["x"] = atomCoordinates[i];
+            variables["v"] = velocities[i];
+            variables["f"] = forces[i];
+            variables["uniform"] = Vec3(SimTKOpenMMUtilities::getUniformlyDistributedRandomNumber(),
+                    SimTKOpenMMUtilities::getUniformlyDistributedRandomNumber(), SimTKOpenMMUtilities::getUniformlyDistributedRandomNumber());
+            variables["gaussian"] = Vec3(SimTKOpenMMUtilities::getNormallyDistributedRandomNumber(),
+                    SimTKOpenMMUtilities::getNormallyDistributedRandomNumber(), SimTKOpenMMUtilities::getNormallyDistributedRandomNumber());
+            for (int j = 0; j < perDof.size(); j++)
+                variables[integrator.getPerDofVariableName(j)] = perDof[j][i];
+            results[i] = expression.evaluate(variables);
+        }
+    }
+}
+
 bool ReferenceCustomDynamics::evaluateCondition(int step) {
     uniform = SimTKOpenMMUtilities::getUniformlyDistributedRandomNumber();
     gaussian = SimTKOpenMMUtilities::getNormallyDistributedRandomNumber();

tests/TestCustomIntegrator.h

@@ -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-2017 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2018 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -102,7 +102,6 @@ void testSingleBond() {
  */
 void testConstraints() {
     const int numParticles = 8;
-    const double temp = 500.0;
     System system;
     CustomIntegrator integrator(0.002);
     integrator.addPerDofVariable("oldx", 0);
@@ -538,9 +537,11 @@ void testPerDofVariables() {
     CustomIntegrator integrator(0.01);
     integrator.addPerDofVariable("temp", 0);
     integrator.addPerDofVariable("pos", 0);
+    integrator.addPerDofVariable("computed", 0);
     integrator.addComputePerDof("v", "v+dt*f/m");
     integrator.addComputePerDof("x", "x+dt*v");
     integrator.addComputePerDof("pos", "x");
+    integrator.addComputePerDof("computed", "step(v)*log(x^2)");
     Context context(system, integrator, platform);
     context.setPositions(positions);
     vector<Vec3> initialValues(numParticles);
@@ -553,13 +554,24 @@ void testPerDofVariables() {
     vector<Vec3> values;
     for (int i = 0; i < 100; ++i) {
         integrator.step(1);
-        State state = context.getState(State::Positions);
+        State state = context.getState(State::Positions | State::Velocities);
         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);
+        integrator.getPerDofVariable(2, values);
+        for (int j = 0; j < numParticles; j++)
+            for (int k = 0; k < 3; k++) {
+                if (state.getVelocities()[j][k] < 0) {
+                    ASSERT(values[j][k] == 0.0);
+                }
+                else {
+                    double v = state.getPositions()[j][k];
+                    ASSERT_EQUAL_TOL(log(v*v), values[j][k], 1e-5);
+                }
+            }
     }
 }
 
@@ -1019,6 +1031,47 @@ void testUpdateContextState() {
     }
 }
 
+/**
+ * Test using expressions that involve vector functions.
+ */
+void testVectorFunctions() {
+    const int numParticles = 8;
+    System system;
+    CustomIntegrator integrator(0.001);
+    integrator.addGlobalVariable("sumy", 0.0);
+    integrator.addPerDofVariable("angular", 0.0);
+    integrator.addPerDofVariable("shuffle", 0.0);
+    integrator.addComputeSum("sumy", "x*vector(0, 1, 0)");
+    integrator.addComputePerDof("angular", "cross(v, x)");
+    integrator.addComputePerDof("shuffle", "dot(vector(_z(x), _x(x), _y(x)), v)");
+    OpenMM_SFMT::SFMT sfmt;
+    init_gen_rand(0, sfmt);
+    vector<Vec3> positions(numParticles);
+    vector<Vec3> velocities(numParticles);
+    for (int i = 0; i < numParticles; ++i) {
+        system.addParticle(1.0);
+        positions[i] = Vec3(genrand_real2(sfmt)-0.5, genrand_real2(sfmt)-0.5, genrand_real2(sfmt)-0.5);
+        velocities[i] = Vec3(genrand_real2(sfmt)-0.5, genrand_real2(sfmt)-0.5, genrand_real2(sfmt)-0.5);
+    }
+    Context context(system, integrator, platform);
+    context.setPositions(positions);
+    context.setVelocities(velocities);
+    integrator.step(1);
+    
+    // See if the expressions were computed correctly.
+    
+    double sumy = 0;
+    vector<Vec3> angular, shuffle;
+    integrator.getPerDofVariable(0, angular);
+    integrator.getPerDofVariable(1, shuffle);
+    for (int i = 0; i < numParticles; i++) {
+        ASSERT_EQUAL_VEC(velocities[i].cross(positions[i]), angular[i], 1e-5);
+        ASSERT_EQUAL_VEC(Vec3(1, 1, 1)*velocities[i].dot(Vec3(positions[i][2], positions[i][0], positions[i][1])), shuffle[i], 1e-5);
+        sumy += positions[i][1];
+    }
+    ASSERT_EQUAL_TOL(sumy, integrator.getGlobalVariable(0), 1e-5);
+}
+
 void runPlatformTests();
 
 int main(int argc, char* argv[]) {
@@ -1044,6 +1097,7 @@ int main(int argc, char* argv[]) {
         testTabulatedFunction();
         testAlternatingGroups();
         testUpdateContextState();
+        testVectorFunctions();
         runPlatformTests();
     }
     catch(const exception& e) {

tests/TestVectorExpression.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) 2018 Stanford University and the Authors.           *
+ * Authors: Robert McGibbon                                                   *
+ * 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/AssertionUtilities.h"
+#include "openmm/internal/VectorExpression.h"
+#include <iostream>
+
+using namespace OpenMM;
+using namespace std;
+
+void verifyEvaluation(const string& expression, Vec3 expectedValue, Vec3 x=Vec3(), Vec3 y=Vec3()) {
+    map<string, Vec3> variables;
+    variables["x"] = x;
+    variables["y"] = y;
+    map<string, Lepton::CustomFunction*> customFunctions;
+    VectorExpression expr(expression, customFunctions);
+    Vec3 value = expr.evaluate(variables);
+    ASSERT_EQUAL_VEC(expectedValue, value, 1e-10);
+}
+
+void testExpressions() {
+    verifyEvaluation("5", Vec3(5, 5, 5));
+    verifyEvaluation("10*x", Vec3(50, 100, 150), Vec3(5, 10, 15));
+    verifyEvaluation("2*dot(x, 3)", Vec3(180, 180, 180), Vec3(5, 10, 15));
+    Vec3 a(1, 1.5, 2);
+    Vec3 b(3, -1, 4);
+    verifyEvaluation("cross(y, x)", b.cross(a), a, b);
+    verifyEvaluation("_x(x)", Vec3(a[0], a[0], a[0]), a);
+    verifyEvaluation("_y(x)", Vec3(a[1], a[1], a[1]), a);
+    verifyEvaluation("_z(x)", Vec3(a[2], a[2], a[2]), a);
+    verifyEvaluation("vector(x, 5, y)", Vec3(a[0], 5, b[2]), a, b);
+}
+
+int main(int argc, char* argv[]) {
+    try {
+        testExpressions();
+    }
+    catch(const exception& e) {
+        cout << "exception: " << e.what() << endl;
+        return 1;
+    }
+    cout << "Done" << endl;
+    return 0;
+}