When computing kinetic energy, make sure the shifted velocities are...

When computing kinetic energy, make sure the shifted velocities are perpendicular to the constraints

platforms/cuda/src/CudaIntegrationUtilities.cpp

@@ -643,6 +643,7 @@ CudaIntegrationUtilities::CudaIntegrationUtilities(CudaContext& context, const S
     vsiteForceKernel = context.getKernel(module, "distributeVirtualSiteForces");
     numVsites = num2Avg+num3Avg+numOutOfPlane;
     randomKernel = context.getKernel(module, "generateRandomNumbers");
+    timeShiftKernel = context.getKernel(module, "timeShiftVelocities");
 }
 
 CudaIntegrationUtilities::~CudaIntegrationUtilities() {
@@ -858,38 +859,46 @@ void CudaIntegrationUtilities::loadCheckpoint(istream& stream) {
 
 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/0x100000000LL;
+    if (timeShift != 0) {
+        float timeShiftFloat = (float) timeShift;
+        void* timeShiftPtr = (context.getUseDoublePrecision() ? (void*) &timeShift : (void*) &timeShiftFloat);
+
+        // Copy the velocities into the posDelta array while we temporarily modify them.
+
+        context.getVelm().copyTo(*posDelta);
+
+        // Apply the time shift.
+
+        void* args[] = {&context.getVelm().getDevicePointer(), &context.getForce().getDevicePointer(), timeShiftPtr};
+        context.executeKernel(timeShiftKernel, args, numParticles);
+        applyConstraints(true, 1e-4);
+    }
+    
+    // Compute the kinetic energy.
+    
     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];
+            if (v.w != 0)
                 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];
+            float4 v = velm[i];
+            if (v.w != 0)
                 energy += (v.x*v.x+v.y*v.y+v.z*v.z)/v.w;
         }
     }
-    }
+    
+    // Restore the velocities.
+    
+    if (timeShift != 0)
+        posDelta->copyTo(context.getVelm());
     return 0.5*energy;
 }

platforms/cuda/src/CudaIntegrationUtilities.h

@@ -122,7 +122,7 @@ private:
     CUfunction ccmaMultiplyKernel;
     CUfunction ccmaUpdateKernel;
     CUfunction vsitePositionKernel, vsiteForceKernel;
-    CUfunction randomKernel;
+    CUfunction randomKernel, timeShiftKernel;
     CudaArray* posDelta;
     CudaArray* settleAtoms;
     CudaArray* settleParams;

platforms/cuda/src/kernels/integrationUtilities.cu

@@ -798,3 +798,19 @@ extern "C" __global__ void distributeVirtualSiteForces(const real4* __restrict__
         addForce(atoms.w, force, fp3);
     }
 }
+
+/**
+ * Apply a time shift to the velocities before computing kinetic energy.
+ */
+extern "C" __global__ void timeShiftVelocities(mixed4* __restrict__ velm, const long long* __restrict__ force, real timeShift) {
+    const mixed scale = timeShift/(mixed) 0x100000000;
+    for (int index = blockIdx.x*blockDim.x+threadIdx.x; index < NUM_ATOMS; index += blockDim.x*gridDim.x) {
+        mixed4 velocity = velm[index];
+        if (velocity.w != 0.0) {
+            velocity.x += scale*force[index]*velocity.w;
+            velocity.y += scale*force[index+PADDED_NUM_ATOMS]*velocity.w;
+            velocity.z += scale*force[index+PADDED_NUM_ATOMS*2]*velocity.w;
+            velm[index] = velocity;
+        }
+    }
+}
\ No newline at end of file

platforms/opencl/src/OpenCLIntegrationUtilities.cpp

@@ -120,6 +120,13 @@ OpenCLIntegrationUtilities::OpenCLIntegrationUtilities(OpenCLContext& context, c
         stepSize->upload(step);
     }
     
+    // Create the time shift kernel for calculating kinetic energy.
+    
+    map<string, string> timeShiftDefines;
+    timeShiftDefines["NUM_ATOMS"] = context.intToString(system.getNumParticles());
+    cl::Program utilitiesProgram = context.createProgram(OpenCLKernelSources::integrationUtilities, timeShiftDefines);
+    timeShiftKernel = cl::Kernel(utilitiesProgram, "timeShiftVelocities");
+
     // Create kernels for enforcing constraints.
 
     map<string, string> velocityDefines;
