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

platforms/reference/src/SimTKReference/ReferenceCustomDynamics.cpp

@@ -59,7 +59,14 @@ ReferenceCustomDynamics::ReferenceCustomDynamics(int numberOfAtoms, const Custom
         string expression;
         integrator.getComputationStep(i, stepType[i], stepVariable[i], expression);
         if (expression.length() > 0)
-            stepExpression[i] = Lepton::Parser::parse(expression).createProgram();
+            stepExpression[i] = Lepton::Parser::parse(expression).optimize().createProgram();
+    }
+    kineticEnergyExpression = Lepton::Parser::parse(integrator.getKineticEnergyExpression()).optimize().createProgram();
+    kineticEnergyNeedsForce = false;
+    for (int i = 0; i < kineticEnergyExpression.getNumOperations(); i++) {
+        const Lepton::Operation& op = kineticEnergyExpression.getOperation(i);
+        if (op.getId() == Lepton::Operation::VARIABLE && op.getName() == "f")
+            kineticEnergyNeedsForce = true;
     }
 }
 
@@ -304,3 +311,35 @@ void ReferenceCustomDynamics::recordChangedParameters(OpenMM::ContextImpl& conte
             context.setParameter(name, globals[name]);
     }
 }
+
+/**---------------------------------------------------------------------------------------
+
+   Compute the kinetic energy of the system.
+
+   @param context             the context this integrator is updating
+   @param numberOfAtoms       number of atoms
+   @param atomCoordinates     atom coordinates
+   @param velocities          velocities
+   @param forces              forces
+   @param masses              atom masses
+   @param globals             a map containing values of global variables
+   @param perDof              the values of per-DOF variables
+   @param forcesAreValid      whether the current forces are valid or need to be recomputed
+
+   --------------------------------------------------------------------------------------- */
+
+double ReferenceCustomDynamics::computeKineticEnergy(OpenMM::ContextImpl& context, int numberOfAtoms, std::vector<OpenMM::RealVec>& atomCoordinates,
+        std::vector<OpenMM::RealVec>& velocities, std::vector<OpenMM::RealVec>& forces, std::vector<RealOpenMM>& masses,
+        std::map<std::string, RealOpenMM>& globals, std::vector<std::vector<OpenMM::RealVec> >& perDof, bool& forcesAreValid) {
+    globals.insert(context.getParameters().begin(), context.getParameters().end());
+    if (kineticEnergyNeedsForce) {
+        energy = context.calcForcesAndEnergy(true, true, -1);
+        forcesAreValid = true;
+    }
+    computePerDof(numberOfAtoms, sumBuffer, atomCoordinates, velocities, forces, masses, globals, perDof, kineticEnergyExpression, "f");
+    RealOpenMM sum = 0.0;
+    for (int j = 0; j < numberOfAtoms; j++)
+        if (masses[j] != 0.0)
+            sum += sumBuffer[j][0]+sumBuffer[j][1]+sumBuffer[j][2];
+    return sum;
+}

platforms/reference/src/SimTKReference/ReferenceCustomDynamics.h

@@ -48,6 +48,8 @@ private:
     std::vector<bool> invalidatesForces, needsForces, needsEnergy;
     std::vector<int> forceGroup;
     RealOpenMM energy;
+    Lepton::ExpressionProgram kineticEnergyExpression;
+    bool kineticEnergyNeedsForce;
     
     void computePerDof(int numberOfAtoms, std::vector<OpenMM::RealVec>& results, const std::vector<OpenMM::RealVec>& atomCoordinates,
                   const std::vector<OpenMM::RealVec>& velocities, const std::vector<OpenMM::RealVec>& forces, const std::vector<RealOpenMM>& masses,
@@ -97,6 +99,25 @@ public:
                   std::vector<OpenMM::RealVec>& velocities, std::vector<OpenMM::RealVec>& forces, std::vector<RealOpenMM>& masses,
                   std::map<std::string, RealOpenMM>& globals, std::vector<std::vector<OpenMM::RealVec> >& perDof, bool& forcesAreValid);
       
