Kinetic energy is computed by the Integrator so it can adjust for leapfrog displacements

olla/include/openmm/kernels.h

@@ -843,6 +843,13 @@ public:
      * @param integrator the VerletIntegrator this kernel is being used for
      */
     virtual void execute(ContextImpl& context, const VerletIntegrator& integrator) = 0;
+    /**
+     * Compute the kinetic energy.
+     * 
+     * @param context    the context in which to execute this kernel
+     * @param integrator the VerletIntegrator this kernel is being used for
+     */
+    virtual double computeKineticEnergy(ContextImpl& context, const VerletIntegrator& integrator) = 0;
 };
 
 /**
@@ -869,6 +876,13 @@ public:
      * @param integrator the LangevinIntegrator this kernel is being used for
      */
     virtual void execute(ContextImpl& context, const LangevinIntegrator& integrator) = 0;
+    /**
+     * Compute the kinetic energy.
+     * 
+     * @param context    the context in which to execute this kernel
+     * @param integrator the LangevinIntegrator this kernel is being used for
+     */
+    virtual double computeKineticEnergy(ContextImpl& context, const LangevinIntegrator& integrator) = 0;
 };
 
 /**
@@ -895,6 +909,13 @@ public:
      * @param integrator the BrownianIntegrator this kernel is being used for
      */
     virtual void execute(ContextImpl& context, const BrownianIntegrator& integrator) = 0;
+    /**
+     * Compute the kinetic energy.
+     * 
+     * @param context    the context in which to execute this kernel
+     * @param integrator the BrownianIntegrator this kernel is being used for
+     */
+    virtual double computeKineticEnergy(ContextImpl& context, const BrownianIntegrator& integrator) = 0;
 };
 
 /**
@@ -918,11 +939,18 @@ public:
      * Execute the kernel.
      *
      * @param context    the context in which to execute this kernel
-     * @param integrator the LangevinIntegrator this kernel is being used for
+     * @param integrator the VariableLangevinIntegrator this kernel is being used for
      * @param maxTime    the maximum time beyond which the simulation should not be advanced
      * @return the size of the step that was taken
      */
     virtual double execute(ContextImpl& context, const VariableLangevinIntegrator& integrator, double maxTime) = 0;
+    /**
+     * Compute the kinetic energy.
+     * 
+     * @param context    the context in which to execute this kernel
+     * @param integrator the VariableLangevinIntegrator this kernel is being used for
+     */
+    virtual double computeKineticEnergy(ContextImpl& context, const VariableLangevinIntegrator& integrator) = 0;
 };
 
 /**
@@ -946,11 +974,18 @@ public:
      * Execute the kernel.
      *
      * @param context    the context in which to execute this kernel
-     * @param integrator the VerletIntegrator this kernel is being used for
+     * @param integrator the VariableVerletIntegrator this kernel is being used for
      * @param maxTime    the maximum time beyond which the simulation should not be advanced
      * @return the size of the step that was taken
      */
     virtual double execute(ContextImpl& context, const VariableVerletIntegrator& integrator, double maxTime) = 0;
+    /**
+     * Compute the kinetic energy.
+     * 
+     * @param context    the context in which to execute this kernel
+     * @param integrator the VariableVerletIntegrator this kernel is being used for
+     */
+    virtual double computeKineticEnergy(ContextImpl& context, const VariableVerletIntegrator& integrator) = 0;
 };
 
 /**
@@ -981,6 +1016,17 @@ public:
      *                       end of the step.
      */
     virtual void execute(ContextImpl& context, CustomIntegrator& integrator, bool& forcesAreValid) = 0;
