Fixed bugs in AndersenThermostat and CMMotionRemover

platforms/cuda/src/CudaKernels.cpp

@@ -607,7 +607,7 @@ void CudaApplyAndersenThermostatKernel::execute(OpenMMContextImpl& context) {
     if (temperature != prevTemp || frequency != prevFrequency || stepSize != prevStepSize) {
         // Initialize the GPU parameters.
         
-        gpuSetAndersenThermostatParameters(gpu, (float) temperature, (float) (frequency*stepSize));
+        gpuSetAndersenThermostatParameters(gpu, (float) temperature, frequency);
         gpuSetConstants(gpu);
         kGenerateRandoms(gpu);
         prevTemp = temperature;

platforms/cuda/src/kernels/cudatypes.h

@@ -336,7 +336,7 @@ struct cudaGmxSimulation {
     float           fix1;                           // Molecular dynamics fix1 constant
     float           oneOverFix1;                    // Molecular dynamics reciprocal of fix1 constant
     float           DOverTauC;                      // Molecular dynamics DOverTauC constant
-    float           collisionProbability;           // Collision probability for Andersen thermostat
+    float           collisionFrequency;             // Collision frequency for Andersen thermostat
     float2*         pObcData;                       // Pointer to fixed Born data
     float2*         pAttr;                          // Pointer to additional atom attributes (sig, eps)
     float2*         pEwaldCosSinSum;                // Pointer to the cos/sin sums (ewald)

platforms/cuda/src/kernels/gpu.cpp

@@ -1351,6 +1351,9 @@ void gpuSetIntegrationParameters(gpuContext gpu, float tau, float deltaT, float
     gpu->sim.X                      = gpu->sim.tau * sqrt(gpu->sim.kT * gpu->sim.C);
     gpu->sim.Yv                     = sqrt(gpu->sim.kT * gpu->sim.B / gpu->sim.C);
     gpu->sim.Yx                     = gpu->sim.tau * sqrt(gpu->sim.kT * gpu->sim.B / (1.0f - gpu->sim.EM));
+    gpu->psStepSize->Download();
+    (*gpu->psStepSize)[0].y = deltaT;
+    gpu->psStepSize->Upload();
 }
 
 extern "C"
@@ -1372,13 +1375,16 @@ void gpuSetBrownianIntegrationParameters(gpuContext gpu, float tau, float deltaT
     gpu->sim.T                      = temperature;
     gpu->sim.kT                     = BOLTZ * gpu->sim.T;
     gpu->sim.Yv = gpu->sim.Yx       = sqrt(2.0f*gpu->sim.kT*deltaT*tau);
+    gpu->psStepSize->Download();
+    (*gpu->psStepSize)[0].y = deltaT;
+    gpu->psStepSize->Upload();
 }
 
 extern "C"
-void gpuSetAndersenThermostatParameters(gpuContext gpu, float temperature, float collisionProbability) {
+void gpuSetAndersenThermostatParameters(gpuContext gpu, float temperature, float collisionFrequency) {
     gpu->sim.T                      = temperature;
     gpu->sim.kT                     = BOLTZ * gpu->sim.T;
-    gpu->sim.collisionProbability   = collisionProbability;
+    gpu->sim.collisionFrequency     = collisionFrequency;
     gpu->sim.Yv = gpu->sim.Yx       = 1.0f;
     gpu->sim.V = gpu->sim.X         = 1.0f;
 }

platforms/cuda/src/kernels/gputypes.h

@@ -216,7 +216,7 @@ extern "C"
 void gpuSetBrownianIntegrationParameters(gpuContext gpu, float tau, float deltaT, float temperature);
 
 extern "C"
-void gpuSetAndersenThermostatParameters(gpuContext gpu, float temperature, float collisionProbability);
+void gpuSetAndersenThermostatParameters(gpuContext gpu, float temperature, float collisionFrequency);
 
 extern "C"
 void gpuShutDown(gpuContext gpu);

platforms/cuda/src/kernels/kCalculateAndersenThermostat.cu

@@ -57,11 +57,12 @@ __global__ void kCalculateAndersenThermostat_kernel()
     unsigned int rpos           = cSim.pRandomPosition[blockIdx.x];
     __syncthreads();
 
+    float collisionProbability = 1.0f-exp(-cSim.collisionFrequency*cSim.pStepSize[0].y);
     while (pos < cSim.atoms)
     {
         float4 velocity         = cSim.pVelm4[pos];
         float4 random4a         = cSim.pRandom4a[rpos + pos];
-        float scale = (random4a.w < cSim.collisionProbability ? 0.0 : 1.0);
+        float scale = (random4a.w < collisionProbability ? 0.0 : 1.0);
         float add = (1.0-scale)*sqrt(cSim.kT*velocity.w);
         velocity.x = scale*velocity.x + add*random4a.x;
         velocity.y = scale*velocity.y + add*random4a.y;

platforms/reference/src/ReferenceKernels.cpp

@@ -967,20 +967,22 @@ void ReferenceRemoveCMMotionKernel::execute(OpenMMContextImpl& context) {
     // Calculate the center of mass momentum.
     
     RealOpenMM momentum[] = {0.0, 0.0, 0.0};
+    RealOpenMM mass = 0.0;
     for (size_t i = 0; i < masses.size(); ++i) {
         momentum[0] += static_cast<RealOpenMM>( masses[i]*velData[i][0] );
         momentum[1] += static_cast<RealOpenMM>( masses[i]*velData[i][1] );
         momentum[2] += static_cast<RealOpenMM>( masses[i]*velData[i][2] );
+        mass += masses[i];
     }
     
     // Adjust the particle velocities.
     
-    momentum[0] /= static_cast<RealOpenMM>( masses.size() );
-    momentum[1] /= static_cast<RealOpenMM>( masses.size() );
-    momentum[2] /= static_cast<RealOpenMM>( masses.size() );
+    momentum[0] /= mass;
+    momentum[1] /= mass;
+    momentum[2] /= mass;
     for (size_t i = 0; i < masses.size(); ++i) {
-        velData[i][0] -= static_cast<RealOpenMM>( momentum[0]/masses[i] );
-        velData[i][1] -= static_cast<RealOpenMM>( momentum[1]/masses[i] );
-        velData[i][2] -= static_cast<RealOpenMM>( momentum[2]/masses[i] );
+        velData[i][0] -= momentum[0];
+        velData[i][1] -= momentum[1];
+        velData[i][2] -= momentum[2];
     }
 }

platforms/reference/src/SimTKReference/ReferenceAndersenThermostat.cpp

@@ -61,7 +61,7 @@
       void ReferenceAndersenThermostat::applyThermostat( int numberOfAtoms, RealOpenMM** atomVelocities, RealOpenMM* atomMasses,
               RealOpenMM temperature, RealOpenMM collisionFrequency, RealOpenMM stepSize ) const {
           
-          const RealOpenMM collisionProbability = collisionFrequency*stepSize;
+          const RealOpenMM collisionProbability = 1.0f - EXP(-collisionFrequency*stepSize);
           for (int i = 0; i < numberOfAtoms; ++i) {
               if (SimTKOpenMMUtilities::getUniformlyDistributedRandomNumber() < collisionProbability) {