Merge pull request #2351 from andysim/nhc

WIP: Nosé-Hoover Thermostat

tests/TestNoseHooverThermostat.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 "openmm/internal/AssertionUtilities.h"
+#include "openmm/NoseHooverChain.h"
+#include "openmm/NoseHooverIntegrator.h"
+#include "openmm/Context.h"
+#include "openmm/State.h"
+#include "openmm/HarmonicBondForce.h"
+#include "openmm/NonbondedForce.h"
+#include "openmm/CustomExternalForce.h"
+#include "openmm/System.h"
+#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;
+
+void testHarmonicOscillator() {
+    const double mass = 1.0;
+    double temperature = 300;
+    double frequency = 1;
+    double mts = 1, ys = 1, chain_length = 3;
+
+    System system;    
+    system.addParticle(mass);
+    vector<Vec3> positions(1);
+    positions[0] = Vec3(0.5,0.5,0.5);
+    vector<Vec3> velocities(1);
+    velocities[0] = Vec3(0, 0, 0);
+    auto harmonic_restraint = new CustomExternalForce("0.5*(x^2+y^2+z^2)");
+    harmonic_restraint->addParticle(0);
+    system.addForce(harmonic_restraint);
+    NoseHooverIntegrator integrator(0.001);
+    integrator.addThermostat(temperature, frequency, chain_length, mts, ys);
+    Context context(system, integrator, platform);
+    context.setPositions(positions);
+    context.setVelocities(velocities);
+
+    double mean_temperature=0;
+    // equilibration
+    integrator.step(2000);
+    for (size_t i=0; i < 2500; i++){
+        integrator.step(10);
+        State state = context.getState(State::Energy | State::Positions | State::Velocities);
+        double kinetic_energy = state.getKineticEnergy();
+        double temp = kinetic_energy/(0.5*3*BOLTZ);
+        mean_temperature = (i*mean_temperature + temp)/(i+1);
+        double PE = state.getPotentialEnergy();
+        double time = state.getTime();
+        double energy = kinetic_energy + PE + integrator.computeHeatBathEnergy();
+    }
+    ASSERT_EQUAL_TOL(temperature, mean_temperature, 0.02);
+}
+
+int makeDimerBox(System& system, std::vector<Vec3>& positions, bool constrain=true, int numMolecules=20, double bondLength=0.1){
+    double boxLength = 2; // nm
+    Vec3 a(boxLength, 0.0, 0.0);
+    Vec3 b(0.0, boxLength, 0.0);
+    Vec3 c(0.0, 0.0, boxLength);
+    double mass = 20;
+    double bondForceConstant = 30000; //0.001;
+    int numDOF = 0;
+    NonbondedForce* forceField = new NonbondedForce();
+    HarmonicBondForce* bondForce = new HarmonicBondForce();
+    for(int molecule = 0; molecule < numMolecules; ++molecule) {
+        int particle1 = system.addParticle(mass);
+        int particle2 = system.addParticle(mass);
+        forceField->addParticle(0.0, 0.1, 1.0);
+        forceField->addParticle(0.0, 0.1, 1.0);
+        forceField->addException(particle1, particle2, 0, 0, 0);
+        bondForce->addBond(particle1, particle2, bondLength, bondForceConstant);
+        numDOF += 6;
+        if (constrain) {
+            system.addConstraint(particle1, particle2, bondLength);
+            numDOF -= 1;
+        }
+    }
+    forceField->setCutoffDistance(.99*boxLength/2);
+    forceField->setSwitchingDistance(.88*boxLength/2);
+    forceField->setUseSwitchingFunction(true);
+    forceField->setUseDispersionCorrection(false);
+    forceField->setNonbondedMethod(NonbondedForce::CutoffPeriodic);
+    system.addForce(forceField);
+    system.addForce(bondForce);
+    system.setDefaultPeriodicBoxVectors(a, b, c);
+    OpenMM_SFMT::SFMT sfmt;
+    init_gen_rand(0, sfmt);
+    for (int i = 0; i < numMolecules; i++) {
+        while (true) {
+            Vec3 pos = Vec3(boxLength*genrand_real2(sfmt), boxLength*genrand_real2(sfmt), boxLength*genrand_real2(sfmt));
+            Vec3 pos1 = pos + Vec3(0,0, bondLength/2);
+            Vec3 pos2 = pos + Vec3(0,0,-bondLength/2);
+            double minDist = 2*boxLength;
+            for (int j = 0; j < i; j++) {
+                Vec3 delta = pos1-positions[j];
+                minDist = std::min(minDist, sqrt(delta.dot(delta)));
+                delta = pos2-positions[j];
+                minDist = std::min(minDist, sqrt(delta.dot(delta)));
+            }
+            if (minDist > 0.15) {
+                positions[2*i+0] = pos1;
+                positions[2*i+1] = pos2;
+                break;
+            }
+        }
+    }
+    return numDOF;
+}
+
+void testDimerBox(bool constrain=true) {
+    // Check conservation of system + bath energy for a harmonic oscillator
+    int numMolecules = 20;
+    double bondLength = 0.1;
+    double bondLengthSquared = bondLength * bondLength;
+    System system;
+    std::vector<Vec3> positions(numMolecules*2);
+    int numDOF = makeDimerBox(system, positions, constrain, numMolecules, bondLength);
+
+    bool simpleConstruct = true;
+    double temperature = 300; // kelvin
+    double collisionFrequency = 200; // 1/ps
+    int numMTS = 3;
+    int numYS = 3;
+    int chainLength = 5;
+    auto integrator = simpleConstruct ? NoseHooverIntegrator(temperature, collisionFrequency, 0.001, chainLength, numMTS, numYS)
+                                      : NoseHooverIntegrator(0.001);
+    if (!simpleConstruct)
+        integrator.addThermostat(temperature, collisionFrequency, chainLength, numMTS, numYS);
+    Context context(system, integrator, platform);
+    context.setPositions(positions);
+    context.setVelocitiesToTemperature(temperature);
+
+    int nSteps = 1500;
+    double mean_temp = 0.0;
+    std::vector<double> energies(nSteps);
+    for (int i = 0; i < nSteps; ++i) {
+        integrator.step(1);
+        State state = context.getState(State::Energy | (constrain ? State::Positions : 0));
+        if (constrain) {
+            auto positions = state.getPositions();
+            for(int i = 0; i < numMolecules; ++i) {
+                Vec3 delta = positions[2*i+1] - positions[2*i];
+                double dR2 = delta.dot(delta);
+                ASSERT_EQUAL_TOL(bondLengthSquared, dR2, 1e-4);
+            }
+        }
+        double KE = state.getKineticEnergy();
+        double PE = state.getPotentialEnergy();
+        double time = state.getTime();
+        double instantaneous_temperature = 2 * KE / (BOLTZ * numDOF);
+        mean_temp = (i*mean_temp + instantaneous_temperature)/(i+1);
+        double energy = KE + PE + integrator.computeHeatBathEnergy();
+        energies[i] = energy;
+    }
+    double sum = std::accumulate(energies.begin(), energies.end(), 0.0);
+    double mean = sum / energies.size();
+    double sq_sum = std::inner_product(energies.begin(), energies.end(), energies.begin(), 0.0);
+    double std = std::sqrt(sq_sum / energies.size() - mean * mean);
+    double relative_std = std / mean;
+    // Check mean temperature
+    ASSERT_USUALLY_EQUAL_TOL(temperature, mean_temp, 1e-2);
+    // Check fluctuation of conserved (total bath + system) energy
+    ASSERT_USUALLY_EQUAL_TOL(relative_std, 0, 5e-4);
+}
+
+void testCheckpoints() {
+    // Create a system with Drude-like particles to be thermostated as a pair, as well as another
+    // particle to be thermostated independently, to test all integrator features.
+    double timeStep = 0.001;
+    NoseHooverIntegrator integrator(timeStep), newIntegrator(timeStep);
+    System system;
+    double mass = 1;
+    system.addParticle(8*mass);
+    system.addParticle(mass);
+    system.addParticle(5*mass);
+    HarmonicBondForce* force = new HarmonicBondForce();
+    force->addBond(0, 1, 0.1, 50.0);
+    force->addBond(0, 2, 0.1, 50.0);
+    system.addForce(force);
+    double kineticEnergy = 1e6;
+    double temperature=300, collisionFrequency=1, chainLength=3, numMTS=3, numYS=3;
+    chainLength = 10;
+    integrator.addSubsystemThermostat(std::vector<int>{2}, std::vector<std::pair<int,int>>{{0,1}},  temperature, collisionFrequency, temperature, collisionFrequency,
+                                      chainLength, numMTS, numYS);
+    newIntegrator.addSubsystemThermostat(std::vector<int>{2}, std::vector<std::pair<int,int>>{{0,1}},  temperature, collisionFrequency, temperature, collisionFrequency,
+                                      chainLength, numMTS, numYS);
+    Context context(system, integrator, platform);
+    Context newContext(system, newIntegrator, platform);
+    std::vector<Vec3> positions(3);
+    std::vector<Vec3> velocities(3);
+    positions[1] = {0.1, 0.0, 0.0};
+    velocities[1] = {0.1,0.2,-0.2};
+    positions[2] = {-0.1, 0.001, 0.001};
+    velocities[2] = {-0.1,0.2,-0.2};
+    context.setPositions(positions);
+    context.setVelocities(velocities);
+
+    // Run a short simulation and checkpoint..
+    integrator.step(500);
+    std::stringstream checkpoint; 
+    context.createCheckpoint(checkpoint);
+
+    // Now continue the simulation
+    integrator.step(5);
+
+    // And try the same, starting from the checkpoint
+    newContext.loadCheckpoint(checkpoint);
+    newIntegrator.step(5);
+
+    State state1 = context.getState(State::Positions | State::Velocities);
+    State state2 = newContext.getState(State::Positions | State::Velocities);
+    ASSERT_EQUAL_VEC(state1.getPositions()[0], state2.getPositions()[0], 1e-6); 
+    ASSERT_EQUAL_VEC(state1.getPositions()[1], state2.getPositions()[1], 1e-6); 
+    ASSERT_EQUAL_VEC(state1.getVelocities()[0], state2.getVelocities()[0], 1e-6); 
+    ASSERT_EQUAL_VEC(state1.getVelocities()[1], state2.getVelocities()[1], 1e-6); 
+}
+
+void testAPIChangeNumParticles() {
+    bool constrain = true;
+    int numMolecules = 20;
+    double bondLength = 0.1;
+    double bondLengthSquared = bondLength * bondLength;
+    System system;
+    std::vector<Vec3> positions(numMolecules*2);
+    int numDOF = makeDimerBox(system, positions, constrain, numMolecules, bondLength);
+
+}
+
+void runPlatformTests();
+
+int main(int argc, char* argv[]) {
+    try {
+        initializeTests(argc, argv);
+        testHarmonicOscillator();
+        bool constrain;
+        constrain = false; testDimerBox(constrain);
+        constrain = true; testDimerBox(constrain);
+        testCheckpoints();
+        runPlatformTests();
+    }
+    catch(const exception& e) {
+        cout << "exception: " << e.what() << endl;
+        return 1;
+    }
+    cout << "Done" << endl;
+    return 0;
+}

