Bug fixes to BAOABLangevinIntegrator

platforms/cuda/include/CudaKernels.h

@@ -1437,7 +1437,7 @@ private:
     CudaContext& cu;
     double prevTemp, prevFriction, prevStepSize;
     CudaArray params, oldDelta;
-    CUfunction kernel1, kernel2, kernel3;
+    CUfunction kernel1, kernel2, kernel3, kernel4;
 };
 
 /**

platforms/cuda/src/CudaKernels.cpp

@@ -7152,6 +7152,7 @@ void CudaIntegrateBAOABStepKernel::initialize(const System& system, const BAOABL
     kernel1 = cu.getKernel(module, "integrateBAOABPart1");
     kernel2 = cu.getKernel(module, "integrateBAOABPart2");
     kernel3 = cu.getKernel(module, "integrateBAOABPart3");
+    kernel4 = cu.getKernel(module, "integrateBAOABPart4");
     if (cu.getUseDoublePrecision() || cu.getUseMixedPrecision()) {
         params.initialize<double>(cu, 3, "baoabParams");
         oldDelta.initialize<double4>(cu, cu.getPaddedNumAtoms(), "oldDelta");
@@ -7202,11 +7203,13 @@ void CudaIntegrateBAOABStepKernel::execute(ContextImpl& context, const BAOABLang
             &oldDelta.getDevicePointer(), &params.getDevicePointer(), &integration.getStepSize().getDevicePointer(), &integration.getRandom().getDevicePointer(), &randomIndex};
     cu.executeKernel(kernel2, args2, numAtoms, 128);
     integration.applyConstraints(integrator.getConstraintTolerance());
-    context.calcForcesAndEnergy(true, false);
-    void* args3[] = {&numAtoms, &paddedNumAtoms, &cu.getPosq().getDevicePointer(), &posCorrection, &cu.getVelm().getDevicePointer(),
-            &cu.getForce().getDevicePointer(), &integration.getPosDelta().getDevicePointer(),
-            &oldDelta.getDevicePointer(), &integration.getStepSize().getDevicePointer()};
+    void* args3[] = {&numAtoms, &cu.getPosq().getDevicePointer(), &posCorrection, &cu.getVelm().getDevicePointer(),
+            &integration.getPosDelta().getDevicePointer(), &oldDelta.getDevicePointer(), &integration.getStepSize().getDevicePointer()};
     cu.executeKernel(kernel3, args3, numAtoms, 128);
+    context.calcForcesAndEnergy(true, false);
+    void* args4[] = {&numAtoms, &paddedNumAtoms, &cu.getVelm().getDevicePointer(),
+            &cu.getForce().getDevicePointer(), &integration.getStepSize().getDevicePointer()};
+    cu.executeKernel(kernel4, args4, numAtoms, 128);
     integration.applyVelocityConstraints(integrator.getConstraintTolerance());
     integration.computeVirtualSites();
 
@@ -7215,6 +7218,8 @@ void CudaIntegrateBAOABStepKernel::execute(ContextImpl& context, const BAOABLang
     cu.setTime(cu.getTime()+stepSize);
     cu.setStepCount(cu.getStepCount()+1);
     cu.reorderAtoms();
+    if (cu.getAtomsWereReordered())
+        forcesAreValid = false;
     
     // Reduce UI lag.
     

platforms/cuda/src/kernels/baoab.cu

@@ -75,18 +75,17 @@ extern "C" __global__ void integrateBAOABPart2(int numAtoms, real4* __restrict__
  * Perform the third step of BAOAB integration.
  */
 
