Fix reference pair integrator, implement hardwall constraint with test

openmmapi/include/openmm/NoseHooverIntegrator.h

@@ -38,6 +38,8 @@
 #include "NoseHooverChain.h"
 #include "internal/windowsExport.h"
 
+#include <tuple>
+
 namespace OpenMM {
 
 class System;
@@ -211,6 +213,24 @@ public:
     bool hasSubsystemThermostats() const {
         return hasSubsystemThermostats_;
     }
+    /**
+     * Gets the maximum distance (in nm) that a connected pair may stray from each other. If zero, there are no
+     * constraints on the intra-pair separation.
+     */
+    double getMaximumPairDistance() const { return maxPairDistance_; }
+    /**
+     * Sets the maximum distance (in nm) that a connected pair may stray from each other, implemented using a hard
+     * wall. If set to zero, the hard wall constraint is omited and the pairs are free to be separated by any distance.
+     */
+    void setMaximumPairDistance(double distance) { maxPairDistance_ = distance; }
+    /**
+     * Get a list of all individual atoms (i.e. not involved in a connected Drude-like pair) in the system.
+     */
+    const std::vector<int> & getAllAtoms() const { return allAtoms; }
+    /**
+     * Get a list of all connected Drude-like pairs, and their target relative temperature, in the system.
+     */
+    const std::vector<std::tuple<int, int, double>> & getAllPairs() const { return allPairs; }
 protected:
    /**
      * Advance any Nose-Hoover chains associated with this integrator and determine
@@ -246,9 +266,12 @@ protected:
     virtual double computeKineticEnergy();
 
     std::vector<NoseHooverChain> noseHooverChains;
+    std::vector<int> allAtoms;
+    std::vector<std::tuple<int, int, double>> allPairs;
     bool forcesAreValid;
     Kernel vvKernel, nhcKernel;
     bool hasSubsystemThermostats_;
+    double maxPairDistance_;
 };
 
 } // namespace OpenMM

openmmapi/src/NoseHooverIntegrator.cpp

@@ -108,10 +108,21 @@ const NoseHooverChain& NoseHooverIntegrator::getNoseHooverThermostat(int chainID
 
 void NoseHooverIntegrator::initializeThermostats(const System &system) {
 
+    allAtoms.clear();
+    allPairs.clear();
+    std::set<int> allAtomsSet;
+    for (int atom = 0; atom < system.getNumParticles(); ++atom) allAtomsSet.insert(atom);
+
     for (auto &thermostat : noseHooverChains) {
         const auto &thermostatedParticles = thermostat.getThermostatedAtoms();
         const auto &thermostatedPairs = thermostat.getThermostatedPairs();
 
+        for( auto& pair : thermostatedPairs ) {
+            allAtomsSet.erase(pair.first);
+            allAtomsSet.erase(pair.second);
+            allPairs.emplace_back(pair.first, pair.second, thermostat.getDefaultRelativeTemperature());
+        }
+
         // figure out the number of DOFs
         int nDOF = 3*(thermostatedParticles.size() + thermostatedPairs.size());
         for (int constraintNum = 0; constraintNum < system.getNumConstraints(); constraintNum++) {
@@ -207,6 +218,8 @@ void NoseHooverIntegrator::initializeThermostats(const System &system) {
             }
         }
     }
+
+    for(const auto& atom : allAtomsSet) allAtoms.push_back(atom);
 }
 
 double NoseHooverIntegrator::getTemperature(int chainID) const {
@@ -278,6 +291,7 @@ void NoseHooverIntegrator::initialize(ContextImpl& contextRef) {
         throw OpenMMException("This Integrator is already bound to a context");
 
     context = &contextRef;
+    const System& system = context->getSystem();
     owner = &contextRef.getOwner();
     vvKernel = context->getPlatform().createKernel(IntegrateVelocityVerletStepKernel::Name(), contextRef);
     vvKernel.getAs<IntegrateVelocityVerletStepKernel>().initialize(contextRef.getSystem(), *this);
@@ -286,7 +300,6 @@ void NoseHooverIntegrator::initialize(ContextImpl& contextRef) {
     forcesAreValid = false;
 
     // check for drude particles and build the Nose-Hoover Chains
-    const System& system = context->getSystem();
     for (auto& thermostat: noseHooverChains){
         // if there are no thermostated particles or pairs in the lists this is a regular thermostat for the whole (non-Drude) system
         if ( (thermostat.getThermostatedAtoms().size() == 0) && (thermostat.getThermostatedPairs().size() == 0) ){

platforms/reference/include/ReferenceVelocityVerletDynamics.h

@@ -71,11 +71,15 @@ class ReferenceVelocityVerletDynamics : public ReferenceDynamics {
          @param masses              atom masses
          @param tolerance           the constraint tolerance
          @param forcesAreValid      whether the forces are valid (duh!)
+         @param allAtoms            a list of all atoms not involved in a Drude-like pair
+         @param allPairs            a list of all Drude-like pairs, and their KT values, in the system
+         @param maxPairDistance     the maximum separation allowed for a Drude-like pair
       
          --------------------------------------------------------------------------------------- */
      
       void update(OpenMM::ContextImpl &context, const OpenMM::System& system, std::vector<OpenMM::Vec3>& atomCoordinates,
-                  std::vector<OpenMM::Vec3>& velocities, std::vector<OpenMM::Vec3>& forces, std::vector<double>& masses, double tolerance, bool &forcesAreValid);
+                  std::vector<OpenMM::Vec3>& velocities, std::vector<OpenMM::Vec3>& forces, std::vector<double>& masses, double tolerance, bool &forcesAreValid,
+                  const std::vector<int> & allAtoms, const std::vector<std::tuple<int, int, double>> & allPairs, double maxPairDistance);
       
 };
 

