Clean up implementation of NHC, add propagator test

olla/include/openmm/kernels.h

@@ -1328,26 +1328,51 @@ public:
  * This kernel is invoked by NoseHooverChainThermostat at the beginning and end of each time step
  * to update the Nose-Hoover chain.
  */
-class PropagateNoseHooverChainKernel : public KernelImpl {
+class NoseHooverChainKernel : public KernelImpl {
 public:
     static std::string Name() {
-        return "PropagateNoseHooverChain";
+        return "NoseHooverChain";
     }
-    PropagateNoseHooverChainKernel(std::string name, const Platform& platform) : KernelImpl(name, platform) {
+    NoseHooverChainKernel(std::string name, const Platform& platform) : KernelImpl(name, platform) {
     }
     /**
      * Initialize the kernel.
-     * 
-     * @param system     the System this kernel will be applied to
-     * @param noseHooverChain the NoseHooverChain this kernel will be used for
      */
-    virtual void initialize(const System& system, const NoseHooverChain& noseHooverChain) = 0;
+    virtual void initialize() = 0;
     /**
-     * Execute the kernel.
+     * Execute the kernel that propagates the Nose Hoover chain and determines the velocity scale factor.
      * 
-     * @param context    the context in which to execute this kernel
+     * @param context  the context in which to execute this kernel
+     * @param noseHooverChain the object describing the chain to be propagated.
+     * @param kineticEnergy the kineticEnergy of the particles being thermostated by this chain.
+     * @param timeStep the time step used by the integrator.
+     * @return the velocity scale factor to apply to the particles associated with this heat bath.
+     */
+    virtual double propagateChain(ContextImpl& context, const NoseHooverChain &noseHooverChain, double kineticEnergy, double timeStep) = 0;
+    /**
+     * Execute the kernal that computes the total (kinetic + potential) heat bath energy.
+     *
+     * @param context the context in which to execute this kernel
+     * @param noseHooverChain the chain whose energy is to be determined.
+     * @return the total heat bath energy.
+     */
+    virtual double computeHeatBathEnergy(ContextImpl& context, const NoseHooverChain &noseHooverChain) = 0;
+    /**
+     * Execute the kernel that computes the kinetic energy for a subset of atoms,
+     * or the relative kinetic energy of Drude particles with respect to their parent atoms
+     *
+     * @param context the context in which to execute this kernel
+     * @param noseHooverChain the chain whose energy is to be determined.
+     */
+    virtual double computeMaskedKineticEnergy(ContextImpl& context, const NoseHooverChain &noseHooverChain) = 0;
+    /**
+     * Execute the kernel that scales the velocities of particles associated with a nose hoover chain
+     *
+     * @param context the context in which to execute this kernel
+     * @param noseHooverChain the chain whose energy is to be determined.
+     * @param scaleFactor the multiplicative factor by which velocities are scaled.
      */
-    virtual void execute(ContextImpl& context, const NoseHooverChain &noseHooverChain, double kineticEnergy, double timeStep) = 0;
+    virtual void scaleVelocities(ContextImpl& context, const NoseHooverChain &noseHooverChain, double scaleFactor) = 0;
 };
 
 /**

openmmapi/include/openmm/NoseHooverChain.h

@@ -35,6 +35,7 @@
 #include "Force.h"
 #include <string>
 #include <vector>
+#include "openmm/OpenMMException.h"
 #include "internal/windowsExport.h"
 
 namespace OpenMM {
@@ -89,18 +90,6 @@ public:
     std::string NumYoshidaSuzukiTimeSteps() const {
         return defaultLabel + "NoseHooverChainNumYoshidaSuzukiTimeSteps";
     }
-    /**
-     * This is the name of the parameter that stores the force acting on the ith bead
-     */
-    std::string Force(int i) const {
-        return defaultLabel + "NoseHooverChainForce" + std::to_string(i);
-    }
-    /**
-     * This is the name of the parameter that stores the mass of the ith bead
-     */
-    std::string Mass(int i) const {
-        return defaultLabel + "NoseHooverChainMass" + std::to_string(i);
-    }
     /**
      * This is the name of the parameter that stores the position of the ith bead
      */
@@ -125,9 +114,12 @@ public:
      * @param defaultNumYoshidaSuzuki   the default number of Yoshida Suzuki steps used to propagate this chain (1, 3, or 5).
      * @param defaultLabel              the default label used to distinguish this Nose-Hoover chain from others that may
      *                                  be used to control a different set of particles, e.g. for Drude oscillators
+     * @param thermostatedAtoms         the list of atoms to be handled by this thermostat
+     * @param parentAtoms               the list of parent atoms, if this thermostat handles Drude particles. Empty vector otherwise.
      */
     NoseHooverChain(double defaultTemperature, double defaultCollisionFrequency, int defaultNumDOFs, int defaultChainLength,
-                    int defaultNumMTS, int defaultNumYoshidaSuzuki, const std::string &defaultLabel);
+                    int defaultNumMTS, int defaultNumYoshidaSuzuki, const std::string &defaultLabel,
+                    const std::vector<int>& thermostatedAtoms, const std::vector<int>& parentAtoms);
     /**
      * Get the default temperature of the heat bath (in Kelvin).
      *
@@ -249,6 +241,44 @@ public:
     void setDefaultLabel(const std::string& label) {
         defaultLabel = label;
     }
+    /**
+     * Get the atom ids of all atoms that are thermostated 
+     *
+     * @returns ids of all atoms that are being handled by this thermostat
+     */
+    const std::vector<int>& getThermostatedAtoms() const {
+        return thermostatedAtoms;
+    }
+    /**
+     * Set list of atoms that are handled by this thermostat
+     *
+     * @param atomIDs 
+     */
+    void setThermostatedAtoms(const std::vector<int>& atomIDs){
+        thermostatedAtoms = atomIDs;
+    }
+    /**
+     * In case this thermostat handles the kinetic energy of Drude particles relative to the
+     * atoms that they are attached to, get the atom ids of all parent atoms.
+     * If this is a regular thermostat, returns an empty vector.
+     *
+     * @returns ids of all parent atoms
+     */
+    const std::vector<int>& getParentAtoms() const {
+        return parentAtoms;
+    }
+    /**
+     * In case this thermostat handles the kinetic energy of Drude particles 
+     * set the atom IDs of all parent atoms. 
+     *
+     * @param parentIDs 
+     */
+    void setParentAtoms(const std::vector<int>& parentIDs){
+        if (not parentIDs.size() == thermostatedAtoms.size()){
+            throw OpenMMException("The number of parent atoms has to be the same as the number of thermostated atoms, or zero.");
+        }
+        parentAtoms = parentIDs;
+    }
     /**
      * Get the default weights used in the Yoshida Suzuki multi time step decomposition (dimensionless) 
      *
@@ -271,6 +301,7 @@ private:
     int defaultNumDOFs, defaultChainLength, defaultNumMTS, defaultNumYS;
     // The suffix used to distinguish NH chains, e.g. for Drude particles vs. regular particles.
     std::string defaultLabel;
+    std::vector<int> thermostatedAtoms, parentAtoms;
 };
 
 } // namespace OpenMM

openmmapi/include/openmm/VelocityVerletIntegrator.h

@@ -67,8 +67,26 @@ public:
      * @param chainID        id of the Nose-Hoover-Chain
      */
     double propagateChain(double kineticEnergy, int chainID=0);
+    /**
+     * Add a Nose-Hoover Chain thermostat to control the temperature of the systeml
+     *
+     * @param system the system to be thermostated.  Note: this must be setup, i.e. all
+     *        particles should have been added, before calling this function.
+     * @param temperature the target temperature for the system.
+     * @param collisionFrequency the frequency of the interaction with the heat bath (in 1/ps).
+     * @param chainLength the number of beads in the Nose-Hoover chain.
+     * @param numMTS the number of step in the  multiple time step chain propagation algorithm.
+     * @param numYoshidaSuzuki the number of terms in the Yoshida-Suzuki multi time step decomposition
+     *        used in the chain propagation algorithm (must be 1, 3, or 5).
+     */
     int addNoseHooverChainThermostat(System& system, double temperature, double collisionFrequency,
                                      int chainLength, int numMTS, int numYoshidaSuzuki);
+    /**
+     * Compute the total (potential + kinetic) heat bath energy for all heat baths
+     * associated with this integrator, at the current time.
+     */
+    double computeHeatBathEnergy();
+
 protected:
     /**
      * This will be called by the Context when it is created.  It informs the Integrator
@@ -90,9 +108,8 @@ protected:
      */
     double computeKineticEnergy();
 private:
-    Kernel vvKernel;
+    Kernel vvKernel, nhcKernel;
     std::vector<NoseHooverChain*> noseHooverChains;
-    std::vector<Kernel> nhcKernels;
 };
 
 } // namespace OpenMM

