Nose-Hoover chain no longer derives from Force

olla/include/openmm/kernels.h

@@ -1081,8 +1081,10 @@ public:
      * 
      * @param context    the context in which to execute this kernel
      * @param integrator the VelocityVerletIntegrator this kernel is being used for
+     * @param forcesAreValid a reference to the parent integrator's boolean for keeping
+     *                       track of the validity of the current forces.
      */
-    virtual void execute(ContextImpl& context, const VelocityVerletIntegrator& integrator) = 0;
+    virtual void execute(ContextImpl& context, const VelocityVerletIntegrator& integrator, bool &forcesAreValid) = 0;
     /**
      * Compute the kinetic energy.
      * 

openmmapi/include/openmm/NoseHooverChain.h

@@ -51,20 +51,8 @@ namespace OpenMM {
  * the number of MTS steps multiplied by the number of terms in the Yoshida-Suzuki decomposition.
  */
 
-class OPENMM_EXPORT NoseHooverChain : public Force {
+class OPENMM_EXPORT NoseHooverChain {
 public:
-    /**
-     * This is the name of the parameter that stores the position of the ith bead
-     */
-    std::string Position(int i) const {
-        return std::to_string(defaultChainID) + "NoseHooverChainPosition" + std::to_string(i);
-    }
-    /**
-     * This is the name of the parameter that stores the velocity of the ith bead
-     */
-    std::string Velocity(int i) const {
-        return std::to_string(defaultChainID) + "NoseHooverChainVelocity" + std::to_string(i);
-    }
     /**
      * Create a NoseHooverChain.
      * 
@@ -257,8 +245,6 @@ public:
     bool usesPeriodicBoundaryConditions() const {
         return false;
     }
-protected:
-    ForceImpl* createImpl() const;
 private:
     double defaultTemp, defaultFreq; //, defaultTimeStep;
     int defaultNumDOFs, defaultChainLength, defaultNumMTS, defaultNumYS;

openmmapi/include/openmm/VelocityVerletIntegrator.h

@@ -136,6 +136,24 @@ public:
      * associated with this integrator, at the current time.
      */
     double computeHeatBathEnergy();
+    /**
+     * Get the number of Nose-Hoover chains registered with this integrator.
+     */
+    int getNumNoseHooverChains() const {
+        return noseHooverChains.size();
+    }
+    /**
+     * This will be called by the Context when the user modifies aspects of the context state, such
+     * as positions, velocities, or parameters.  This gives the Integrator a chance to discard cached
+     * information.  This is <i>only</i> called when the user modifies information using methods of the Context
+     * object.  It is <i>not</i> called when a ForceImpl object modifies state information in its updateContextState()
+     * method (unless the ForceImpl calls a Context method to perform the modification).
+     * 
+     * @param changed     this specifies what aspect of the Context was changed
+     */
+    virtual void stateChanged(State::DataType changed) {
+       if (State::Positions == changed) forcesAreValid = false;
+    }
 
 protected:
     /**
@@ -157,7 +175,8 @@ protected:
      * Compute the kinetic energy of the system at the current time.
      */
     double computeKineticEnergy();
-    std::map<int, NoseHooverChain*> noseHooverChains;
+    std::vector<NoseHooverChain> noseHooverChains;
+    bool forcesAreValid;
 private:
     Kernel vvKernel, nhcKernel;
 };

openmmapi/include/openmm/internal/NoseHooverChainImpl.h

-#ifndef OPENMM_NOSEHOOVERCHAINIMPL_H_
-#define OPENMM_NOSEHOOVERCHAINIMPL_H_
-
-/* -------------------------------------------------------------------------- *
- *                                   OpenMM                                   *
- * -------------------------------------------------------------------------- *
- * This is part of the OpenMM molecular simulation toolkit originating from   *
- * Simbios, the NIH National Center for Physics-Based Simulation of           *
- * Biological Structures at Stanford, funded under the NIH Roadmap for        *
- * Medical Research, grant U54 GM072970. See https://simtk.org.               *
- *                                                                            *
- * Portions copyright (c) 2019 Stanford University and the Authors.           *
- * Authors: Andreas Krämer and Andrew C. Simmonett                            *
- * Contributors:                                                              *
- *                                                                            *
- * Permission is hereby granted, free of charge, to any person obtaining a    *
- * copy of this software and associated documentation files (the "Software"), *
- * to deal in the Software without restriction, including without limitation  *
- * the rights to use, copy, modify, merge, publish, distribute, sublicense,   *
- * and/or sell copies of the Software, and to permit persons to whom the      *
- * Software is furnished to do so, subject to the following conditions:       *
- *                                                                            *
- * The above copyright notice and this permission notice shall be included in *
- * all copies or substantial portions of the Software.                        *
- *                                                                            *
- * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR *
- * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,   *
- * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL    *
- * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,    *
- * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR      *
- * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE  *
- * USE OR OTHER DEALINGS IN THE SOFTWARE.                                     *
- * -------------------------------------------------------------------------- */
-
-#include "ForceImpl.h"
-#include "openmm/NoseHooverChain.h"
-#include "openmm/Kernel.h"
-#include <string>
-
-namespace OpenMM {
-
-class System;
-
-/**
- * This is the internal implementation of NoseHooverChain.
- */
-
-class OPENMM_EXPORT NoseHooverChainImpl : public ForceImpl {
-public:
-    NoseHooverChainImpl(const NoseHooverChain& owner);
-    void initialize(ContextImpl& context);
-    const NoseHooverChain& getOwner() const {
-        return owner;
-    }
-    void updateContextState(ContextImpl& context, bool& forcesInvalid);
-    double calcForcesAndEnergy(ContextImpl& context, bool includeForces, bool includeEnergy, int groups);
-    std::map<std::string, double> getDefaultParameters();
-    std::vector<std::string> getKernelNames();
-private:
-    const NoseHooverChain& owner;
-};
-
-} // namespace OpenMM
-
-#endif /*OPENMM_NOSEHOOVERCHAINIMPL_H_*/

openmmapi/src/ContextImpl.cpp

@@ -486,6 +486,10 @@ void ContextImpl::loadCheckpoint(istream& stream) {
     }
     updateStateDataKernel.getAs<UpdateStateDataKernel>().loadCheckpoint(*this, stream);
     hasSetPositions = true;
+    integrator.stateChanged(State::Positions);
+    integrator.stateChanged(State::Velocities);
+    integrator.stateChanged(State::Parameters);
+    integrator.stateChanged(State::Energy);
 }
 
 void ContextImpl::systemChanged() {

openmmapi/src/NoseHooverChain.cpp

@@ -30,7 +30,6 @@
  * -------------------------------------------------------------------------- */
 
 #include "openmm/NoseHooverChain.h"
-#include "openmm/internal/NoseHooverChainImpl.h"
 #include "openmm/OpenMMException.h"
 
 using namespace OpenMM;
@@ -44,9 +43,6 @@ NoseHooverChain::NoseHooverChain(double defaultTemperature, double defaultCollis
         defaultChainID(defaultChainID), thermostatedAtoms(thermostatedAtoms), parentAtoms(parentAtoms) 
 {}
 
-ForceImpl* NoseHooverChain::createImpl() const {
-    return new NoseHooverChainImpl(*this);
-}
 
 std::vector<double> NoseHooverChain::getDefaultYoshidaSuzukiWeights() const {
     switch (defaultNumYS) {

openmmapi/src/NoseHooverChainImpl.cpp

-/* -------------------------------------------------------------------------- *
- *                                   OpenMM                                   *
- * -------------------------------------------------------------------------- *
- * This is part of the OpenMM molecular simulation toolkit originating from   *
- * Simbios, the NIH National Center for Physics-Based Simulation of           *
- * Biological Structures at Stanford, funded under the NIH Roadmap for        *
- * Medical Research, grant U54 GM072970. See https://simtk.org.               *
- *                                                                            *
- * Portions copyright (c) 2008-2010 Stanford University and the Authors.      *
- * Authors: Andreas Krämer and Andrew C. Simmonett                            *
- * Contributors:                                                              *
- *                                                                            *
- * Permission is hereby granted, free of charge, to any person obtaining a    *
- * copy of this software and associated documentation files (the "Software"), *
- * to deal in the Software without restriction, including without limitation  *
- * the rights to use, copy, modify, merge, publish, distribute, sublicense,   *
- * and/or sell copies of the Software, and to permit persons to whom the      *
- * Software is furnished to do so, subject to the following conditions:       *
- *                                                                            *
- * The above copyright notice and this permission notice shall be included in *
- * all copies or substantial portions of the Software.                        *
- *                                                                            *
- * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR *
- * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,   *
- * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL    *
- * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,    *
- * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR      *
- * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE  *
- * USE OR OTHER DEALINGS IN THE SOFTWARE.                                     *
- * -------------------------------------------------------------------------- */
-
-#include "openmm/internal/NoseHooverChainImpl.h"
-#include "openmm/internal/ContextImpl.h"
-#include "openmm/Integrator.h"
-#include "openmm/System.h"
-#include "openmm/kernels.h"
-#include "SimTKOpenMMRealType.h"
-#include "sfmt/SFMT.h"
-#include <string>
-#include <vector>
-#include <iostream>
-
-using namespace OpenMM;
-using std::vector;
-
-NoseHooverChainImpl::NoseHooverChainImpl(const NoseHooverChain& owner) : owner(owner) {
-}
-
-void NoseHooverChainImpl::initialize(ContextImpl& context) {
-}
-
-void NoseHooverChainImpl::updateContextState(ContextImpl& context, bool& forcesInvalid) {
-}
-
-double NoseHooverChainImpl::calcForcesAndEnergy(ContextImpl& context, bool includeForces, bool includeEnergy, int groups) {
-    return 0;
-}
-
-std::map<std::string, double> NoseHooverChainImpl::getDefaultParameters() {
-    std::map<std::string, double> parameters;
-    const auto &owner = getOwner();
-    int chainLength = owner.getDefaultChainLength();
-    for(int i = 0; i < chainLength; ++i) {
-        parameters[owner.Position(i)] = 0;
-        parameters[owner.Velocity(i)] = 0;
-    }
-    return parameters;
-}
-
-std::vector<std::string> NoseHooverChainImpl::getKernelNames() {
-    return std::vector<std::string>();
-}

openmmapi/src/VelocityVerletIntegrator.cpp

@@ -46,7 +46,9 @@ using namespace OpenMM;
 using std::string;
 using std::vector;
 
-VelocityVerletIntegrator::VelocityVerletIntegrator(double stepSize) {
+VelocityVerletIntegrator::VelocityVerletIntegrator(double stepSize):
+    forcesAreValid(false)
+{
     setStepSize(stepSize);
     setConstraintTolerance(1e-5);
 }
@@ -54,7 +56,7 @@ VelocityVerletIntegrator::VelocityVerletIntegrator(double stepSize) {
 VelocityVerletIntegrator::~VelocityVerletIntegrator() {}
 
 double VelocityVerletIntegrator::propagateChain(double kineticEnergy, int chainID) {
-    return nhcKernel.getAs<NoseHooverChainKernel>().propagateChain(*context, *noseHooverChains.at(chainID), kineticEnergy, getStepSize());
+    return nhcKernel.getAs<NoseHooverChainKernel>().propagateChain(*context, noseHooverChains.at(chainID), kineticEnergy, getStepSize());
 }
 
 int VelocityVerletIntegrator::addNoseHooverChainThermostat(System& system, double temperature, double collisionFrequency,
@@ -112,42 +114,25 @@ int VelocityVerletIntegrator::addMaskedNoseHooverChainThermostat(System& system,
         }
     }
 
-    // figure out the chain ID
-    std::vector<int> takenIDs;
-    for (int forceNum = 0; forceNum < numForces; ++forceNum){
-        const auto & other_nhc = dynamic_cast<NoseHooverChain*>(&system.getForce(forceNum));
-        if (other_nhc){
-            takenIDs.push_back(other_nhc->getDefaultChainID());
-        }
-    }
-    int chainID = 0;
-    while (true) {
-        bool isChainIDTaken = (std::find(takenIDs.begin(), takenIDs.end(), chainID) != takenIDs.end());
-        bool isChainIDRegistered = (noseHooverChains.find(chainID) != noseHooverChains.end());
-        if (isChainIDTaken || isChainIDRegistered) {
-            chainID ++; 
-        } else {
-            break;
-        }
-    }
-    auto nhcForce = new NoseHooverChain(temperature, collisionFrequency, nDOF, chainLength,
-                                        numMTS, numYoshidaSuzuki, chainID,
-                                        thermostatedParticles, parentParticles);
-
     // make sure that thermostats do not overlap 
     for (auto &other_nhc: noseHooverChains) {
         for (auto &particle: mask){
-            bool isParticleInOtherChain = (std::find(other_nhc.second->getThermostatedAtoms().begin(), 
-                                                    other_nhc.second->getThermostatedAtoms().end(),
-                                                    particle) == other_nhc.second->getThermostatedAtoms().begin());
+            bool isParticleInOtherChain = (std::find(other_nhc.getThermostatedAtoms().begin(), 
+                                                    other_nhc.getThermostatedAtoms().end(),
+                                                    particle) == other_nhc.getThermostatedAtoms().begin());
             if (isParticleInOtherChain){
                 throw OpenMMException("Found particle " + std::to_string(particle) + "in a different NoseHooverChain, "
                                       "but particles can only be thermostated by one thermostat.");
             }
         }
     }
-    system.addForce(nhcForce);
-    noseHooverChains[chainID] = nhcForce;
+
+    // create and add new chain 
+    int chainID = noseHooverChains.size();
+    auto nhc = NoseHooverChain(temperature, collisionFrequency, nDOF, chainLength,
+                                    numMTS, numYoshidaSuzuki, chainID,
+                                    thermostatedParticles, parentParticles);
+    noseHooverChains.push_back(nhc);
     return chainID;
 }
 
@@ -157,7 +142,7 @@ double VelocityVerletIntegrator::getTemperature(int chainID) const {
                 + ". Only " + std::to_string(noseHooverChains.size()) + " have been registered with this integrator."
         );
     }
-    return noseHooverChains.at(chainID)->getDefaultTemperature();
+    return noseHooverChains.at(chainID).getDefaultTemperature();
 }
 
 void VelocityVerletIntegrator::setTemperature(double temperature, int chainID){
@@ -166,7 +151,7 @@ void VelocityVerletIntegrator::setTemperature(double temperature, int chainID){
                 + ". Only " + std::to_string(noseHooverChains.size()) + " have been registered with this integrator."
         );
     }
-    noseHooverChains.at(chainID)->setDefaultTemperature(temperature);
+    noseHooverChains.at(chainID).setDefaultTemperature(temperature);
 
 }
 
@@ -176,7 +161,7 @@ double VelocityVerletIntegrator::getCollisionFrequency(int chainID) const {
                 + ". Only " + std::to_string(noseHooverChains.size()) + " have been registered with this integrator."
         );
     }
-    return noseHooverChains.at(chainID)->getDefaultCollisionFrequency();
+    return noseHooverChains.at(chainID).getDefaultCollisionFrequency();
 } 
 void VelocityVerletIntegrator::setCollisionFrequency(double frequency, int chainID){
     if (chainID >= noseHooverChains.size()) {
@@ -184,7 +169,7 @@ void VelocityVerletIntegrator::setCollisionFrequency(double frequency, int chain
                 + ". Only " + std::to_string(noseHooverChains.size()) + " have been registered with this integrator."
         );
     }
-    noseHooverChains.at(chainID)->setDefaultCollisionFrequency(frequency);
+    noseHooverChains.at(chainID).setDefaultCollisionFrequency(frequency);
 }
 
 double VelocityVerletIntegrator::computeKineticEnergy() {
@@ -195,7 +180,7 @@ double VelocityVerletIntegrator::computeKineticEnergy() {
 double VelocityVerletIntegrator::computeHeatBathEnergy() {
     double energy = 0;
     for(auto &nhc : noseHooverChains) {
-        energy += nhcKernel.getAs<NoseHooverChainKernel>().computeHeatBathEnergy(*context, *nhc.second);
+        energy += nhcKernel.getAs<NoseHooverChainKernel>().computeHeatBathEnergy(*context, nhc);
     }
     return energy;
 }
@@ -209,6 +194,7 @@ void VelocityVerletIntegrator::initialize(ContextImpl& contextRef) {
     vvKernel.getAs<IntegrateVelocityVerletStepKernel>().initialize(contextRef.getSystem(), *this);
     nhcKernel = context->getPlatform().createKernel(NoseHooverChainKernel::Name(), contextRef);
     nhcKernel.getAs<NoseHooverChainKernel>().initialize();
+    forcesAreValid = false;
 }
 
 void VelocityVerletIntegrator::cleanup() {
@@ -230,15 +216,15 @@ void VelocityVerletIntegrator::step(int steps) {
     for (int i = 0; i < steps; ++i) {
         context->updateContextState();
         for(auto &nhc : noseHooverChains) {
-            kineticEnergy = nhcKernel.getAs<NoseHooverChainKernel>().computeMaskedKineticEnergy(*context, *nhc.second);
-            scale = nhcKernel.getAs<NoseHooverChainKernel>().propagateChain(*context, *nhc.second, kineticEnergy, getStepSize());
-            nhcKernel.getAs<NoseHooverChainKernel>().scaleVelocities(*context, *nhc.second, scale);
+            kineticEnergy = nhcKernel.getAs<NoseHooverChainKernel>().computeMaskedKineticEnergy(*context, nhc);
+            scale = nhcKernel.getAs<NoseHooverChainKernel>().propagateChain(*context, nhc, kineticEnergy, getStepSize());
+            nhcKernel.getAs<NoseHooverChainKernel>().scaleVelocities(*context, nhc, scale);
         }
-        vvKernel.getAs<IntegrateVelocityVerletStepKernel>().execute(*context, *this);
+        vvKernel.getAs<IntegrateVelocityVerletStepKernel>().execute(*context, *this, forcesAreValid);
         for(auto &nhc : noseHooverChains) {
-            kineticEnergy = nhcKernel.getAs<NoseHooverChainKernel>().computeMaskedKineticEnergy(*context, *nhc.second);
-            scale = nhcKernel.getAs<NoseHooverChainKernel>().propagateChain(*context, *nhc.second, kineticEnergy, getStepSize());
-            nhcKernel.getAs<NoseHooverChainKernel>().scaleVelocities(*context, *nhc.second, scale);
+            kineticEnergy = nhcKernel.getAs<NoseHooverChainKernel>().computeMaskedKineticEnergy(*context, nhc);
+            scale = nhcKernel.getAs<NoseHooverChainKernel>().propagateChain(*context, nhc, kineticEnergy, getStepSize());
+            nhcKernel.getAs<NoseHooverChainKernel>().scaleVelocities(*context, nhc, scale);
         }
     }
 }

platforms/cuda/include/CudaArray.h

@@ -71,6 +71,10 @@ public:
      * @param name              the name of the array
      */
     CudaArray(CudaContext& context, int size, int elementSize, const std::string& name);
+    CudaArray(const CudaArray &other) noexcept;
+    CudaArray& operator=(CudaArray &&other) noexcept;
+    CudaArray& operator=(const CudaArray& other) noexcept;
+    CudaArray(CudaArray &&other) noexcept ; 
     ~CudaArray();
     /**
      * Initialize this object.

platforms/cuda/include/CudaKernels.h

@@ -1386,8 +1386,10 @@ public:
      * 
      * @param context    the context in which to execute this kernel
      * @param integrator the VerletIntegrator this kernel is being used for
+     * @param forcesAreValid a reference to the parent integrator's boolean for keeping
+     *                       track of the validity of the current forces.
      */
-    void execute(ContextImpl& context, const VelocityVerletIntegrator& integrator);
+    void execute(ContextImpl& context, const VelocityVerletIntegrator& integrator, bool &forcesAreValid);
     /**
      * Compute the kinetic energy.
      * 

platforms/cuda/src/CudaArray.cpp

@@ -35,6 +35,47 @@ using namespace OpenMM;
 CudaArray::CudaArray() : pointer(0), ownsMemory(false) {
 }
 
+CudaArray::CudaArray(const CudaArray &other) noexcept {
+    context = other.context;
+    pointer = other.pointer;
+    size = other.size;
+    elementSize = other.elementSize;
+    ownsMemory = false;
+    name = other.name;
+}
+
+CudaArray::CudaArray(CudaArray &&other) noexcept :
+    context(std::move(other.context)),
+    pointer(std::move(other.pointer)),
+    size(std::move(other.size)),
+    elementSize(std::move(other.elementSize)),
+    ownsMemory(false),
+    name(std::move(other.name)) { }
+
+CudaArray& CudaArray::operator =(const CudaArray& other) noexcept{
+    if(this != &other) {
+        context = other.context;
+        pointer = other.pointer;
+        size = other.size;
+        elementSize = other.elementSize;
+        ownsMemory = false;
+        name = other.name;
+    }
+    return *this;
+}
+
+CudaArray& CudaArray::operator =(CudaArray &&other) noexcept{
+    if(this != &other) {
+        context = other.context;
+        pointer = other.pointer;
+        size = other.size;
+        elementSize = other.elementSize;
+        ownsMemory = false;
+        name = other.name;
+    }
+    return *this;
+}
+
 CudaArray::CudaArray(CudaContext& context, int size, int elementSize, const std::string& name) : pointer(0) {
     initialize(context, size, elementSize, name);
 }

platforms/cuda/src/CudaKernels.cpp

@@ -37,7 +37,6 @@
 #include "openmm/internal/CustomManyParticleForceImpl.h"
 #include "openmm/internal/CustomNonbondedForceImpl.h"
 #include "openmm/internal/NonbondedForceImpl.h"
-#include "openmm/internal/NoseHooverChainImpl.h"
 #include "openmm/internal/OSRngSeed.h"
 #include "CudaBondedUtilities.h"
 #include "CudaExpressionUtilities.h"
@@ -7087,7 +7086,7 @@ void CudaIntegrateVelocityVerletStepKernel::initialize(const System& system, con
     kernel3 = cu.getKernel(module, "integrateVelocityVerletPart3");
 }
 
-void CudaIntegrateVelocityVerletStepKernel::execute(ContextImpl& context, const VelocityVerletIntegrator& integrator) {
+void CudaIntegrateVelocityVerletStepKernel::execute(ContextImpl& context, const VelocityVerletIntegrator& integrator, bool &forcesAreValid) {
     cu.setAsCurrent();
     CudaIntegrationUtilities& integration = cu.getIntegrationUtilities();
     int numAtoms = cu.getNumAtoms();
@@ -7095,6 +7094,8 @@ void CudaIntegrateVelocityVerletStepKernel::execute(ContextImpl& context, const
     double dt = integrator.getStepSize();
     cu.getIntegrationUtilities().setNextStepSize(dt);
 
+    if( !forcesAreValid ) context.calcForcesAndEnergy(true, false);
+
     //// Call the first integration kernel.
 
     CUdeviceptr posCorrection = (cu.getUseMixedPrecision() ? cu.getPosqCorrection().getDevicePointer() : 0);
@@ -7116,6 +7117,7 @@ void CudaIntegrateVelocityVerletStepKernel::execute(ContextImpl& context, const
 
     //// Update forces
     context.calcForcesAndEnergy(true, false);
+    forcesAreValid = true;
 
     //// Call the third integration kernel.
 
@@ -8390,16 +8392,16 @@ double CudaNoseHooverChainKernel::propagateChain(ContextImpl& context, const Nos
 
     auto & chainState = cu.getIntegrationUtilities().getNoseHooverChainState();
     if (chainID >= chainState.size()) chainState.resize(chainID+1);
-    if (!chainState[chainID].isInitialized() || chainState[chainID].getSize() != chainLength) {
+    if (!chainState.at(chainID).isInitialized() || chainState.at(chainID).getSize() != chainLength) {
         // We need to upload the CUDA array
         if(useDouble){
-            chainState[chainID].initialize<double2>(cu, chainLength, "chainState" + std::to_string(chainID));
+            chainState.at(chainID).initialize<double2>(cu, chainLength, "chainState" + std::to_string(chainID));
             std::vector<double2> zeros(chainLength, make_double2(0, 0));
-            chainState[chainID].upload(zeros.data());
+            chainState.at(chainID).upload(zeros.data());
         } else {
-            chainState[chainID].initialize<float2>(cu, chainLength, "chainState" + std::to_string(chainID));
+            chainState.at(chainID).initialize<float2>(cu, chainLength, "chainState" + std::to_string(chainID));
             std::vector<float2> zeros(chainLength, make_float2(0, 0));
-            chainState[chainID].upload(zeros.data());
+            chainState.at(chainID).upload(zeros.data());
         }
     }
 
@@ -8415,21 +8417,22 @@ double CudaNoseHooverChainKernel::propagateChain(ContextImpl& context, const Nos
         }
     }
 
-    if (!chainForces.isInitialized() || !chainMasses.isInitialized() || chainForces.getSize() < chainLength || chainMasses.getSize() < chainLength) {
-        if(useDouble){
     std::vector<double> zeros(chainLength,0);
+    if (!chainForces.isInitialized() || !chainMasses.isInitialized() ){
+        if(useDouble){
             chainMasses.initialize<double>(cu, chainLength, "chainMasses");
             chainForces.initialize<double>(cu, chainLength, "chainForces");
             chainMasses.upload(zeros);
             chainForces.upload(zeros);
         } else {
-            std::vector<float> zeros(chainLength,0);
             chainMasses.initialize<float>(cu, chainLength, "chainMasses");
             chainForces.initialize<float>(cu, chainLength, "chainForces");
             chainMasses.upload(zeros);
             chainForces.upload(zeros);
         }
     }
+    if (chainForces.getSize() < chainLength) chainMasses.resize(chainLength);
+    if (chainMasses.getSize() < chainLength) chainMasses.resize(chainLength);
     double kT = BOLTZ * temperature;
     float kTfloat = (float) kT;
     float timeStepFloat = (float) timeStep;
@@ -8511,7 +8514,7 @@ double CudaNoseHooverChainKernel::computeMaskedKineticEnergy(ContextImpl& contex
     const auto & thermostatedAtoms = nhc.getThermostatedAtoms();
     if (chainID >= masks.size()) masks.resize(chainID+1);
     auto nAtoms = cu.getPaddedNumAtoms();
-    if (!masks[chainID].isInitialized()) { 
+    if (!masks.at(chainID).isInitialized()) { 
         // We need to upload the CUDA array
         std::vector<int> maskVec(nAtoms);
         for (int i = 0; i < nAtoms; ++i) maskVec[i] = i;
@@ -8531,8 +8534,8 @@ double CudaNoseHooverChainKernel::computeMaskedKineticEnergy(ContextImpl& contex
         for(int i = thermostatedAtoms.size(); i < nAtoms; ++i){
             maskVec[i] = i;
         }
-        masks[chainID].initialize<int>(cu, nAtoms, "mask" + std::to_string(chainID));
-        masks[chainID].upload(maskVec);
+        masks.at(chainID).initialize<int>(cu, nAtoms, "mask" + std::to_string(chainID));
+        masks.at(chainID).upload(maskVec);
     }
     if (masks[chainID].getSize() != nAtoms) {
         throw OpenMMException("Number of atoms changed. Cannot be handled by the same Nose-Hoover thermostat.");

platforms/reference/include/ReferenceKernels.h

@@ -1156,8 +1156,10 @@ public:
      * 
      * @param context    the context in which to execute this kernel
      * @param integrator the VerletIntegrator this kernel is being used for
+     * @param forcesAreValid a reference to the parent integrator's boolean for keeping
+     *                       track of the validity of the current forces.
      */
-    void execute(ContextImpl& context, const VelocityVerletIntegrator& integrator);
+    void execute(ContextImpl& context, const VelocityVerletIntegrator& integrator, bool &forcesAreValid);
     /**
      * Compute the kinetic energy.
      * 

platforms/reference/include/ReferencePlatform.h

@@ -72,6 +72,8 @@ public:
     Vec3* periodicBoxVectors;
     ReferenceConstraints* constraints;
     std::map<std::string, double>* energyParameterDerivatives;
+    std::vector<std::vector<double>>* noseHooverPositions;
+    std::vector<std::vector<double>>* noseHooverVelocities;
 };
 } // namespace OpenMM
 

platforms/reference/include/ReferenceVelocityVerletDynamics.h

@@ -70,11 +70,12 @@ class ReferenceVelocityVerletDynamics : public ReferenceDynamics {
          @param forces              forces
          @param masses              atom masses
          @param tolerance           the constraint tolerance
+         @param forcesAreValid      whether the forces are valid (duh!)
       
          --------------------------------------------------------------------------------------- */
      
       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);
+                  std::vector<OpenMM::Vec3>& velocities, std::vector<OpenMM::Vec3>& forces, std::vector<double>& masses, double tolerance, bool &forcesAreValid);
       
 };
 