-extern "C" __global__ void integrateBAOABPart3(int numAtoms, int paddedNumAtoms, real4* __restrict__ posq, real4* __restrict__ posqCorrection, mixed4* __restrict__ velm,
-        const long long* __restrict__ force, mixed4* __restrict__ posDelta, mixed4* __restrict__ oldDelta, const mixed2* __restrict__ dt) {
+extern "C" __global__ void integrateBAOABPart3(int numAtoms, real4* __restrict__ posq, real4* __restrict__ posqCorrection, mixed4* __restrict__ velm,
+        mixed4* __restrict__ posDelta, mixed4* __restrict__ oldDelta, const mixed2* __restrict__ dt) {
     mixed halfdt = 0.5*dt[0].y;
     mixed invHalfdt = 1/halfdt;
-    mixed fscale = halfdt/(mixed) 0x100000000;
     for (int index = blockIdx.x*blockDim.x+threadIdx.x; index < numAtoms; index += blockDim.x*gridDim.x) {
         mixed4 velocity = velm[index];
         if (velocity.w != 0.0) {
             mixed4 delta = posDelta[index];
-            velocity.x += (delta.x-oldDelta[index].x)*invHalfdt + fscale*velocity.w*force[index];
-            velocity.y += (delta.y-oldDelta[index].y)*invHalfdt + fscale*velocity.w*force[index+paddedNumAtoms];
-            velocity.z += (delta.z-oldDelta[index].z)*invHalfdt + fscale*velocity.w*force[index+paddedNumAtoms*2];
+            velocity.x += (delta.x-oldDelta[index].x)*invHalfdt;
+            velocity.y += (delta.y-oldDelta[index].y)*invHalfdt;
+            velocity.z += (delta.z-oldDelta[index].z)*invHalfdt;
             velm[index] = velocity;
 #ifdef USE_MIXED_PRECISION
             real4 pos1 = posq[index];
@@ -107,3 +106,22 @@ extern "C" __global__ void integrateBAOABPart3(int numAtoms, int paddedNumAtoms,
         }
     }
 }
+
+/**
+ * Perform the fourth step of BAOAB integration.
+ */
+
+extern "C" __global__ void integrateBAOABPart4(int numAtoms, int paddedNumAtoms, mixed4* __restrict__ velm,
+        const long long* __restrict__ force, const mixed2* __restrict__ dt) {
+    mixed halfdt = 0.5*dt[0].y;
+    mixed fscale = halfdt/(mixed) 0x100000000;
+    for (int index = blockIdx.x*blockDim.x+threadIdx.x; index < numAtoms; index += blockDim.x*gridDim.x) {
+        mixed4 velocity = velm[index];
+        if (velocity.w != 0.0) {
+            velocity.x += fscale*velocity.w*force[index];
+            velocity.y += fscale*velocity.w*force[index+paddedNumAtoms];
+            velocity.z += fscale*velocity.w*force[index+paddedNumAtoms*2];
+            velm[index] = velocity;
+        }
+    }
+}

platforms/opencl/include/OpenCLKernels.h

@@ -1421,7 +1421,7 @@ private:
     double prevTemp, prevFriction, prevStepSize;
     bool hasInitializedKernels;
     OpenCLArray params, oldDelta;
-    cl::Kernel kernel1, kernel2, kernel3;
+    cl::Kernel kernel1, kernel2, kernel3, kernel4;
 };
 
 /**

platforms/opencl/src/OpenCLKernels.cpp

@@ -7483,6 +7483,7 @@ void OpenCLIntegrateBAOABStepKernel::initialize(const System& system, const BAOA
     kernel1 = cl::Kernel(program, "integrateBAOABPart1");
     kernel2 = cl::Kernel(program, "integrateBAOABPart2");
     kernel3 = cl::Kernel(program, "integrateBAOABPart3");
+    kernel4 = cl::Kernel(program, "integrateBAOABPart4");
     if (cl.getUseDoublePrecision() || cl.getUseMixedPrecision()) {
         params.initialize<cl_double>(cl, 2, "baoabParams");
         oldDelta.initialize<mm_double4>(cl, cl.getPaddedNumAtoms(), "oldDelta");
@@ -7515,10 +7516,12 @@ void OpenCLIntegrateBAOABStepKernel::execute(ContextImpl& context, const BAOABLa
         kernel3.setArg<cl::Buffer>(0, cl.getPosq().getDeviceBuffer());
         setPosqCorrectionArg(cl, kernel3, 1);
         kernel3.setArg<cl::Buffer>(2, cl.getVelm().getDeviceBuffer());
-        kernel3.setArg<cl::Buffer>(3, cl.getForce().getDeviceBuffer());
-        kernel3.setArg<cl::Buffer>(4, integration.getPosDelta().getDeviceBuffer());
-        kernel3.setArg<cl::Buffer>(5, oldDelta.getDeviceBuffer());
-        kernel3.setArg<cl::Buffer>(6, integration.getStepSize().getDeviceBuffer());
+        kernel3.setArg<cl::Buffer>(3, integration.getPosDelta().getDeviceBuffer());
+        kernel3.setArg<cl::Buffer>(4, oldDelta.getDeviceBuffer());
+        kernel3.setArg<cl::Buffer>(5, integration.getStepSize().getDeviceBuffer());
+        kernel4.setArg<cl::Buffer>(0, cl.getVelm().getDeviceBuffer());
+        kernel4.setArg<cl::Buffer>(1, cl.getForce().getDeviceBuffer());
+        kernel4.setArg<cl::Buffer>(2, integration.getStepSize().getDeviceBuffer());
     }
     if (!forcesAreValid) {
         context.calcForcesAndEnergy(true, false);
@@ -7550,8 +7553,9 @@ void OpenCLIntegrateBAOABStepKernel::execute(ContextImpl& context, const BAOABLa
     integration.applyConstraints(integrator.getConstraintTolerance());
     cl.executeKernel(kernel2, numAtoms);
     integration.applyConstraints(integrator.getConstraintTolerance());
-    context.calcForcesAndEnergy(true, false);
     cl.executeKernel(kernel3, numAtoms);
+    context.calcForcesAndEnergy(true, false);
+    cl.executeKernel(kernel4, numAtoms);
     integration.applyVelocityConstraints(integrator.getConstraintTolerance());
     integration.computeVirtualSites();
 
@@ -7560,6 +7564,8 @@ void OpenCLIntegrateBAOABStepKernel::execute(ContextImpl& context, const BAOABLa
     cl.setTime(cl.getTime()+stepSize);
     cl.setStepCount(cl.getStepCount()+1);
     cl.reorderAtoms();
+    if (cl.getAtomsWereReordered())
+        forcesAreValid = false;
     
     // Reduce UI lag.
     

platforms/opencl/src/kernels/baoab.cl

@@ -71,16 +71,16 @@ __kernel void integrateBAOABPart2(__global real4* restrict posq, __global real4*
  */
 
 __kernel void integrateBAOABPart3(__global real4* restrict posq, __global real4* restrict posqCorrection, __global mixed4* restrict velm,
-        __global const real4* restrict force, __global mixed4* restrict posDelta, __global mixed4* restrict oldDelta, __global const mixed2* restrict dt) {
+         __global mixed4* restrict posDelta, __global mixed4* restrict oldDelta, __global const mixed2* restrict dt) {
     mixed halfdt = 0.5*dt[0].y;
     mixed invHalfdt = 1/halfdt;
     for (int index = get_global_id(0); index < NUM_ATOMS; index += get_global_size(0)) {
         mixed4 velocity = velm[index];
         if (velocity.w != 0.0) {
             mixed4 delta = posDelta[index];
-            velocity.x += (delta.x-oldDelta[index].x)*invHalfdt + halfdt*velocity.w*force[index].x;
-            velocity.y += (delta.y-oldDelta[index].y)*invHalfdt + halfdt*velocity.w*force[index].y;
-            velocity.z += (delta.z-oldDelta[index].z)*invHalfdt + halfdt*velocity.w*force[index].z;
+            velocity.x += (delta.x-oldDelta[index].x)*invHalfdt;
+            velocity.y += (delta.y-oldDelta[index].y)*invHalfdt;
+            velocity.z += (delta.z-oldDelta[index].z)*invHalfdt;
             velm[index] = velocity;
 #ifdef USE_MIXED_PRECISION
             real4 pos1 = posq[index];
@@ -99,3 +99,20 @@ __kernel void integrateBAOABPart3(__global real4* restrict posq, __global real4*
         }
     }
 }
+
+/**
+ * Perform the fourth step of BAOAB integration.
+ */
+
+__kernel void integrateBAOABPart4(__global mixed4* restrict velm, __global const real4* restrict force, __global const mixed2* restrict dt) {
+    mixed halfdt = 0.5*dt[0].y;
+    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.x += halfdt*velocity.w*force[index].x;
+            velocity.y += halfdt*velocity.w*force[index].y;
+            velocity.z += halfdt*velocity.w*force[index].z;
+            velm[index] = velocity;
+        }
+    }
+}