+    /**
+     * Compute the kinetic energy.
+     * 
+     * @param context    the context in which to execute this kernel
+     * @param integrator the CustomIntegrator this kernel is being used for
+     * @param forcesAreValid if the context has been modified since the last time step, this will be
+     *                       false to show that cached forces are invalid and must be recalculated.
+     *                       On exit, this should specify whether the cached forces are valid at the
+     *                       end of the step.
+     */
+    virtual double computeKineticEnergy(ContextImpl& context, CustomIntegrator& integrator, bool& forcesAreValid) = 0;
     /**
      * Get the values of all global variables.
      *
@@ -1074,30 +1120,6 @@ public:
     virtual void restoreCoordinates(ContextImpl& context) = 0;
 };
 
-/**
- * This kernel is invoked to calculate the kinetic energy of the system.
- */
-class CalcKineticEnergyKernel : public KernelImpl {
-public:
-    static std::string Name() {
-        return "CalcKineticEnergy";
-    }
-    CalcKineticEnergyKernel(std::string name, const Platform& platform) : KernelImpl(name, platform) {
-    }
-    /**
-     * Initialize the kernel.
-     * 
-     * @param system     the System this kernel will be applied to
-     */
-    virtual void initialize(const System& system) = 0;
-    /**
-     * Execute the kernel.
-     * 
-     * @param context    the context in which to execute this kernel
-     */
-    virtual double execute(ContextImpl& context) = 0;
-};
-
 /**
  * This kernel is invoked to remove center of mass motion from the system.
  */

openmmapi/include/openmm/BrownianIntegrator.h

@@ -125,6 +125,10 @@ protected:
      * Get the names of all Kernels used by this Integrator.
      */
     std::vector<std::string> getKernelNames();
+    /**
+     * Compute the kinetic energy of the system at the current time.
+     */
+    double computeKineticEnergy();
 private:
     double temperature, friction;
     int randomNumberSeed;

openmmapi/include/openmm/CustomIntegrator.h

@@ -177,6 +177,31 @@ namespace OpenMM {
  * integrator.addComputePerDof("v", "v+0.5*dt*f1/m");
  * </pre></tt>
  * 
+ * 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
+ * freedom may not give the correct answer.  For example, in a leapfrog integrator
+ * the velocities are "delayed" by half a time step, so the above formula would
+ * give the kinetic energy half a time step ago, not at the current time.
+ * 
+ * Call setKineticEnergyExpression() to set an expression for the kinetic energy.
+ * It is computed for every degree of freedom (excluding ones whose mass is 0) and
+ * the result is summed.  The default expression is "m*v*v/2", which is correct
+ * for many integrators.
+ * 
+ * As example, the following line defines the correct way to compute kinetic energy
+ * when using a leapfrog algorithm:
+ * 
+ * <tt><pre>
+ * integrator.setKineticEnergyExpression("m*v1*v1/2; v1=v+0.5*dt*f/m");
+ * </pre></tt>
+ * 
+ * The kinetic energy expression may depend on the following pre-defined variables:
+ * x, v, f, m, dt.  It also may depend on user-defined global and per-DOF variables,
+ * and on the values of adjustable parameters defined  in the integrator's Context.
+ * It may <i>not</i> depend on any other variable, such as the potential energy,
+ * the force from a single force group, or a random number.
+ * 
  * Expressions may involve the operators + (add), - (subtract), * (multiply), / (divide), and ^ (power), and the following
  * functions: sqrt, exp, log, sin, cos, sec, csc, tan, cot, asin, acos, atan, sinh, cosh, tanh, erf, erfc, min, max, abs, step, delta.  All trigonometric functions
  * 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.  An expression
@@ -378,6 +403,18 @@ public:
      *                   evaluate an expression, this will be an empty string.
      */
     void getComputationStep(int index, ComputationType& type, std::string& variable, std::string& expression) const;
+    /**
+     * Get the expression to use for computing the kinetic energy.  The expression is evaluated
+     * for every degree of freedom.  Those values are then added together, and the sum
+     * is reported as the current kinetic energy.
+     */
+    const std::string& getKineticEnergyExpression() const;
+    /**
+     * Set the expression to use for computing the kinetic energy.  The expression is evaluated
+     * for every degree of freedom.  Those values are then added together, and the sum
+     * is reported as the current kinetic energy.
+     */
+    void setKineticEnergyExpression(const std::string& expression);
     /**
      * Get the random number seed.  See setRandomNumberSeed() for details.
      */
@@ -421,6 +458,10 @@ protected:
      * Get the names of all Kernels used by this Integrator.
      */
     std::vector<std::string> getKernelNames();
+    /**
+     * Compute the kinetic energy of the system at the current time.
+     */
+    double computeKineticEnergy();
 private:
     class ComputationInfo;
     std::vector<std::string> globalNames;
@@ -428,6 +469,7 @@ private:
     mutable std::vector<double> globalValues;
     std::vector<std::vector<Vec3> > perDofValues;
     std::vector<ComputationInfo> computations;
+    std::string kineticEnergy;
     mutable bool globalsAreCurrent;
     int randomNumberSeed;
     bool forcesAreValid;

openmmapi/include/openmm/Integrator.h

@@ -119,6 +119,12 @@ protected:
      */
     virtual void stateChanged(State::DataType changed) {
     }
+    /**
+     * 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.
+     */
+    virtual double computeKineticEnergy() = 0;
 private:
     double stepSize, constraintTol;
 };

