AndersenThermostat works correctly with constraints (see bug 1319)

openmmapi/include/openmm/internal/AndersenThermostatImpl.h

@@ -9,7 +9,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2008 Stanford University and the Authors.           *
+ * Portions copyright (c) 2008-2010 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -39,6 +39,8 @@
 
 namespace OpenMM {
 
+class System;
+
 /**
  * This is the internal implementation of AndersenThermostat.
  */
@@ -57,7 +59,13 @@ public:
     }
     std::map<std::string, double> getDefaultParameters();
     std::vector<std::string> getKernelNames();
+    /**
+     * This is a utility routine that computes the groups of particles the thermostat should be
+     * applied to.
+     */
+    static std::vector<std::vector<int> > calcParticleGroups(const System& system);
 private:
+    static void tagParticlesInGroup(int particle, int group, std::vector<int>& particleGroup, std::vector<std::vector<int> >& particleConstraints);
     AndersenThermostat& owner;
     Kernel kernel;
 };

openmmapi/src/AndersenThermostatImpl.cpp

@@ -6,7 +6,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2008 Stanford University and the Authors.           *
+ * Portions copyright (c) 2008-2010 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -63,3 +63,38 @@ std::vector<std::string> AndersenThermostatImpl::getKernelNames() {
     names.push_back(ApplyAndersenThermostatKernel::Name());
     return names;
 }
+
+vector<vector<int> > AndersenThermostatImpl::calcParticleGroups(const System& system) {
+    // First make a list of every other particle to which each particle is connected by a constraint.
+
+    int numParticles = system.getNumParticles();
+    vector<vector<int> > particleConstraints(numParticles);
+    for (int i = 0; i < system.getNumConstraints(); i++) {
+        int particle1, particle2;
+        double distance;
+        system.getConstraintParameters(i, particle1, particle2, distance);
+        particleConstraints[particle1].push_back(particle2);
+        particleConstraints[particle2].push_back(particle1);
+    }
+
+    // Now tag particles by which molecule they belong to.
+
+    vector<int> particleGroup(numParticles, -1);
+    int numGroups = 0;
+    for (int i = 0; i < numParticles; i++)
+        if (particleGroup[i] == -1)
+            tagParticlesInGroup(i, numGroups++, particleGroup, particleConstraints);
+    vector<vector<int> > particleIndices(numGroups);
+    for (int i = 0; i < numParticles; i++)
+        particleIndices[particleGroup[i]].push_back(i);
+    return particleIndices;
+}
+
+void AndersenThermostatImpl::tagParticlesInGroup(int particle, int group, vector<int>& particleGroup, vector<vector<int> >& particleConstraints) {
+    // Recursively tag particles as belonging to a particular group.
+
+    particleGroup[particle] = group;
+    for (int i = 0; i < (int) particleConstraints[particle].size(); i++)
+        if (particleGroup[particleConstraints[particle][i]] == -1)
+            tagParticlesInGroup(particleConstraints[particle][i], group, particleGroup, particleConstraints);
+}

platforms/cuda/src/CudaKernels.cpp

@@ -28,6 +28,7 @@
 #include "openmm/LangevinIntegrator.h"
 #include "openmm/Context.h"
 #include "openmm/OpenMMException.h"
+#include "openmm/internal/AndersenThermostatImpl.h"
 #include "openmm/internal/CMAPTorsionForceImpl.h"
 #include "openmm/internal/ContextImpl.h"
 #include "openmm/internal/NonbondedForceImpl.h"