platforms/reference/src/ReferenceKernels.cpp

@@ -126,6 +126,15 @@ static map<string, double>& extractEnergyParameterDerivatives(ContextImpl& conte
     return *data->energyParameterDerivatives;
 }
 
+static vector<vector<double> >& extractNoseHooverPositions(ContextImpl& context) {
+    ReferencePlatform::PlatformData *data = reinterpret_cast<ReferencePlatform::PlatformData*>(context.getPlatformData());  
+    return *((vector<vector<double> >*) data->noseHooverPositions);
+}
+
+static vector<vector<double> >& extractNoseHooverVelocities(ContextImpl& context) {
+    ReferencePlatform::PlatformData *data = reinterpret_cast<ReferencePlatform::PlatformData*>(context.getPlatformData());  
+    return *((vector<vector<double> >*) data->noseHooverVelocities);
+}
 /**
  * Make sure an expression doesn't use any undefined variables.
  */
@@ -278,7 +287,7 @@ void ReferenceUpdateStateDataKernel::setPeriodicBoxVectors(ContextImpl& context,
 }
 
 void ReferenceUpdateStateDataKernel::createCheckpoint(ContextImpl& context, ostream& stream) {
-    int version = 2;
+    int version = 3;
     stream.write((char*) &version, sizeof(int));
     stream.write((char*) &data.time, sizeof(data.time));
     vector<Vec3>& posData = extractPositions(context);
@@ -287,13 +296,27 @@ void ReferenceUpdateStateDataKernel::createCheckpoint(ContextImpl& context, ostr
     stream.write((char*) &velData[0], sizeof(Vec3)*velData.size());
     Vec3* vectors = extractBoxVectors(context);
     stream.write((char*) vectors, 3*sizeof(Vec3));
+    auto& allNoseHooverPositions = extractNoseHooverPositions(context);
+    auto& allNoseHooverVelocities = extractNoseHooverVelocities(context);
+    size_t numChains = allNoseHooverPositions.size();
+    assert(numChains == allNoseHooverVelocities.size());
+    stream.write((char*) &numChains, sizeof(size_t));
+    for (size_t i=0; i<numChains; i++){
+        auto & noseHooverPositions = allNoseHooverPositions.at(i);
+        auto & noseHooverVelocities = allNoseHooverVelocities.at(i);
+        size_t numBeads = noseHooverPositions.size();
+        assert(numBeads == noseHooverVelocities.size());
+        stream.write((char*) &numBeads, sizeof(size_t));
+        stream.write((char*) noseHooverPositions.data(), sizeof(double)*numBeads);
+        stream.write((char*) noseHooverVelocities.data(), sizeof(double)*numBeads);
+    }
     SimTKOpenMMUtilities::createCheckpoint(stream);
 }
 
 void ReferenceUpdateStateDataKernel::loadCheckpoint(ContextImpl& context, istream& stream) {
     int version;
     stream.read((char*) &version, sizeof(int));
-    if (version != 2)
+    if (version != 3)
         throw OpenMMException("Checkpoint was created with a different version of OpenMM");
     stream.read((char*) &data.time, sizeof(data.time));
     vector<Vec3>& posData = extractPositions(context);
@@ -302,6 +325,21 @@ void ReferenceUpdateStateDataKernel::loadCheckpoint(ContextImpl& context, istrea
     stream.read((char*) &velData[0], sizeof(Vec3)*velData.size());
     Vec3* vectors = extractBoxVectors(context);
     stream.read((char*) vectors, 3*sizeof(Vec3));
+    size_t numChains, numBeads;
+    auto& allNoseHooverPositions = extractNoseHooverPositions(context);
+    auto& allNoseHooverVelocities = extractNoseHooverVelocities(context);
+    stream.read((char*) &numChains, sizeof(size_t));
+    allNoseHooverPositions.clear();
+    allNoseHooverVelocities.clear();
+    for (size_t i=0; i<numChains; i++){
+        stream.read((char*) &numBeads, sizeof(size_t));
+        std::vector<double> noseHooverPositions(numBeads);
+        std::vector<double> noseHooverVelocities(numBeads);
+        stream.read((char*) &noseHooverPositions[0], sizeof(double)*numBeads);
+        stream.read((char*) &noseHooverVelocities[0], sizeof(double)*numBeads);
+        allNoseHooverPositions.push_back(noseHooverPositions);
+        allNoseHooverVelocities.push_back(noseHooverVelocities);
+    }
     SimTKOpenMMUtilities::loadCheckpoint(stream);
 }
 
@@ -2119,7 +2157,7 @@ void ReferenceIntegrateVelocityVerletStepKernel::initialize(const System& system
         masses[i] = system.getParticleMass(i);
 }
 
-void ReferenceIntegrateVelocityVerletStepKernel::execute(ContextImpl& context, const VelocityVerletIntegrator& integrator) {
+void ReferenceIntegrateVelocityVerletStepKernel::execute(ContextImpl& context, const VelocityVerletIntegrator& integrator, bool &forcesAreValid) {
     double stepSize = integrator.getStepSize();
     vector<Vec3>& posData = extractPositions(context);
     vector<Vec3>& velData = extractVelocities(context);
@@ -2132,7 +2170,7 @@ void ReferenceIntegrateVelocityVerletStepKernel::execute(ContextImpl& context, c
         dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
         prevStepSize = stepSize;
     }
-    dynamics->update(context, context.getSystem(), posData, velData, forceData, masses, integrator.getConstraintTolerance());
+    dynamics->update(context, context.getSystem(), posData, velData, forceData, masses, integrator.getConstraintTolerance(), forcesAreValid);
     data.time += stepSize;
     data.stepCount++;
 }
@@ -2436,27 +2474,32 @@ double ReferenceNoseHooverChainKernel::propagateChain(ContextImpl& context, cons
     if (kineticEnergy < 1e-8) return 1.0;  // (catches the problem of zero velocities in the first dynamics step, where we have nothing to scale)
     // Get the variables describing the NHC
     int chainLength = nhc.getDefaultChainLength();
+    int chainID = nhc.getDefaultChainID();
     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));
+    auto& allChainPositions = extractNoseHooverPositions(context);
+    auto& allChainVelocities = extractNoseHooverVelocities(context);
+    if (allChainPositions.size() < chainID+1){
+        allChainPositions.resize(chainID+1);
+    }
+    if (allChainVelocities.size() < chainID+1){
+        allChainVelocities.resize(chainID+1);
+    }
+    auto& chainPositions = allChainPositions.at(chainID);
+    auto& chainVelocities = allChainVelocities.at(chainID);
+    if (chainPositions.size() < chainLength){
+        chainPositions.resize(chainLength, 0);
+    }
+    if (chainVelocities.size() < chainLength){
+        chainVelocities.resize(chainLength, 0);
     }
     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;
 }
 
@@ -2464,18 +2507,21 @@ double ReferenceNoseHooverChainKernel::computeHeatBathEnergy(ContextImpl& contex
     double potentialEnergy = 0;
     double kineticEnergy = 0;
     int chainLength = nhc.getDefaultChainLength();
+    int chainID = nhc.getDefaultChainID();
     double temperature = nhc.getDefaultTemperature();
     double collisionFrequency = nhc.getDefaultCollisionFrequency();
     double kT = temperature * BOLTZ;
     int numDOFs = nhc.getDefaultNumDegreesOfFreedom();
+    auto& nhcPositions = extractNoseHooverPositions(context);
+    auto& nhcVelocities = extractNoseHooverVelocities(context);
     for(int i = 0; i < chainLength; ++i) {
         double prefac = i ? 1 : numDOFs;
         double mass = prefac * kT / (collisionFrequency * collisionFrequency);
-        double velocity = context.getParameter(nhc.Velocity(i));
+        double velocity = nhcVelocities[chainID][i];
         // The kinetic energy of this bead
         kineticEnergy += 0.5 * mass * velocity * velocity;
         // The potential energy of this bead
-        double position = context.getParameter(nhc.Position(i));
+        double position = nhcPositions[chainID][i];
         potentialEnergy += prefac * kT * position;
     }
     return kineticEnergy + potentialEnergy;
@@ -2485,7 +2531,6 @@ double ReferenceNoseHooverChainKernel::computeMaskedKineticEnergy(ContextImpl& c
     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);

platforms/reference/src/ReferencePlatform.cpp

@@ -103,6 +103,8 @@ ReferencePlatform::PlatformData::PlatformData(const System& system) : time(0.0),
     periodicBoxVectors = new Vec3[3];
     constraints = new ReferenceConstraints(system);
     energyParameterDerivatives = new map<string, double>();
+    noseHooverPositions = new vector<vector<double> >();
+    noseHooverVelocities = new vector<vector<double> >();
 }
 
 ReferencePlatform::PlatformData::~PlatformData() {
@@ -113,4 +115,6 @@ ReferencePlatform::PlatformData::~PlatformData() {
     delete[] periodicBoxVectors;
     delete constraints;
     delete energyParameterDerivatives;
+    delete noseHooverPositions;
+    delete noseHooverVelocities;
 }

platforms/reference/src/SimTKReference/ReferenceNoseHooverChain.cpp

@@ -26,7 +26,6 @@
 #include <string.h>
 #include <sstream>
 #include <exception>
-#include <iostream>
 #include "SimTKOpenMMUtilities.h"
 #include "ReferenceNoseHooverChain.h"
 

platforms/reference/src/SimTKReference/ReferenceVelocityVerletDynamics.cpp

@@ -77,8 +77,10 @@ ReferenceVelocityVerletDynamics::~ReferenceVelocityVerletDynamics() {
 
 void ReferenceVelocityVerletDynamics::update(OpenMM::ContextImpl &context, const OpenMM::System& system, vector<Vec3>& atomCoordinates,
                                           vector<Vec3>& velocities,
-                                          vector<Vec3>& forces, vector<double>& masses, double tolerance) {
+                                          vector<Vec3>& forces, vector<double>& masses, double tolerance, bool &forcesAreValid) {
+
     // first-time-through initialization
+    if (!forcesAreValid) context.calcForcesAndEnergy(true, false);
 
     int numberOfAtoms = system.getNumParticles();
     if (getTimeStep() == 0) {
@@ -110,8 +112,8 @@ void ReferenceVelocityVerletDynamics::update(OpenMM::ContextImpl &context, const
        referenceConstraintAlgorithm->apply(xPrime, atomCoordinates, inverseMasses, tolerance);
 
     ReferenceVirtualSites::computePositions(system, atomCoordinates);
-
     context.calcForcesAndEnergy(true, false);
+    forcesAreValid = true;
 
     for (int i = 0; i < numberOfAtoms; ++i) {
         if (masses[i] != 0.0)

tests/TestNoseHooverThermostat.h

@@ -33,6 +33,7 @@
 #include "openmm/NoseHooverChain.h"
 #include "openmm/VelocityVerletIntegrator.h"
 #include "openmm/Context.h"
+#include "openmm/State.h"
 #include "openmm/HarmonicBondForce.h"
 #include "openmm/NonbondedForce.h"
 #include "openmm/CustomExternalForce.h"
@@ -40,9 +41,11 @@
 #include "SimTKOpenMMRealType.h"
 #include "sfmt/SFMT.h"
 #include <iostream>
+#include <sstream>
 #include <iomanip>
 #include <vector>
 #include <algorithm>
+#include <numeric>
 
 using namespace OpenMM;
 using namespace std;
@@ -260,6 +263,67 @@ void testPropagateChainConsistentWithPythonReference() {
     ASSERT_EQUAL_TOL(0.9674732261005896, scale, 1e-5)
 }
 
+
+void testCheckpoints() {
+    VelocityVerletIntegrator integrator(0.001);
+    System system;
+    double mass = 1;
+    system.addParticle(mass);
+    system.addParticle(mass);
+    NonbondedForce* force = new NonbondedForce();
+    force->addParticle(0, 0.1, 1.0);
+    force->addParticle(0, 0.1, 1.0);
+    system.addForce(force);
+    double kineticEnergy = 1e6;
+    double temperature=300, collisionFrequency=1, chainLength=3, numMTS=3, numYS=3;
+    integrator.addMaskedNoseHooverChainThermostat(system, std::vector<int>(1,0), std::vector<int>(), temperature, collisionFrequency,
+                                                                  chainLength, numMTS, numYS);
+    chainLength = 10;
+    integrator.addMaskedNoseHooverChainThermostat(system, std::vector<int>(1,1),  std::vector<int>(1,0), 
+                                                                  temperature, collisionFrequency,
+                                                                  chainLength, numMTS, numYS);
+    Context context(system, integrator, platform);
+    std::vector<Vec3> positions(2);
+    positions[1] = {0.1,0.1,0.1};
+    context.setPositions(positions);
+    integrator.step(300);
+#if DEBUG
+    std::cout << std::endl << std::endl << std::endl << "writing checkpoint";
+#endif
+    std::stringstream checkpoint; 
+    context.createCheckpoint(checkpoint);
+    
+    State state = context.getState(State::Positions | State::Velocities);
+    for (size_t i=0; i<100; i++){
+        state = context.getState(State::Positions | State::Velocities);
+#if DEBUG
+        std::cout << "posvel" << state.getPositions()[0] << " " << state.getVelocities()[0] << std::endl;
+#endif
+        integrator.step(1);
+    }
+#if DEBUG
+    std::cout << std::endl << std::endl << "loading checkpoint" << std::endl;
+#endif
+    context.loadCheckpoint(checkpoint);
+
+    
+    State state2 = context.getState(State::Positions | State::Velocities);
+    for (size_t i=0; i<100; i++){
+        state2 = context.getState(State::Positions | State::Velocities);
+#if DEBUG
+        std::cout << "posvel" << state2.getPositions()[0] << " " << state2.getVelocities()[0] << std::endl;
+#endif
+        integrator.step(1);
+    }
+
+    for (int i=0; i<3;i++){
+        ASSERT_EQUAL_VEC(state.getPositions()[0], state2.getPositions()[0], 1e-6); 
+        ASSERT_EQUAL_VEC(state.getPositions()[1], state2.getPositions()[1], 1e-6); 
+        ASSERT_EQUAL_VEC(state.getVelocities()[0], state2.getVelocities()[0], 1e-6); 
+        ASSERT_EQUAL_VEC(state.getVelocities()[1], state2.getVelocities()[1], 1e-6); 
+    }
+}
+
 void runPlatformTests();
 
 int main(int argc, char* argv[]) {
@@ -269,6 +333,7 @@ int main(int argc, char* argv[]) {
         bool constrain;
         constrain = false; testDimerBox(constrain);
         constrain = true; testDimerBox(constrain);
+        testCheckpoints();
         runPlatformTests();
     }
     catch(const exception& e) {