platforms/reference/include/ReferenceBAOABDynamics.h

@@ -121,6 +121,7 @@ class OPENMM_EXPORT ReferenceBAOABDynamics : public ReferenceDynamics {
       
          Third update
       
+         @param context             the context this integrator is updating
          @param numberOfAtoms       number of atoms
          @param atomCoordinates     atom coordinates
          @param velocities          velocities
@@ -129,7 +130,7 @@ class OPENMM_EXPORT ReferenceBAOABDynamics : public ReferenceDynamics {
       
          --------------------------------------------------------------------------------------- */
       
-      virtual void updatePart3(int numberOfAtoms, std::vector<OpenMM::Vec3>& atomCoordinates, std::vector<OpenMM::Vec3>& velocities,
+      virtual void updatePart3(OpenMM::ContextImpl& context, int numberOfAtoms, std::vector<OpenMM::Vec3>& atomCoordinates, std::vector<OpenMM::Vec3>& velocities,
                        std::vector<OpenMM::Vec3>& forces, std::vector<double>& masses, std::vector<OpenMM::Vec3>& xPrime);
 };
 

platforms/reference/src/SimTKReference/ReferenceBAOABDynamics.cpp

@@ -87,17 +87,21 @@ void ReferenceBAOABDynamics::updatePart2(int numberOfAtoms, vector<Vec3>& atomCo
     }
 }
 
-void ReferenceBAOABDynamics::updatePart3(int numberOfAtoms, vector<Vec3>& atomCoordinates,
+void ReferenceBAOABDynamics::updatePart3(OpenMM::ContextImpl& context, int numberOfAtoms, vector<Vec3>& atomCoordinates,
                                               vector<Vec3>& velocities, vector<Vec3>& forces, vector<double>& inverseMasses,
                                               vector<Vec3>& xPrime) {
     const double halfdt = 0.5*getDeltaT();
     for (int i = 0; i < numberOfAtoms; i++) {
         if (inverseMasses[i] != 0.0) {
             velocities[i] += (xPrime[i]-oldx[i])/halfdt;
-            velocities[i] += (halfdt*inverseMasses[i])*forces[i];
             atomCoordinates[i] = xPrime[i];
         }
     }
+    context.calcForcesAndEnergy(true, false);
+    for (int i = 0; i < numberOfAtoms; i++) {
+        if (inverseMasses[i] != 0.0)
+            velocities[i] += (halfdt*inverseMasses[i])*forces[i];
+    }
 }
 
 void ReferenceBAOABDynamics::update(ContextImpl& context, vector<Vec3>& atomCoordinates,
@@ -135,8 +139,7 @@ void ReferenceBAOABDynamics::update(ContextImpl& context, vector<Vec3>& atomCoor
 
     // 3rd update
 
-    context.calcForcesAndEnergy(true, false);
-    updatePart3(numberOfAtoms, atomCoordinates, velocities, forces, inverseMasses, xPrime);
+    updatePart3(context, numberOfAtoms, atomCoordinates, velocities, forces, inverseMasses, xPrime);
     if (referenceConstraintAlgorithm)
         referenceConstraintAlgorithm->applyToVelocities(atomCoordinates, velocities, inverseMasses, tolerance);