@@ -1027,6 +1028,8 @@ void CudaIntegrateVariableLangevinStepKernel::execute(ContextImpl& context, cons
 }
 
 CudaApplyAndersenThermostatKernel::~CudaApplyAndersenThermostatKernel() {
+    if (atomGroups != NULL)
+        delete atomGroups;
 }
 
 void CudaApplyAndersenThermostatKernel::initialize(const System& system, const AndersenThermostat& thermostat) {
@@ -1036,6 +1039,16 @@ void CudaApplyAndersenThermostatKernel::initialize(const System& system, const A
     prevTemp = -1.0;
     prevFrequency = -1.0;
     prevStepSize = -1.0;
+
+    // Create the arrays with the group definitions.
+
+    vector<vector<int> > groups = AndersenThermostatImpl::calcParticleGroups(system);
+    atomGroups = new CUDAStream<int>(system.getNumParticles(), 1, "atomGroups");
+    for (int i = 0; i < (int) groups.size(); i++) {
+        for (int j = 0; j < (int) groups[i].size(); j++)
+            (*atomGroups)[groups[i][j]] = i;
+    }
+    atomGroups->Upload();
 }
 
 void CudaApplyAndersenThermostatKernel::execute(ContextImpl& context) {
@@ -1053,7 +1066,7 @@ void CudaApplyAndersenThermostatKernel::execute(ContextImpl& context) {
         prevFrequency = frequency;
         prevStepSize = stepSize;
     }
-    kCalculateAndersenThermostat(gpu);
+    kCalculateAndersenThermostat(gpu, *atomGroups);
 }
 
 CudaApplyMonteCarloBarostatKernel::~CudaApplyMonteCarloBarostatKernel() {

platforms/cuda/src/CudaKernels.h

@@ -737,7 +737,8 @@ private:
  */
 class CudaApplyAndersenThermostatKernel : public ApplyAndersenThermostatKernel {
 public:
-    CudaApplyAndersenThermostatKernel(std::string name, const Platform& platform, CudaPlatform::PlatformData& data) : ApplyAndersenThermostatKernel(name, platform), data(data) {
+    CudaApplyAndersenThermostatKernel(std::string name, const Platform& platform, CudaPlatform::PlatformData& data) : ApplyAndersenThermostatKernel(name, platform),
+            data(data), atomGroups(NULL) {
     }
     ~CudaApplyAndersenThermostatKernel();
     /**
@@ -756,6 +757,7 @@ public:
 private:
     CudaPlatform::PlatformData& data;
     double prevTemp, prevFrequency, prevStepSize;
+    CUDAStream<int>* atomGroups;
 };
 
 /**

platforms/cuda/src/kernels/cudaKernels.h

@@ -52,7 +52,7 @@ extern void kReduceObcGbsaBornForces(gpuContext gpu);
 extern void OPENMMCUDA_EXPORT kCalculateObcGbsaForces2(gpuContext gpu);
 extern void kCalculateGBVIForces2(gpuContext gpu);
 extern void kCalculateLocalForces(gpuContext gpu);
-extern void kCalculateAndersenThermostat(gpuContext gpu);
+extern void kCalculateAndersenThermostat(gpuContext gpu, CUDAStream<int>& atomGroups);
 extern void kReduceBornSumAndForces(gpuContext gpu);
 extern void kApplyFirstShake(gpuContext gpu);
 extern void kApplyFirstCCMA(gpuContext gpu);

platforms/cuda/src/kernels/kCalculateAndersenThermostat.cu

@@ -51,7 +51,7 @@ void GetCalculateAndersenThermostatSim(gpuContext gpu)
     RTERROR(status, "cudaMemcpyFromSymbol: SetSim copy from cSim failed");
 }
 
-__global__ void kCalculateAndersenThermostat_kernel()
+__global__ void kCalculateAndersenThermostat_kernel(int* atomGroups)
 {
     unsigned int pos            = threadIdx.x + blockIdx.x * blockDim.x;
     unsigned int rpos           = cSim.pRandomPosition[blockIdx.x];
@@ -62,15 +62,16 @@ __global__ void kCalculateAndersenThermostat_kernel()
     while (pos < cSim.atoms)
     {
         float4 velocity         = cSim.pVelm4[pos];
-        float4 random4a         = cSim.pRandom4[rpos + pos];
-        float scale = (random4a.w > -randomRange && random4a.w < randomRange ? 0.0f : 1.0f);
+        float4 selectRand       = cSim.pRandom4[rpos + atomGroups[pos]];
+        float4 velRand          = cSim.pRandom4[rpos + pos];
+        float scale = (selectRand.w > -randomRange && selectRand.w < randomRange ? 0.0f : 1.0f);
         float add = (1.0f-scale)*sqrt(cSim.kT*velocity.w);
-        velocity.x = scale*velocity.x + add*random4a.x;
-        velocity.y = scale*velocity.y + add*random4a.y;
-        velocity.z = scale*velocity.z + add*random4a.z;
+        velocity.x = scale*velocity.x + add*velRand.x;
+        velocity.y = scale*velocity.y + add*velRand.y;
+        velocity.z = scale*velocity.z + add*velRand.z;
         cSim.pVelm4[pos]        = velocity;
-         
-        pos                    += blockDim.x * gridDim.x;    
+
+        pos                    += blockDim.x * gridDim.x;
     }
 
     // Update random position pointer
@@ -84,10 +85,10 @@ __global__ void kCalculateAndersenThermostat_kernel()
 }
 
 extern void kGenerateRandoms(gpuContext gpu);
-void kCalculateAndersenThermostat(gpuContext gpu)
+void kCalculateAndersenThermostat(gpuContext gpu, CUDAStream<int>& atomGroups)
 {
 //    printf("kCalculateAndersenThermostat\n");
-    kCalculateAndersenThermostat_kernel<<<gpu->sim.blocks, gpu->sim.update_threads_per_block>>>();
+    kCalculateAndersenThermostat_kernel<<<gpu->sim.blocks, gpu->sim.update_threads_per_block>>>(atomGroups._pDevData);
     LAUNCHERROR("kCalculateAndersenThermostat");
     
     // Update randoms if necessary

platforms/cuda/tests/TestCudaAndersenThermostat.cpp

@@ -52,9 +52,10 @@ void testTemperature() {
     const int numParticles = 8;
     const double temp = 100.0;
     const double collisionFreq = 10.0;
+    const int numSteps = 10000;
     CudaPlatform platform;
     System system;
-    VerletIntegrator integrator(0.01);
+    VerletIntegrator integrator(0.005);
     NonbondedForce* forceField = new NonbondedForce();
     for (int i = 0; i < numParticles; ++i) {
         system.addParticle(2.0);
@@ -68,22 +69,75 @@ void testTemperature() {
     for (int i = 0; i < numParticles; ++i)
         positions[i] = Vec3((i%2 == 0 ? 2 : -2), (i%4 < 2 ? 2 : -2), (i < 4 ? 2 : -2));
     context.setPositions(positions);
-    
+
     // Let it equilibrate.
-    
+
     integrator.step(10000);
-    
+
     // Now run it for a while and see if the temperature is correct.
-    
+
     double ke = 0.0;
-    for (int i = 0; i < 1000; ++i) {
+    for (int i = 0; i < numSteps; ++i) {
         State state = context.getState(State::Energy);
         ke += state.getKineticEnergy();
         integrator.step(1);
     }
-    ke /= 1000;
+    ke /= numSteps;
     double expected = 0.5*numParticles*3*BOLTZ*temp;
-    ASSERT_EQUAL_TOL(expected, ke, 3*expected/std::sqrt(1000.0));
+    ASSERT_EQUAL_TOL(expected, ke, 6/std::sqrt((double) numSteps));
+}
+
+void testConstraints() {
+    const int numParticles = 8;
+    const double temp = 100.0;
+    const double collisionFreq = 10.0;
+    const int numSteps = 10000;
+    CudaPlatform platform;
+    System system;
+    VerletIntegrator integrator(0.005);
+    NonbondedForce* forceField = new NonbondedForce();
+    for (int i = 0; i < numParticles; ++i) {
+        system.addParticle(2.0);
+        forceField->addParticle((i%2 == 0 ? 1.0 : -1.0), 1.0, 5.0);
+    }
+    system.addForce(forceField);
+    system.addConstraint(0, 1, 1);
+    system.addConstraint(1, 2, 1);
+    system.addConstraint(2, 3, 1);
+    system.addConstraint(3, 0, 1);
+    system.addConstraint(4, 5, 1);
+    system.addConstraint(5, 6, 1);
+    system.addConstraint(6, 7, 1);
+    system.addConstraint(7, 4, 1);
+    AndersenThermostat* thermstat = new AndersenThermostat(temp, collisionFreq);
+    system.addForce(thermstat);
+    Context context(system, integrator, platform);
+    vector<Vec3> positions(numParticles);
+    positions[0] = Vec3(0, 0, 0);
+    positions[1] = Vec3(1, 0, 0);
+    positions[2] = Vec3(1, 1, 0);
+    positions[3] = Vec3(0, 1, 0);
+    positions[4] = Vec3(1, 0, 1);
+    positions[5] = Vec3(1, 1, 1);
+    positions[6] = Vec3(0, 1, 1);
+    positions[7] = Vec3(0, 0, 1);
+    context.setPositions(positions);
+
+    // Let it equilibrate.
+
+    integrator.step(10000);
+
+    // Now run it for a while and see if the temperature is correct.
+
+    double ke = 0.0;
+    for (int i = 0; i < numSteps; ++i) {
+        State state = context.getState(State::Energy);
+        ke += state.getKineticEnergy();
+        integrator.step(1);
+    }
+    ke /= numSteps;
+    double expected = 0.5*(numParticles*3-system.getNumConstraints())*BOLTZ*temp;
+    ASSERT_EQUAL_TOL(expected, ke, 6/std::sqrt((double) numSteps));
 }
 
 void testRandomSeed() {
@@ -150,6 +204,7 @@ void testRandomSeed() {
 int main() {
     try {
         testTemperature();
+        testConstraints();
         testRandomSeed();
     }
     catch(const exception& e) {

platforms/opencl/src/OpenCLKernels.cpp

@@ -28,6 +28,7 @@
 #include "OpenCLForceInfo.h"
 #include "openmm/LangevinIntegrator.h"
 #include "openmm/Context.h"
+#include "openmm/internal/AndersenThermostatImpl.h"
 #include "openmm/internal/CMAPTorsionForceImpl.h"
 #include "openmm/internal/ContextImpl.h"
 #include "openmm/internal/CustomHbondForceImpl.h"
@@ -3534,15 +3535,27 @@ void OpenCLIntegrateVariableLangevinStepKernel::execute(ContextImpl& context, co
 }
 
 OpenCLApplyAndersenThermostatKernel::~OpenCLApplyAndersenThermostatKernel() {
+    if (atomGroups != NULL)
+        delete atomGroups;
 }
 
 void OpenCLApplyAndersenThermostatKernel::initialize(const System& system, const AndersenThermostat& thermostat) {
     randomSeed = thermostat.getRandomNumberSeed();
     map<string, string> defines;
     defines["NUM_ATOMS"] = intToString(cl.getNumAtoms());
-    defines["PADDED_NUM_ATOMS"] = intToString(cl.getPaddedNumAtoms());
     cl::Program program = cl.createProgram(OpenCLKernelSources::andersenThermostat, defines);
     kernel = cl::Kernel(program, "applyAndersenThermostat");
+
+    // Create the arrays with the group definitions.
+
+    vector<vector<int> > groups = AndersenThermostatImpl::calcParticleGroups(system);
+    atomGroups = new OpenCLArray<int>(cl, cl.getNumAtoms(), "atomGroups");
+    vector<int> atoms(atomGroups->getSize());
+    for (int i = 0; i < (int) groups.size(); i++) {
+        for (int j = 0; j < (int) groups[i].size(); j++)
+            atoms[groups[i][j]] = i;
+    }
+    atomGroups->upload(atoms);
 }
 
 void OpenCLApplyAndersenThermostatKernel::execute(ContextImpl& context) {
@@ -3552,6 +3565,7 @@ void OpenCLApplyAndersenThermostatKernel::execute(ContextImpl& context) {
         kernel.setArg<cl::Buffer>(2, cl.getVelm().getDeviceBuffer());
         kernel.setArg<cl::Buffer>(3, cl.getIntegrationUtilities().getStepSize().getDeviceBuffer());
         kernel.setArg<cl::Buffer>(4, cl.getIntegrationUtilities().getRandom().getDeviceBuffer());
+        kernel.setArg<cl::Buffer>(6, atomGroups->getDeviceBuffer());
     }
     kernel.setArg<cl_float>(0, (cl_float) context.getParameter(AndersenThermostat::CollisionFrequency()));
     kernel.setArg<cl_float>(1, (cl_float) (BOLTZ*context.getParameter(AndersenThermostat::Temperature())));

platforms/opencl/src/OpenCLKernels.h

@@ -895,7 +895,7 @@ private:
 class OpenCLApplyAndersenThermostatKernel : public ApplyAndersenThermostatKernel {
 public:
     OpenCLApplyAndersenThermostatKernel(std::string name, const Platform& platform, OpenCLContext& cl) : ApplyAndersenThermostatKernel(name, platform), cl(cl),
-            hasInitializedKernels(false) {
+            hasInitializedKernels(false), atomGroups(NULL) {
     }
     ~OpenCLApplyAndersenThermostatKernel();
     /**
@@ -915,6 +915,7 @@ private:
     OpenCLContext& cl;
     bool hasInitializedKernels;
     int randomSeed;
+    OpenCLArray<cl_int>* atomGroups;
     cl::Kernel kernel;
 };
 

platforms/opencl/src/kernels/andersenThermostat.cl

@@ -2,17 +2,17 @@
  * Apply the Andersen thermostat to adjust particle velocities.
  */
 
-__kernel void applyAndersenThermostat(float collisionFrequency, float kT, __global float4* velm, __global float2* stepSize, __global float4* random, unsigned int randomIndex) {
-    randomIndex += get_global_id(0);
+__kernel void applyAndersenThermostat(float collisionFrequency, float kT, __global float4* velm, __global float2* stepSize, __global float4* random,
+        unsigned int randomIndex, __global int* atomGroups) {
     float collisionProbability = 1.0f-exp(-collisionFrequency*stepSize[0].y);
     float randomRange = erf(collisionProbability/sqrt(2.0f));
     for (int index = get_global_id(0); index < NUM_ATOMS; index += get_global_size(0)) {
         float4 velocity = velm[index];
-        float4 rand = random[randomIndex];
-        float scale = (rand.w > -randomRange && rand.w < randomRange ? 0.0f : 1.0f);
+        float4 selectRand = random[randomIndex+atomGroups[index]];
+        float4 velRand = random[randomIndex+index];
+        float scale = (selectRand.w > -randomRange && selectRand.w < randomRange ? 0.0f : 1.0f);
         float add = (1.0f-scale)*sqrt(kT*velocity.w);
-        velocity.xyz = scale*velocity.xyz + add*rand.xyz;
+        velocity.xyz = scale*velocity.xyz + add*velRand.xyz;
         velm[index] = velocity;
-        randomIndex += get_global_size(0);
     }
 }

platforms/opencl/tests/TestOpenCLAndersenThermostat.cpp

@@ -52,9 +52,10 @@ void testTemperature() {
     const int numParticles = 8;
     const double temp = 100.0;
     const double collisionFreq = 10.0;
+    const int numSteps = 10000;
     OpenCLPlatform platform;
     System system;
-    VerletIntegrator integrator(0.01);
+    VerletIntegrator integrator(0.005);
     NonbondedForce* forceField = new NonbondedForce();
     for (int i = 0; i < numParticles; ++i) {
         system.addParticle(2.0);
@@ -76,14 +77,67 @@ void testTemperature() {
     // Now run it for a while and see if the temperature is correct.
 
     double ke = 0.0;
-    for (int i = 0; i < 1000; ++i) {
+    for (int i = 0; i < numSteps; ++i) {
         State state = context.getState(State::Energy);
         ke += state.getKineticEnergy();
         integrator.step(1);
     }
-    ke /= 1000;
+    ke /= numSteps;
     double expected = 0.5*numParticles*3*BOLTZ*temp;
-    ASSERT_EQUAL_TOL(expected, ke, 3*expected/std::sqrt(1000.0));
+    ASSERT_EQUAL_TOL(expected, ke, 6/std::sqrt((double) numSteps));
+}
+
+void testConstraints() {
+    const int numParticles = 8;
+    const double temp = 100.0;
+    const double collisionFreq = 10.0;
+    const int numSteps = 10000;
+    OpenCLPlatform platform;
+    System system;
+    VerletIntegrator integrator(0.005);
+    NonbondedForce* forceField = new NonbondedForce();
+    for (int i = 0; i < numParticles; ++i) {
+        system.addParticle(2.0);
+        forceField->addParticle((i%2 == 0 ? 1.0 : -1.0), 1.0, 5.0);
+    }
+    system.addForce(forceField);
+    system.addConstraint(0, 1, 1);
+    system.addConstraint(1, 2, 1);
+    system.addConstraint(2, 3, 1);
+    system.addConstraint(3, 0, 1);
+    system.addConstraint(4, 5, 1);
+    system.addConstraint(5, 6, 1);
+    system.addConstraint(6, 7, 1);
+    system.addConstraint(7, 4, 1);
+    AndersenThermostat* thermstat = new AndersenThermostat(temp, collisionFreq);
+    system.addForce(thermstat);
+    Context context(system, integrator, platform);
+    vector<Vec3> positions(numParticles);
+    positions[0] = Vec3(0, 0, 0);
+    positions[1] = Vec3(1, 0, 0);
+    positions[2] = Vec3(1, 1, 0);
+    positions[3] = Vec3(0, 1, 0);
+    positions[4] = Vec3(1, 0, 1);
+    positions[5] = Vec3(1, 1, 1);
+    positions[6] = Vec3(0, 1, 1);
+    positions[7] = Vec3(0, 0, 1);
+    context.setPositions(positions);
+
+    // Let it equilibrate.
+
+    integrator.step(10000);
+
+    // Now run it for a while and see if the temperature is correct.
+
+    double ke = 0.0;
+    for (int i = 0; i < numSteps; ++i) {
+        State state = context.getState(State::Energy);
+        ke += state.getKineticEnergy();
+        integrator.step(1);
+    }
+    ke /= numSteps;
+    double expected = 0.5*(numParticles*3-system.getNumConstraints())*BOLTZ*temp;
+    ASSERT_EQUAL_TOL(expected, ke, 6/std::sqrt((double) numSteps));
 }
 
 void testRandomSeed() {
@@ -150,6 +204,7 @@ void testRandomSeed() {
 int main() {
     try {
         testTemperature();
+        testConstraints();
         testRandomSeed();
     }
     catch(const exception& e) {

platforms/reference/src/ReferenceKernels.cpp

@@ -58,6 +58,7 @@
 #include "openmm/CMMotionRemover.h"
 #include "openmm/Context.h"
 #include "openmm/System.h"
+#include "openmm/internal/AndersenThermostatImpl.h"
 #include "openmm/internal/ContextImpl.h"
 #include "openmm/internal/CustomHbondForceImpl.h"
 #include "openmm/internal/CMAPTorsionForceImpl.h"
@@ -90,7 +91,7 @@ static RealOpenMM** allocateRealArray(int length, int width) {
 
 static int** copyToArray(const vector<vector<int> > vec) {
     if (vec.size() == 0)
-        return new int*[0];
+        return new int*[1];
     int** array = allocateIntArray(vec.size(), vec[0].size());
     for (size_t i = 0; i < vec.size(); ++i)
         for (size_t j = 0; j < vec[i].size(); ++j)
@@ -100,7 +101,7 @@ static int** copyToArray(const vector<vector<int> > vec) {
 
 static RealOpenMM** copyToArray(const vector<vector<double> > vec) {
     if (vec.size() == 0)
-        return new RealOpenMM*[0];
+        return new RealOpenMM*[1];
     RealOpenMM** array = allocateRealArray(vec.size(), vec[0].size());
     for (size_t i = 0; i < vec.size(); ++i)
         for (size_t j = 0; j < vec[i].size(); ++j)
@@ -1566,17 +1567,15 @@ void ReferenceApplyAndersenThermostatKernel::initialize(const System& system, co
         masses[i] = static_cast<RealOpenMM>(system.getParticleMass(i));
     this->thermostat = new ReferenceAndersenThermostat();
     SimTKOpenMMUtilities::setRandomNumberSeed((unsigned int) thermostat.getRandomNumberSeed());
+    particleGroups = AndersenThermostatImpl::calcParticleGroups(system);
 }
 
 void ReferenceApplyAndersenThermostatKernel::execute(ContextImpl& context) {
     RealOpenMM** velData = extractVelocities(context);
-    thermostat->applyThermostat(
-			context.getSystem().getNumParticles(),
-			velData, 
-			masses, 
-			static_cast<RealOpenMM>(context.getParameter(AndersenThermostat::Temperature())), 
-			static_cast<RealOpenMM>(context.getParameter(AndersenThermostat::CollisionFrequency())), 
-			static_cast<RealOpenMM>(context.getIntegrator().getStepSize()) );
+    thermostat->applyThermostat(particleGroups, velData, masses,
+        static_cast<RealOpenMM>(context.getParameter(AndersenThermostat::Temperature())),
+        static_cast<RealOpenMM>(context.getParameter(AndersenThermostat::CollisionFrequency())),
+        static_cast<RealOpenMM>(context.getIntegrator().getStepSize()));
 }
 
 ReferenceApplyMonteCarloBarostatKernel::~ReferenceApplyMonteCarloBarostatKernel() {

platforms/reference/src/ReferenceKernels.h

@@ -880,6 +880,7 @@ public:
     void execute(ContextImpl& context);
 private:
     ReferenceAndersenThermostat* thermostat;
+    std::vector<std::vector<int> > particleGroups;
     RealOpenMM* masses;
 };
 

platforms/reference/src/SimTKReference/ReferenceAndersenThermostat.cpp

 
-/* Portions copyright (c) 2008 Stanford University and Simbios.
+/* Portions copyright (c) 2008-2010 Stanford University and Simbios.
  * Contributors: Peter Eastman
  *
  * Permission is hereby granted, free of charge, to any person obtaining
@@ -28,6 +28,8 @@
 #include "../SimTKUtilities/SimTKOpenMMUtilities.h"
 #include "ReferenceAndersenThermostat.h"
 
+using std::vector;
+
       /**---------------------------------------------------------------------------------------
       
          Constructor
@@ -50,7 +52,7 @@
       
          Apply the thermostat at the start of a time step.
       
-         @param numberOfAtoms      number of atoms
+         @param atomGroups         the groups of atoms to apply the thermostat to
          @param atomVelocities     atom velocities
          @param temperature        thermostat temperature in Kelvin
          @param collisionFrequency collision frequency for each atom in fs^-1
@@ -58,19 +60,22 @@
                   
          --------------------------------------------------------------------------------------- */
           
-      void ReferenceAndersenThermostat::applyThermostat( int numberOfAtoms, RealOpenMM** atomVelocities, RealOpenMM* atomMasses,
+      void ReferenceAndersenThermostat::applyThermostat( const vector<vector<int> >& atomGroups, RealOpenMM** atomVelocities, RealOpenMM* atomMasses,
               RealOpenMM temperature, RealOpenMM collisionFrequency, RealOpenMM stepSize ) const {
           
           const RealOpenMM collisionProbability = 1.0f - EXP(-collisionFrequency*stepSize);
-          for (int i = 0; i < numberOfAtoms; ++i) {
+          for (int i = 0; i < (int) atomGroups.size(); ++i) {
               if (SimTKOpenMMUtilities::getUniformlyDistributedRandomNumber() < collisionProbability) {
                   
-                  // A collision occurred, so set the velocity to a new value chosen from a Boltzmann distribution.
-                  
-                  const RealOpenMM velocityScale = static_cast<RealOpenMM>( sqrt(BOLTZ*temperature/atomMasses[i]) );
-                  atomVelocities[i][0] = velocityScale*SimTKOpenMMUtilities::getNormallyDistributedRandomNumber();
-                  atomVelocities[i][1] = velocityScale*SimTKOpenMMUtilities::getNormallyDistributedRandomNumber();
-                  atomVelocities[i][2] = velocityScale*SimTKOpenMMUtilities::getNormallyDistributedRandomNumber();
+                  // A collision occurred, so set the velocities to new values chosen from a Boltzmann distribution.
+
+                  for (int j = 0; j < (int) atomGroups[i].size(); j++) {
+                      int atom = atomGroups[i][j];
+                      const RealOpenMM velocityScale = static_cast<RealOpenMM>(sqrt(BOLTZ*temperature/atomMasses[atom]));
+                      atomVelocities[atom][0] = velocityScale*SimTKOpenMMUtilities::getNormallyDistributedRandomNumber();
+                      atomVelocities[atom][1] = velocityScale*SimTKOpenMMUtilities::getNormallyDistributedRandomNumber();
+                      atomVelocities[atom][2] = velocityScale*SimTKOpenMMUtilities::getNormallyDistributedRandomNumber();
+                  }
               }
           }
           

platforms/reference/src/SimTKReference/ReferenceAndersenThermostat.h

 
-/* Portions copyright (c) 2008 Stanford University and Simbios.
+/* Portions copyright (c) 2008-2010 Stanford University and Simbios.
  * Contributors: Peter Eastman
  *
  * Permission is hereby granted, free of charge, to any person obtaining
@@ -26,6 +26,7 @@
 #define __ReferenceAndersenThermostat_H__
 
 #include "../SimTKUtilities/SimTKOpenMMCommon.h"
+#include <vector>
 
 // ---------------------------------------------------------------------------------------
 
@@ -55,7 +56,7 @@ class ReferenceAndersenThermostat {
       
          Apply the thermostat at the start of a time step.
       
-         @param numberOfAtoms      number of atoms
+         @param atomGroups         the groups of atoms to apply the thermostat to
          @param atomVelocities     atom velocities
          @param atomMasses         atom masses
          @param temperature        thermostat temperature in Kelvin
@@ -64,7 +65,7 @@ class ReferenceAndersenThermostat {
                   
          --------------------------------------------------------------------------------------- */
           
-      void applyThermostat( int numberOfAtoms, RealOpenMM** atomVelocities, RealOpenMM* atomMasses,
+      void applyThermostat( const std::vector<std::vector<int> >& atomGroups, RealOpenMM** atomVelocities, RealOpenMM* atomMasses,
               RealOpenMM temperature, RealOpenMM collisionFrequency, RealOpenMM stepSize ) const;
       
 };

platforms/reference/tests/TestReferenceAndersenThermostat.cpp

@@ -52,9 +52,10 @@ void testTemperature() {
     const int numParticles = 8;
     const double temp = 100.0;
     const double collisionFreq = 10.0;
+    const int numSteps = 10000;
     ReferencePlatform platform;
     System system;
-    VerletIntegrator integrator(0.01);
+    VerletIntegrator integrator(0.005);
     NonbondedForce* forceField = new NonbondedForce();
     for (int i = 0; i < numParticles; ++i) {
         system.addParticle(2.0);
@@ -76,14 +77,67 @@ void testTemperature() {
     // Now run it for a while and see if the temperature is correct.
     
     double ke = 0.0;
-    for (int i = 0; i < 1000; ++i) {
+    for (int i = 0; i < numSteps; ++i) {
         State state = context.getState(State::Energy);
         ke += state.getKineticEnergy();
         integrator.step(1);
     }
-    ke /= 1000;
+    ke /= numSteps;
     double expected = 0.5*numParticles*3*BOLTZ*temp;
-    ASSERT_EQUAL_TOL(expected, ke, 3*expected/std::sqrt(1000.0));
+    ASSERT_EQUAL_TOL(expected, ke, 6/std::sqrt((double) numSteps));
+}
+
+void testConstraints() {
+    const int numParticles = 8;
+    const double temp = 100.0;
+    const double collisionFreq = 10.0;
+    const int numSteps = 10000;
+    ReferencePlatform platform;
+    System system;
+    VerletIntegrator integrator(0.005);
+    NonbondedForce* forceField = new NonbondedForce();
+    for (int i = 0; i < numParticles; ++i) {
+        system.addParticle(2.0);
+        forceField->addParticle((i%2 == 0 ? 1.0 : -1.0), 1.0, 5.0);
+    }
+    system.addForce(forceField);
+    system.addConstraint(0, 1, 1);
+    system.addConstraint(1, 2, 1);
+    system.addConstraint(2, 3, 1);
+    system.addConstraint(3, 0, 1);
+    system.addConstraint(4, 5, 1);
+    system.addConstraint(5, 6, 1);
+    system.addConstraint(6, 7, 1);
+    system.addConstraint(7, 4, 1);
+    AndersenThermostat* thermstat = new AndersenThermostat(temp, collisionFreq);
+    system.addForce(thermstat);
+    Context context(system, integrator, platform);
+    vector<Vec3> positions(numParticles);
+    positions[0] = Vec3(0, 0, 0);
+    positions[1] = Vec3(1, 0, 0);
+    positions[2] = Vec3(1, 1, 0);
+    positions[3] = Vec3(0, 1, 0);
+    positions[4] = Vec3(1, 0, 1);
+    positions[5] = Vec3(1, 1, 1);
+    positions[6] = Vec3(0, 1, 1);
+    positions[7] = Vec3(0, 0, 1);
+    context.setPositions(positions);
+
+    // Let it equilibrate.
+
+    integrator.step(10000);
+
+    // Now run it for a while and see if the temperature is correct.
+
+    double ke = 0.0;
+    for (int i = 0; i < numSteps; ++i) {
+        State state = context.getState(State::Energy);
+        ke += state.getKineticEnergy();
+        integrator.step(1);
+    }
+    ke /= numSteps;
+    double expected = 0.5*(numParticles*3-system.getNumConstraints())*BOLTZ*temp;
+    ASSERT_EQUAL_TOL(expected, ke, 6/std::sqrt((double) numSteps));
 }
 
 void testRandomSeed() {
@@ -150,6 +204,7 @@ void testRandomSeed() {
 int main() {
     try {
         testTemperature();
+        testConstraints();
         testRandomSeed();
     }
     catch(const exception& e) {