openmmapi/include/openmm/LangevinIntegrator.h

@@ -125,6 +125,10 @@ protected:
      * Get the names of all Kernels used by this Integrator.
      */
     std::vector<std::string> getKernelNames();
+    /**
+     * Compute the kinetic energy of the system at the current time.
+     */
+    double computeKineticEnergy();
 private:
     double temperature, friction;
     int randomNumberSeed;

openmmapi/include/openmm/State.h

@@ -79,6 +79,11 @@ public:
     const std::vector<Vec3>& getForces() const;
     /**
      * Get the total kinetic energy of the system.  If this State does not contain energies, this will throw an exception.
+     * 
+     * Note that this may be different from simply mv<sup>2</sup>/2 summed over all particles.  For example, a leapfrog
+     * integrator will store velocities offset by half a step, so they must be adjusted before computing the kinetic energy.
+     * This routine returns the kinetic energy at the current time, computed in a way that is appropriate for whatever
+     * Integrator is being used.
      */
     double getKineticEnergy() const;
     /**

openmmapi/include/openmm/VariableLangevinIntegrator.h

@@ -156,6 +156,10 @@ protected:
      * Get the names of all Kernels used by this Integrator.
      */
     std::vector<std::string> getKernelNames();
+    /**
+     * Compute the kinetic energy of the system at the current time.
+     */
+    double computeKineticEnergy();
 private:
     double temperature, friction, errorTol;
     int randomNumberSeed;

openmmapi/include/openmm/VariableVerletIntegrator.h

@@ -109,6 +109,10 @@ protected:
      * Get the names of all Kernels used by this Integrator.
      */
     std::vector<std::string> getKernelNames();
+    /**
+     * Compute the kinetic energy of the system at the current time.
+     */
+    double computeKineticEnergy();
 private:
     double errorTol;
     ContextImpl* context;

openmmapi/include/openmm/VerletIntegrator.h

@@ -72,6 +72,10 @@ protected:
      * Get the names of all Kernels used by this Integrator.
      */
     std::vector<std::string> getKernelNames();
+    /**
+     * Compute the kinetic energy of the system at the current time.
+     */
+    double computeKineticEnergy();
 private:
     ContextImpl* context;
     Context* owner;

openmmapi/include/openmm/internal/ContextImpl.h

@@ -249,7 +249,7 @@ private:
     bool hasInitializedForces, hasSetPositions;
     int lastForceGroups;
     Platform* platform;
-    Kernel initializeForcesKernel, kineticEnergyKernel, updateStateDataKernel, applyConstraintsKernel, virtualSitesKernel;
+    Kernel initializeForcesKernel, updateStateDataKernel, applyConstraintsKernel, virtualSitesKernel;
     void* platformData;
 };
 

openmmapi/src/BrownianIntegrator.cpp

@@ -68,6 +68,10 @@ vector<string> BrownianIntegrator::getKernelNames() {
     return names;
 }
 