openmmapi/include/openmm/internal/NoseHooverChainImpl.h

@@ -58,7 +58,6 @@ public:
     std::vector<std::string> getKernelNames();
 private:
     const NoseHooverChain& owner;
-    Kernel kernel;
 };
 
 } // namespace OpenMM

openmmapi/src/ContextImpl.cpp

@@ -52,6 +52,11 @@ using namespace OpenMM;
 using namespace std;
 const static char CHECKPOINT_MAGIC_BYTES[] = "OpenMM Binary Checkpoint\n";
 
+void printvars(std::map<std::string, double> & map) {
+    for(auto & v: map) {
+        std::cout << v.first << "  " << v.second << std::endl;
+    }
+}
 
 ContextImpl::ContextImpl(Context& owner, const System& system, Integrator& integrator, Platform* platform, const map<string, string>& properties, ContextImpl* originalContext) :
         owner(owner), system(system), integrator(integrator), hasInitializedForces(false), hasSetPositions(false), integratorIsDeleted(false),

openmmapi/src/NoseHooverChain.cpp

@@ -37,10 +37,12 @@ using namespace OpenMM;
 
 NoseHooverChain::NoseHooverChain(double defaultTemperature, double defaultCollisionFrequency, 
                                  int defaultNumDOFs, int defaultChainLength, int defaultNumMTS,
-                                 int defaultNumYoshidaSuzuki, const std::string &defaultLabel) :
+                                 int defaultNumYoshidaSuzuki, const std::string &defaultLabel,
+                                 const std::vector<int>& thermostatedAtoms, const std::vector<int>& parentAtoms):
         defaultTemp(defaultTemperature), defaultFreq(defaultCollisionFrequency), defaultNumDOFs(defaultNumDOFs),
