Merge pull request #2447 from peastman/energy

Avoid unnecessary force evaluation when computing kinetic energy

openmmapi/include/openmm/BAOABLangevinIntegrator.h

@@ -140,6 +140,10 @@ protected:
      * Compute the kinetic energy of the system at the current time.
      */
     double computeKineticEnergy();
+    /**
+     * Computing kinetic energy for this integrator does not require forces.
+     */
+    bool kineticEnergyRequiresForce() const;
 private:
     double temperature, friction;
     int randomNumberSeed;

openmmapi/include/openmm/BrownianIntegrator.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) 2008-2012 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2019 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -133,6 +133,10 @@ protected:
      * Compute the kinetic energy of the system at the current time.
      */
     double computeKineticEnergy();
+    /**
+     * Computing kinetic energy for this integrator does not require forces.
+     */
+    bool kineticEnergyRequiresForce() const;
 private:
     double temperature, friction;
     int randomNumberSeed;

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-2018 Stanford University and the Authors.      *
+ * Portions copyright (c) 2011-2019 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -627,6 +627,10 @@ protected:
      * Compute the kinetic energy of the system at the current time.
      */
     double computeKineticEnergy();
+    /**
+     * Get whether computeKineticEnergy() expects forces to have been computed.
+     */
+    bool kineticEnergyRequiresForce() const;
 private:
     class ComputationInfo;
     class FunctionInfo;
@@ -639,7 +643,7 @@ private:
     std::string kineticEnergy;
     mutable bool globalsAreCurrent;
     int randomNumberSeed;
-    bool forcesAreValid;
+    bool forcesAreValid, keNeedsForce;
     Kernel kernel;
 };
 

openmmapi/include/openmm/Integrator.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) 2008-2015 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2019 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -120,8 +120,19 @@ protected:
      * Compute the kinetic energy of the system at the current time.  This may be different from simply
      * mv<sup>2</sup>/2.  For example, a leapfrog integrator will store velocities offset by half a step,
      * but the kinetic energy should be computed at the current time, not delayed by half a step.
+     * 
+     * If kineticEnergyRequiresForce() returns true, this method can assume that valid forces
+     * have already been computed.
      */
     virtual double computeKineticEnergy() = 0;
+    /**
+     * Get whether computeKineticEnergy() expects forces to have been computed.  The default
+     * implementation returns true to be safe.  Non-leapfrog integrators can override this to
+     * return false, which makes calling getState() to query the energy less expensive.
+     */
+    virtual bool kineticEnergyRequiresForce() const {
+        return true;
+    }
 private:
     double stepSize, constraintTol;
 };

openmmapi/src/BAOABLangevinIntegrator.cpp

@@ -76,6 +76,10 @@ double BAOABLangevinIntegrator::computeKineticEnergy() {
     return kernel.getAs<IntegrateBAOABStepKernel>().computeKineticEnergy(*context, *this);
 }
 