platforms/reference/src/ReferenceKernels.cpp

@@ -2170,7 +2170,8 @@ void ReferenceIntegrateVelocityVerletStepKernel::execute(ContextImpl& context, c
         dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
         prevStepSize = stepSize;
     }
-    dynamics->update(context, context.getSystem(), posData, velData, forceData, masses, integrator.getConstraintTolerance(), forcesAreValid);
+    dynamics->update(context, context.getSystem(), posData, velData, forceData, masses, integrator.getConstraintTolerance(), forcesAreValid,
+                     integrator.getAllAtoms(), integrator.getAllPairs(), integrator.getMaximumPairDistance());
     data.time += stepSize;
     data.stepCount++;
 }

platforms/reference/src/SimTKReference/ReferenceVelocityVerletDynamics.cpp

@@ -26,6 +26,7 @@
 #include <iostream>
 #include <sstream>
 
+#include "openmm/OpenMMException.h"
 #include "SimTKOpenMMUtilities.h"
 #include "openmm/internal/ContextImpl.h"
 #include "ReferenceVelocityVerletDynamics.h"
@@ -72,12 +73,17 @@ ReferenceVelocityVerletDynamics::~ReferenceVelocityVerletDynamics() {
    @param velocities          velocities
    @param forces              forces
    @param masses              atom masses
+   @param atomList            list of all atoms not involved in a Drude-like pair
+   @param pairList            list of all Drude-like pairs
+   @param maxPairDistance     the maximum separation of any Drude-like pairs
 
    --------------------------------------------------------------------------------------- */
 
 void ReferenceVelocityVerletDynamics::update(OpenMM::ContextImpl &context, const OpenMM::System& system, vector<Vec3>& atomCoordinates,
                                           vector<Vec3>& velocities,
-                                          vector<Vec3>& forces, vector<double>& masses, double tolerance, bool &forcesAreValid) {
+                                          vector<Vec3>& forces, vector<double>& masses, double tolerance, bool &forcesAreValid,
+                                          const std::vector<int> & atomList, const std::vector<std::tuple<int, int, double>> &pairList,
+                                          double maxPairDistance) {
 
     // first-time-through initialization
     if (!forcesAreValid) context.calcForcesAndEnergy(true, false);
@@ -94,23 +100,125 @@ void ReferenceVelocityVerletDynamics::update(OpenMM::ContextImpl &context, const
        }
     }
 
-    // Perform the integration.
+    //// Perform the integration.
 
-    for (int i = 0; i < numberOfAtoms; ++i) {
-        if (masses[i] != 0.0)
-            for (int j = 0; j < 3; ++j) {
-                velocities[i][j] += 0.5 * inverseMasses[i]*forces[i][j]*getDeltaT();
-                xPrime[i][j] = atomCoordinates[i][j];
-                atomCoordinates[i][j] += velocities[i][j]*getDeltaT();
+    // Regular atoms
+    for (const auto &atom : atomList) {
+        if (masses[atom] != 0.0) {
+            for (int xyz = 0; xyz < 3; ++xyz) {
+                velocities[atom][xyz] += 0.5 * inverseMasses[atom]*forces[atom][xyz]*getDeltaT();
+                xPrime[atom][xyz] = atomCoordinates[atom][xyz];
+                atomCoordinates[atom][xyz] += velocities[atom][xyz]*getDeltaT();
+            }
+        }
+    }
+    // Connected particles
+    for (const auto &pair : pairList) {
+        const auto &atom1 = std::get<0>(pair);
+        const auto &atom2 = std::get<1>(pair);
+        double m1 = system.getParticleMass(atom1);
+        double m2 = system.getParticleMass(atom2);
+        double mass1fract = m1 / (m1 + m2);
+        double mass2fract = m2 / (m1 + m2);
+        double invRedMass = (m1 * m2 != 0.0) ? (m1 + m2)/(m1 * m2) : 0.0;
+        double invTotMass = (m1 + m2 != 0.0) ? 1.0 /(m1 + m2) : 0.0;
+        Vec3 comVel = velocities[atom1]*mass1fract + velocities[atom2]*mass2fract;
+        Vec3 relVel = velocities[atom2] - velocities[atom1];
+        Vec3 comForce = forces[atom1] + forces[atom2];
+        Vec3 relForce = mass1fract*forces[atom2] - mass2fract*forces[atom1];
+        comVel += 0.5 * comForce * getDeltaT() * invTotMass;
+        relVel += 0.5 * relForce * getDeltaT() * invRedMass; 
+        if (m1 != 0.0) {
+            for (int xyz = 0; xyz < 3; ++xyz) {
+                velocities[atom1][xyz] = comVel[xyz] - relVel[xyz]*mass2fract;
+                xPrime[atom1][xyz] = atomCoordinates[atom1][xyz];
+                atomCoordinates[atom1][xyz] += velocities[atom1][xyz]*getDeltaT();
+            }
+        }
+        if (m2 != 0.0) {
+            for (int xyz = 0; xyz < 3; ++xyz) {
+                velocities[atom2][xyz] = comVel[xyz] + relVel[xyz]*mass1fract;
+                xPrime[atom2][xyz] = atomCoordinates[atom2][xyz];
+                atomCoordinates[atom2][xyz] += velocities[atom2][xyz]*getDeltaT();
             }
+        }
     }
-   
+
     // 
 
     ReferenceConstraintAlgorithm* referenceConstraintAlgorithm = getReferenceConstraintAlgorithm();
     if (referenceConstraintAlgorithm)
        referenceConstraintAlgorithm->apply(xPrime, atomCoordinates, inverseMasses, tolerance);
 
+
+    // Apply hard wall constraints.
+    if (maxPairDistance > 0) {
+        for (const auto & pair : pairList) {
+            const int atom1 = std::get<0>(pair);
+            const int atom2 = std::get<1>(pair);
+            const double hardWallScale = sqrt(std::get<2>(pair)*BOLTZ);
+            Vec3 delta = atomCoordinates[atom1]-atomCoordinates[atom2];
+            double r = sqrt(delta.dot(delta));
+            double rInv = 1/r;
+            if (rInv*maxPairDistance < 1.0) {
+                // The constraint has been violated, so make the inter-particle distance "bounce"
+                // off the hard wall.
+                if (rInv*maxPairDistance < 0.5)
+                    throw OpenMMException("Drude particle moved too far beyond hard wall constraint");
+                Vec3 bondDir = delta*rInv;
+                Vec3 vel1 = velocities[atom1];
+                Vec3 vel2 = velocities[atom2];
+                double m1 = masses[atom1];
+                double m2 = masses[atom2];
+                double invTotMass = (m1 + m2 != 0.0) ? 1.0 /(m1 + m2) : 0.0;
+                double deltaR = r-maxPairDistance;
+                double deltaT = getDeltaT();
+                double dt = getDeltaT();
+
+                double dotvr1 = vel1.dot(bondDir);
+                Vec3 vb1 = bondDir*dotvr1;
+                Vec3 vp1 = vel1-vb1;
+                if (m2 == 0) {
+                    // The parent particle is massless, so move only the Drude particle.
+
+                    if (dotvr1 != 0.0)
+                        deltaT = deltaR/std::abs(dotvr1);
+                    if (deltaT > getDeltaT())
+                        deltaT = getDeltaT();
+                    dotvr1 = -dotvr1*hardWallScale/(std::abs(dotvr1)*sqrt(m1));
+                    double dr = -deltaR + deltaT*dotvr1;
+                    atomCoordinates[atom1] += bondDir*dr;
+                    velocities[atom1] = vp1 + bondDir*dotvr1;
+                }
+                else {
+                    // Move both particles.
+                    double dotvr2 = vel2.dot(bondDir);
+                    Vec3 vb2 = bondDir*dotvr2;
+                    Vec3 vp2 = vel2-vb2;
+                    double vbCMass = (m1*dotvr1 + m2*dotvr2)*invTotMass;
+                    dotvr1 -= vbCMass;
+                    dotvr2 -= vbCMass;
+                    if (dotvr1 != dotvr2)
+                        deltaT = deltaR/std::abs(dotvr1-dotvr2);
+                    if (deltaT > dt)
+                        deltaT = dt;
+                    double vBond = hardWallScale/sqrt(m1);
+                    dotvr1 = -dotvr1*vBond*m2*invTotMass/std::abs(dotvr1);
+                    dotvr2 = -dotvr2*vBond*m1*invTotMass/std::abs(dotvr2);
+                    double dr1 = -deltaR*m2*invTotMass + deltaT*dotvr1;
+                    double dr2 = deltaR*m1*invTotMass + deltaT*dotvr2;
+                    dotvr1 += vbCMass;
+                    dotvr2 += vbCMass;
+                    atomCoordinates[atom1] += bondDir*dr1;
+                    atomCoordinates[atom2] += bondDir*dr2;
+                    velocities[atom1] = vp1 + bondDir*dotvr1;
+                    velocities[atom2] = vp2 + bondDir*dotvr2;
+                }
+            }
+        }
+    } /* end of hard wall constraint part */
+
+
     ReferenceVirtualSites::computePositions(system, atomCoordinates);
     context.calcForcesAndEnergy(true, false);
     forcesAreValid = true;
@@ -119,16 +227,44 @@ void ReferenceVelocityVerletDynamics::update(OpenMM::ContextImpl &context, const
         if (masses[i] != 0.0)
             for (int j = 0; j < 3; ++j) {
                 xPrime[i][j] += velocities[i][j]*getDeltaT();
-            }
-    }
+        }
+   } 
 
     // Update the positions and velocities.
-
-    for (int i = 0; i < numberOfAtoms; ++i) {
-        if (masses[i] != 0.0)
-            for (int j = 0; j < 3; ++j) {
-                velocities[i][j] += 0.5*inverseMasses[i]*forces[i][j]*getDeltaT() + (atomCoordinates[i][j] - xPrime[i][j]) / getDeltaT();
+    // Regular atoms
+    for (const auto &atom : atomList) {
+        if (masses[atom] != 0.0) {
+            for (int xyz = 0; xyz < 3; ++xyz) {
+                velocities[atom][xyz] += 0.5 * inverseMasses[atom]*forces[atom][xyz]*getDeltaT() + (atomCoordinates[atom][xyz] - xPrime[atom][xyz])/getDeltaT();
+            }
+        }
+    }
+    // Connected particles
+    for (const auto &pair : pairList) {
+        const auto &atom1 = std::get<0>(pair);
+        const auto &atom2 = std::get<1>(pair);
+        double m1 = system.getParticleMass(atom1);
+        double m2 = system.getParticleMass(atom2);
+        double mass1fract = m1 / (m1 + m2);
+        double mass2fract = m2 / (m1 + m2);
+        double invRedMass = (m1 * m2 != 0.0) ? (m1 + m2)/(m1 * m2) : 0.0;
+        double invTotMass = (m1 + m2 != 0.0) ? 1.0 /(m1 + m2) : 0.0;
+        Vec3 comVel = velocities[atom1]*mass1fract + velocities[atom2]*mass2fract;
+        Vec3 relVel = velocities[atom2] - velocities[atom1];
+        Vec3 comForce = forces[atom1] + forces[atom2];
+        Vec3 relForce = mass1fract*forces[atom2] - mass2fract*forces[atom1];
+        comVel += 0.5 * comForce * getDeltaT() * invTotMass;
+        relVel += 0.5 * relForce * getDeltaT() * invRedMass; 
+        if (m1 != 0.0) {
+            for (int xyz = 0; xyz < 3; ++xyz) {
+                velocities[atom1][xyz] = comVel[xyz] - relVel[xyz]*mass2fract + (atomCoordinates[atom1][xyz] - xPrime[atom1][xyz])/getDeltaT();
+            }
+        }
+        if (m2 != 0.0) {
+            for (int xyz = 0; xyz < 3; ++xyz) {
+                velocities[atom2][xyz] = comVel[xyz] + relVel[xyz]*mass1fract + (atomCoordinates[atom2][xyz] - xPrime[atom2][xyz])/getDeltaT();
             }
+        }
     }
     if (referenceConstraintAlgorithm)
        referenceConstraintAlgorithm->applyToVelocities(atomCoordinates, velocities, inverseMasses, tolerance);

plugins/drude/openmmapi/include/openmm/DrudeNoseHooverIntegrator.h

@@ -75,6 +75,16 @@ public:
      * It will also get called again if the application calls reinitialize() on the Context.
      */
     void initialize(ContextImpl& context);
+    /**
+     * Get the maximum distance a Drude particle can ever move from its parent particle, measured in nm.  This is implemented
+     * with a hard wall constraint.  If this distance is set to 0 (the default), the hard wall constraint is omitted.
+     */
+    double getMaxDrudeDistance() const;
+    /**
+     * Set the maximum distance a Drude particle can ever move from its parent particle, measured in nm.  This is implemented
+     * with a hard wall constraint.  If this distance is set to 0 (the default), the hard wall constraint is omitted.
+     */
+    void setMaxDrudeDistance(double distance);
     /**
      * Compute the kinetic energy of the drude particles at the current time.
      */

plugins/drude/openmmapi/src/DrudeNoseHooverIntegrator.cpp

@@ -51,6 +51,7 @@ DrudeNoseHooverIntegrator::DrudeNoseHooverIntegrator(double temperature, double
                                                      double stepSize, int chainLength, int numMTS, int numYoshidaSuzuki) :
     NoseHooverIntegrator(stepSize) {
 
+    setMaxDrudeDistance(0);
     hasSubsystemThermostats_ = false;
     addSubsystemThermostat(std::vector<int>(), std::vector<std::pair<int, int>>(), temperature,
                            collisionFrequency, drudeTemperature, drudeCollisionFrequency,
@@ -59,6 +60,15 @@ DrudeNoseHooverIntegrator::DrudeNoseHooverIntegrator(double temperature, double
 
 DrudeNoseHooverIntegrator::~DrudeNoseHooverIntegrator() { }
 
+double DrudeNoseHooverIntegrator::getMaxDrudeDistance() const {
+    return maxPairDistance_;
+}
+
+void DrudeNoseHooverIntegrator::setMaxDrudeDistance(double distance) {
+    if (distance < 0)
+        throw OpenMMException("setMaxDrudeDistance: Distance cannot be negative");
+    maxPairDistance_ = distance;
+}
 
 void DrudeNoseHooverIntegrator::initialize(ContextImpl& contextRef) {
 

plugins/drude/tests/TestDrudeNoseHoover.h

@@ -56,16 +56,12 @@ using namespace std;
 extern "C" OPENMM_EXPORT void registerKernelFactories();
 Platform& initializePlatform(int argc, char* argv[]);
 
-void testWaterBox(Platform& platform){
+void build_waterbox(System &system, int gridSize, double polarizability, vector<Vec3> & positions) {
     // Create a box of SWM4-NDP water molecules.  This involves constraints, virtual sites,
     // and Drude particles.
-    const int gridSize = 3;
     const int numMolecules = gridSize*gridSize*gridSize;
     const double spacing = 0.8;
     const double boxSize = spacing*(gridSize+1);
-    const double temperature = 300.0;
-    const double temperatureDrude = 10.0;
-    System system;
     NonbondedForce* nonbonded = new NonbondedForce();
     DrudeForce* drude = new DrudeForce();
     CMMotionRemover* cmm = new CMMotionRemover(1);
@@ -95,11 +91,10 @@ void testWaterBox(Platform& platform){
         system.addConstraint(startIndex, startIndex+3, 0.09572);
         system.addConstraint(startIndex+2, startIndex+3, 0.15139);
         system.setVirtualSite(startIndex+4, new ThreeParticleAverageSite(startIndex, startIndex+2, startIndex+3, 0.786646558, 0.106676721, 0.106676721));
-        drude->addParticle(startIndex+1, startIndex, -1, -1, -1, -1.71636, ONE_4PI_EPS0*1.71636*1.71636/(100000*4.184), 1, 1);
+        drude->addParticle(startIndex+1, startIndex, -1, -1, -1, -1.71636, polarizability, 1, 1);
     }
-    vector<Vec3> positions;
-    for (int i = 0; i < gridSize; i++)
-        for (int j = 0; j < gridSize; j++)
+    for (int i = 0; i < gridSize; i++) {
+        for (int j = 0; j < gridSize; j++) {
             for (int k = 0; k < gridSize; k++) {
                 Vec3 pos(i*spacing, j*spacing, k*spacing);
                 positions.push_back(pos);
@@ -108,17 +103,33 @@ void testWaterBox(Platform& platform){
                 positions.push_back(pos+Vec3(-0.023999, 0.092663, 0));
                 positions.push_back(pos);
             }
-    
+        }
+    }
+}
+
+void testWaterBox(Platform& platform) {
+    // Create a box of SWM4-NDP water molecules.  This involves constraints, virtual sites,
+    // and Drude particles.
+    System system;
+    const int gridSize = 3;
+    vector<Vec3> positions;
+    double polarizability = ONE_4PI_EPS0*1.71636*1.71636/(100000*4.184);
+
+    build_waterbox(system, gridSize, polarizability, positions);
+
+    const int numMolecules = gridSize*gridSize*gridSize;
     int numStandardDof = 3*3*numMolecules - system.getNumConstraints();
     int numDrudeDof = 3*numMolecules;
     int numDof = numStandardDof+numDrudeDof;
+    const double temperature = 300.0;
+    const double temperatureDrude = 10.0;
+
     // Simulate it and check the temperature.
-    
     int chainLength = 4;
     int numMTS = 3;
     int numYS = 3;
-    double frequency = 100.0;
-    double frequencyDrude = 80.0;
+    double frequency = 200.0;
+    double frequencyDrude = 2000.0;
     int randomSeed = 100;
     DrudeNoseHooverIntegrator integ(temperature, frequency, 
                                     temperatureDrude, frequencyDrude, 0.0005, 
@@ -139,12 +150,12 @@ void testWaterBox(Platform& platform){
     context.applyConstraints(1e-6);
     // Equilibrate.
     
-    integ.step(800);
+    integ.step(500);
     
     // Compute the internal and center of mass temperatures.
     
     double totalKE = 0;
-    const int numSteps = 800;
+    const int numSteps = 2000;
     double meanTemp = 0.0;
     double meanDrudeTemp = 0.0;
     double meanConserved = 0.0;
@@ -179,16 +190,112 @@ void testWaterBox(Platform& platform){
         ASSERT(fabs(meanConserved - conserved) < 0.2);
     }
     totalKE /= numSteps;
-    ASSERT_USUALLY_EQUAL_TOL(temperature, meanTemp, 0.03);
-    ASSERT_USUALLY_EQUAL_TOL(temperatureDrude, meanDrudeTemp, 0.04);
+    ASSERT_USUALLY_EQUAL_TOL(temperature, meanTemp, 0.001);
+    ASSERT_USUALLY_EQUAL_TOL(temperatureDrude, meanDrudeTemp, 0.001);
+}
+
+
+double testWaterBoxWithHardWallConstraint(Platform& platform, double hardWallConstraint){
+    // Create a box of SWM4-NDP water molecules.  This involves constraints, virtual sites,
+    // and Drude particles.
+    System system;
+    const int gridSize = 3;
+    vector<Vec3> positions;
+
+    double polarizability = 1e-2;
+    build_waterbox(system, gridSize, polarizability, positions);
+
+    const int numMolecules = gridSize*gridSize*gridSize;
+    int numStandardDof = 3*3*numMolecules - system.getNumConstraints();
+    int numDrudeDof = 3*numMolecules;
+    int numDof = numStandardDof+numDrudeDof;
+    const double temperature = 300.0;
+    const double temperatureDrude = 10.0;
+
+    // Simulate it and check the temperature.
+    int chainLength = 4;
+    int numMTS = 3;
+    int numYS = 3;
+    double frequency = 25.0;
+    double frequencyDrude = 25.0;
+    int randomSeed = 100;
+    DrudeNoseHooverIntegrator integ(temperature, frequency, 
+                                    temperatureDrude, frequencyDrude, 0.0005, 
+                                    chainLength, numMTS, numYS);
+    integ.setMaxDrudeDistance(hardWallConstraint);
+    Context context(system, integ, platform);
+    context.setPositions(positions);
+    context.setVelocitiesToTemperature(temperature, randomSeed);
+    std::vector<Vec3> velocities = context.getState(State::Velocities).getVelocities();
+    for (int i = 0; i < numMolecules; i++){
+        Vec3 noize;
+        for (int j = 0; j < 3; j++){
+            noize[j] = float(((i+18311)*(j+18253) * 313419097822414) % 18313) / float(18313);
+            noize[j] *= sqrt(3 * BOLTZ * temperatureDrude / 0.4);
+        }
+        velocities[5*i+1] = velocities[5*i] + noize;
+    }
+    context.setVelocities(velocities);
+    context.applyConstraints(1e-6);
+    // Equilibrate.
+    
+    integ.step(10);
+    
+    // Compute the internal and center of mass temperatures.
+    
+    double totalKE = 0;
+    const int numSteps = 500;
+    double meanTemp = 0.0;
+    double meanDrudeTemp = 0.0;
+    double meanConserved = 0.0;
+    double maxR = 0.0;
+    for (int i = 0; i < numSteps; i++) {
+        integ.step(1);
+        State state = context.getState(State::Energy | State::Positions);
+        double KE = state.getKineticEnergy();
+        double PE = state.getPotentialEnergy();
+        double fullKE = integ.computeTotalKineticEnergy();
+        double drudeKE = integ.computeDrudeKineticEnergy();
+        double temp = KE/(0.5*numStandardDof*BOLTZ);
+        double drudeTemp = drudeKE/(0.5*numDrudeDof*BOLTZ);
+        meanTemp = (i*meanTemp + temp)/(i+1);
+        meanDrudeTemp = (i*meanDrudeTemp + drudeTemp)/(i+1);
+        double heatBathEnergy = integ.computeHeatBathEnergy();
+        double conserved = PE + fullKE + heatBathEnergy;
+        meanConserved = (i*meanConserved + conserved)/(i+1);
+        const auto & positions = state.getPositions();
+        for(int mol = 0; mol < gridSize*gridSize*gridSize; ++mol) {
+            auto dR = positions[5*mol+1] - positions[5*mol];
+            maxR = std::max(maxR, std::sqrt(dR.dot(dR)));
+        }
+#if DEBUG
+        if(i%1 == 0)
+        std::cout << std::setw(6) << i
+                  << std::setprecision(8) << std::setw(16) << KE
+                  << std::setprecision(8) << std::setw(16) << drudeKE
+                  << std::setprecision(8) << std::setw(16) << meanTemp
+                  << std::setprecision(8) << std::setw(16) << meanDrudeTemp
+                  << std::setprecision(8) << std::setw(16) << heatBathEnergy
+                  << std::setprecision(8) << std::setw(16) << fullKE
+                  << std::setprecision(8) << std::setw(16) << conserved
+                  << std::setprecision(8) << std::setw(16) << maxR
+                  << std::endl;
+#endif
+        totalKE += KE;
+    }
+    totalKE /= numSteps;
+    return maxR;
 }
 
 int main(int argc, char* argv[]) {
     try {
         Platform& platform = initializePlatform(argc, argv);
-        //registerDrudeReferenceKernelFactories();
-        //registerKernelFactories();
         testWaterBox(platform);
+        double maxDrudeDistance = 0.005;
+        double observedDrudeDistance = testWaterBoxWithHardWallConstraint(platform, 0.0);
+        ASSERT(observedDrudeDistance > maxDrudeDistance);
+        observedDrudeDistance = testWaterBoxWithHardWallConstraint(platform, maxDrudeDistance);
+        ASSERT(observedDrudeDistance < maxDrudeDistance);
     }
     catch(const exception& e) {
         cout << "exception: " << e.what() << endl;