-        defaultChainLength(defaultChainLength), defaultNumMTS(defaultNumMTS), defaultNumYS(defaultNumYS),
-        defaultLabel(defaultLabel) {}
+        defaultChainLength(defaultChainLength), defaultNumMTS(defaultNumMTS), defaultNumYS(defaultNumYoshidaSuzuki),
+        defaultLabel(defaultLabel), thermostatedAtoms(thermostatedAtoms), parentAtoms(parentAtoms) 
+{}
 
 ForceImpl* NoseHooverChain::createImpl() const {
     return new NoseHooverChainImpl(*this);

openmmapi/src/NoseHooverChainImpl.cpp

@@ -38,6 +38,7 @@
 #include "sfmt/SFMT.h"
 #include <string>
 #include <vector>
+#include <iostream>
 
 using namespace OpenMM;
 using std::vector;
@@ -46,76 +47,26 @@ NoseHooverChainImpl::NoseHooverChainImpl(const NoseHooverChain& owner) : owner(o
 }
 
 void NoseHooverChainImpl::initialize(ContextImpl& context) {
-    kernel = context.getPlatform().createKernel(PropagateNoseHooverChainKernel::Name(), context);
-    kernel.getAs<PropagateNoseHooverChainKernel>().initialize(context.getSystem(), owner);
 }
 
 void NoseHooverChainImpl::updateContextState(ContextImpl& context, bool& forcesInvalid) {
-    // This kernel updates the Nose-Hoover particles when invoked explicitly
-    // via its propagate(), which is called from the integrator.
 }
 
 double NoseHooverChainImpl::calcForcesAndEnergy(ContextImpl& context, bool includeForces, bool includeEnergy, int groups) {
-    // This kernel doesn't compute an energy or force.  Instead it is invoked
-    // directly by the integrator via it's propagate() method.
     return 0;
 }
 
 std::map<std::string, double> NoseHooverChainImpl::getDefaultParameters() {
     std::map<std::string, double> parameters;
     const auto &owner = getOwner();
-    double frequency = owner.getDefaultCollisionFrequency();
     int chainLength = owner.getDefaultChainLength();
-    double T = owner.getDefaultTemperature();
-    double kT = BOLTZ * T;
-    int DOFs = owner.getDefaultNumDegreesOfFreedom();
-    parameters[owner.Temperature()] = T;
-    parameters[owner.CollisionFrequency()] = frequency;
-    parameters[owner.ChainLength()] = chainLength;
-    parameters[owner.NumYoshidaSuzukiTimeSteps()] = owner.getDefaultNumYoshidaSuzukiTimeSteps();
-    parameters[owner.NumMultiTimeSteps()] = owner.getDefaultNumMultiTimeSteps();
-    parameters[owner.NumDegreesOfFreedom()] = DOFs;
     for(int i = 0; i < chainLength; ++i) {
-        parameters[owner.Force(i)] = 0;
         parameters[owner.Position(i)] = 0;
-        parameters[owner.Mass(i)] = kT / (frequency * frequency);
+        parameters[owner.Velocity(i)] = 0;
     }
-    parameters[owner.Mass(0)] *= DOFs;
-
-    // Set the velocities to the appropriate Boltzmann distribution;
-    // this is copied from Context::setVelocitiesToTemperature()
-    OpenMM_SFMT::SFMT sfmt;
-    int randomSeed = 0; //TODO figure out where / how this should be handled
-    init_gen_rand(randomSeed, sfmt);
-    vector<double> randoms;
-    while (randoms.size() < chainLength) {
-        double x, y, r2;
-        do {
-            x = 2.0*genrand_real2(sfmt)-1.0;
-            y = 2.0*genrand_real2(sfmt)-1.0;
-            r2 = x*x + y*y;
-        } while (r2 >= 1.0 || r2 == 0.0);
-        double multiplier = sqrt((-2.0*log(r2))/r2);
-        randoms.push_back(x*multiplier);
-        randoms.push_back(y*multiplier);
-    }
-    int nextRandom = 0;
-    for (int i = 0; i < chainLength; i++) {
-        double velocityScale = 1 / frequency; // = sqrt(kT / Q) = sqrt(kT / [ kT frequency^-2])
-        parameters[owner.Velocity(i)] = velocityScale * randoms[nextRandom++];
-    }
-    // N.B. the zeroth entry is computed as a function of the instantaneous KE at the start of propagate
-    for(int i = 1; i < chainLength-1; ++i) {
-        const double & v = parameters[owner.Velocity(i)];
-        parameters[owner.Force(i+1)] = (parameters[owner.Mass(i)] * v * v - kT) / parameters[owner.Mass(i+1)];
-    }
-
-
     return parameters;
 }
 
 std::vector<std::string> NoseHooverChainImpl::getKernelNames() {
-    std::vector<std::string> names;
-    names.push_back(PropagateNoseHooverChainKernel::Name());
-    return names;
+    return std::vector<std::string>();
 }

openmmapi/src/VelocityVerletIntegrator.cpp

@@ -39,6 +39,7 @@
 #include "openmm/internal/ContextImpl.h"
 #include "openmm/kernels.h"
 #include <string>
+#include <iostream>
 
 using namespace OpenMM;
 using std::string;
@@ -50,6 +51,7 @@ VelocityVerletIntegrator::VelocityVerletIntegrator(double stepSize) {
 }
 
 double VelocityVerletIntegrator::propagateChain(double kineticEnergy, int chainID) {
+    return nhcKernel.getAs<NoseHooverChainKernel>().propagateChain(*context, *noseHooverChains[chainID], kineticEnergy, getStepSize());
 }
 
 /*int VelocityVerletIntegrator::addMaskedNoseHooverChain(System& system, std::vector<int> mask, std::vector<int> parents, double temperature,
@@ -62,8 +64,13 @@ int VelocityVerletIntegrator::addNoseHooverChainThermostat(System& system, doubl
     std::vector<int> mask, parents;
     int nDOF = 0;
     int numForces = system.getNumForces();
+    vector<int> thermostatedParticles;
+    vector<int> parentParticles;
     for(int particle = 0; particle < system.getNumParticles(); ++particle) {
-        if(system.getParticleMass(particle) > 0) nDOF += 3;
+        if(system.getParticleMass(particle) > 0) {
+            nDOF += 3;
+            thermostatedParticles.push_back(particle);
+        }
     }
     nDOF -= system.getNumConstraints();
     for (int forceNum = 0; forceNum < numForces; ++forceNum) {
@@ -71,15 +78,25 @@ int VelocityVerletIntegrator::addNoseHooverChainThermostat(System& system, doubl
     }
 
     auto nhcForce = new NoseHooverChain(temperature, collisionFrequency, nDOF, chainLength,
-                                        numMTS, numYoshidaSuzuki, std::to_string(noseHooverChains.size()));
+                                        numMTS, numYoshidaSuzuki, std::to_string(noseHooverChains.size()),
+                                        thermostatedParticles, parentParticles);
     system.addForce(nhcForce);
     noseHooverChains.push_back(nhcForce);
+    return noseHooverChains.size() - 1;
 }
 
 double VelocityVerletIntegrator::computeKineticEnergy() {
     return vvKernel.getAs<IntegrateVelocityVerletStepKernel>().computeKineticEnergy(*context, *this);
 }
 
+double VelocityVerletIntegrator::computeHeatBathEnergy() {
+    double energy = 0;
+    for(auto &nhc : noseHooverChains) {
+        energy += nhcKernel.getAs<NoseHooverChainKernel>().computeHeatBathEnergy(*context, *nhc);
+    }
+    return energy;
+}
+
 void VelocityVerletIntegrator::initialize(ContextImpl& contextRef) {
     if (owner != NULL && &contextRef.getOwner() != owner)
         throw OpenMMException("This Integrator is already bound to a context");
@@ -87,20 +104,18 @@ void VelocityVerletIntegrator::initialize(ContextImpl& contextRef) {
     owner = &contextRef.getOwner();
     vvKernel = context->getPlatform().createKernel(IntegrateVelocityVerletStepKernel::Name(), contextRef);
     vvKernel.getAs<IntegrateVelocityVerletStepKernel>().initialize(contextRef.getSystem(), *this);
-    for (auto * nhc: noseHooverChains){
-        nhcKernels.push_back(context->getPlatform().createKernel(PropagateNoseHooverChainKernel::Name(), contextRef));
-        nhcKernels.back().getAs<PropagateNoseHooverChainKernel>().initialize(contextRef.getSystem(), *nhc);
-    }
+    nhcKernel = context->getPlatform().createKernel(NoseHooverChainKernel::Name(), contextRef);
+    nhcKernel.getAs<NoseHooverChainKernel>().initialize();
 }
 
 void VelocityVerletIntegrator::cleanup() {
     vvKernel = Kernel();
-    nhcKernels.clear();
+    nhcKernel = Kernel();
 }
 
 vector<string> VelocityVerletIntegrator::getKernelNames() {
     std::vector<std::string> names;
-    names.push_back(PropagateNoseHooverChainKernel::Name());
+    names.push_back(NoseHooverChainKernel::Name());
     names.push_back(IntegrateVelocityVerletStepKernel::Name());
     return names;
 }
@@ -108,8 +123,23 @@ vector<string> VelocityVerletIntegrator::getKernelNames() {
 void VelocityVerletIntegrator::step(int steps) {
     if (context == NULL)
         throw OpenMMException("This Integrator is not bound to a context!");
+    double scale, kineticEnergy;
     for (int i = 0; i < steps; ++i) {
         context->updateContextState();
+        for(auto &nhc : noseHooverChains) {
+            
+            kineticEnergy = nhcKernel.getAs<NoseHooverChainKernel>().computeMaskedKineticEnergy(*context, *nhc);
+            //std::cout << "ke " << kineticEnergy << std::endl;
+            scale = nhcKernel.getAs<NoseHooverChainKernel>().propagateChain(*context, *nhc, kineticEnergy, getStepSize());
+            //std::cout << "scl " << scale << std::endl;
+            nhcKernel.getAs<NoseHooverChainKernel>().scaleVelocities(*context, *nhc, scale);
+        }
         vvKernel.getAs<IntegrateVelocityVerletStepKernel>().execute(*context, *this);
+        for(auto &nhc : noseHooverChains) {
+            kineticEnergy = nhcKernel.getAs<NoseHooverChainKernel>().computeMaskedKineticEnergy(*context, *nhc);
+            scale = nhcKernel.getAs<NoseHooverChainKernel>().propagateChain(*context, *nhc, kineticEnergy, getStepSize());
+            //std::cout << "ke " << kineticEnergy << " scale " << scale << std::endl;
+            nhcKernel.getAs<NoseHooverChainKernel>().scaleVelocities(*context, *nhc, scale);
+        }
     }
 }

platforms/reference/include/ReferenceKernels.h

@@ -1430,24 +1430,52 @@ private:
  * This kernel is invoked by NoseHooverChain at the start of each time step to adjust the thermostat
  * and update the associated particle velocities.
  */
-class ReferencePropagateNoseHooverChainKernel : public PropagateNoseHooverChainKernel {
+class ReferenceNoseHooverChainKernel : public NoseHooverChainKernel {
 public:
-    ReferencePropagateNoseHooverChainKernel(std::string name, const Platform& platform) : PropagateNoseHooverChainKernel(name, platform), chainPropagator(0) {
+    ReferenceNoseHooverChainKernel(std::string name, const Platform& platform) : NoseHooverChainKernel(name, platform), chainPropagator(0) {
     }
-    ~ReferencePropagateNoseHooverChainKernel();
+    ~ReferenceNoseHooverChainKernel();
     /**
      * Initialize the kernel.
-     * 
-     * @param system     the System this kernel will be applied to
-     * @param thermostat the NoseHooverChain this kernel will be used for
      */
-    void initialize(const System& system, const NoseHooverChain& chain);
+    virtual void initialize();
     /**
-     * Execute the kernel.
+     * Execute the kernel that propagates the Nose Hoover chain and determines the velocity scale factor.
      * 
-     * @param context    the context in which to execute this kernel
+     * @param context  the context in which to execute this kernel
+     * @param noseHooverChain the object describing the chain to be propagated.
+     * @param kineticEnergy the kineticEnergy of the particles being thermostated by this chain.
+     * @param timeStep the time step used by the integrator.
+     * @return the velocity scale factor to apply to the particles associated with this heat bath.
+     */
+    virtual double propagateChain(ContextImpl& context, const NoseHooverChain &nhc, double kineticEnergy, double timeStep);
+    /**
+     * Execute the kernal that computes the total (kinetic + potential) heat bath energy.
+     *
+     * @param context the context in which to execute this kernel
+     * @param noseHooverChain the chain whose energy is to be determined.
+     * @return the total heat bath energy.
+     */
+    virtual double computeHeatBathEnergy(ContextImpl& context, const NoseHooverChain &nhc);
+    /**
+     * Execute the kernel that computes the kinetic energy for a subset of atoms,
+     * or the relative kinetic energy of Drude particles with respect to their parent atoms
+     *
+     * @param context the context in which to execute this kernel
+     * @param noseHooverChain the chain whose energy is to be determined.
+     */
+     virtual double computeMaskedKineticEnergy(ContextImpl& context, const NoseHooverChain &noseHooverChain);
+
+    /**
+     * Execute the kernel that scales the velocities of particles associated with a nose hoover chain
+     *
+     * @param context the context in which to execute this kernel
+     * @param noseHooverChain the chain whose energy is to be determined.
+     * @param scaleFactor the multiplicative factor by which velocities are scaled.
      */
-    void execute(ContextImpl& context, const NoseHooverChain &nhc, double kineticEnergy, double timeStep);
+    virtual void scaleVelocities(ContextImpl& context, const NoseHooverChain &noseHooverChain, double scaleFactor);
+
+
 private:
     ReferenceNoseHooverChain* chainPropagator;
 };

platforms/reference/include/ReferenceNoseHooverChain.h

@@ -59,10 +59,8 @@ class ReferenceNoseHooverChain {
          Propagate the Nose-Hoover chain a half timestep and find the appropriate velocity scaling
 
          @param kineticEnergy      the instantaneous kinetic energy of the particles being thermostated
-         @param chainMasses        the "masses" assigned to each thermostat bead in ps^2 kJ / mol
          @param chainVelocities    the velocities of the chain's beads in nm / ps
          @param chainPositions     the positions of the chains's beads in nm
-         @param chainForces        the forces on each bead in the chain
          @param numDOFs            the number of degrees of freedom in the system that this chain thermostats
          @param temperature        thermostat temperature in Kelvin
          @param collisionFrequency collision frequency for each atom in ps^-1
@@ -70,9 +68,8 @@ class ReferenceNoseHooverChain {
          @param numMTS             number of multi timestep increments used in the Trotter expansion
          @param YSWeights          vector of weights used in the Yoshida-Suzuki multi-timestepping.
          --------------------------------------------------------------------------------------- */
-      double propagate(double kineticEnergy, const vector<double>& chainMasses,
-                       vector<double>& chainVelocities, vector<double>& chainPositions,
-                       vector<double>& chainForces, int numDOFs,
+      double propagate(double kineticEnergy, vector<double>& chainVelocities,
+                       vector<double>& chainPositions, int numDOFs,
                        double temperature, double collisionFrequency, double timeStep,
                        int numMTS, const vector<double>& YSWeights) const;
 };

platforms/reference/src/ReferenceKernelFactory.cpp

@@ -90,8 +90,8 @@ KernelImpl* ReferenceKernelFactory::createKernelImpl(std::string name, const Pla
         return new ReferenceIntegrateVerletStepKernel(name, platform, data);
     if (name == IntegrateVelocityVerletStepKernel::Name())
         return new ReferenceIntegrateVelocityVerletStepKernel(name, platform, data);
-    if (name == PropagateNoseHooverChainKernel::Name())
-        return new ReferencePropagateNoseHooverChainKernel(name, platform);
+    if (name == NoseHooverChainKernel::Name())
+        return new ReferenceNoseHooverChainKernel(name, platform);
     if (name == IntegrateLangevinStepKernel::Name())
         return new ReferenceIntegrateLangevinStepKernel(name, platform, data);
     if (name == IntegrateBrownianStepKernel::Name())

platforms/reference/src/ReferenceKernels.cpp

@@ -2422,39 +2422,118 @@ void ReferenceApplyAndersenThermostatKernel::execute(ContextImpl& context) {
         context.getIntegrator().getStepSize());
 }
 
-ReferencePropagateNoseHooverChainKernel::~ReferencePropagateNoseHooverChainKernel() {
+ReferenceNoseHooverChainKernel::~ReferenceNoseHooverChainKernel() {
     if (chainPropagator)
         delete chainPropagator;
 }
 
-void ReferencePropagateNoseHooverChainKernel::initialize(const System& system, const NoseHooverChain& nhc) {
+void ReferenceNoseHooverChainKernel::initialize() {
     this->chainPropagator = new ReferenceNoseHooverChain();
     //SimTKOpenMMUtilities::setRandomNumberSeed((unsigned int) thermostat.getRandomNumberSeed());
 }
 
-void ReferencePropagateNoseHooverChainKernel::execute(ContextImpl& context, const NoseHooverChain &nhc, double kineticEnergy, double timeStep) {
-    int chainLength = context.getParameter(nhc.ChainLength());
-    std::vector<double> chainPositions(chainLength), chainVelocities(chainLength), chainForces(chainLength),
-                        chainMasses(chainLength);
-    double temperature = context.getParameter(nhc.Temperature());
-    double collisionFrequency = context.getParameter(nhc.CollisionFrequency());
-    int numDOFs = context.getParameter(nhc.NumDegreesOfFreedom());
-    int numMTS = context.getParameter(nhc.NumMultiTimeSteps());
+double ReferenceNoseHooverChainKernel::propagateChain(ContextImpl& context, const NoseHooverChain &nhc, double kineticEnergy, double timeStep) {
+    // Get the variables describing the NHC
+    int chainLength = nhc.getDefaultChainLength();
+    double temperature = nhc.getDefaultTemperature();
+    double collisionFrequency = nhc.getDefaultCollisionFrequency();
+    int numDOFs = nhc.getDefaultNumDegreesOfFreedom();
+    int numMTS = nhc.getDefaultNumMultiTimeSteps();
+
+    // Get the state of the NHC from the context
+    std::vector<double> chainPositions(chainLength), chainVelocities(chainLength);
     for(int i = 0; i < chainLength; ++i) {
         chainPositions[i] = context.getParameter(nhc.Position(i));
         chainVelocities[i] = context.getParameter(nhc.Velocity(i));
-        chainMasses[i] = context.getParameter(nhc.Mass(i));
-        chainForces[i] = context.getParameter(nhc.Force(i));
     }
 
-    chainPropagator->propagate(kineticEnergy, chainMasses, chainVelocities, chainPositions, chainForces, numDOFs,
-                               temperature, collisionFrequency, timeStep, numMTS, nhc.getDefaultYoshidaSuzukiWeights());
-    //chain->applyThermostat(particleGroups, velData, masses,
-    //    context.getParameter(AndersenThermostat::Temperature()),
-    //    context.getParameter(AndersenThermostat::CollisionFrequency()),
-    //    context.getIntegrator().getStepSize());
+    double scale = chainPropagator->propagate(kineticEnergy, chainVelocities, chainPositions, numDOFs,
+                                              temperature, collisionFrequency, timeStep,
+                                              numMTS, nhc.getDefaultYoshidaSuzukiWeights());
+
+    // Update the state of the NHC in the context
+    for(int i = 0; i < chainLength; ++i) {
+        context.setParameter(nhc.Position(i), chainPositions[i]);
+        context.setParameter(nhc.Velocity(i), chainVelocities[i]);
+    }
+
+    return scale;
+}
+
+double ReferenceNoseHooverChainKernel::computeHeatBathEnergy(ContextImpl& context, const NoseHooverChain &nhc) {
+    double energy = 0;
+    int chainLength = nhc.getDefaultChainLength();
+    double temperature = nhc.getDefaultTemperature();
+    double collisionFrequency = nhc.getDefaultCollisionFrequency();
+    double kT = temperature * BOLTZ;
+    int numDOFs = nhc.getDefaultNumDegreesOfFreedom();
+    for(int i = 0; i < chainLength; ++i) {
+         double mass = kT / (collisionFrequency * collisionFrequency);
+         mass *= i ? 1 : numDOFs;
+         double velocity = context.getParameter(nhc.Velocity(i));
+         // The kinetic energy of this bead
+         energy += 0.5 * mass * velocity * velocity;
+         // The potential energy of this bead
+         double position = context.getParameter(nhc.Position(i));
+         energy += numDOFs * kT * position;
+    }
+    return energy;
+}
+
+double ReferenceNoseHooverChainKernel::computeMaskedKineticEnergy(ContextImpl& context, const NoseHooverChain &noseHooverChain) {
+    const std::vector<int>& mask = noseHooverChain.getThermostatedAtoms();
+    const std::vector<int>& parents = noseHooverChain.getParentAtoms();
+    std::vector<Vec3>& velocities = extractVelocities(context);
+    // TODO move these bad boys into the class
+    const System& system = context.getSystem();
+    int numParticles = system.getNumParticles();
+    std::vector<double> masses(numParticles);
+    for (int i = 0; i < numParticles; ++i)
+        masses[i] = system.getParticleMass(i);
+
+    double ke = 0;
+    if (parents.size() == 0){
+        // scale absolute velocities
+        for (auto m: mask){ 
+            ke += 0.5 * masses[m] * velocities[m].dot(velocities[m]);
+        }
+    } else {
+        // scale velocities relative to parent
+        assert(parents.size() == mask.size()); 
+        Vec3 velocity;
+        double mass;
+        for (int i=0; i < mask.size(); i++){
+            velocity = velocities[mask[i]] - velocities[parents[i]];
+            mass = masses[mask[i]] * masses[parents[i]] / (masses[mask[i]] + masses[parents[i]]);
+            ke += 0.5 * mass * velocity.dot(velocity);
+        }
+    }
+    return ke;
+}
+
+
+void ReferenceNoseHooverChainKernel::scaleVelocities(ContextImpl& context, const NoseHooverChain &noseHooverChain, double scaleFactor) {
+    const std::vector<int>& mask = noseHooverChain.getThermostatedAtoms();
+    const std::vector<int>& parents = noseHooverChain.getParentAtoms();
+    std::vector<Vec3>& velocities = extractVelocities(context);
+
+    if (parents.size() == 0){
+        // scale absolute velocities
+        for (auto m: mask){
+            //std::cout << m << " " << velocities[m] << " " << scaleFactor << std::endl;
+            velocities[m] *= scaleFactor;
+        }
+    } else {
+        std::cout << "OUCHH" << std::endl; 
+        // scale velocities relative to parent
+        assert(parents.size() == mask.size());
+        for (int i=0; i < mask.size(); i++){
+            velocities[mask[i]] = velocities[parents[i]] + scaleFactor * (velocities[mask[i]] - velocities[parents[i]]);
+        }
+    }
 }
 
+
 ReferenceApplyMonteCarloBarostatKernel::~ReferenceApplyMonteCarloBarostatKernel() {
     if (barostat)
         delete barostat;

platforms/reference/src/ReferencePlatform.cpp

@@ -67,7 +67,7 @@ ReferencePlatform::ReferencePlatform() {
     registerKernelFactory(CalcGayBerneForceKernel::Name(), factory);
     registerKernelFactory(IntegrateVerletStepKernel::Name(), factory);
     registerKernelFactory(IntegrateVelocityVerletStepKernel::Name(), factory);
-    registerKernelFactory(PropagateNoseHooverChainKernel::Name(), factory);
+    registerKernelFactory(NoseHooverChainKernel::Name(), factory);
     registerKernelFactory(IntegrateLangevinStepKernel::Name(), factory);
     registerKernelFactory(IntegrateBrownianStepKernel::Name(), factory);
     registerKernelFactory(IntegrateVariableLangevinStepKernel::Name(), factory);

platforms/reference/src/SimTKReference/ReferenceNoseHooverChain.cpp

@@ -25,7 +25,9 @@
 #include <cmath>
 #include <string.h>
 #include <sstream>
-
+#include <exception>
+#include <iostream>
+#include <fenv.h>
 #include "SimTKOpenMMUtilities.h"
 #include "ReferenceNoseHooverChain.h"
 
@@ -50,16 +52,22 @@ using namespace OpenMM;
  ReferenceNoseHooverChain::~ReferenceNoseHooverChain() {
  }
 
-double ReferenceNoseHooverChain::propagate(double kineticEnergy, const vector<double>& chainMasses,
-                                           vector<double>& chainVelocities, vector<double>& chainPositions,
-                                           vector<double>& chainForces, int numDOFs,
+double ReferenceNoseHooverChain::propagate(double kineticEnergy, vector<double>& chainVelocities,
+                                           vector<double>& chainPositions, int numDOFs,
                                            double temperature, double collisionFrequency, double timeStep,
                                            int numMTS, const vector<double>& YSWeights) const {
      double scale = 1;
-     double KE2 = 2 * kineticEnergy;
      const double kT = BOLTZ * temperature;
-     const size_t chainLength = chainMasses.size();
+     const size_t chainLength = chainPositions.size();
+     std::vector<double> chainForces(chainLength, 0);
+     std::vector<double> chainMasses(chainLength, kT/(collisionFrequency*collisionFrequency));
+     chainMasses[0] *= numDOFs;
+      feenableexcept(FE_DIVBYZERO | FE_INVALID | FE_OVERFLOW);
+     double KE2 = 2 * kineticEnergy;
      chainForces[0] = (KE2 - numDOFs * kT) / chainMasses[0];
+     for (int bead = 0; bead < chainLength - 1; ++bead) {
+         chainForces[bead + 1] = (chainMasses[bead] * chainVelocities[bead] * chainVelocities[bead] - kT) / chainMasses[bead + 1];
+     }
      for (int mts = 0; mts < numMTS; ++mts) {
          for (const auto &ys : YSWeights) {
              double wdt = ys * timeStep / numMTS;
@@ -83,7 +91,10 @@ double ReferenceNoseHooverChain::propagate(double kineticEnergy, const vector<do
                  chainVelocities[bead] = aa * (aa * chainVelocities[bead] + 0.25 * wdt * chainForces[bead]);
                  chainForces[bead + 1] = (chainMasses[bead] * chainVelocities[bead] * chainVelocities[bead] - kT) / chainMasses[bead + 1];
              }
-             chainVelocities.back() += 0.25 * wdt * chainForces.back();
+             chainVelocities[chainLength-1] += 0.25 * wdt * chainForces.back();
+             //  std::cout << "P " << chainPositions.back() << std::endl;
+             //  std::cout << "V " << chainVelocities.back() << std::endl;
+             //  std::cout << "F " << chainForces.back() << std::endl;
          }  // YS loop
      } // MTS loop
      return scale;

tests/TestNoseHooverThermostat.h

@@ -39,16 +39,26 @@
 #include "SimTKOpenMMRealType.h"
 #include "sfmt/SFMT.h"
 #include <iostream>
+#include <iomanip>
 #include <vector>
 
 using namespace OpenMM;
 using namespace std;
 
 void testSingleBond() {
+    // Check conservation of system + bath energy for a harmonic oscillator
     System system;
     system.addParticle(2.0);
     system.addParticle(2.0);
-    VelocityVerletIntegrator integrator(0.01);
+    VelocityVerletIntegrator integrator(0.001);
+    double temperature = 300; // kelvin
+    double collisionFrequency = 0.1; // 1/ps
+    int numMTS = 3;
+    int numYS = 3;
+    int chainLength = 5;
+    int numDOF = 6;
+    integrator.addNoseHooverChainThermostat(system, temperature, collisionFrequency,
+                                            chainLength, numMTS, numYS);
     HarmonicBondForce* forceField = new HarmonicBondForce();
     forceField->addBond(0, 1, 1.5, 1);
     system.addForce(forceField);
@@ -57,26 +67,25 @@ void testSingleBond() {
     positions[0] = Vec3(-1, 0, 0);
     positions[1] = Vec3(1, 0, 0);
     context.setPositions(positions);
-    
-    // This is simply a harmonic oscillator, so compare it to the analytical solution.
-    
-    const double freq = 1.0;
-    State state = context.getState(State::Energy);
-    const double initialEnergy = state.getKineticEnergy()+state.getPotentialEnergy();
-    for (int i = 0; i < 1000; ++i) {
-        state = context.getState(State::Positions | State::Velocities | State::Energy);
+
+    integrator.step(10000);
+
+    double mean_temp = 0.0;
+    for (int i = 0; i < 10000000; ++i) {
+        State state = context.getState(State::Energy);
+        double KE = state.getKineticEnergy();
+        double PE = state.getPotentialEnergy();
         double time = state.getTime();
-        double expectedDist = 1.5+0.5*std::cos(freq*time);
-        ASSERT_EQUAL_VEC(Vec3(-0.5*expectedDist, 0, 0), state.getPositions()[0], 0.02);
-        ASSERT_EQUAL_VEC(Vec3(0.5*expectedDist, 0, 0), state.getPositions()[1], 0.02);
-        double expectedSpeed = -0.5*freq*std::sin(freq*time);
-        ASSERT_EQUAL_VEC(Vec3(-0.5*expectedSpeed, 0, 0), state.getVelocities()[0], 0.02);
-        ASSERT_EQUAL_VEC(Vec3(0.5*expectedSpeed, 0, 0), state.getVelocities()[1], 0.02);
-        double energy = state.getKineticEnergy()+state.getPotentialEnergy();
-        ASSERT_EQUAL_TOL(initialEnergy, energy, 0.01);
+        double temperature = 2 * KE / (BOLTZ * numDOF);
+        mean_temp = (i*mean_temp + temperature)/(i+1);
+        double energy = KE + PE + integrator.computeHeatBathEnergy();
+        if(i % 100 == 0)
+        std::cout << " Time: " << std::setw(4) << time 
+                  << " Temperature: " << std::setw(8) << std::setprecision(3) << temperature
+                  << " Mean Temperature: " << std::setw(8) << std::setprecision(3) << mean_temp
+                  << " Total Energy: " << std::setw(12) << std::setprecision(8) << energy << std::endl;
         integrator.step(1);
     }
-    ASSERT_EQUAL_TOL(10.0, context.getState(0).getTime(), 1e-5);
 }
 
 void testConstraints() {
@@ -132,15 +141,49 @@ void testConstraints() {
     }
 }
 void testNHCPropagation() {
-    double temperature = 300; // kelvin
-    double collisionFrequency = 20; // 1/ps
-    int numDOFs = 1;
-    int chainLength = 4;
-    int numMTS = 3;
-    int numYS = 1;
-    std::string label;
-    NoseHooverChain nhc(temperature, collisionFrequency, numDOFs, chainLength, numMTS, numYS, label);
-    
+    // test if the velocity scale factor goes to one for a single particle
+    // with no forces in the system
+    for (int numMTS = 1; numMTS < 5; numMTS++){
+        for (int numYS = 1; numYS <=5; numYS+=2){
+            for (int chainLength = 2; chainLength < 6; chainLength += 2){
+            double temperature = 300; // kelvin
+            double collisionFrequency = 0.01; // 1/ps
+            VelocityVerletIntegrator integrator(0.001);
+            // make system
+            System system;
+            double mass = 1;
+            system.addParticle(mass);
+            std::vector<int> particleList {0};
+            int chainID = integrator.addNoseHooverChainThermostat(system, temperature, collisionFrequency,
+                                                                  chainLength, numMTS, numYS);
+            Context context(system, integrator, platform);
+            // propagate the chain
+            double velocity = 0.1;
+            double scale, kineticEnergy;
+            for(int i = 0; i < 100; ++i) {
+                kineticEnergy = 3 * 0.5 * mass * velocity * velocity;
+                scale = integrator.propagateChain(kineticEnergy, chainID);
+                velocity *= scale;
+            }
+            ASSERT_EQUAL_TOL(1, scale,  1e-3);
+            }
+        }
+    }
+}
+
+void testPropagateChainConsistentWithPythonReference() {
+    VelocityVerletIntegrator integrator(0.001);
+    System system;
+    double mass = 1;
+    system.addParticle(mass);
+    double kineticEnergy = 1e6;
+    double temperature=300, collisionFrequency=1, chainLength=3, numMTS=3, numYS=3;
+    int chainID = integrator.addNoseHooverChainThermostat(system, temperature, collisionFrequency,
+                                                                  chainLength, numMTS, numYS);
+    Context context(system, integrator, platform);
+    double scale = integrator.propagateChain(kineticEnergy, chainID);
+    std::cout << std::setw(12) << std::setprecision(10) << scale << std::endl;
+    ASSERT_EQUAL_TOL(0.9674732261005896, scale, 1e-5)
 }
 
 void runPlatformTests();
@@ -148,9 +191,11 @@ void runPlatformTests();
 int main(int argc, char* argv[]) {
     try {
         initializeTests(argc, argv);
-        testSingleBond();
-        testConstraints();
-        runPlatformTests();
+        //testNHCPropagation();
+        testPropagateChainConsistentWithPythonReference();
+        //testSingleBond();
+        //testConstraints();
+        //runPlatformTests();
     }
     catch(const exception& e) {
         cout << "exception: " << e.what() << endl;

wrappers/python/src/swig_doxygen/swigInputConfig.py

@@ -67,7 +67,7 @@ SKIP_METHODS = [('State', 'getPositions'),
                 ('IntegrateVariableVerletStepKernel',),
                 ('IntegrateVerletStepKernel',),
                 ('IntegrateVelocityVerletStepKernel',),
-                ('PropagateNoseHooverChainKernel',),
+                ('NoseHooverChainKernel',),
                 ('IntegrateCustomStepKernel',),
                 ('Kernel',),
                 ('KernelFactory',),
@@ -468,7 +468,9 @@ UNITS = {
 ("MonteCarloMembraneBarostat", "getXYMode") : (None, ()),
 ("MonteCarloMembraneBarostat", "getZMode") : (None, ()),
 ("DrudeLangevinIntegrator", "getDrudeFriction") : ("1/unit.picosecond", ()),
-("NoseHooverChain", "getDefaultYoshidaSuzukiWeights") : ("unit.dimensionless", ()),
+("NoseHooverChain", "getParentAtoms") : (None, ()),
+("NoseHooverChain", "getThermostatedAtoms") : (None, ()),
+("NoseHooverChain", "getDefaultYoshidaSuzukiWeights") : (None, ()),
 ("DrudeSCFIntegrator", "getMinimizationErrorTolerance") : ("unit.kilojoules_per_mole/unit.nanometer", ()),
 ("RPMDIntegrator", "getContractions") : (None, ()),
 ("RPMDIntegrator", "getTotalEnergy") : ("unit.kilojoules_per_mole", ()),