wrappers/generateWrappers.py

@@ -76,7 +76,12 @@ class WrapperGenerator:
                             'static std::vector<std::string> OpenMM::Platform::loadPluginsFromDirectory',
                             'Vec3 OpenMM::LocalCoordinatesSite::getOriginWeights',
                             'Vec3 OpenMM::LocalCoordinatesSite::getXWeights',
-                            'Vec3 OpenMM::LocalCoordinatesSite::getYWeights']
+                            'Vec3 OpenMM::LocalCoordinatesSite::getYWeights',
+                            'std::vector<double> OpenMM::NoseHooverChain::getYoshidaSuzukiWeights',
+                            'const std::vector<int>& OpenMM::NoseHooverIntegrator::getAllThermostatedIndividualParticles',
+                            'const std::vector<std::tuple<int, int, double> >& OpenMM::NoseHooverIntegrator::getAllThermostatedPairs',
+                            'virtual void OpenMM::NoseHooverIntegrator::stateChanged'
+                           ] 
         self.hideClasses = ['Kernel', 'KernelImpl', 'KernelFactory', 'ContextImpl', 'SerializationNode', 'SerializationProxy']
         self.nodeByID={}
 
@@ -165,6 +170,7 @@ class CHeaderGenerator(WrapperGenerator):
                                  'std::vector< std::string >': 'OpenMM_StringArray',
                                  'std::vector< Vec3 >': 'OpenMM_Vec3Array',
                                  'std::vector< std::pair< int, int > >': 'OpenMM_BondArray',