@@ -961,54 +968,48 @@ void OpenCLIntegrationUtilities::loadCheckpoint(istream& stream) {
 
 double OpenCLIntegrationUtilities::computeKineticEnergy(double timeShift) {
     int numParticles = context.getNumAtoms();
-    double energy = 0.0;
-    if (context.getUseDoublePrecision()) {
-        mm_double4* force = (mm_double4*) context.getPinnedBuffer();
-        context.getForce().download(force);
-        vector<mm_double4> velm;
-        context.getVelm().download(velm);
-        for (int i = 0; i < numParticles; i++) {
-            mm_double4 v = velm[i];
-            if (v.w != 0) {
-                double scale = timeShift*v.w;
-                v.x += scale*force[i].x;
-                v.y += scale*force[i].y;
-                v.z += scale*force[i].z;
-                energy += (v.x*v.x+v.y*v.y+v.z*v.z)/v.w;
-            }
-        }
+    if (timeShift != 0) {
+        // Copy the velocities into the posDelta array while we temporarily modify them.
+
+        context.getVelm().copyTo(*posDelta);
+
+        // Apply the time shift.
+
+        timeShiftKernel.setArg<cl::Buffer>(0, context.getVelm().getDeviceBuffer());
+        timeShiftKernel.setArg<cl::Buffer>(1, context.getForce().getDeviceBuffer());
+        if (context.getUseDoublePrecision() || context.getUseMixedPrecision())
+            timeShiftKernel.setArg<cl_double>(2, timeShift);
+        else
+            timeShiftKernel.setArg<cl_float>(2, (cl_float) timeShift);
+        context.executeKernel(timeShiftKernel, numParticles);
+        applyConstraints(true, 1e-4);
     }
-    else if (context.getUseMixedPrecision()) {
-        mm_float4* force = (mm_float4*) context.getPinnedBuffer();
-        context.getForce().download(force);
+    
+    // Compute the kinetic energy.
+    
+    double energy = 0.0;
+    if (context.getUseDoublePrecision() || context.getUseMixedPrecision()) {
         vector<mm_double4> velm;
         context.getVelm().download(velm);
         for (int i = 0; i < numParticles; i++) {
             mm_double4 v = velm[i];
-            if (v.w != 0) {
-                double scale = timeShift*v.w;
-                v.x += scale*force[i].x;
-                v.y += scale*force[i].y;
-                v.z += scale*force[i].z;
+            if (v.w != 0)
                 energy += (v.x*v.x+v.y*v.y+v.z*v.z)/v.w;
         }
     }
-    }
     else {
-        mm_float4* force = (mm_float4*) context.getPinnedBuffer();
-        context.getForce().download(force);
         vector<mm_float4> velm;
         context.getVelm().download(velm);
         for (int i = 0; i < numParticles; i++) {
-            mm_double4 v = mm_double4(velm[i].x, velm[i].y, velm[i].z, velm[i].w);
-            if (v.w != 0) {
-                double scale = timeShift*v.w;
-                v.x += scale*force[i].x;
-                v.y += scale*force[i].y;
-                v.z += scale*force[i].z;
+            mm_float4 v = velm[i];
+            if (v.w != 0)
                 energy += (v.x*v.x+v.y*v.y+v.z*v.z)/v.w;
         }
     }
-    }
+    
+    // Restore the velocities.
+    
+    if (timeShift != 0)
+        posDelta->copyTo(context.getVelm());
     return 0.5*energy;
 }

platforms/opencl/src/OpenCLIntegrationUtilities.h

@@ -122,7 +122,7 @@ private:
     cl::Kernel ccmaMultiplyKernel;
     cl::Kernel ccmaPosUpdateKernel, ccmaVelUpdateKernel;
     cl::Kernel vsitePositionKernel, vsiteForceKernel;
-    cl::Kernel randomKernel;
+    cl::Kernel randomKernel, timeShiftKernel;
     OpenCLArray* posDelta;
     OpenCLArray* settleAtoms;
     OpenCLArray* settleParams;

platforms/opencl/src/kernels/integrationUtilities.cl

+/**
+ * Apply a time shift to the velocities before computing kinetic energy.
+ */
+__kernel void timeShiftVelocities(__global mixed4* restrict velm, __global const real4* restrict force, real timeShift) {
+    for (int index = get_global_id(0); index < NUM_ATOMS; index += get_global_size(0)) {
+        mixed4 velocity = velm[index];
+        if (velocity.w != 0.0) {
+            velocity.xyz += timeShift*force[index].xyz*velocity.w;
+            velm[index] = velocity;
+        }
+    }
+}
\ No newline at end of file

platforms/reference/src/ReferenceKernels.cpp

@@ -149,16 +149,38 @@ static void findAnglesForCCMA(const System& system, vector<ReferenceCCMAAlgorith
  * Compute the kinetic energy of the system, possibly shifting the velocities in time to account
  * for a leapfrog integrator.
  */
-static double computeShiftedKineticEnergy(ContextImpl& context, vector<double>& masses, double timeShift) {
+static double computeShiftedKineticEnergy(ContextImpl& context, vector<double>& masses, double timeShift, ReferenceConstraintAlgorithm* constraints) {
+    vector<RealVec>& posData = extractPositions(context);
     vector<RealVec>& velData = extractVelocities(context);
     vector<RealVec>& forceData = extractForces(context);
-    double energy = 0.0;
-    for (int i = 0; i < (int) masses.size(); ++i) {
-        if (masses[i] > 0) {
-            RealVec v = velData[i]+forceData[i]*(timeShift/masses[i]);
-            energy += masses[i]*(v.dot(v));
+    int numParticles = context.getSystem().getNumParticles();
+    
+    // Compute the shifted velocities.
+    
+    vector<RealVec> shiftedVel(numParticles);
+    for (int i = 0; i < numParticles; ++i) {
+        if (masses[i] > 0)
+            shiftedVel[i] = velData[i]+forceData[i]*(timeShift/masses[i]);
+        else
+            shiftedVel[i] = velData[i];
     }
+    
+    // Apply constraints to them.
+    
+    if (constraints != NULL) {
+        vector<double> inverseMasses(numParticles);
+        for (int i = 0; i < numParticles; i++)
+            inverseMasses[i] = (masses[i] == 0 ? 0 : 1/masses[i]);
+        constraints->setTolerance(1e-4);
+        constraints->applyToVelocities(numParticles, posData, shiftedVel, inverseMasses);
     }
+    
+    // Compute the kinetic energy.
+    
+    double energy = 0.0;
+    for (int i = 0; i < numParticles; ++i)
+        if (masses[i] > 0)
+            energy += masses[i]*(shiftedVel[i].dot(shiftedVel[i]));
     return 0.5*energy;
 }
 
@@ -1707,7 +1729,7 @@ void ReferenceIntegrateVerletStepKernel::execute(ContextImpl& context, const Ver
 }
 
 double ReferenceIntegrateVerletStepKernel::computeKineticEnergy(ContextImpl& context, const VerletIntegrator& integrator) {
-    return computeShiftedKineticEnergy(context, masses, 0.5*integrator.getStepSize());
+    return computeShiftedKineticEnergy(context, masses, 0.5*integrator.getStepSize(), constraints);
 }
 
 ReferenceIntegrateLangevinStepKernel::~ReferenceIntegrateLangevinStepKernel() {
@@ -1775,7 +1797,7 @@ void ReferenceIntegrateLangevinStepKernel::execute(ContextImpl& context, const L
 }
 
 double ReferenceIntegrateLangevinStepKernel::computeKineticEnergy(ContextImpl& context, const LangevinIntegrator& integrator) {
-    return computeShiftedKineticEnergy(context, masses, 0.5*integrator.getStepSize());
+    return computeShiftedKineticEnergy(context, masses, 0.5*integrator.getStepSize(), constraints);
 }
 
 ReferenceIntegrateBrownianStepKernel::~ReferenceIntegrateBrownianStepKernel() {
@@ -1842,7 +1864,7 @@ void ReferenceIntegrateBrownianStepKernel::execute(ContextImpl& context, const B
 }
 
 double ReferenceIntegrateBrownianStepKernel::computeKineticEnergy(ContextImpl& context, const BrownianIntegrator& integrator) {
-    return computeShiftedKineticEnergy(context, masses, 0);
+    return computeShiftedKineticEnergy(context, masses, 0, constraints);
 }
 
 ReferenceIntegrateVariableLangevinStepKernel::~ReferenceIntegrateVariableLangevinStepKernel() {
@@ -1910,7 +1932,7 @@ double ReferenceIntegrateVariableLangevinStepKernel::execute(ContextImpl& contex
 }
 
 double ReferenceIntegrateVariableLangevinStepKernel::computeKineticEnergy(ContextImpl& context, const VariableLangevinIntegrator& integrator) {
-    return computeShiftedKineticEnergy(context, masses, 0.5*integrator.getStepSize());
+    return computeShiftedKineticEnergy(context, masses, 0.5*integrator.getStepSize(), constraints);
 }
 
 ReferenceIntegrateVariableVerletStepKernel::~ReferenceIntegrateVariableVerletStepKernel() {
@@ -1972,7 +1994,7 @@ double ReferenceIntegrateVariableVerletStepKernel::execute(ContextImpl& context,
 }
 
 double ReferenceIntegrateVariableVerletStepKernel::computeKineticEnergy(ContextImpl& context, const VariableVerletIntegrator& integrator) {
-    return computeShiftedKineticEnergy(context, masses, 0.5*integrator.getStepSize());
+    return computeShiftedKineticEnergy(context, masses, 0.5*integrator.getStepSize(), constraints);
 }
 
 ReferenceIntegrateCustomStepKernel::~ReferenceIntegrateCustomStepKernel() {