+bool BAOABLangevinIntegrator::kineticEnergyRequiresForce() const {
+    return false;
+}
+
 void BAOABLangevinIntegrator::step(int steps) {
     if (context == NULL)
         throw OpenMMException("This Integrator is not bound to a context!");  

openmmapi/src/BrownianIntegrator.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-2012 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2019 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -71,6 +71,10 @@ double BrownianIntegrator::computeKineticEnergy() {
     return kernel.getAs<IntegrateBrownianStepKernel>().computeKineticEnergy(*context, *this);
 }
 
+bool BrownianIntegrator::kineticEnergyRequiresForce() const {
+    return false;
+}
+
 void BrownianIntegrator::step(int steps) {
     if (context == NULL)
         throw OpenMMException("This Integrator is not bound to a context!");  

openmmapi/src/Context.cpp

@@ -96,8 +96,9 @@ State Context::getState(int types, bool enforcePeriodicBox, int groups) const {
     bool includeForces = types&State::Forces;
     bool includeEnergy = types&State::Energy;
     bool includeParameterDerivs = types&State::ParameterDerivatives;
+    bool needForcesForEnergy = (includeEnergy && getIntegrator().kineticEnergyRequiresForce());
     if (includeForces || includeEnergy || includeParameterDerivs) {
-        double energy = impl->calcForcesAndEnergy(includeForces || includeEnergy || includeParameterDerivs, includeEnergy, groups);
+        double energy = impl->calcForcesAndEnergy(includeForces || needForcesForEnergy || includeParameterDerivs, includeEnergy, groups);
         if (includeEnergy)
             builder.setEnergy(impl->calcKineticEnergy(), energy);
         if (includeForces) {

openmmapi/src/CustomIntegrator.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) 2011-2017 Stanford University and the Authors.      *
+ * Portions copyright (c) 2011-2019 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -35,6 +35,9 @@
 #include "openmm/internal/AssertionUtilities.h"
 #include "openmm/internal/ContextImpl.h"
 #include "openmm/kernels.h"
+#include "lepton/CompiledExpression.h"
+#include "lepton/ParsedExpression.h"
+#include "lepton/Parser.h"
 #include <set>
 #include <string>
 
@@ -45,7 +48,7 @@ CustomIntegrator::CustomIntegrator(double stepSize) : globalsAreCurrent(true), f
     setStepSize(stepSize);
     setConstraintTolerance(1e-5);
     setRandomNumberSeed(0);
-    kineticEnergy = "m*v*v/2";
+    setKineticEnergyExpression("m*v*v/2");
 }
 
 CustomIntegrator::~CustomIntegrator() {
@@ -103,9 +106,14 @@ vector<string> CustomIntegrator::getKernelNames() {
 }
 
 double CustomIntegrator::computeKineticEnergy() {
+    forcesAreValid = keNeedsForce;
     return kernel.getAs<IntegrateCustomStepKernel>().computeKineticEnergy(*context, *this, forcesAreValid);
 }
 
+bool CustomIntegrator::kineticEnergyRequiresForce() const {
+    return keNeedsForce;
+}
+
 void CustomIntegrator::step(int steps) {
     if (context == NULL)
         throw OpenMMException("This Integrator is not bound to a context!");  
@@ -312,4 +320,6 @@ const string& CustomIntegrator::getKineticEnergyExpression() const {
 
 void CustomIntegrator::setKineticEnergyExpression(const string& expression) {
     kineticEnergy = expression;
+    Lepton::CompiledExpression expr = Lepton::Parser::parse(kineticEnergy).createCompiledExpression();
+    keNeedsForce = (expr.getVariables().find("f") != expr.getVariables().end());
 }

platforms/cuda/include/CudaKernels.h

@@ -1633,7 +1633,7 @@ private:
     double energy;
     float energyFloat;
     int numGlobalVariables, sumWorkGroupSize;
-    bool hasInitializedKernels, deviceGlobalsAreCurrent, modifiesParameters, keNeedsForce, hasAnyConstraints, needsEnergyParamDerivs;
+    bool hasInitializedKernels, deviceGlobalsAreCurrent, modifiesParameters, hasAnyConstraints, needsEnergyParamDerivs;
     std::vector<bool> deviceValuesAreCurrent;
     mutable std::vector<bool> localValuesAreCurrent; 
     CudaArray globalValues;

platforms/cuda/src/CudaKernels.cpp

@@ -7957,7 +7957,6 @@ void CudaIntegrateCustomStepKernel::prepareForComputation(ContextImpl& context,
             kineticEnergyArgs.push_back(&array.getDevicePointer());
         for (auto& array : tabulatedFunctions)
             kineticEnergyArgs.push_back(&array.getDevicePointer());
-        keNeedsForce = usesVariable(keExpression, "f");
 
         // Create a second kernel to sum the values.
 
@@ -8213,17 +8212,6 @@ bool CudaIntegrateCustomStepKernel::evaluateCondition(int step) {
 
 double CudaIntegrateCustomStepKernel::computeKineticEnergy(ContextImpl& context, CustomIntegrator& integrator, bool& forcesAreValid) {
     prepareForComputation(context, integrator, forcesAreValid);
-    if (keNeedsForce && !forcesAreValid) {
-        // Compute the force.  We want to then mark that forces are valid, which means also computing
-        // potential energy if any steps will expect it to be valid too.
-        
-        bool willNeedEnergy = false;
-        for (int i = 0; i < integrator.getNumComputations(); i++)
-            willNeedEnergy |= needsEnergy[i];
-        energy = context.calcForcesAndEnergy(true, willNeedEnergy, -1);
-        energyFloat = (float) energy;
-        forcesAreValid = true;
-    }
     CUdeviceptr posCorrection = (cu.getUseMixedPrecision() ? cu.getPosqCorrection().getDevicePointer() : 0);
     int randomIndex = 0;
     kineticEnergyArgs[1] = &posCorrection;

platforms/opencl/include/OpenCLKernels.h

@@ -1625,7 +1625,7 @@ private:
     double energy;
     float energyFloat;
     int numGlobalVariables, sumWorkGroupSize;
-    bool hasInitializedKernels, deviceGlobalsAreCurrent, modifiesParameters, keNeedsForce, hasAnyConstraints, needsEnergyParamDerivs;
+    bool hasInitializedKernels, deviceGlobalsAreCurrent, modifiesParameters, hasAnyConstraints, needsEnergyParamDerivs;
     std::vector<bool> deviceValuesAreCurrent;
     mutable std::vector<bool> localValuesAreCurrent;
     OpenCLArray globalValues;

platforms/opencl/src/OpenCLKernels.cpp

@@ -8364,7 +8364,6 @@ void OpenCLIntegrateCustomStepKernel::prepareForComputation(ContextImpl& context
             kineticEnergyKernel.setArg<cl::Buffer>(index++, array.getDeviceBuffer());
         for (auto& array : tabulatedFunctions)
             kineticEnergyKernel.setArg<cl::Buffer>(index++, array.getDeviceBuffer());
-        keNeedsForce = usesVariable(keExpression, "f");
 
         // Create a second kernel to sum the values.
 
@@ -8627,16 +8626,6 @@ bool OpenCLIntegrateCustomStepKernel::evaluateCondition(int step) {
 
 double OpenCLIntegrateCustomStepKernel::computeKineticEnergy(ContextImpl& context, CustomIntegrator& integrator, bool& forcesAreValid) {
     prepareForComputation(context, integrator, forcesAreValid);
-    if (keNeedsForce && !forcesAreValid) {
-        // Compute the force.  We want to then mark that forces are valid, which means also computing
-        // potential energy if any steps will expect it to be valid too.
-        
-        bool willNeedEnergy = false;
-        for (int i = 0; i < integrator.getNumComputations(); i++)
-            willNeedEnergy |= needsEnergy[i];
-        energy = context.calcForcesAndEnergy(true, willNeedEnergy, -1);
-        forcesAreValid = true;
-    }
     cl.clearBuffer(sumBuffer);
     kineticEnergyKernel.setArg<cl::Buffer>(8, cl.getIntegrationUtilities().getRandom().getDeviceBuffer());
     kineticEnergyKernel.setArg<cl_uint>(9, 0);

platforms/reference/src/SimTKReference/ReferenceCustomDynamics.cpp

@@ -463,10 +463,6 @@ double ReferenceCustomDynamics::computeKineticEnergy(OpenMM::ContextImpl& contex
     globals.insert(context.getParameters().begin(), context.getParameters().end());
     for (auto& global : globals)
         expressionSet.setVariable(expressionSet.getVariableIndex(global.first), global.second);
-    if (kineticEnergyNeedsForce) {
-        energy = context.calcForcesAndEnergy(true, true, -1);
-        forcesAreValid = true;
-    }
     computePerDof(numberOfAtoms, sumBuffer, atomCoordinates, velocities, forces, masses, perDof, kineticEnergyExpression);
     double sum = 0.0;
     for (int j = 0; j < numberOfAtoms; j++)