Adaptive quantum thermal bath (#4995)

* Began implementing QTBIntegrator

* Adaptation and deconvolution

* Continuing reference implementation

* Continuing to implement QTBIntegrator

* Use common thread pool

* More tests, documentation, and threading

* Fix segfault

* Serialize adapted friction when creating a State

* Beginning of GPU implementation

* Added missing files

* Bug fixes

* Fixed inverse FFT

* Continuing GPU implementation

* Checkpointing

* Bug fixes

* Test cases run faster

* Changes needed for latest main branch

* Minor optimizations

* Documentation

* Fixed atom reordering

* Added parahydrogen test case

* Workaround for bug in Microsoft's compiler

* Added a Python test

* Normalize kernel in deconvolution

* Minor documentation improvements

platforms/opencl/src/OpenCLPlatform.cpp

@@ -96,6 +96,7 @@ OpenCLPlatform::OpenCLPlatform() {
     registerKernelFactory(IntegrateVariableLangevinStepKernel::Name(), factory);
     registerKernelFactory(IntegrateCustomStepKernel::Name(), factory);
     registerKernelFactory(IntegrateDPDStepKernel::Name(), factory);
+    registerKernelFactory(IntegrateQTBStepKernel::Name(), factory);
     registerKernelFactory(ApplyAndersenThermostatKernel::Name(), factory);
     registerKernelFactory(ApplyMonteCarloBarostatKernel::Name(), factory);
     registerKernelFactory(RemoveCMMotionKernel::Name(), factory);

platforms/opencl/tests/TestOpenCLQTBIntegrator.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) 2025 Stanford University and the Authors.           *
+ * Authors: Peter Eastman                                                     *
+ * 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 "OpenCLTests.h"
+#include "TestQTBIntegrator.h"
+
+void runPlatformTests() {
+}

platforms/reference/include/ReferenceKernels.h

@@ -69,6 +69,7 @@ class ReferenceVariableVerletDynamics;
 class ReferenceVerletDynamics;
 class ReferenceCustomDynamics;
 class ReferenceDPDDynamics;
+class ReferenceQTBDynamics;
 
 /**
  * This kernel is invoked at the beginning and end of force and energy computations.  It gives the
@@ -1546,6 +1547,76 @@ private:
     std::vector<double> masses;
 };
 
+/**
+ * This kernel is invoked by QTBIntegrator to take one time step.
+ */
+class ReferenceIntegrateQTBStepKernel : public IntegrateQTBStepKernel {
+public:
+    ReferenceIntegrateQTBStepKernel(std::string name, const Platform& platform, ReferencePlatform::PlatformData& data) : IntegrateQTBStepKernel(name, platform),
+        data(data), dynamics(NULL), hasInitialized(false) {
+    }
+    ~ReferenceIntegrateQTBStepKernel();
+    /**
+     * Initialize the kernel.
+     * 
+     * @param system     the System this kernel will be applied to
+     * @param integrator the QTBIntegrator this kernel will be used for
+     */
+    void initialize(const System& system, const QTBIntegrator& integrator);
+    /**
+     * Execute the kernel.
+     * 
+     * @param context    the context in which to execute this kernel
+     * @param integrator the QTBIntegrator this kernel is being used for
+     */
+    void execute(ContextImpl& context, const QTBIntegrator& integrator);
+    /**
+     * Compute the kinetic energy.
+     * 
+     * @param context    the context in which to execute this kernel
+     * @param integrator the QTBIntegrator this kernel is being used for
+     */
+    double computeKineticEnergy(ContextImpl& context, const QTBIntegrator& integrator);
+    /**
+     * Get the adapted friction coefficients for a particle.
+     * 
+     * @param context    the context in which to execute this kernel
+     * @param particle   the index of the particle for which to get the friction
+     * @param friction   the adapted friction coefficients used in generating the
+     *                   random force
+     */
+    void getAdaptedFriction(ContextImpl& context, int particle, std::vector<double>& friction) const;
+    /**
+     * Set the adapted friction coefficients for a particle.  This affects the
+     * specified particle, and all others that have the same type.
+     * 
+     * @param context    the context in which to execute this kernel
+     * @param particle   the index of the particle for which to get the friction
+     * @param friction   the adapted friction coefficients used in generating the
+     *                   random force
+     */
+    void setAdaptedFriction(ContextImpl& context, int particle, const std::vector<double>& friction);
+    /**
+     * Write the adapted friction to a checkpoint.
+     * 
+     * @param context    the context in which to execute this kernel
+     * @param stream     the stream to write the checkpoint to
+     */
+    void createCheckpoint(ContextImpl& context, std::ostream& stream) const;
+    /**
+     * Load the adapted friction from a checkpoint.
+     * 
+     * @param context    the context in which to execute this kernel
+     * @param stream     the stream to read the checkpoint from
+     */
+    void loadCheckpoint(ContextImpl& context, std::istream& stream);
+private:
+    ReferencePlatform::PlatformData& data;
+    ReferenceQTBDynamics* dynamics;
+    std::vector<double> masses;
+    bool hasInitialized;
+};
+
 /**
  * This kernel is invoked by AndersenThermostat at the start of each time step to adjust the particle velocities.
  */

platforms/reference/include/ReferenceQTBDynamics.h

+/* Portions copyright (c) 2025 Stanford University and the Authors.
+ * Authors: Peter Eastman
+ *
+ * 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.
+ */
+
+#ifndef __ReferenceQTBDynamics_H__
+#define __ReferenceQTBDynamics_H__
+
+#include "ReferenceDynamics.h"
+#include "openmm/QTBIntegrator.h"
+#include "openmm/internal/ContextImpl.h"
+#include "openmm/internal/ThreadPool.h"
+#include "openmm/internal/windowsExport.h"
+#include <complex>
+
+namespace OpenMM {
+
+class OPENMM_EXPORT ReferenceQTBDynamics : public ReferenceDynamics {
+protected:
+    double friction, lastTemperature;
+    int segmentLength, stepIndex, numFreq;
+    std::vector<OpenMM::Vec3> xPrime, oldx, randomForce, segmentVelocity;
+    std::vector<double> inverseMasses, typeAdaptationRate, typeMass;
+    std::vector<double> noise, theta, thetad, cutoffFunction;
+    std::vector<int> particleType;
+    std::vector<std::vector<int> > typeParticles;
+    std::vector<std::vector<double> > adaptedFriction;
+
+public:
+    /**
+     * Constructor
+     * 
+     * @param system       the system to simulate
+     * @param integrator   the integrator being used
+     */
+    ReferenceQTBDynamics(const System& system, const QTBIntegrator& integrator);
+
+    /**
+     * Destructor
+     */
+    ~ReferenceQTBDynamics();
+
+    /**
+     * Perform a time step, updating the positions and velocities.
+     * 
+     * @param context             the context this integrator is updating
+     * @param system              the System to be integrated
+     * @param atomCoordinates     atom coordinates
+     * @param velocities          velocities
+     * @param masses              atom masses
+     * @param tolerance           the constraint tolerance
+     * @param boxVectors          the current periodic box vectors
+     * @param threads      a ThreadPool to use for parallelization
+     */
+    void update(OpenMM::ContextImpl& context, std::vector<OpenMM::Vec3>& atomCoordinates,
+                std::vector<OpenMM::Vec3>& velocities, std::vector<double>& masses, double tolerance, const Vec3* boxVectors, ThreadPool& threads);
+
+    /**
+     * The first stage of the update algorithm.
+     */
+    virtual void updatePart1(int numParticles, std::vector<OpenMM::Vec3>& velocities, std::vector<OpenMM::Vec3>& forces);
+
+    /**
+     * The second stage of the update algorithm.
+     */
+    virtual void updatePart2(int numParticles, std::vector<OpenMM::Vec3>& atomCoordinates, std::vector<OpenMM::Vec3>& velocities,
+                             std::vector<OpenMM::Vec3>& xPrime);
+
+    /**
+     * The third stage of the update algorithm.
+     */
+    virtual void updatePart3(OpenMM::ContextImpl& context, int numParticles, std::vector<OpenMM::Vec3>& atomCoordinates,
+                             std::vector<OpenMM::Vec3>& velocities, std::vector<OpenMM::Vec3>& xPrime);
+    /**
+     * Get the adapted friction coefficients for a particle.
+     * 
+     * @param particle   the index of the particle for which to get the friction
+     * @param friction   the adapted friction coefficients used in generating the
+     *                   random force.
+     */
+    void getAdaptedFriction(int particle, std::vector<double>& friction) const;
+    /**
+     * Set the adapted friction coefficients for a particle.  This affects the
+     * specified particle, and all others that have the same type.
+     * 
+     * @param particle   the index of the particle for which to get the friction
+     * @param friction   the adapted friction coefficients used in generating the
+     *                   random force.
+     */
+    void setAdaptedFriction(int particle, const std::vector<double>& friction);
+    /**
+     * Write the adapted friction to a checkpoint.
+     */
+    void createCheckpoint(std::ostream& stream) const;
+    /**
+     * Load the adapted friction from a checkpoint.
+     */
+    void loadCheckpoint(std::istream& stream);
+
+private:
+    /**
+     * Generate noise for the next segment.
+     */
+    void generateNoise(int numParticles, std::vector<double>& masses, ThreadPool& threads);
+    /**
+     * Update the friction rates used for generating noise.
+     */
+    void adaptFriction(ThreadPool& threads);
+};
+
+} // namespace OpenMM
+
+#endif // __ReferenceQTBDynamics_H__

platforms/reference/src/ReferenceKernelFactory.cpp

@@ -106,6 +106,8 @@ KernelImpl* ReferenceKernelFactory::createKernelImpl(std::string name, const Pla
         return new ReferenceIntegrateCustomStepKernel(name, platform, data);
     if (name == IntegrateDPDStepKernel::Name())
         return new ReferenceIntegrateDPDStepKernel(name, platform, data);
+    if (name == IntegrateQTBStepKernel::Name())
+        return new ReferenceIntegrateQTBStepKernel(name, platform, data);
     if (name == ApplyAndersenThermostatKernel::Name())
         return new ReferenceApplyAndersenThermostatKernel(name, platform);
     if (name == ApplyMonteCarloBarostatKernel::Name())

platforms/reference/src/ReferenceKernels.cpp

@@ -60,6 +60,7 @@
 #include "ReferenceNoseHooverDynamics.h"
 #include "ReferencePointFunctions.h"
 #include "ReferenceProperDihedralBond.h"
+#include "ReferenceQTBDynamics.h"
 #include "ReferenceRbDihedralBond.h"
 #include "ReferenceRMSDForce.h"
 #include "ReferenceTabulatedFunction.h"
@@ -2899,6 +2900,55 @@ double ReferenceIntegrateDPDStepKernel::computeKineticEnergy(ContextImpl& contex
     return computeShiftedKineticEnergy(context, masses, 0.0);
 }
 
+ReferenceIntegrateQTBStepKernel::~ReferenceIntegrateQTBStepKernel() {
+    if (dynamics != NULL)
+        delete dynamics;
+}
+
+void ReferenceIntegrateQTBStepKernel::initialize(const System& system, const QTBIntegrator& integrator) {
+    int numParticles = system.getNumParticles();
+    masses.resize(numParticles);
+    for (int i = 0; i < numParticles; ++i)
+        masses[i] = system.getParticleMass(i);
+    SimTKOpenMMUtilities::setRandomNumberSeed((unsigned int) integrator.getRandomNumberSeed());
+    dynamics = new ReferenceQTBDynamics(system, integrator);
+}
+
+void ReferenceIntegrateQTBStepKernel::execute(ContextImpl& context, const QTBIntegrator& integrator) {
+    double stepSize = integrator.getStepSize();
+    vector<Vec3>& posData = extractPositions(context);
+    vector<Vec3>& velData = extractVelocities(context);
+    if (!hasInitialized) {
+        hasInitialized = true;
+        dynamics->setReferenceConstraintAlgorithm(&extractConstraints(context));
+        dynamics->setVirtualSites(extractVirtualSites(context));
+    }
+    dynamics->setTemperature(integrator.getTemperature());
+    dynamics->update(context, posData, velData, masses, integrator.getConstraintTolerance(), extractBoxVectors(context), extractThreadPool(context));
+    data.time += stepSize;
+    data.stepCount++;
+}
+
+double ReferenceIntegrateQTBStepKernel::computeKineticEnergy(ContextImpl& context, const QTBIntegrator& integrator) {
+    return computeShiftedKineticEnergy(context, masses, 0.0);
+}
+
+void ReferenceIntegrateQTBStepKernel::getAdaptedFriction(ContextImpl& context, int particle, std::vector<double>& friction) const {
+    dynamics->getAdaptedFriction(particle, friction);
+}
+
+void ReferenceIntegrateQTBStepKernel::setAdaptedFriction(ContextImpl& context, int particle, const std::vector<double>& friction) {
+    dynamics->setAdaptedFriction(particle, friction);
+}
+
+void ReferenceIntegrateQTBStepKernel::createCheckpoint(ContextImpl& context, ostream& stream) const {
+    dynamics->createCheckpoint(stream);
+}
+
+void ReferenceIntegrateQTBStepKernel::loadCheckpoint(ContextImpl& context, istream& stream) {
+    dynamics->loadCheckpoint(stream);
+}
+
 ReferenceApplyAndersenThermostatKernel::~ReferenceApplyAndersenThermostatKernel() {
     if (thermostat)
         delete thermostat;

platforms/reference/src/ReferencePlatform.cpp

@@ -76,6 +76,7 @@ ReferencePlatform::ReferencePlatform() {
     registerKernelFactory(IntegrateVariableVerletStepKernel::Name(), factory);
     registerKernelFactory(IntegrateCustomStepKernel::Name(), factory);
     registerKernelFactory(IntegrateDPDStepKernel::Name(), factory);
+    registerKernelFactory(IntegrateQTBStepKernel::Name(), factory);
     registerKernelFactory(ApplyAndersenThermostatKernel::Name(), factory);
     registerKernelFactory(ApplyMonteCarloBarostatKernel::Name(), factory);
     registerKernelFactory(RemoveCMMotionKernel::Name(), factory);

platforms/reference/src/SimTKReference/ReferenceQTBDynamics.cpp

+/* Portions copyright (c) 2025 Stanford University and the Authors.
+ * Authors: Peter Eastman
+ *
+ * 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 "SimTKOpenMMUtilities.h"
+#include "ReferenceQTBDynamics.h"
+#include "ReferencePlatform.h"
+#include "ReferenceVirtualSites.h"
+#include "openmm/OpenMMException.h"
+#include "openmm/internal/ContextImpl.h"
+#include "openmm/internal/QTBIntegratorUtilities.h"
+#include <algorithm>
+#include <cmath>
+#include <map>
+
+#ifdef _MSC_VER
+  #define POCKETFFT_NO_VECTORS
+#endif
+#include "pocketfft_hdronly.h"
+
+using namespace OpenMM;
+using namespace std;
+
+ReferenceQTBDynamics::ReferenceQTBDynamics(const System& system, const QTBIntegrator& integrator) :
+            ReferenceDynamics(system.getNumParticles(), integrator.getStepSize(), integrator.getTemperature()), friction(integrator.getFriction()), stepIndex(0) {
+    segmentLength = (int) ceil(integrator.getSegmentLength()/integrator.getStepSize());
+    int numParticles = system.getNumParticles();
+    xPrime.resize(numParticles);
+    oldx.resize(numParticles);
+    noise.resize(3*3*segmentLength*numParticles);
+    randomForce.resize(segmentLength*numParticles);
+    segmentVelocity.resize(segmentLength*numParticles);
+    for (int i = 0; i < noise.size(); i++)
+        noise[i] = SimTKOpenMMUtilities::getNormallyDistributedRandomNumber();
+    QTBIntegratorUtilities::findTypes(system, integrator, particleType, typeParticles, typeMass, typeAdaptationRate);
+
+    // Calculate the target energy distribution.
+
+    numFreq = (3*segmentLength+1)/2;
+    cutoffFunction.resize(numFreq);
+    double cutoff = integrator.getCutoffFrequency();
+    double cutoffWidth = cutoff/100;
+    for (int i = 0; i < numFreq; i++) {
+        double w = M_PI*i/(numFreq*getDeltaT());
+        cutoffFunction[i] = 1.0/(1.0+exp((w-cutoff)/cutoffWidth));
+    }
+
+    // Allocate space for adaptation.
+
+    int numTypes = typeParticles.size();
+    adaptedFriction.resize(numTypes, vector<double>(numFreq, friction));
+}
+
+ReferenceQTBDynamics::~ReferenceQTBDynamics() {
+}
+
+void ReferenceQTBDynamics::updatePart1(int numParticles, vector<Vec3>& velocities, vector<Vec3>& forces) {
+    for (int i = 0; i < numParticles; i++) {
+        segmentVelocity[i*segmentLength+stepIndex] = velocities[i];
+        if (inverseMasses[i] != 0.0)
+            velocities[i] += (getDeltaT()*inverseMasses[i])*forces[i];
+    }
+}
+
+void ReferenceQTBDynamics::updatePart2(int numParticles, vector<Vec3>& atomCoordinates,
+                                         vector<Vec3>& velocities, vector<Vec3>& xPrime) {
+    const double dt = getDeltaT();
+    const double halfdt = 0.5*dt;
+    const double vscale = exp(-dt*friction);
+
+    for (int i = 0; i < numParticles; i++) {
+        if (inverseMasses[i] != 0.0) {
+            xPrime[i] = atomCoordinates[i] + velocities[i]*halfdt;
+            velocities[i] = vscale*velocities[i] + inverseMasses[i]*dt*randomForce[segmentLength*i+stepIndex];
+            xPrime[i] = xPrime[i] + velocities[i]*halfdt;
+            oldx[i] = xPrime[i];
+        }
+    }
+}
+
+void ReferenceQTBDynamics::updatePart3(OpenMM::ContextImpl& context, int numParticles, vector<Vec3>& atomCoordinates,
+                                         vector<Vec3>& velocities, vector<Vec3>& xPrime) {
+    for (int i = 0; i < numParticles; i++) {
+        if (inverseMasses[i] != 0.0) {
+            velocities[i] += (xPrime[i]-oldx[i])/getDeltaT();
+            atomCoordinates[i] = xPrime[i];
+        }
+    }
+}
+
+void ReferenceQTBDynamics::update(ContextImpl& context, vector<Vec3>& atomCoordinates, vector<Vec3>& velocities,
+            vector<double>& masses, double tolerance, const Vec3* boxVectors, ThreadPool& threads) {
+    int numParticles = context.getSystem().getNumParticles();
+    ReferenceConstraintAlgorithm* referenceConstraintAlgorithm = getReferenceConstraintAlgorithm();
+    if (inverseMasses.size() == 0) {
+        // Invert masses
+
+        inverseMasses.resize(numParticles);
+        for (int ii = 0; ii < numParticles; ii++) {
+            if (masses[ii] == 0.0)
+                inverseMasses[ii] = 0.0;
+            else
+                inverseMasses[ii] = 1.0/masses[ii];
+        }
+    }
+    ReferencePlatform::PlatformData* data = reinterpret_cast<ReferencePlatform::PlatformData*>(context.getPlatformData());
+    vector<Vec3>& forces = *data->forces;
+    if (stepIndex%segmentLength == 0) {
+        if (lastTemperature != getTemperature() || thetad.size() == 0) {
+            QTBIntegratorUtilities::calculateSpectrum(getTemperature(), friction, getDeltaT(), numFreq, theta, thetad, threads);
+            lastTemperature = getTemperature();
+        }
+        adaptFriction(threads);
+        generateNoise(numParticles, masses, threads);
+        stepIndex = 0;
+    }
+
+    // 1st update
+
+    updatePart1(numParticles, velocities, forces);
+    if (referenceConstraintAlgorithm)
+        referenceConstraintAlgorithm->applyToVelocities(atomCoordinates, velocities, inverseMasses, tolerance);
+
+    // 2nd update
+
+    updatePart2(numParticles, atomCoordinates, velocities, xPrime);
+    if (referenceConstraintAlgorithm)
+        referenceConstraintAlgorithm->apply(atomCoordinates, xPrime, inverseMasses, tolerance);
+
+    // 3rd update
+
+    updatePart3(context, numParticles, atomCoordinates, velocities, xPrime);
+
+    getVirtualSites().computePositions(context.getSystem(), atomCoordinates, boxVectors);
+    incrementTimeStep();
+    stepIndex += 1;
+}
+
+void ReferenceQTBDynamics::generateNoise(int numParticles, vector<double>& masses, ThreadPool& threads) {
+    // Update the buffer of white noise.
+
+    for (int base = 0; base < noise.size(); base += 3*segmentLength) {
+        for (int i = 0; i < 2*segmentLength; i++)
+            noise[base+i] = noise[base+i+segmentLength];
+        for (int i = 0; i < segmentLength; i++)
+            noise[base+2*segmentLength+i] = SimTKOpenMMUtilities::getNormallyDistributedRandomNumber();
+    }
+
+    // Generate the random force for the next segment.
+
+    double dt = getDeltaT();
+    vector<ptrdiff_t> realStride = {(ptrdiff_t) sizeof(double)};
+    vector<ptrdiff_t> complexStride = {(ptrdiff_t) sizeof(complex<double>)};
+    vector<size_t> shape = {(size_t) 3*segmentLength}, axes = {0};
+    threads.execute([&] (ThreadPool& threads, int threadIndex) {
+        vector<complex<double> > recipData(3*segmentLength);
+        vector<double> force(3*segmentLength);
+        for (int particle = threadIndex; particle < numParticles; particle += threads.getNumThreads()) {
+            int type = particleType[particle];
+            for (int axis = 0; axis < 3; axis++) {
+                double* data = &noise[(3*particle+axis)*3*segmentLength];
+                pocketfft::r2c(shape, realStride, complexStride, axes, true, data, recipData.data(), 1.0, 1);
+                for (int i = 0; i < numFreq; i++) {
+                    double w = M_PI*i/(numFreq*dt);
+                    double gamma = adaptedFriction[type][i];
+                    double cw = (1 - 2*exp(-dt*friction)*cos(w*dt) + exp(-2*friction*dt)) / ((friction*friction+w*w)*dt*dt);
+                    recipData[i] *= sqrt(cutoffFunction[i]*thetad[i]*cw*gamma/friction);
+                }
+                pocketfft::c2r(shape, complexStride, realStride, axes, false, recipData.data(), force.data(), 1.0/(3*segmentLength), 1);
+                for (int i = 0; i < segmentLength; i++)
+                    randomForce[particle*segmentLength+i][axis] = sqrt(2*masses[particle]*friction/dt)*force[segmentLength+i];
+            }
+        }
+    });
+    threads.waitForThreads();
+}
+
+void ReferenceQTBDynamics::adaptFriction(ThreadPool& threads) {
+    vector<ptrdiff_t> realStride = {(ptrdiff_t) sizeof(double)};
+    vector<ptrdiff_t> complexStride = {(ptrdiff_t) sizeof(complex<double>)};
+    vector<size_t> shape = {(size_t) 3*segmentLength}, axes = {0};
+    threads.execute([&] (ThreadPool& threads, int threadIndex) {
+        vector<double> vel(3*segmentLength, 0.0), force(3*segmentLength, 0.0);
+        vector<complex<double> > recipVel(3*segmentLength), recipForce(3*segmentLength);
+        vector<double> dfdt(numFreq);
+        for (int type = threadIndex; type < typeParticles.size(); type += threads.getNumThreads()) {
+            for (int i = 0; i < dfdt.size(); i++)
+                dfdt[i] = 0;
+            for (int particle : typeParticles[type]) {
+                for (int axis = 0; axis < 3; axis++) {
+                    // Compute the Fourier transformed velocity and force.
+
+                    for (int i = 0; i < segmentLength; i++) {
+                        vel[i] = segmentVelocity[particle*segmentLength+i][axis];
+                        force[i] = randomForce[particle*segmentLength+i][axis];
+                    }
+                    pocketfft::r2c(shape, realStride, complexStride, axes, true, vel.data(), recipVel.data(), 1.0, 1);
+                    pocketfft::r2c(shape, realStride, complexStride, axes, true, force.data(), recipForce.data(), 1.0, 1);
+
+                    // Compute the error in the fluctuation dissipation theorem.
+
+                    double mass = typeMass[type];
+                    for (int i = 0; i < numFreq; i++) {
+                        double cvv = norm(recipVel[i]);
+                        complex<double> cvf = recipVel[i]*conj(recipForce[i]);
+                        dfdt[i] += mass*adaptedFriction[type][i]*cvv - cvf.real();
+                    }
+                }
+            }
+
+            // Average over particles and axes, and update the friction.
+
+            double scale = getDeltaT()/(3*typeParticles[type].size()*segmentLength);
+            for (int i = 0; i < adaptedFriction[type].size(); i++)
+                adaptedFriction[type][i] = max(0.0, adaptedFriction[type][i]-scale*typeAdaptationRate[type]*dfdt[i]);
+        }
+    });
+    threads.waitForThreads();
+}
+
+void ReferenceQTBDynamics::getAdaptedFriction(int particle, vector<double>& friction) const {
+    friction = adaptedFriction[particleType[particle]];
+}
+
+void ReferenceQTBDynamics::setAdaptedFriction(int particle, const std::vector<double>& friction) {
+    adaptedFriction[particleType[particle]] = friction;
+}
+
+void ReferenceQTBDynamics::createCheckpoint(std::ostream& stream) const {
+    stream.write((char*) &stepIndex, sizeof(int));
+    stream.write((char*) noise.data(), sizeof(double)*noise.size());
+    stream.write((char*) randomForce.data(), sizeof(Vec3)*randomForce.size());
+    stream.write((char*) segmentVelocity.data(), sizeof(Vec3)*segmentVelocity.size());
+    for (int i = 0; i < adaptedFriction.size(); i++)
+        stream.write((char*) adaptedFriction[i].data(), sizeof(double)*adaptedFriction[i].size());
+}
+
+void ReferenceQTBDynamics::loadCheckpoint(std::istream& stream) {
+    stream.read((char*) &stepIndex, sizeof(int));
+    stream.read((char*) noise.data(), sizeof(double)*noise.size());
+    stream.read((char*) randomForce.data(), sizeof(Vec3)*randomForce.size());
+    stream.read((char*) segmentVelocity.data(), sizeof(Vec3)*segmentVelocity.size());
+    for (int i = 0; i < adaptedFriction.size(); i++)
+        stream.read((char*) adaptedFriction[i].data(), sizeof(double)*adaptedFriction[i].size());
+}

platforms/reference/tests/TestReferenceQTBIntegrator.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) 2025 Stanford University and the Authors.           *
+ * Authors: Peter Eastman                                                     *
+ * 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 "ReferenceTests.h"
+#include "TestQTBIntegrator.h"
+
+void runPlatformTests() {
+}

serialization/include/openmm/serialization/QTBIntegratorProxy.h

+#ifndef OPENMM_QTB_INTEGRATOR_PROXY_H_
+#define OPENMM_QTB_INTEGRATOR_PROXY_H_
+
+#include "openmm/serialization/XmlSerializer.h"
+
+namespace OpenMM {
+
+class QTBIntegratorProxy : public SerializationProxy {
+public:
+    QTBIntegratorProxy();
+    void serialize(const void* object, SerializationNode& node) const;
+    void* deserialize(const SerializationNode& node) const;
+};
+
+}
+
+#endif /*OPENMM_QTB_INTEGRATOR_PROXY_H_*/
\ No newline at end of file

serialization/src/QTBIntegratorProxy.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) 2025 Stanford University and the Authors.           *
+ * Authors: Peter Eastman                                                     *
+ * 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/serialization/QTBIntegratorProxy.h"
+#include <OpenMM.h>
+
+using namespace std;
+using namespace OpenMM;
+
+QTBIntegratorProxy::QTBIntegratorProxy() : SerializationProxy("QTBIntegrator") {
+}
+
+void QTBIntegratorProxy::serialize(const void* object, SerializationNode& node) const {
+    node.setIntProperty("version", 1);
+    const QTBIntegrator& integrator = *reinterpret_cast<const QTBIntegrator*>(object);
+    node.setDoubleProperty("stepSize", integrator.getStepSize());
+    node.setDoubleProperty("constraintTolerance", integrator.getConstraintTolerance());
+    node.setDoubleProperty("temperature", integrator.getTemperature());
+    node.setDoubleProperty("friction", integrator.getFriction());
+    node.setDoubleProperty("segmentLength", integrator.getSegmentLength());
+    node.setDoubleProperty("cutoffFrequency", integrator.getCutoffFrequency());
+    node.setDoubleProperty("defaultAdaptationRate", integrator.getDefaultAdaptationRate());
+    node.setIntProperty("randomSeed", integrator.getRandomNumberSeed());
+    node.setIntProperty("integrationForceGroups", integrator.getIntegrationForceGroups());
+    SerializationNode& particleTypes = node.createChildNode("ParticleTypes");
+    for (auto type : integrator.getParticleTypes())
+        particleTypes.createChildNode("Particle").setIntProperty("index", type.first).setIntProperty("type", type.second);
+    SerializationNode& adaptationRates = node.createChildNode("TypeAdaptationRates");
+    for (auto type : integrator.getTypeAdaptationRates())
+        adaptationRates.createChildNode("Type").setIntProperty("index", type.first).setDoubleProperty("rate", type.second);
+}
+
+void* QTBIntegratorProxy::deserialize(const SerializationNode& node) const {
+    if (node.getIntProperty("version") != 1)
+        throw OpenMMException("Unsupported version number");
+    QTBIntegrator *integrator = new QTBIntegrator(node.getDoubleProperty("temperature"),
+                                                  node.getDoubleProperty("friction"),
+                                                  node.getDoubleProperty("stepSize"));
+    integrator->setConstraintTolerance(node.getDoubleProperty("constraintTolerance"));
+    integrator->setSegmentLength(node.getDoubleProperty("segmentLength"));
+    integrator->setCutoffFrequency(node.getDoubleProperty("cutoffFrequency"));
+    integrator->setDefaultAdaptationRate(node.getDoubleProperty("defaultAdaptationRate"));
+    integrator->setRandomNumberSeed(node.getIntProperty("randomSeed"));
+    integrator->setIntegrationForceGroups(node.getIntProperty("integrationForceGroups", 0xFFFFFFFF));
+    const SerializationNode& particleTypes = node.getChildNode("ParticleTypes");
+    for (auto& particle : particleTypes.getChildren())
+        integrator->setParticleType(particle.getIntProperty("index"), particle.getIntProperty("type"));
+    const SerializationNode& adaptationRates = node.getChildNode("TypeAdaptationRates");
+    for (auto& type : adaptationRates.getChildren())
+        integrator->setTypeAdaptationRate(type.getIntProperty("index"), type.getDoubleProperty("rate"));
+    return integrator;
+}
\ No newline at end of file

serialization/src/SerializationProxyRegistration.cpp

@@ -62,6 +62,7 @@
 #include "openmm/NonbondedForce.h"
 #include "openmm/NoseHooverIntegrator.h"
 #include "openmm/PeriodicTorsionForce.h"
+#include "openmm/QTBIntegrator.h"
 #include "openmm/RBTorsionForce.h"
 #include "openmm/RMSDForce.h"
 #include "openmm/System.h"
@@ -104,6 +105,7 @@
 #include "openmm/serialization/NonbondedForceProxy.h"
 #include "openmm/serialization/NoseHooverIntegratorProxy.h"
 #include "openmm/serialization/PeriodicTorsionForceProxy.h"
+#include "openmm/serialization/QTBIntegratorProxy.h"
 #include "openmm/serialization/RBTorsionForceProxy.h"
 #include "openmm/serialization/RMSDForceProxy.h"
 #include "openmm/serialization/StateProxy.h"
@@ -168,6 +170,7 @@ extern "C" void registerSerializationProxies() {
     SerializationProxy::registerProxy(typeid(NoseHooverIntegrator), new NoseHooverIntegratorProxy());
     SerializationProxy::registerProxy(typeid(PeriodicTorsionForce), new PeriodicTorsionForceProxy());
     SerializationProxy::registerProxy(typeid(RBTorsionForce), new RBTorsionForceProxy());
+    SerializationProxy::registerProxy(typeid(QTBIntegrator), new QTBIntegratorProxy());
     SerializationProxy::registerProxy(typeid(RMSDForce), new RMSDForceProxy());
     SerializationProxy::registerProxy(typeid(System), new SystemProxy());
     SerializationProxy::registerProxy(typeid(State), new StateProxy());

serialization/tests/TestSerializeQTBIntegrator.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) 2025 Stanford University and the Authors.           *
+ * Authors: Peter Eastman                                                     *
+ * 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/QTBIntegrator.h"
+#include "openmm/serialization/XmlSerializer.h"
+#include <iostream>
+#include <sstream>
+#include <stdlib.h>
+
+#include <fstream>
+
+using namespace OpenMM;
+using namespace std;
+
+void testSerialization() {
+    QTBIntegrator integ(310.0, 10.0, 0.002);
+    integ.setSegmentLength(0.5);
+    integ.setCutoffFrequency(600.0);
+    integ.setDefaultAdaptationRate(0.02);
+    integ.setRandomNumberSeed(10);
+    integ.setIntegrationForceGroups(3);
+    for (int i = 0; i < 5; i++)
+        integ.setParticleType(i, i%3);
+    for (int i = 0; i < 3; i++)
+        integ.setTypeAdaptationRate(i, 0.1*i);
+    stringstream ss;
+    XmlSerializer::serialize<Integrator>(&integ, "QTBIntegrator", ss);
+    QTBIntegrator *copy = dynamic_cast<QTBIntegrator*>(XmlSerializer::deserialize<Integrator>(ss));
+    QTBIntegrator& integ2 = *copy;
+    ASSERT_EQUAL(integ.getConstraintTolerance(), integ2.getConstraintTolerance());
+    ASSERT_EQUAL(integ.getStepSize(), integ2.getStepSize());
+    ASSERT_EQUAL(integ.getTemperature(), integ2.getTemperature());
+    ASSERT_EQUAL(integ.getFriction(), integ2.getFriction());
+    ASSERT_EQUAL(integ.getSegmentLength(), integ2.getSegmentLength());
+    ASSERT_EQUAL(integ.getCutoffFrequency(), integ2.getCutoffFrequency());
+    ASSERT_EQUAL(integ.getDefaultAdaptationRate(), integ2.getDefaultAdaptationRate());
+    ASSERT_EQUAL(integ.getRandomNumberSeed(), integ2.getRandomNumberSeed());
+    ASSERT_EQUAL(integ.getIntegrationForceGroups(), integ2.getIntegrationForceGroups());
+    ASSERT_EQUAL_CONTAINERS(integ.getParticleTypes(), integ2.getParticleTypes());
+    ASSERT_EQUAL_CONTAINERS(integ.getTypeAdaptationRates(), integ2.getTypeAdaptationRates());
+    delete copy;
+}
+
+int main() {
+    try {
+        testSerialization();
+    }
+    catch(const exception& e) {
+        cout << "exception: " << e.what() << endl;
+        return 1;
+    }
+    cout << "Done" << endl;
+    return 0;
+}

tests/TestQTBIntegrator.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) 2025 Stanford University and the Authors.           *
+ * Authors: Peter Eastman                                                     *
+ * 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/Context.h"
+#include "openmm/CustomBondForce.h"
+#include "openmm/CustomExternalForce.h"
+#include "openmm/CustomNonbondedForce.h"
+#include "openmm/HarmonicBondForce.h"
+#include "openmm/NonbondedForce.h"
+#include "openmm/System.h"
+#include "openmm/QTBIntegrator.h"
+#include "SimTKOpenMMRealType.h"
+#include "sfmt/SFMT.h"
+#include <iostream>
+#include <vector>
+
+using namespace OpenMM;
+using namespace std;
+
+void testHarmonic() {
+    // Create a collection of uncoupled harmonic oscillators.
+
+    int numFrequencies = 10;
+    int numReplicas = 10;
+    int numParticles = numFrequencies*numReplicas;
+    double temperature = 300.0;
+    double mass = 1.0;
+    System system;
+    vector<Vec3> positions(numParticles);
+    vector<double> k;
+    CustomExternalForce* force = new CustomExternalForce("0.5*k*(x*x+y*y+z*z)");
+    system.addForce(force);
+    force->addPerParticleParameter("k");
+    QTBIntegrator integrator(temperature, 10.0, 0.001);
+    integrator.setCutoffFrequency(500.0);
+    for (int j = 0; j < numReplicas; j++) {
+        int base = system.getNumParticles();
+        for (int i = 0; i < numFrequencies; i++) {
+            system.addParticle(mass);
+            k.push_back(1000.0*(i+1)*(i+1));
+            force->addParticle(i+base, {k[i]});
+            integrator.setParticleType(i+base, i);
+        }
+    }
+    Context context(system, integrator, platform);
+    context.setPositions(positions);
+    context.setVelocitiesToTemperature(temperature);
+
+    // Compute the average energy of each particle over a simulation.
+
+    integrator.step(10000);
+    vector<double> energy(numParticles, 0.0);
+    int numSteps = 10000;
+    for (int i = 0; i < numSteps; i++) {
+        integrator.step(10);
+        State state = context.getState(State::Positions);
+        for (int j = 0; j < numParticles; j++) {
+            Vec3 p = state.getPositions()[j];
+            energy[j%numFrequencies] += 0.5*k[j]*p.dot(p);
+        }
+    }
+    for (int i = 0; i < numParticles; i++)
+        energy[i] /= numSteps*numReplicas;
+
+    // Compare to the expected distribution.
+
+    for (int i = 0; i < numFrequencies; i++) {
+        double w = sqrt(k[i]/mass);
+        double hbar = 1.054571628e-34*AVOGADRO/(1000*1e-12);
+        double kT = BOLTZ*temperature;
+        double expected = 1.5*hbar*w*(0.5+1/(exp(hbar*w/kT)-1));
+        ASSERT_USUALLY_EQUAL_TOL(expected, energy[i], 0.07);
+    }
+}
+
+void testCoupledHarmonic() {
+    // Create a collection of weakly coupled harmonic oscillators.
+
+    int numFrequencies = 8;
+    int numReplicas = 10;
+    int numParticles = numFrequencies*numReplicas;
+    double temperature = 10.0;
+    double mass = 1.0;
+    System system;
+    vector<Vec3> positions(numParticles);
+    vector<double> k;
+    CustomExternalForce* force = new CustomExternalForce("0.5*k*(x*x+y*y+z*z)");
+    system.addForce(force);
+    force->addPerParticleParameter("k");
+    CustomBondForce* bonds = new CustomBondForce("sin(100*r)");
+    system.addForce(bonds);
+    QTBIntegrator integrator(temperature, 50.0, 0.001);
+    integrator.setCutoffFrequency(500.0);
+    integrator.setDefaultAdaptationRate(1.0);
+    for (int j = 0; j < numReplicas; j++) {
+        int base = system.getNumParticles();
+        for (int i = 0; i < numFrequencies; i++) {
+            system.addParticle(mass);
+            k.push_back(1000.0*(i+1)*(i+1));
+            force->addParticle(i+base, {k[i]});
+            integrator.setParticleType(i+base, i);
+        }
+        for (int i = 0; i < numFrequencies; i++)
+            for (int k = 0; k < i; k++)
+                bonds->addBond(i+base, k+base);
+    }
+    Context context(system, integrator, platform);
+    context.setPositions(positions);
+    context.setVelocitiesToTemperature(temperature);
+
+    // Equilibrate with a high adaptation rate to let the spectrum converge.
+
+    for (int i = 0; i < 15; i++) {
+        integrator.step(10000);
+        // OpenCL on Mac hangs if we go too long without downloading anything.  I don't know why.
+        context.getState(State::Positions);
+    }
+    integrator.setDefaultAdaptationRate(0.05);
+    context.reinitialize(true);
+
+    // Compute the average energy of each particle over a simulation.
+
+    vector<double> energy(numFrequencies, 0.0);
+    int numSteps = 20000;
+    for (int i = 0; i < numSteps; i++) {
+        integrator.step(10);
+        State state = context.getState(State::Positions);
+        for (int j = 0; j < numParticles; j++) {
+            Vec3 p = state.getPositions()[j];
+            energy[j%numFrequencies] += 0.5*k[j]*p.dot(p);
+        }
+    }
+    for (int i = 0; i < numFrequencies; i++)
+        energy[i] /= numSteps*numReplicas;
+
+    // Compare to the expected distribution.
+
+    for (int i = 0; i < numFrequencies; i++) {
+        double w = sqrt(k[i]/mass);
+        double hbar = 1.054571628e-34*AVOGADRO/(1000*1e-12);
+        double kT = BOLTZ*temperature;
+        double expected = 1.5*hbar*w*(0.5+1/(exp(hbar*w/kT)-1));
+        ASSERT_USUALLY_EQUAL_TOL(expected, energy[i], 0.15);
+    }
+}
+
+void testParaHydrogen() {
+    const int numParticles = 32;
+    const double temperature = 25.0;
+    const double mass = 2.0;
+    const double boxSize = 1.1896;
+    const int numSteps = 2000;
+    const int numBins = 200;
+    const double reference[] = {
+        0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
+        0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
+        0.0, 0.0, 0.0, 4.932814042206152e-5, 1.244331241336431e-4, 4.052316284060125e-4,
+        1.544810863683946e-3, 4.376197806690222e-3, 1.025847561714293e-2, 2.286702037465422e-2,
+        4.371052180263602e-2, 7.518538770734748e-2, 0.122351534531647, 0.185758975626622,
+        0.266399984652322, 0.363380262153250, 0.473696401293219, 0.595312098494172,
+        0.726049519422861, 0.862264551954547, 0.991102029379444, 1.1147503922535,
+        1.23587006992066, 1.33495411932817, 1.42208208736987, 1.49273884004107,
+        1.54633319690403, 1.58714702233941, 1.60439217751355, 1.61804190608902,
+        1.60680198476058, 1.58892222973695, 1.56387607986781, 1.52629494593350,
+        1.48421439018970, 1.43656176771959, 1.38752775598872, 1.33310695719931,
+        1.28363477223121, 1.23465642750248, 1.18874848666326, 1.14350496170519,
+        1.10292486009936, 1.06107270157688, 1.02348927970441, 0.989729345271297,
+        0.959273446941802, 0.932264875865758, 0.908818658748942, 0.890946420768315,
+        0.869332737718165, 0.856401736350349, 0.842370069917020, 0.834386614237393,
+        0.826268072171045, 0.821547250199453, 0.818786865315836, 0.819441757028076,
+        0.819156933383128, 0.822275325148621, 0.828919078023881, 0.837233720599450,
+        0.846961908186718, 0.855656955481099, 0.864520333201247, 0.876082425547566,
+        0.886950044046000, 0.900275658318995
+    };
+
+    // Create a box of para-hydrogen.
+
+    System system;
+    for (int i = 0; i < numParticles; i++)
+        system.addParticle(mass);
+    system.setDefaultPeriodicBoxVectors(Vec3(boxSize,0,0), Vec3(0,boxSize,0), Vec3(0,0,boxSize));
+    CustomNonbondedForce* nb = new CustomNonbondedForce("2625.49963*(exp(1.713-1.5671*p-0.00993*p*p)-(12.14/p^6+215.2/p^8-143.1/p^9+4813.9/p^10)*(step(rc-p)*exp(-(rc/p-1)^2)+1-step(rc-p))); p=r/0.05291772108; rc=8.32");
+    nb->setNonbondedMethod(CustomNonbondedForce::CutoffPeriodic);
+    nb->setCutoffDistance(boxSize/2);
+    vector<double> params;
+    for (int i = 0; i < numParticles; i++)
+        nb->addParticle(params);
+    system.addForce(nb);
+    QTBIntegrator integ(temperature, 40.0, 0.001);
+    for (int i = 0; i < numParticles; i++)
+        integ.setParticleType(i, 0);
+    integ.setDefaultAdaptationRate(0.5);
+    integ.setSegmentLength(0.5);
+    Context context(system, integ, platform);
+    OpenMM_SFMT::SFMT sfmt;
+    init_gen_rand(0, sfmt);
+    vector<Vec3> positions(numParticles);
+    for (int i = 0; i < numParticles; i++)
+        positions[i] = Vec3(boxSize*genrand_real2(sfmt), boxSize*genrand_real2(sfmt), boxSize*genrand_real2(sfmt));
+    context.setPositions(positions);
+    integ.step(50000);
+
+    // Simulate it.
+
+    vector<int> counts(numBins, 0);
+    const double invBoxSize = 1.0/boxSize;
+    for (int step = 0; step < numSteps; step++) {
+        integ.step(20);
+        State state = context.getState(State::Positions);
+
+        // Record the radial distribution function.
+
+        const vector<Vec3>& pos = state.getPositions();
+        for (int j = 0; j < numParticles; j++)
+            for (int k = 0; k < j; k++) {
+                Vec3 delta = pos[j]-pos[k];
+                delta[0] -= floor(delta[0]*invBoxSize+0.5)*boxSize;
+                delta[1] -= floor(delta[1]*invBoxSize+0.5)*boxSize;
+                delta[2] -= floor(delta[2]*invBoxSize+0.5)*boxSize;
+                double dist = sqrt(delta.dot(delta));
+                int bin = (int) (numBins*(dist/boxSize));
+                counts[bin]++;
+            }
+    }
+
+    // Check against expected values.
+
+    double scale = (boxSize*boxSize*boxSize)/(numSteps*0.5*numParticles*numParticles);
+    for (int i = 0; i < numBins/2; i++) {
+        double r1 = i*boxSize/numBins;
+        double r2 = (i+1)*boxSize/numBins;
+        double volume = (4.0/3.0)*M_PI*(r2*r2*r2-r1*r1*r1);
+        ASSERT_USUALLY_EQUAL_TOL(reference[i], scale*counts[i]/volume, 0.2);
+    }
+}
+
+void testConstraints() {
+    const int numParticles = 8;
+    const int numConstraints = 5;
+    const double temp = 100.0;
+    System system;
+    QTBIntegrator integrator(temp, 2.0, 0.01);
+    integrator.setConstraintTolerance(1e-5);
+    NonbondedForce* forceField = new NonbondedForce();
+    for (int i = 0; i < numParticles; ++i) {
+        system.addParticle(10.0);
+        forceField->addParticle((i%2 == 0 ? 0.2 : -0.2), 0.5, 5.0);
+    }
+    system.addConstraint(0, 1, 1.0);
+    system.addConstraint(1, 2, 1.0);
+    system.addConstraint(2, 3, 1.0);
+    system.addConstraint(4, 5, 1.0);
+    system.addConstraint(6, 7, 1.0);
+    system.addForce(forceField);
+    Context context(system, integrator, platform);
+    vector<Vec3> positions(numParticles);
+    vector<Vec3> velocities(numParticles);
+    OpenMM_SFMT::SFMT sfmt;
+    init_gen_rand(0, sfmt);
+
+    for (int i = 0; i < numParticles; ++i) {
+        positions[i] = Vec3(i/2, (i+1)/2, 0);
+        velocities[i] = Vec3(genrand_real2(sfmt)-0.5, genrand_real2(sfmt)-0.5, genrand_real2(sfmt)-0.5);
+    }
+    context.setPositions(positions);
+    context.setVelocities(velocities);
+
+    // Simulate it and see whether the constraints remain satisfied.
+
+    for (int i = 0; i < 1000; ++i) {
+        State state = context.getState(State::Positions);
+        for (int j = 0; j < numConstraints; ++j) {
+            int particle1, particle2;
+            double distance;
+            system.getConstraintParameters(j, particle1, particle2, distance);
+            Vec3 p1 = state.getPositions()[particle1];
+            Vec3 p2 = state.getPositions()[particle2];
+            double dist = std::sqrt((p1[0]-p2[0])*(p1[0]-p2[0])+(p1[1]-p2[1])*(p1[1]-p2[1])+(p1[2]-p2[2])*(p1[2]-p2[2]));
+            ASSERT_EQUAL_TOL(distance, dist, 1e-4);
+        }
+        integrator.step(1);
+    }
+}
+
+void testConstrainedMasslessParticles() {
+    System system;
+    system.addParticle(0.0);
+    system.addParticle(1.0);
+    system.addConstraint(0, 1, 1.5);
+    vector<Vec3> positions(2);
+    positions[0] = Vec3(-1, 0, 0);
+    positions[1] = Vec3(1, 0, 0);
+    QTBIntegrator integrator(300.0, 2.0, 0.01);
+    bool failed = false;
+    try {
+        // This should throw an exception.
+        
+        Context context(system, integrator, platform);
+    }
+    catch (exception& ex) {
+        failed = true;
+    }
+    ASSERT(failed);
+    
+    // Now make both particles massless, which should work.
+    
+    system.setParticleMass(1, 0.0);
+    Context context(system, integrator, platform);
+    context.setPositions(positions);
+    context.setVelocitiesToTemperature(300.0);
+    integrator.step(1);
+    State state = context.getState(State::Velocities);
+    ASSERT_EQUAL(0.0, state.getVelocities()[0][0]);
+}
+
+void testRandomSeed() {
+    const int numParticles = 8;
+    const double temp = 100.0;
+    System system;
+    QTBIntegrator integrator(temp, 2.0, 0.01);
+    NonbondedForce* forceField = new NonbondedForce();
+    for (int i = 0; i < numParticles; ++i) {
+        system.addParticle(2.0);
+        forceField->addParticle((i%2 == 0 ? 1.0 : -1.0), 1.0, 5.0);
+    }
+    system.addForce(forceField);
+    vector<Vec3> positions(numParticles);
+    vector<Vec3> velocities(numParticles);
+    for (int i = 0; i < numParticles; ++i) {
+        positions[i] = Vec3((i%2 == 0 ? 2 : -2), (i%4 < 2 ? 2 : -2), (i < 4 ? 2 : -2));
+        velocities[i] = Vec3(0, 0, 0);
+    }
+
+    // Try twice with the same random seed.
+
+    integrator.setRandomNumberSeed(5);
+    Context context(system, integrator, platform);
+    context.setPositions(positions);
+    context.setVelocities(velocities);
+    integrator.step(10);
+    State state1 = context.getState(State::Positions);
+    context.reinitialize();
+    context.setPositions(positions);
+    context.setVelocities(velocities);
+    integrator.step(10);
+    State state2 = context.getState(State::Positions);
+
+    // Try twice with a different random seed.
+
+    integrator.setRandomNumberSeed(10);
+    context.reinitialize();
+    context.setPositions(positions);
+    context.setVelocities(velocities);
+    integrator.step(10);
+    State state3 = context.getState(State::Positions);
+    context.reinitialize();
+    context.setPositions(positions);
+    context.setVelocities(velocities);
+    integrator.step(10);
+    State state4 = context.getState(State::Positions);
+
+    // Compare the results.
+
+    for (int i = 0; i < numParticles; i++) {
+        for (int j = 0; j < 3; j++) {
+            ASSERT_EQUAL_TOL(state1.getPositions()[i][j], state2.getPositions()[i][j], 1e-6);
+            ASSERT_EQUAL_TOL(state3.getPositions()[i][j], state4.getPositions()[i][j], 1e-6);
+            ASSERT(state1.getPositions()[i][j] != state3.getPositions()[i][j]);
+        }
+    }
+}
+
+void testInitialTemperature() {
+    // Check temperature initialization for a collection of randomly placed particles
+    const int numParticles = 50000;
+    const int nDoF = 3 * numParticles;
+    const double targetTemperature = 300;
+    System system;
+    OpenMM_SFMT::SFMT sfmt;
+    init_gen_rand(0, sfmt);
+    std::vector<Vec3> positions(numParticles);
+
+    for (int i = 0; i < numParticles; i++) {
+        system.addParticle(1.0);
+        positions[i][0] = genrand_real2(sfmt);
+        positions[i][1] = genrand_real2(sfmt);
+        positions[i][2] = genrand_real2(sfmt);
+    }
+
+    QTBIntegrator integrator(300, 25, 0.001);
+    Context context(system, integrator, platform);
+    context.setPositions(positions);
+    context.setVelocitiesToTemperature(targetTemperature);
+    auto velocities = context.getState(State::Velocities).getVelocities();
+    double kineticEnergy = 0;
+    for(const auto &v : velocities) kineticEnergy += 0.5 * v.dot(v);
+    double temperature = (2*kineticEnergy / (nDoF*BOLTZ));
+    ASSERT_USUALLY_EQUAL_TOL(targetTemperature, temperature, 0.01);
+}
+
+void testSerializeParameters() {
+    // Test serializing the integrator's parameters.
+
+    int numParticles = 10;
+    double temperature = 300.0;
+    double mass = 2.0;
+    System system;
+    vector<Vec3> positions(numParticles);
+    CustomExternalForce* force = new CustomExternalForce("0.5*k*(x*x+y*y+z*z)");
+    system.addForce(force);
+    force->addPerParticleParameter("k");
+    QTBIntegrator integrator(temperature, 10.0, 0.001);
+    integrator.setCutoffFrequency(500.0);
+    integrator.setDefaultAdaptationRate(0.1);
+    for (int i = 0; i < numParticles; i++) {
+        system.addParticle(mass);
+        force->addParticle(i, {1000.0*(i+1)*(i+1)});
+        if (i < 5)
+            integrator.setParticleType(i, i%3);
+    }
+    Context context(system, integrator, platform);
+    context.setPositions(positions);
+    context.setVelocitiesToTemperature(temperature);
+
+    // Run for a little while, then record a State.
+
+    integrator.step(5000);
+    State state = context.getState(State::IntegratorParameters);
+
+    // Create a new Integrator and Context, set the State, and see if the adapted
+    // friction coefficients were set correctly.
+
+    QTBIntegrator integrator2(temperature, 10.0, 0.001);
+    for (auto type : integrator.getParticleTypes())
+        integrator2.setParticleType(type.first, type.second);
+    Context context2(system, integrator2, platform);
+    context2.setPositions(positions);
+    context2.setState(state);
+    for (int i = 0; i < numParticles; i++) {
+        vector<double> f1, f2;
+        integrator.getAdaptedFriction(i, f1);
+        integrator2.getAdaptedFriction(i, f2);
+        ASSERT_EQUAL_CONTAINERS(f1, f2);
+    }
+}
+
+void runPlatformTests();
+
+int main(int argc, char* argv[]) {
+    try {
+        initializeTests(argc, argv);
+        testHarmonic();
+        testCoupledHarmonic();
+        testParaHydrogen();
+        testConstraints();
+        testConstrainedMasslessParticles();
+        testRandomSeed();
+        testInitialTemperature();
+        testSerializeParameters();
+        runPlatformTests();
+    }
+    catch(const exception& e) {
+        cout << "exception: " << e.what() << endl;
+        return 1;
+    }
+    cout << "Done" << endl;
+    return 0;
+}

wrappers/generateWrappers.py

@@ -95,7 +95,9 @@ class WrapperGenerator:
                             'bool OpenMM::Discrete1DFunction::operator!=',
                             'bool OpenMM::Discrete2DFunction::operator!=',
                             'bool OpenMM::Discrete3DFunction::operator!=',
-                            'const std::map<int, int>& OpenMM::DPDIntegrator::getParticleTypes'
+                            'const std::map<int, int>& OpenMM::DPDIntegrator::getParticleTypes',
+                            'const std::map<int, int>& OpenMM::QTBIntegrator::getParticleTypes',
+                            'const std::map<int, double>& OpenMM::QTBIntegrator::getTypeAdaptationRates'
                            ]
         self.skipMethods = [s.replace(' ', '') for s in self.skipMethods]
         self.hideClasses = ['Kernel', 'KernelImpl', 'KernelFactory', 'ContextImpl', 'SerializationNode', 'SerializationProxy']

wrappers/python/src/swig_doxygen/swigInputConfig.py

@@ -512,6 +512,14 @@ UNITS = {
 ("DPDIntegrator", "getDefaultCutoff") : ("unit.nanometer", ()),
 ("DPDIntegrator", "setDefaultCutoff") : (None, ("unit.nanometer",)),
 ("DPDIntegrator", "getParticleTypes") : (None, ()),
+("QTBIntegrator", "getSegmentLength") : ("unit.picosecond", ()),
+("QTBIntegrator", "setSegmentLength") : (None, ("unit.picosecond",)),
+("QTBIntegrator", "getDefaultAdaptationRate") : (None, ()),
+("QTBIntegrator", "setDefaultAdaptationRate") : (None, (None,)),
+("QTBIntegrator", "getTypeAdaptationRates") : (None, ()),
+("QTBIntegrator", "getParticleTypes") : (None, ()),
+("QTBIntegrator", "getCutoffFrequency") : ("unit.picosecond**-1", ()),
+("QTBIntegrator", "setCutoffFrequency") : (None, ("unit.picosecond**-1",)),
 ("LocalEnergyMinimizer", "minimize") : (None, (None, "unit.kilojoules_per_mole/unit.nanometer", None)),
 ("ATMForce", "getForce") : (None, ()),
 ("ATMForce", "getPerturbationEnergy") :  ('unit.kilojoule_per_mole', ()),

wrappers/python/tests/TestAPIUnits.py

@@ -1089,5 +1089,22 @@ class TestAPIUnits(unittest.TestCase):
         integrator = CustomIntegrator(1.0*femtoseconds)
         self.assertEqual(integrator.getStepSize(), 1.0*femtoseconds)
 
+    def testQTBIntegrator(self):
+        """ Tests the LangevinIntegrator API features """
+        integrator = QTBIntegrator(300, 0.1, 0.001)
+        self.assertEqual(integrator.getTemperature(), 300*kelvin)
+        self.assertEqual(integrator.getFriction(), 0.1/picosecond)
+        self.assertEqual(integrator.getStepSize(), 0.001*picosecond)
+
+        integrator = QTBIntegrator(300*kelvin, 0.1/microsecond, 1*femtosecond)
+        self.assertEqual(integrator.getTemperature(), 300*kelvin)
+        self.assertAlmostEqualUnit(integrator.getFriction(), 0.1/microsecond)
+        self.assertEqual(integrator.getStepSize(), 1*femtosecond)
+
+        integrator.setSegmentLength(0.5)
+        self.assertEqual(integrator.getSegmentLength(), 0.5*picosecond)
+        integrator.setSegmentLength(100*femtosecond)
+        self.assertEqual(integrator.getSegmentLength(), 0.1*picosecond)
+
 if __name__ == '__main__':
     unittest.main()