+      /**---------------------------------------------------------------------------------------
+      
+         Compute the kinetic energy of the system.
+      
+         @param context             the context this integrator is updating
+         @param numberOfAtoms       number of atoms
+         @param atomCoordinates     atom coordinates
+         @param velocities          velocities
+         @param forces              forces
+         @param masses              atom masses
+         @param globals             a map containing values of global variables
+         @param perDof              the values of per-DOF variables
+         @param forcesAreValid      whether the current forces are valid or need to be recomputed
+
+         --------------------------------------------------------------------------------------- */
+
+       double computeKineticEnergy(OpenMM::ContextImpl& context, int numberOfAtoms, std::vector<OpenMM::RealVec>& atomCoordinates,
+                                   std::vector<OpenMM::RealVec>& velocities, std::vector<OpenMM::RealVec>& forces, std::vector<RealOpenMM>& masses,
+                                   std::map<std::string, RealOpenMM>& globals, std::vector<std::vector<OpenMM::RealVec> >& perDof, bool& forcesAreValid);
 };
 
 // ---------------------------------------------------------------------------------------

platforms/reference/tests/TestReferenceCustomIntegrator.cpp

@@ -59,9 +59,11 @@ void testSingleBond() {
     System system;
     system.addParticle(2.0);
     system.addParticle(2.0);
-    CustomIntegrator integrator(0.01);
+    const double dt = 0.01;
+    CustomIntegrator integrator(dt);
     integrator.addComputePerDof("v", "v+dt*f/m");
     integrator.addComputePerDof("x", "x+dt*v");
+    integrator.setKineticEnergyExpression("m*v1*v1/2; v1=v+0.5*dt*f/m");
     HarmonicBondForce* forceField = new HarmonicBondForce();
     forceField->addBond(0, 1, 1.5, 1);
     system.addForce(forceField);
@@ -70,23 +72,25 @@ void testSingleBond() {
     positions[0] = Vec3(-1, 0, 0);
     positions[1] = Vec3(1, 0, 0);
     context.setPositions(positions);
+    vector<Vec3> velocities(2);
+    velocities[0] = Vec3(-0.5*dt*0.5*0.5, 0, 0);
+    velocities[1] = Vec3(0.5*dt*0.5*0.5, 0, 0);
+    context.setVelocities(velocities);
     
     // This is simply a harmonic oscillator, so compare it to the analytical solution.
     
     const double freq = 1.0;;
-    State state = context.getState(State::Energy);
-    const double initialEnergy = state.getKineticEnergy()+state.getPotentialEnergy();
     for (int i = 0; i < 1000; ++i) {
-        state = context.getState(State::Positions | State::Velocities | State::Energy);
+        State state = context.getState(State::Positions | State::Velocities | State::Energy);
         double time = state.getTime();
         double expectedDist = 1.5+0.5*std::cos(freq*time);
-        ASSERT_EQUAL_VEC(Vec3(-0.5*expectedDist, 0, 0), state.getPositions()[0], 0.02);
-        ASSERT_EQUAL_VEC(Vec3(0.5*expectedDist, 0, 0), state.getPositions()[1], 0.02);
-        double expectedSpeed = -0.5*freq*std::sin(freq*time);
-        ASSERT_EQUAL_VEC(Vec3(-0.5*expectedSpeed, 0, 0), state.getVelocities()[0], 0.02);
-        ASSERT_EQUAL_VEC(Vec3(0.5*expectedSpeed, 0, 0), state.getVelocities()[1], 0.02);
+        ASSERT_EQUAL_VEC(Vec3(-0.5*expectedDist, 0, 0), state.getPositions()[0], 1e-4);
+        ASSERT_EQUAL_VEC(Vec3(0.5*expectedDist, 0, 0), state.getPositions()[1], 1e-4);
+        double expectedSpeed = -0.5*freq*std::sin(freq*(time-dt/2));
+        ASSERT_EQUAL_VEC(Vec3(-0.5*expectedSpeed, 0, 0), state.getVelocities()[0], 1e-4);
+        ASSERT_EQUAL_VEC(Vec3(0.5*expectedSpeed, 0, 0), state.getVelocities()[1], 1e-4);
         double energy = state.getKineticEnergy()+state.getPotentialEnergy();
-        ASSERT_EQUAL_TOL(initialEnergy, energy, 0.01);
+        ASSERT_EQUAL_TOL(0.5*0.5*0.5, energy, 1e-4);
         integrator.step(1);
     }
 }

platforms/reference/tests/TestReferenceVariableVerletIntegrator.cpp

@@ -50,23 +50,6 @@ using namespace std;
 
 const double TOL = 1e-5;
 
-/**
- * Compute the energy of a state, taking into account the half step offset between
- * positions and velocities.
- */
-
-static double computeEnergy(const State& state, const System& system, double dt) {
-    const vector<Vec3>& v = state.getVelocities();
-    const vector<Vec3>& f = state.getForces();
-    double energy = 0.0;
-    for (int i = 0; i < system.getNumParticles(); i++) {
-        double m = system.getParticleMass(i);
-        Vec3 vel = v[i]+f[i]*(0.5*dt/m);
-        energy += 0.5*m*vel.dot(vel);
-    }
-    return energy+state.getPotentialEnergy();
-}
-
 void testSingleBond() {
     ReferencePlatform platform;
     System system;
@@ -145,7 +128,7 @@ void testConstraints() {
             double dist = std::sqrt((p1[0]-p2[0])*(p1[0]-p2[0])+(p1[1]-p2[1])*(p1[1]-p2[1])+(p1[2]-p2[2])*(p1[2]-p2[2]));
             ASSERT_EQUAL_TOL(distance, dist, 2e-5);
         }
-        double energy = computeEnergy(state, system, integrator.getStepSize());
+        double energy = state.getKineticEnergy()+state.getPotentialEnergy();
         if (i == 1)
             initialEnergy = energy;
         else if (i > 1)
@@ -217,7 +200,7 @@ void testConstrainedClusters() {
             double dist = std::sqrt((p1[0]-p2[0])*(p1[0]-p2[0])+(p1[1]-p2[1])*(p1[1]-p2[1])+(p1[2]-p2[2])*(p1[2]-p2[2]));
             ASSERT_EQUAL_TOL(distance, dist, 2e-5);
         }
-        double energy = computeEnergy(state, system, integrator.getStepSize());
+        double energy = state.getKineticEnergy()+state.getPotentialEnergy();
         if (i == 1)
             initialEnergy = energy;
         else if (i > 1)

platforms/reference/tests/TestReferenceVerletIntegrator.cpp

@@ -50,23 +50,6 @@ using namespace std;
 
 const double TOL = 1e-5;
 
-/**
- * Compute the energy of a state, taking into account the half step offset between
- * positions and velocities.
- */
-
-static double computeEnergy(const State& state, const System& system, double dt) {
-    const vector<Vec3>& v = state.getVelocities();
-    const vector<Vec3>& f = state.getForces();
-    double energy = 0.0;
-    for (int i = 0; i < system.getNumParticles(); i++) {
-        double m = system.getParticleMass(i);
-        Vec3 vel = v[i]+f[i]*(0.5*dt/m);
-        energy += 0.5*m*vel.dot(vel);
-    }
-    return energy+state.getPotentialEnergy();
-}
-
 void testSingleBond() {
     ReferencePlatform platform;
     System system;
@@ -144,7 +127,7 @@ void testConstraints() {
             double dist = std::sqrt((p1[0]-p2[0])*(p1[0]-p2[0])+(p1[1]-p2[1])*(p1[1]-p2[1])+(p1[2]-p2[2])*(p1[2]-p2[2]));
             ASSERT_EQUAL_TOL(distance, dist, 2e-5);
         }
-        double energy = computeEnergy(state, system, integrator.getStepSize());
+        double energy = state.getPotentialEnergy()+state.getKineticEnergy();
         if (i == 1)
             initialEnergy = energy;
         else if (i > 1)
@@ -208,7 +191,7 @@ void testConstrainedClusters() {
             double dist = std::sqrt((p1[0]-p2[0])*(p1[0]-p2[0])+(p1[1]-p2[1])*(p1[1]-p2[1])+(p1[2]-p2[2])*(p1[2]-p2[2]));
             ASSERT_EQUAL_TOL(distance, dist, 2e-5);
         }
-        double energy = computeEnergy(state, system, integrator.getStepSize());
+        double energy = state.getPotentialEnergy()+state.getKineticEnergy();
         if (i == 1)
             initialEnergy = energy;
         else if (i > 1)

plugins/rpmd/openmmapi/include/openmm/RPMDIntegrator.h

@@ -175,6 +175,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 numCopies, randomNumberSeed;

plugins/rpmd/openmmapi/include/openmm/RpmdKernels.h

@@ -80,6 +80,10 @@ public:
      * Copy positions and velocities for one copy into the context.
      */
     virtual void copyToContext(int copy, ContextImpl& context) = 0;
+    /**
+     * Compute the kinetic energy.
+     */
+    virtual double computeKineticEnergy(ContextImpl& context, const RPMDIntegrator& integrator) = 0;
 };
 
 } // namespace OpenMM

plugins/rpmd/openmmapi/src/RPMDIntegrator.cpp

@@ -90,6 +90,10 @@ State RPMDIntegrator::getState(int copy, int types, bool enforcePeriodicBox, int
     return context->getOwner().getState(types, enforcePeriodicBox, groups);
 }
 
+double RPMDIntegrator::computeKineticEnergy() {
+    kernel.getAs<IntegrateRPMDStepKernel>().computeKineticEnergy(*context, *this);
+}
+
 void RPMDIntegrator::step(int steps) {
     if (!hasSetPosition) {
         // Initialize the positions from the context.

plugins/rpmd/platforms/opencl/src/OpenCLRpmdKernels.cpp

@@ -185,6 +185,10 @@ void OpenCLIntegrateRPMDStepKernel::computeForces(ContextImpl& context) {
     }
 }
 
+double OpenCLIntegrateRPMDStepKernel::computeKineticEnergy(ContextImpl& context, const RPMDIntegrator& integrator) {
+    return cl.getIntegrationUtilities().computeKineticEnergy(0);
+}
+
 void OpenCLIntegrateRPMDStepKernel::setPositions(int copy, const vector<Vec3>& pos) {
     if (positions == NULL)
         throw OpenMMException("RPMDIntegrator: Cannot set positions before the integrator is added to a Context");

plugins/rpmd/platforms/opencl/src/OpenCLRpmdKernels.h

@@ -64,6 +64,13 @@ public:
      *                       false to show that cached forces are invalid and must be recalculated
      */
     void execute(ContextImpl& context, const RPMDIntegrator& integrator, bool forcesAreValid);
+    /**
+     * Compute the kinetic energy.
+     * 
+     * @param context        the context in which to execute this kernel
+     * @param integrator     the RPMDIntegrator this kernel is being used for
+     */
+    double computeKineticEnergy(ContextImpl& context, const RPMDIntegrator& integrator);
     /**
      * Get the positions of all particles in one copy of the system.
      */

plugins/rpmd/platforms/reference/src/ReferenceRpmdKernels.cpp

@@ -320,6 +320,21 @@ void ReferenceIntegrateRPMDStepKernel::execute(ContextImpl& context, const RPMDI
 //    }
 //}
 
+double ReferenceIntegrateRPMDStepKernel::computeKineticEnergy(ContextImpl& context, const RPMDIntegrator& integrator) {
+    System& system = context.getSystem();
+    int numParticles = system.getNumParticles();
+    vector<RealVec>& velData = extractVelocities(context);
+    double energy = 0.0;
+    for (int i = 0; i < numParticles; ++i) {
+        double mass = system.getParticleMass(i);
+        if (mass > 0) {
+            RealVec v = velData[i];
+            energy += mass*(v.dot(v));
+        }
+    }
+    return 0.5*energy;
+}
+
 void ReferenceIntegrateRPMDStepKernel::setPositions(int copy, const vector<Vec3>& pos) {
     int numParticles = positions[copy].size();
     for (int i = 0; i < numParticles; i++)

plugins/rpmd/platforms/reference/src/ReferenceRpmdKernels.h

@@ -65,6 +65,13 @@ public:
      *                       false to show that cached forces are invalid and must be recalculated
      */
     void execute(ContextImpl& context, const RPMDIntegrator& integrator, bool forcesAreValid);
+    /**
+     * Compute the kinetic energy.
+     * 
+     * @param context    the context in which to execute this kernel
+     * @param integrator     the RPMDIntegrator this kernel is being used for
+     */
+    double computeKineticEnergy(ContextImpl& context, const RPMDIntegrator& integrator);
     /**
      * Get the positions of all particles in one copy of the system.
      */