+                                 'const std::vector< std::pair< int, int > >': 'OpenMM_BondArray',
                                  'std::map< std::string, double >': 'OpenMM_ParameterArray',
                                  'std::map< std::string, std::string >': 'OpenMM_PropertyArray',
                                  'std::vector< double >': 'OpenMM_DoubleArray',

wrappers/python/setup.py

@@ -184,7 +184,7 @@ def buildKeywordDictionary(major_version_num=MAJOR_VERSION_NUM,
     if not openmm_lib_path:
         reportError("Set OPENMM_LIB_PATH to point to the lib directory for OpenMM")
 
-    extra_compile_args=[]
+    extra_compile_args=['-std=c++11']
     extra_link_args=[]
     if platform.system() == "Windows":
         define_macros.append( ('WIN32', None) )

wrappers/python/src/swig_doxygen/swigInputBuilder.py

@@ -21,6 +21,9 @@ except ImportError:
 INDENT = "   "
 docTags = {'emphasis':'i', 'bold':'b', 'itemizedlist':'ul', 'listitem':'li', 'preformatted':'pre', 'computeroutput':'tt', 'subscript':'sub'}
 
+def is_method_abstract(argstring):
+    return argstring.split(")")[-1].find("=0") >= 0
+
 def striphtmltags(s):
     """Strip a couple html tags used inside docstrings in the C++ source
     to produce something more easily read as plain text.
@@ -511,7 +514,7 @@ class SwigInputBuilder:
                 mArgsstring = getText("argsstring", memberNode)
             if self.fOutPythonprepend and \
                len(paramList) and \
-               mArgsstring.find('=0') < 0:
+               not is_method_abstract(mArgsstring):
                 text = '''
 %pythonprepend OpenMM::{shortClassName}::{methName}{mArgsstring} %{{{{{{0}}
 %}}}}'''.format(shortClassName=shortClassName, methName=methName, mArgsstring=mArgsstring)
@@ -529,6 +532,7 @@ class SwigInputBuilder:
              % self.__class__.__name__)
         raise Exception(s)'''.format(argName=argName)
 
+
                 # Convert input arguments to the proper units, if specified.
                 if key not in methodsWithOutputArgs:
                     if key in self.configModule.UNITS:
@@ -563,9 +567,10 @@ class SwigInputBuilder:
 
             # write pythonappend blocks
             if self.fOutPythonappend \
-               and mArgsstring.find('=0') < 0:
+               and not is_method_abstract(mArgsstring):
                 key = (shortClassName, methName)
-                #print "key %s %s \n" % (shortClassName, methName)
+                #sys.stdout.write("key %s %s \n" % (shortClassName, methName))
+
                 addText=''
                 returnType = getText("type", memberNode)
 
@@ -619,8 +624,15 @@ class SwigInputBuilder:
                             pType = getText('type', pNode)
                         except IndexError:
                             pType = getText('type/ref', pNode)
+                        # parse default arguments
+                        try:
+                            defaultValue = getText('defval', pNode)
+                        except:
+                            defaultValue = ""
+                        if defaultValue != "":
+                            defaultValue = "=%s" %defaultValue
                         pName = getText('declname', pNode)
-                        self.fOutPythonappend.write("%s%s %s" % (sepChar, pType, pName))
+                        self.fOutPythonappend.write("%s%s %s%s" % (sepChar, pType, pName, defaultValue))
                         sepChar=', '
 
                         if pType.find('&')>=0 and \

wrappers/python/src/swig_doxygen/swigInputConfig.py

@@ -66,6 +66,8 @@ SKIP_METHODS = [('State', 'getPositions'),
                 ('IntegrateVariableLangevinStepKernel',),
                 ('IntegrateVariableVerletStepKernel',),
                 ('IntegrateVerletStepKernel',),
+                ('IntegrateVelocityVerletStepKernel',),
+                ('NoseHooverChainKernel',),
                 ('IntegrateCustomStepKernel',),
                 ('Kernel',),
                 ('KernelFactory',),
@@ -112,6 +114,8 @@ SKIP_METHODS = [('State', 'getPositions'),
                 ('LocalCoordinatesSite',  'getOriginWeights', 0),
                 ('LocalCoordinatesSite',  'getXWeights', 0),
                 ('LocalCoordinatesSite',  'getYWeights', 0),
+                ("NoseHooverIntegrator", "getAllThermostatedIndividualParticles"),
+                ("NoseHooverIntegrator", "getAllThermostatedPairs"),
 ]
 
 # The build script assumes method args that are non-const references are
@@ -180,6 +184,7 @@ UNITS = {
 ("*", "setDefaultSurfaceTension") : (None, ("unit.bar*unit.nanometer",)),
 ("*", "getDefaultTemperature") : ("unit.kelvin", ()),
 ("*", "setDefaultTemperature") : (None, ("unit.kelvin",)),
+("*", "getRelativeTemperature") : ("unit.kelvin", ()),
 ("*", "getErrorTolerance") : (None, ()),
 ("*", "getEwaldErrorTolerance") : (None, ()),
 ("*", "getFriction") : ("1/unit.picosecond", ()),
@@ -213,6 +218,9 @@ UNITS = {
 ("*", "getTabulatedFunction") : (None, ()),
 ("*", "getUseDispersionCorrection") : (None, ()),
 ("*", "getTemperature") : ("unit.kelvin", ()),
+("*", "getRelativeTemperature") : ("unit.kelvin", ()),
+("*", "getCollisionFrequency") : ( "1/unit.picosecond", ()),
+("*", "getRelativeCollisionFrequency") : ( "1/unit.picosecond", ()),
 ("*", "getUseDispersionCorrection") : (None, ()),
 ("*", "getWeight") : (None, ()),
 ("*", "getWeight12") : (None, ()),
@@ -339,9 +347,11 @@ UNITS = {
 ("HippoNonbondedForce",                 "getInducedDipoles")                             :  ( None, ()),
 ("HippoNonbondedForce",                 "getLabFramePermanentDipoles")                   :  ( None, ()),
 
+
 ("Context", "getParameter") : (None, ()),
 ("Context", "getParameters") : (None, ()),
 ("Context", "getMolecules") : (None, ()),
+("Context", "getState") : (None, (None, None, None)),
 ("CMAPTorsionForce", "getMapParameters") : (None, (None, 'unit.kilojoule_per_mole')),
 ("CMAPTorsionForce", "getTorsionParameters") : (None, ()),
 ("CMMotionRemover", "getFrequency") : (None, ()),
@@ -465,11 +475,28 @@ UNITS = {
 ("MonteCarloMembraneBarostat", "MonteCarloMembraneBarostat") : (None, ("unit.bar", "unit.bar*unit.nanometer", "unit.kelvin", None, None, None)),
 ("MonteCarloMembraneBarostat", "getXYMode") : (None, ()),
 ("MonteCarloMembraneBarostat", "getZMode") : (None, ()),
-("DrudeLangevinIntegrator", "getDrudeFriction") : ("1/unit.picosecond", ()),
+("DrudeLangevinIntegrator", "getDrudeFriction") : ("unit.picosecond**-1", ()),
 ("DrudeSCFIntegrator", "getMinimizationErrorTolerance") : ("unit.kilojoules_per_mole/unit.nanometer", ()),
 ("RPMDIntegrator", "getContractions") : (None, ()),
 ("RPMDIntegrator", "getTotalEnergy") : ("unit.kilojoules_per_mole", ()),
+("RPMDIntegrator", "getState"): (None,(None, None, None, None)),
 ("RMSDForce", "getReferencePositions") : ("unit.nanometer", ()),
 ("RMSDForce", "getParticles") : (None, ()),
+("NoseHooverChain", "getThermostatedPairs") : (None, ()),
+("NoseHooverChain", "getThermostatedAtoms") : (None, ()),
+("NoseHooverChain", "getYoshidaSuzukiWeights") : (None, ()),
+("NoseHooverIntegrator", "setTemperature") : (None, ("unit.kelvin", None)),
+("NoseHooverIntegrator", "setRelativeTemperature") : (None, ("unit.kelvin", None) ),
+("NoseHooverIntegrator", "setCollisionFrequency") : (None, ("unit.picosecond**-1", None)),
+("NoseHooverIntegrator", "setRelativeCollisionFrequency") : (None, ("unit.picosecond**-1", None)),
+("NoseHooverIntegrator", "computeHeatBathEnergy") : ( "unit.kilojoules_per_mole", ()),
+("NoseHooverIntegrator", "addThermostat"): (None, ("unit.kelvin", "unit.picosecond**-1", None, None, None)),
+("NoseHooverIntegrator", "addSubsystemThermostat"):
+    (None, (None, None, "unit.kelvin", "unit.picosecond**-1", "unit.kelvin", "unit.picosecond**-1", None, None, None)),
+("NoseHooverIntegrator", "getNumThermostats") : (None, ()),
+("NoseHooverIntegrator", "getThermostat") : (None, ()),
+("NoseHooverIntegrator", "getMaximumPairDistance") : ("unit.nanometer", ()),
+("NoseHooverIntegrator", "setMaximumPairDistance") : (None, ("unit.nanometer",)),
+("DrudeNoseHooverIntegrator", "getMaxDrudeDistance") : ("unit.nanometer", ()),
+("DrudeNoseHooverIntegrator", "setMaxDrudeDistance") : (None, ("unit.nanometer",)),
 }
-

wrappers/python/tests/TestAPIUnits.py

@@ -682,7 +682,7 @@ class TestAPIUnits(unittest.TestCase):
         force.setSolventDielectric(80)
         self.assertEqual(force.getSolventDielectric(), 80)
 
-        self.assertEqual(force.getSurfaceAreaFactor(),
+        self.assertAlmostEqualUnit(force.getSurfaceAreaFactor(),
                 -170.35173066268223*kilojoule_per_mole/nanometer**2) # default
         force.setSurfaceAreaFactor(-1.0*kilocalorie_per_mole/angstrom**2)
         self.assertAlmostEqualUnit(force.getSurfaceAreaFactor(),
@@ -1165,6 +1165,60 @@ class TestAPIUnits(unittest.TestCase):
         self.assertAlmostEqualUnit(integrator.getFriction(), 0.1/microsecond)
         self.assertEqual(integrator.getStepSize(), 1*femtosecond)
 
+    def testNoseHooverIntegrator(self):
+        """ Tests the NoseHooverIntegrator API features """
+        integrator = NoseHooverIntegrator(300, 0.1, 1)
+        self.assertEqual(integrator.getTemperature(0), 300*kelvin)
+        self.assertEqual(integrator.getCollisionFrequency(), 0.1/picosecond)
+        self.assertEqual(integrator.getStepSize(), 1*picosecond)
+
+        integrator = NoseHooverIntegrator(300*kelvin, 0.1/microsecond, 1*femtosecond)
+        self.assertEqual(integrator.getTemperature(), 300*kelvin)
+        self.assertAlmostEqualUnit(integrator.getCollisionFrequency(), 0.1/microsecond)
+        self.assertEqual(integrator.getStepSize(), 1*femtosecond)
+        self.assertEqual(integrator.computeHeatBathEnergy(), 0.0*kilojoule_per_mole)
+
+        # Test setters
+        integrator.setTemperature(200*kelvin)
+        self.assertEqual(integrator.getTemperature(), 200*kelvin)
+        integrator.setCollisionFrequency(0.1/picosecond)
+        self.assertEqual(integrator.getCollisionFrequency(), 0.1/picosecond)
+        integrator.setRelativeTemperature(200*kelvin)
+        self.assertEqual(integrator.getRelativeTemperature(), 200*kelvin)
+        integrator.setRelativeCollisionFrequency(0.1/picosecond)
+        self.assertEqual(integrator.getRelativeCollisionFrequency(), 0.1/picosecond)
+
+        # Test bare consructor and addThermostat
+        integrator = NoseHooverIntegrator(1*femtosecond)
+        self.assertEqual(integrator.getStepSize(), 1*femtosecond)
+
+        integrator.addThermostat(300*kelvin, 0.1/microsecond, 3, 3, 3)
+        self.assertAlmostEqualUnit(integrator.getTemperature(), 300*kelvin)
+        self.assertAlmostEqualUnit(integrator.getCollisionFrequency(), 0.1/microsecond)
+
+        integrator = NoseHooverIntegrator(1*femtosecond)
+        integrator.addSubsystemThermostat([0], [], 300*kelvin, 0.1/microsecond, 1.0*kelvin, 1.0/microsecond, 3, 3, 3)
+        self.assertAlmostEqualUnit(integrator.getTemperature(), 300*kelvin)
+        self.assertAlmostEqualUnit(integrator.getCollisionFrequency(), 0.1/microsecond)
+        self.assertAlmostEqualUnit(integrator.getRelativeTemperature(), 1.0*kelvin)
+        self.assertAlmostEqualUnit(integrator.getRelativeCollisionFrequency(), 1.0/microsecond)
+
+    def testDrudeNoseHooverIntegrator(self):
+        """ Tests the DrudeNoseHooverIntegrator API features """
+        integrator = DrudeNoseHooverIntegrator(300, 0.1, 1.0, 1.0, 1)
+        self.assertEqual(integrator.getTemperature(0), 300*kelvin)
+        self.assertEqual(integrator.getCollisionFrequency(), 0.1/picosecond)
+        self.assertEqual(integrator.getRelativeTemperature(0), 1.0*kelvin)
+        self.assertEqual(integrator.getRelativeCollisionFrequency(), 1.0/picosecond)
+        self.assertEqual(integrator.getStepSize(), 1*picosecond)
+
+        integrator = DrudeNoseHooverIntegrator(300*kelvin, 0.1/microsecond, 5.0*kelvin, 0.01/microsecond, 1*femtosecond)
+        self.assertEqual(integrator.getTemperature(), 300*kelvin)
+        self.assertAlmostEqualUnit(integrator.getCollisionFrequency(), 0.1/microsecond)
+        self.assertEqual(integrator.getRelativeTemperature(), 5.0*kelvin)
+        self.assertAlmostEqualUnit(integrator.getRelativeCollisionFrequency(), 0.01/microsecond)
+        self.assertEqual(integrator.getStepSize(), 1*femtosecond)
+
     def testAndersenThermostat(self):
         """ Tests the AndersenThermostat API features """
         force = AndersenThermostat(300*kelvin, 1/microsecond)

wrappers/python/tests/TestIntegrators.py

 import unittest
+import warnings
 import tempfile
 from datetime import datetime, timedelta
 from simtk.openmm import *
@@ -120,5 +121,36 @@ class TestIntegrators(unittest.TestCase):
             context.setPositions(pdb.positions)
             integrator.step(10)
 
+    def testNoseHooverIntegrator(self):
+        """Test partial thermostating in the NoseHooverIntegrator (only API)"""
+        pdb = PDBFile('systems/alanine-dipeptide-explicit.pdb')
+        ff = ForceField('amber99sbildn.xml', 'tip3p.xml')
+        system = ff.createSystem(pdb.topology, cutoffMethod=PME)
+
+        integrator = NoseHooverIntegrator(1.0*femtosecond) 
+        integrator.addSubsystemThermostat(list(range(5)), [], 200*kelvin, 1/picosecond, 200*kelvin, 1/picosecond, 3,3,3)
+        con = Context(system, integrator)
+        con.setPositions(pdb.positions)
+
+        integrator.step(5)
+        self.assertNotEqual(integrator.computeHeatBathEnergy(), 0.0*kilojoule_per_mole)
+
+    def testDrudeNoseHooverIntegrator(self):
+        """Test the DrudeNoseHooverIntegrator"""
+        warnings.filterwarnings('ignore', category=CharmmPSFWarning)
+        psf = CharmmPsfFile('systems/ala3_solv_drude.psf')
+        crd = CharmmCrdFile('systems/ala3_solv_drude.crd')
+        params = CharmmParameterSet('systems/toppar_drude_master_protein_2013e.str')
+        # Box dimensions (cubic box)
+        psf.setBox(33.2*angstroms, 33.2*angstroms, 33.2*angstroms)
+
+        system = psf.createSystem(params, nonbondedMethod=PME, ewaldErrorTolerance=0.0005)
+        integrator = DrudeNoseHooverIntegrator(300*kelvin, 1.0/picosecond, 1*kelvin, 10/picosecond, 0.001*picoseconds)
+        con = Context(system, integrator)
+        con.setPositions(crd.positions)
+
+        integrator.step(5)
+        self.assertNotEqual(integrator.computeHeatBathEnergy(), 0.0*kilojoule_per_mole)
+
 if __name__ == '__main__':
     unittest.main()