+double BrownianIntegrator::computeKineticEnergy() {
+    return kernel.getAs<IntegrateBrownianStepKernel>().computeKineticEnergy(*context, *this);
+}
+
 void BrownianIntegrator::step(int steps) {
     for (int i = 0; i < steps; ++i) {
         context->updateContextState();

openmmapi/src/Context.cpp

@@ -87,7 +87,7 @@ State Context::getState(int types, bool enforcePeriodicBox, int groups) const {
     bool includeForces = types&State::Forces;
     bool includeEnergy = types&State::Energy;
     if (includeForces || includeEnergy) {
-        double energy = impl->calcForcesAndEnergy(includeForces, includeEnergy, groups);
+        double energy = impl->calcForcesAndEnergy(includeForces || includeEnergy, includeEnergy, groups);
         if (includeEnergy)
             state.setEnergy(impl->calcKineticEnergy(), energy);
         if (includeForces)

openmmapi/src/ContextImpl.cpp

@@ -75,7 +75,6 @@ ContextImpl::ContextImpl(Context& owner, System& system, Integrator& integrator,
     // Find the list of kernels required.
     
     vector<string> kernelNames;
-    kernelNames.push_back(CalcKineticEnergyKernel::Name());
     kernelNames.push_back(CalcForcesAndEnergyKernel::Name());
     kernelNames.push_back(UpdateStateDataKernel::Name());
     for (int i = 0; i < system.getNumForces(); ++i) {
@@ -98,8 +97,6 @@ ContextImpl::ContextImpl(Context& owner, System& system, Integrator& integrator,
     platform->contextCreated(*this, properties);
     initializeForcesKernel = platform->createKernel(CalcForcesAndEnergyKernel::Name(), *this);
     initializeForcesKernel.getAs<CalcForcesAndEnergyKernel>().initialize(system);
-    kineticEnergyKernel = platform->createKernel(CalcKineticEnergyKernel::Name(), *this);
-    kineticEnergyKernel.getAs<CalcKineticEnergyKernel>().initialize(system);
     updateStateDataKernel = platform->createKernel(UpdateStateDataKernel::Name(), *this);
     updateStateDataKernel.getAs<UpdateStateDataKernel>().initialize(system);
     applyConstraintsKernel = platform->createKernel(ApplyConstraintsKernel::Name(), *this);
@@ -122,7 +119,6 @@ ContextImpl::~ContextImpl() {
     // Make sure all kernels get properly deleted before contextDestroyed() is called.
     
     initializeForcesKernel = Kernel();
-    kineticEnergyKernel = Kernel();
     updateStateDataKernel = Kernel();
     applyConstraintsKernel = Kernel();
     virtualSitesKernel = Kernel();
@@ -218,7 +214,7 @@ int ContextImpl::getLastForceGroups() const {
 }
 
 double ContextImpl::calcKineticEnergy() {
-    return kineticEnergyKernel.getAs<CalcKineticEnergyKernel>().execute(*this);
+    return integrator.computeKineticEnergy();
 }
 
 void ContextImpl::updateContextState() {

openmmapi/src/CustomIntegrator.cpp

@@ -46,6 +46,7 @@ CustomIntegrator::CustomIntegrator(double stepSize) : owner(NULL), globalsAreCur
     setStepSize(stepSize);
     setConstraintTolerance(1e-4);
     setRandomNumberSeed((int) time(NULL));
+    kineticEnergy = "m*v*v/2";
 }
 
 void CustomIntegrator::initialize(ContextImpl& contextRef) {
@@ -77,6 +78,10 @@ vector<string> CustomIntegrator::getKernelNames() {
     return names;
 }
 
+double CustomIntegrator::computeKineticEnergy() {
+    return kernel.getAs<IntegrateCustomStepKernel>().computeKineticEnergy(*context, *this, forcesAreValid);
+}
+
 void CustomIntegrator::step(int steps) {
     globalsAreCurrent = false;
     for (int i = 0; i < steps; ++i) {
@@ -214,3 +219,11 @@ void CustomIntegrator::getComputationStep(int index, ComputationType& type, stri
     variable = computations[index].variable;
     expression = computations[index].expression;
 }
+
+const string& CustomIntegrator::getKineticEnergyExpression() const {
+    return kineticEnergy;
+}
+
+void CustomIntegrator::setKineticEnergyExpression(const string& expression) {
+    kineticEnergy = expression;
+}

openmmapi/src/LangevinIntegrator.cpp

@@ -68,6 +68,10 @@ vector<string> LangevinIntegrator::getKernelNames() {
     return names;
 }
 
+double LangevinIntegrator::computeKineticEnergy() {
+    return kernel.getAs<IntegrateLangevinStepKernel>().computeKineticEnergy(*context, *this);
+}
+
 void LangevinIntegrator::step(int steps) {
     for (int i = 0; i < steps; ++i) {
         context->updateContextState();

openmmapi/src/VariableLangevinIntegrator.cpp

@@ -69,6 +69,10 @@ vector<string> VariableLangevinIntegrator::getKernelNames() {
     return names;
 }
 
+double VariableLangevinIntegrator::computeKineticEnergy() {
+    return kernel.getAs<IntegrateVariableLangevinStepKernel>().computeKineticEnergy(*context, *this);
+}
+
 void VariableLangevinIntegrator::step(int steps) {
     for (int i = 0; i < steps; ++i) {
         context->updateContextState();

openmmapi/src/VariableVerletIntegrator.cpp

@@ -64,6 +64,10 @@ vector<string> VariableVerletIntegrator::getKernelNames() {
     return names;
 }
 
+double VariableVerletIntegrator::computeKineticEnergy() {
+    return kernel.getAs<IntegrateVariableVerletStepKernel>().computeKineticEnergy(*context, *this);
+}
+
 void VariableVerletIntegrator::step(int steps) {
     for (int i = 0; i < steps; ++i) {
         context->updateContextState();

openmmapi/src/VerletIntegrator.cpp

@@ -64,6 +64,10 @@ vector<string> VerletIntegrator::getKernelNames() {
     return names;
 }
 
+double VerletIntegrator::computeKineticEnergy() {
+    return kernel.getAs<IntegrateVerletStepKernel>().computeKineticEnergy(*context, *this);
+}
+
 void VerletIntegrator::step(int steps) {
     for (int i = 0; i < steps; ++i) {
         context->updateContextState();

platforms/cuda/src/CudaIntegrationUtilities.cpp

@@ -851,3 +851,41 @@ void CudaIntegrationUtilities::loadCheckpoint(istream& stream) {
     stream.read((char*) &randomSeedVec[0], sizeof(int4)*randomSeed->getSize());
     randomSeed->upload(randomSeedVec);
 }
+
+double CudaIntegrationUtilities::computeKineticEnergy(double timeShift) {
+    int numParticles = context.getNumAtoms();
+    int paddedNumParticles = context.getPaddedNumAtoms();
+    long long* force = (long long*) context.getPinnedBuffer();
+    context.getForce().download(force);
+    double forceScale = timeShift/0xFFFFFFFF;
+    double energy = 0.0;
+    if (context.getUseDoublePrecision() || context.getUseMixedPrecision()) {
+        vector<double4> velm;
+        context.getVelm().download(velm);
+        for (int i = 0; i < numParticles; i++) {
+            double4 v = velm[i];
+            if (v.w != 0) {
+                double scale = forceScale*v.w;
+                v.x += scale*force[i];
+                v.y += scale*force[i+paddedNumParticles];
+                v.z += scale*force[i+paddedNumParticles*2];
+                energy += (v.x*v.x+v.y*v.y+v.z*v.z)/v.w;
+            }
+        }
+    }
+    else {
+        vector<float4> velm;
+        context.getVelm().download(velm);
+        for (int i = 0; i < numParticles; i++) {
+            double4 v = make_double4(velm[i].x, velm[i].y, velm[i].z, velm[i].w);
+            if (v.w != 0) {
+                double scale = forceScale*v.w;
+                v.x += scale*force[i];
+                v.y += scale*force[i+paddedNumParticles];
+                v.z += scale*force[i+paddedNumParticles*2];
+                energy += (v.x*v.x+v.y*v.y+v.z*v.z)/v.w;
+            }
+        }
+    }
+    return 0.5*energy;
+}

platforms/cuda/src/CudaIntegrationUtilities.h

@@ -105,6 +105,13 @@ public:
      * @param stream    an input stream the checkpoint data should be read from
      */
     void loadCheckpoint(std::istream& stream);
+    /**
+     * Compute the kinetic energy of the system, possibly shifting the velocities in time to account
+     * for a leapfrog integrator.
+     * 
+     * @param timeShift   the amount by which to shift the velocities in time
+     */
+    double computeKineticEnergy(double timeShift);
 private:
     void applyConstraints(bool constrainVelocities, double tol);
     CudaContext& context;