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

docs-source/api-c++/integrators.rst

@@ -57,6 +57,17 @@ The RPMD integrator implements Ring Polymer MD.
     generated/RPMDIntegrator
 
 
+Quantum Thermal Bath integrators
+================================
+
+The QTB integrator implements the adaptive quantum thermal bath (adQTB) method.
+
+.. toctree::
+    :maxdepth: 2
+
+    generated/QTBIntegrator
+
+
 Dissipative Particle Dynamics integrators
 =========================================
 

docs-source/usersguide/references.bib

@@ -346,6 +346,17 @@
    type = {Journal Article}
 }
 
+@article{Mangaud2019
+    author = {Mangaud, Etienne and Huppert, Simon and Pl{\'e}, Thomas and Depondt, Philippe and Bonella, Sara and Finocchi, Fabio},
+    title = {The Fluctuation–Dissipation Theorem as a Diagnosis and Cure for Zero-Point Energy Leakage in Quantum Thermal Bath Simulations},
+    journal = {Journal of Chemical Theory and Computation},
+    volume = {15},
+    number = {5},
+    pages = {2863-2880},
+    year = {2019},
+    doi = {10.1021/acs.jctc.8b01164},
+}
+
 @article{Marinari1992,
     doi = {10.1209/0295-5075/19/6/002},
     year = 1992,
@@ -378,6 +389,17 @@
    type = {Journal Article}
 }
 
+@article{Mauger2021
+    author = {Mauger, Nastasia and Pl{\'e}, Thomas and Lagardère, Louis and Bonella, Sara and Mangaud, {\'E}tienne and Piquemal, Jean-Philip and Huppert, Simon},
+    title = {Nuclear Quantum Effects in Liquid Water at Near Classical Computational Cost Using the Adaptive Quantum Thermal Bath},
+    journal = {The Journal of Physical Chemistry Letters},
+    volume = {12},
+    number = {34},
+    pages = {8285-8291},
+    year = {2021},
+    doi = {10.1021/acs.jpclett.1c01722},
+}
+
 @article{Mongan2007
    author = {Mongan, John and Simmerling, Carlos and McCammon, J. Andrew and Case, David A. and Onufriev, Alexey},
    title = {Generalized {Born} model with a simple, robust molecular volume correction},

docs-source/usersguide/theory/04_integrators.rst

@@ -294,3 +294,102 @@ on the types of the particles.
 
 The integration is done using the same LFMiddle discretization used for
 LangevinIntegrator. :cite:`Zhang2019`
+
+QTBIntegrator
+*************
+
+This integrator implements the Adaptive Quantum Thermal Bath (adQTB) algorithm.
+:cite:`Mangaud2019`  This is a fast method for approximating nuclear quantum
+effects by applying a Langevin thermostat whose random force varies with
+frequency to match the expected energy of a quantum harmonic oscillator.
+
+It integrates the Langevin equation
+
+.. math::
+   m_i\frac{d\mathbf{v}_i}{dt}=\mathbf{f}_i-\gamma_f m_i \mathbf{v}_i+\mathbf{R}_i
+
+Unlike LangevinIntegrator, for which the random noise force :math:`\mathbf{R}_i`
+is uncorrelated white noise, in this case its cross correlation function is given
+by
+
+.. math::
+   C(\omega) = 2 m_i \gamma_r \theta(\omega, T)
+
+where
+
+.. math::
+   \theta(\omega, T) = \hbar \omega \left( \frac{1}{2} + \frac{1}{e^{\hbar \omega / k_B T} - 1} \right)
+
+In the limit of high temperature, :math:`\theta(\omega, T) \approx k_B T`,
+reproducing the classical behavior.  In the limit of low temperature,
+:math:`\theta(\omega, T) \approx \hbar \omega / 2`, which is the ground state
+energy of a quantum harmonic oscillator with frequency :math:`\omega`.
+
+A problem that can arise with this method is zero-point energy leakage.  The
+thermostat drives the system toward the desired quantum energy distribution, but
+the classical dynamics of the system causes it to continuously relax toward the
+classical distribution.  The result is too much energy in the low frequency modes
+and too little energy in the high frequency modes.  A solution is to allow the
+two friction coefficients to differ.  The coefficient :math:`\gamma_f` appearing
+in the friction term of the Langevin equation is fixed as a constant.  The
+coefficient :math:`\gamma_r` that determines the magnitude of the random force
+becomes frequency dependent, :math:`\gamma_r(\omega)`.  Its spectrum is dynamically
+adjusted to produce the correct distribution of energy between modes.
+
+In practice, the simulation is divided into short segments, typically on the
+order of 1000 time steps.  At the end of each segment, we calculate the velocity
+autocorrelation function :math:`C_{vv}(\omega)` of each degree of freedom, as
+well as the correlation :math:`C_{vR}(\omega)` between the velocity and random
+force :math:`\mathbf{R}`.  From these we can calculate the amount by which the
+fluctuation-dissipation theorem is violated at each frequency:
+
+.. math::
+   \Delta_{FDT}(\omega) = \mathrm{Re}[C_{vR}(\omega)] - m \gamma_r(\omega) C_{vv}(\omega)
+
+We then select new friction coefficients according to
+
+.. math::
+   \gamma_r^{k+1}(\omega) = \gamma_r^k(\omega) - A \Delta_{FDT}^k(\omega)
+
+where :math:`k` is the index of the segment and the adaptation rate :math:`A`
+may vary between degrees of freedom.  Random noise for the next segment is
+generated based on the new spectrum, and the simulation continues.  The value of
+:math:`A` generally needs to be determined by trial and error to find a value
+that produces fast adaptation and good convergence.
+
+This process is much more robust if data is pooled over all degrees of freedom
+that are expected to be equivalent.  One specifies a type index for each
+particle.  :math:`\Delta_{FDT}(\omega)` is computed as the average over all
+three coordinates of all particles with the same type.  The more particles that
+are averaged over, the larger :math:`A` can be, leading to faster adaptation.
+
+When running simulations with adQTB, it is critical to equilibrate the system
+long enough for :math:`\gamma_r(\omega)` to converge.  Until it does, the
+simulation does not explore the correct distribution of states.
+
+Simulations with adQTB tend to require a larger friction coefficient
+:math:`\gamma_f` than is usual in fully classical simulations, typically 10
+ps\ :sup:`-1` or more.  If the friction is too low, it is impossible to fully
+compensate for zero-point energy leakage.  Even reducing the random force to
+zero may still leave too much energy in low frequency modes.  It is important
+to check the converged spectrum of :math:`\gamma_r(\omega)`.  If it is close to
+zero at some frequencies, that suggests :math:`\gamma_f` needs to be increased.
+
+A possible consequence of the large friction coefficient is unwanted broadening
+of spectral peaks, leading to a less accurate energy spectrum.  This is
+compensated through a deconvolution procedure. :cite:`Mauger2021`  The target
+spectrum :math:`\theta(\omega, T)` is replaced with a deconvolved version
+:math:`\tilde{\theta}(\omega, T)` related to it by
+
+.. math::
+   \frac{\theta(\omega_0, T)}{2} = \int \frac{d\omega}{\pi} \frac{\gamma \omega_0^2}{(\omega^2-\omega_0^2)^2 + \gamma^2\omega^2} \tilde{\theta}(\omega, T)
+
+An iterative procedure is used to solve for :math:`\tilde{\theta}(\omega, T)`,
+which is then used in place of :math:`\theta(\omega, T)` when computing the
+random force.
+
+One must be very careful when trying to compute velocity dependent thermodynamic
+quantities, such as the instantaneous temperature or instantaneous pressure.
+The standard calculations for these quantities assume the velocities follow a
+classical distribution.  They do not produce correct results for an adQTB
+simulation, in which the velocities follow a quantum distribution.

olla/include/openmm/kernels.h

@@ -57,6 +57,7 @@
 #include "openmm/KernelImpl.h"
 #include "openmm/MonteCarloBarostat.h"
 #include "openmm/PeriodicTorsionForce.h"
+#include "openmm/QTBIntegrator.h"
 #include "openmm/RBTorsionForce.h"
 #include "openmm/RMSDForce.h"
 #include "openmm/NonbondedForce.h"
@@ -1435,6 +1436,72 @@ public:
     virtual double computeKineticEnergy(ContextImpl& context, const DPDIntegrator& integrator) = 0;
 };
 
+/**
+ * This kernel is invoked by QTBIntegrator to take one time step.
+ */
+class IntegrateQTBStepKernel : public KernelImpl {
+public:
+    static std::string Name() {
+        return "IntegrateQTBStep";
+    }
+    IntegrateQTBStepKernel(std::string name, const Platform& platform) : KernelImpl(name, platform) {
+    }
+    /**
+     * Initialize the kernel.
+     * 
+     * @param system     the System this kernel will be applied to
+     * @param integrator the QTBIntegrator this kernel will be used for
+     */
+    virtual void initialize(const System& system, const QTBIntegrator& integrator) = 0;
+    /**
+     * Execute the kernel.
+     * 
+     * @param context    the context in which to execute this kernel
+     * @param integrator the QTBIntegrator this kernel is being used for
+     */
+    virtual void execute(ContextImpl& context, const QTBIntegrator& integrator) = 0;
+    /**
+     * Compute the kinetic energy.
+     * 
+     * @param context    the context in which to execute this kernel
+     * @param integrator the QTBIntegrator this kernel is being used for
+     */
+    virtual double computeKineticEnergy(ContextImpl& context, const QTBIntegrator& integrator) = 0;
+    /**
+     * 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
+     */
+    virtual void getAdaptedFriction(ContextImpl& context, int particle, std::vector<double>& friction) const = 0;
+    /**
+     * 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
+     */
+    virtual void setAdaptedFriction(ContextImpl& context, int particle, const std::vector<double>& friction) = 0;
+    /**
+     * 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
+     */
+    virtual void createCheckpoint(ContextImpl& context, std::ostream& stream) const = 0;
+    /**
+     * 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
+     */
+    virtual void loadCheckpoint(ContextImpl& context, std::istream& stream) = 0;
+};
+
 /**
  * This kernel is invoked by AndersenThermostat at the start of each time step to adjust the particle velocities.
  */

openmmapi/include/OpenMM.h

@@ -68,6 +68,7 @@
 #include "openmm/Context.h"
 #include "openmm/OpenMMException.h"
 #include "openmm/PeriodicTorsionForce.h"
+#include "openmm/QTBIntegrator.h"
 #include "openmm/RBTorsionForce.h"
 #include "openmm/RMSDForce.h"
 #include "openmm/State.h"

openmmapi/include/openmm/QTBIntegrator.h

+#ifndef OPENMM_QTBINTEGRATOR_H_
+#define OPENMM_QTBINTEGRATOR_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 "Integrator.h"
+#include "openmm/Kernel.h"
+#include "internal/windowsExport.h"
+
+namespace OpenMM {
+
+/**
+ * This integrator implements the Adaptive Quantum Thermal Bath (adQTB) algorithm
+ * as described in https://doi.org/10.1021/acs.jctc.8b01164.  This is a fast method
+ * for approximating nuclear quantum effects by applying a Langevin thermostat whose
+ * random force varies with frequency to match the expected energy of a quantum
+ * harmonic oscillator.
+ * 
+ * To compensate for zero point energy leakage, the spectrum of the random force
+ * is adjusted automatically.  The trajectory is divided into short segments.  At
+ * the end of each segment, the distribution of energy across frequencies is evaluated,
+ * and the friction coefficients used in generating the random force are adjusted
+ * to better match the target distribution.  You can monitor this process by calling
+ * getAdaptedFriction().  It is important to equilibrate the simulation long enough
+ * for the friction coefficients to converge before beginning the production part of
+ * the simulation.
+ * 
+ * To make this process more robust, it is recommended to average the data over all
+ * particles that are expected to behave identically.  To do this, you can optionally
+ * call setParticleType() to define a particle as having a particular type.  The
+ * data for all particles of the same type is averaged and they are adjusted together.
+ * You also can set a different adaptation rate for each particle type.  The more
+ * particles that are being averaged over, the higher the adaptation rate can reasonably
+ * be set, leading to faster adaptation.
+ * 
+ * Most properties of this integrator are fixed at Context creation time and can
+ * only be changed after that by reinitializing the Context.  That includes the
+ * step size, friction coefficient, segment length, cutoff frequency, particle types,
+ * and adaptation rates.  The only property of the integrator that can be freely
+ * changed in the middle of a simulation is the temperature, and even that requires
+ * some care.  The new temperature will not take effect until the beginning of the
+ * next segment.  Furthermore, changing the temperature is a potentially expensive
+ * operation, since it requires performing a calculation whose cost scales as the
+ * cube of the number of time steps in a segment.
+ * 
+ * One must be very careful when trying to compute velocity dependent thermodynamic
+ * quantities, such as the instantaneous temperature or instantaneous pressure.
+ * The standard calculations for these quantities assume the velocities follow a
+ * classical distribution.  They do not produce correct results for an adQTB
+ * simulation, in which the velocities follow a quantum distribution.
+ */
+class OPENMM_EXPORT QTBIntegrator : public Integrator {
+public:
+    /**
+     * Create a QTBIntegrator.
+     * 
+     * @param temperature    the temperature of the heat bath (in Kelvin)
+     * @param frictionCoeff  the friction coefficient which couples the system to the heat bath (in inverse picoseconds)
+     * @param stepSize       the step size with which to integrate the system (in picoseconds)
+     */
+    QTBIntegrator(double temperature, double frictionCoeff, double stepSize);
+    /**
+     * Get the temperature of the heat bath (in Kelvin).
+     *
+     * @return the temperature of the heat bath, measured in Kelvin
+     */
+    double getTemperature() const {
+        return temperature;
+    }
+    /**
+     * Set the temperature of the heat bath (in Kelvin).
+     *
+     * @param temp    the temperature of the heat bath, measured in Kelvin
+     */
+    void setTemperature(double temp);
+    /**
+     * Get the friction coefficient which determines how strongly the system is coupled to
+     * the heat bath (in inverse ps).
+     *
+     * @return the friction coefficient, measured in 1/ps
+     */
+    double getFriction() const {
+        return friction;
+    }
+    /**
+     * Set the friction coefficient which determines how strongly the system is coupled to
+     * the heat bath (in inverse ps).
+     *
+     * @param coeff    the friction coefficient, measured in 1/ps
+     */
+    void setFriction(double coeff);
+    /**
+     * Get the length of each segment used for adapting the noise spectrum.
+     * 
+     * @return the segment length, measured in ps
+     */
+    double getSegmentLength() const;
+    /**
+     * Set the length of each segment used for adapting the noise spectrum.
+     * 
+     * @param length    the segment length, measured in ps
+     */
+    void setSegmentLength(double length);
+    /**
+     * Get the cutoff frequency applied to the colored noise.
+     * 
+     * @return the cutoff frequency, measured in 1/ps
+     */
+    double getCutoffFrequency() const;
+    /**
+     * Set the cutoff frequency applied to the colored noise.
+     * 
+     * @param cutoff    the cutoff frequency, measured in 1/ps
+     */
+    void setCutoffFrequency(double cutoff);
+    /**
+     * Get a map whose keys are particle indices and whose values are particle types.  This
+     * contains only the particles that have been specifically set with setParticleType().
+     */
+    const std::map<int, int>& getParticleTypes() const;
+    /**
+     * Set the type of a particle.  This is an arbitrary integer.
+     */
+    void setParticleType(int index, int type);
+    /**
+     * Get the default adaptation rate.  This is the rate used for particles whose
+     * rate has not been otherwise specified with setTypeAdaptationRate().  It
+     * determines how much the noise spectrum changes after each segment.
+     */
+    double getDefaultAdaptationRate() const;
+    /**
+     * Set the default adaptation rate.  This is the rate used for particles whose
+     * rate has not been otherwise specified with setTypeAdaptationRate().  It
+     * determines how much the noise spectrum changes after each segment.
+     */
+    void setDefaultAdaptationRate(double rate);
+    /**
+     * Get a map whose keys are particle types and whose values are adaptation rates.
+     * These are the rates used for particles whose rates have been specified with
+     * setTypeAdaptationRate().  The rate determines how much the noise spectrum
+     * changes after each segment.
+     */
+    const std::map<int, double>& getTypeAdaptationRates() const;
+    /**
+     * Set the adaptation rate to use for particles of a given type.  The rate
+     * determines how much the noise spectrum changes after each segment.
+     */
+    void setTypeAdaptationRate(int type, double rate);
+    /**
+     * Get the adapted friction coefficients for a particle.  The return value is
+     * a vector of length numFreq = (3*n+1)/2, where n is the number of time steps
+     * in a segment.  Element i is the friction coefficient used in generating
+     * the random force with frequency i*pi/(numFreq*stepSize).
+     * 
+     * This method is guaranteed to return identical results for particles 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 getAdaptedFriction(int particle, std::vector<double>& friction) const;
+    /**
+     * Get the random number seed.  See setRandomNumberSeed() for details.
+     */
+    int getRandomNumberSeed() const {
+        return randomNumberSeed;
+    }
+    /**
+     * Set the random number seed.  The precise meaning of this parameter is undefined, and is left up
+     * to each Platform to interpret in an appropriate way.  It is guaranteed that if two simulations
+     * are run with different random number seeds, the sequence of random forces will be different.  On
+     * the other hand, no guarantees are made about the behavior of simulations that use the same seed.
+     * In particular, Platforms are permitted to use non-deterministic algorithms which produce different
+     * results on successive runs, even if those runs were initialized identically.
+     *
+     * If seed is set to 0 (which is the default value assigned), a unique seed is chosen when a Context
+     * is created from this Integrator. This is done to ensure that each Context receives unique random seeds
+     * without you needing to set them explicitly.
+     */
+    void setRandomNumberSeed(int seed) {
+        randomNumberSeed = seed;
+    }
+    /**
+     * Advance a simulation through time by taking a series of time steps.
+     * 
+     * @param steps   the number of time steps to take
+     */
+    void step(int steps);
+protected:
+    /**
+     * This will be called by the Context when it is created.  It informs the Integrator
+     * of what context it will be integrating, and gives it a chance to do any necessary initialization.
+     * It will also get called again if the application calls reinitialize() on the Context.
+     */
+    void initialize(ContextImpl& context);
+    /**
+     * This will be called by the Context when it is destroyed to let the Integrator do any necessary
+     * cleanup.  It will also get called again if the application calls reinitialize() on the Context.
+     */
+    void cleanup();
+    /**
+     * Get the names of all Kernels used by this Integrator.
+     */
+    std::vector<std::string> getKernelNames();
+    /**
+     * Compute the kinetic energy of the system at the current time.
+     */
+    double computeKineticEnergy();
+    /**
+     * Computing kinetic energy for this integrator does not require forces.
+     */
+    bool kineticEnergyRequiresForce() const;
+    /**
+     * Get the time interval by which velocities are offset from positions.  This is used to
+     * adjust velocities when setVelocitiesToTemperature() is called on a Context.
+     */
+    double getVelocityTimeOffset() const {
+        return getStepSize()/2;
+    }
+    /**
+     * This is called while writing checkpoints.  It gives the integrator a chance to write
+     * its own data.
+     */
+    void createCheckpoint(std::ostream& stream) const;
+    /**
+     * This is called while loading a checkpoint.  The integrator should read in whatever
+     * data it wrote in createCheckpoint() and update its internal state accordingly.
+     */
+    void loadCheckpoint(std::istream& stream);
+    /**
+     * This is called while creating a State.  The Integrator should store the values
+     * of all time-varying parameters into the SerializationNode so they can be saved
+     * as part of the state.
+     */
+    void serializeParameters(SerializationNode& node) const;
+    /**
+     * This is called when loading a previously saved State.  The Integrator should
+     * load the values of all time-varying parameters from the SerializationNode.  If
+     * the node contains parameters that are not defined for this Integrator, it should
+     * throw an exception.
+     */
+    void deserializeParameters(const SerializationNode& node);
+private:
+    double temperature, friction, segmentLength, defaultAdaptationRate, cutoffFrequency;
+    std::map<int, int> particleType;
+    std::map<int, double> typeAdaptationRates;
+    int randomNumberSeed;
+    Kernel kernel;
+};
+
+} // namespace OpenMM
+
+#endif /*OPENMM_QTBINTEGRATOR_H_*/

openmmapi/include/openmm/internal/QTBIntegratorUtilities.h

+#ifndef OPENMM_QTBINTEGRATORUTILITIES_H_
+#define OPENMM_QTBINTEGRATORUTILITIES_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/QTBIntegrator.h"
+#include "openmm/System.h"
+#include "openmm/internal/ThreadPool.h"
+#include <vector>
+
+namespace OpenMM {
+
+/**
+ * This class defines a set of utility functions that are useful in implementing QTBIntegrator.
+ */
+
+class OPENMM_EXPORT QTBIntegratorUtilities {
+public:
+    /**
+     * Identify unique particle types, and record the sets of particles with the same type.
+     * 
+     * @param system              the System being simulated
+     * @param integrator          the integrator defining the particle types
+     * @param particleType        on exit, element i is the type index of particle i
+     * @param typeParticles       on exit, element i contains the indices of all particles with type i
+     * @param typeMass            on exit, element i contains the mass of particles with type i
+     * @param typeAdaptationRate  on exit, element i contains the adaptation rate of particles with type i
+     */
+    static void findTypes(const System& system, const QTBIntegrator& integrator, std::vector<int>& particleType,
+            std::vector<std::vector<int> >& typeParticles, std::vector<double>& typeMass, std::vector<double>& typeAdaptationRate);
+    /**
+     * Calculate the target noise spectrum.
+     * 
+     * @param temperature  the target temperature, in Kelvin
+     * @param friction     the friction coefficient in 1/ps
+     * @param dt           the integration step size, in ps
+     * @param numFreq      the number of frequency bins in the spectrum
+     * @param theta        the standard version of the spectrum is stored into this
+     * @param thetad       the deconvolved version of the spectrum is stored into this
+     * @param threads      a ThreadPool to use for parallelization
+     */
+    static void calculateSpectrum(double temperature, double friction, double dt, int numFreq, std::vector<double>& theta, std::vector<double>& thetad, ThreadPool& threads);
+};
+
+} // namespace OpenMM
+
+#endif /*OPENMM_QTBINTEGRATORUTILITIES_H_*/

openmmapi/src/QTBIntegrator.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/QTBIntegrator.h"
+#include "openmm/Context.h"
+#include "openmm/OpenMMException.h"
+#include "openmm/internal/ContextImpl.h"
+#include "openmm/kernels.h"
+#include <set>
+#include <string>
+
+using namespace OpenMM;
+using std::map;
+using std::set;
+using std::string;
+using std::vector;
+
+QTBIntegrator::QTBIntegrator(double temperature, double frictionCoeff, double stepSize) {
+    setTemperature(temperature);
+    setFriction(frictionCoeff);
+    setSegmentLength(1.0);
+    setCutoffFrequency(500.0);
+    setDefaultAdaptationRate(0.01);
+    setStepSize(stepSize);
+    setConstraintTolerance(1e-5);
+    setRandomNumberSeed(0);
+}
+
+void QTBIntegrator::initialize(ContextImpl& contextRef) {
+    if (owner != NULL && &contextRef.getOwner() != owner)
+        throw OpenMMException("This Integrator is already bound to a context");
+    context = &contextRef;
+    owner = &contextRef.getOwner();
+    kernel = context->getPlatform().createKernel(IntegrateQTBStepKernel::Name(), contextRef);
+    kernel.getAs<IntegrateQTBStepKernel>().initialize(contextRef.getSystem(), *this);
+}
+
+void QTBIntegrator::setTemperature(double temp) {
+    if (temp < 0)
+        throw OpenMMException("Temperature cannot be negative");
+    temperature = temp;
+}
+
+void QTBIntegrator::setFriction(double coeff) {
+    if (coeff < 0)
+        throw OpenMMException("Friction cannot be negative");
+    friction = coeff;
+}
+
+double QTBIntegrator::getSegmentLength() const {
+    return segmentLength;
+}
+
+void QTBIntegrator::setSegmentLength(double length) {
+    segmentLength = length;
+}
+
+double QTBIntegrator::getCutoffFrequency() const {
+    return cutoffFrequency;
+}
+
+void QTBIntegrator::setCutoffFrequency(double cutoff) {
+    cutoffFrequency = cutoff;
+}
+
+const map<int, int>& QTBIntegrator::getParticleTypes() const {
+    return particleType;
+}
+
+void QTBIntegrator::setParticleType(int index, int type) {
+    particleType[index] = type;
+}
+
+double QTBIntegrator::getDefaultAdaptationRate() const {
+    return defaultAdaptationRate;
+}
+
+void QTBIntegrator::setDefaultAdaptationRate(double rate) {
+    defaultAdaptationRate = rate;
+}
+
+const map<int, double>& QTBIntegrator::getTypeAdaptationRates() const {
+    return typeAdaptationRates;
+}
+
+void QTBIntegrator::setTypeAdaptationRate(int type, double rate) {
+    typeAdaptationRates[type] = rate;
+}
+
+void QTBIntegrator::getAdaptedFriction(int particle, vector<double>& friction) const {
+    kernel.getAs<IntegrateQTBStepKernel>().getAdaptedFriction(*context, particle, friction);
+}
+
+void QTBIntegrator::cleanup() {
+    kernel = Kernel();
+}
+
+vector<string> QTBIntegrator::getKernelNames() {
+    vector<std::string> names;
+    names.push_back(IntegrateQTBStepKernel::Name());
+    return names;
+}
+
+double QTBIntegrator::computeKineticEnergy() {
+    return kernel.getAs<IntegrateQTBStepKernel>().computeKineticEnergy(*context, *this);
+}
+
+bool QTBIntegrator::kineticEnergyRequiresForce() const {
+    return false;
+}
+
+void QTBIntegrator::step(int steps) {
+    if (context == NULL)
+        throw OpenMMException("This Integrator is not bound to a context!");
+    for (int i = 0; i < steps; ++i) {
+        context->updateContextState();
+        context->calcForcesAndEnergy(true, false, getIntegrationForceGroups());
+        kernel.getAs<IntegrateQTBStepKernel>().execute(*context, *this);
+    }
+}
+
+void QTBIntegrator::createCheckpoint(std::ostream& stream) const {
+    kernel.getAs<IntegrateQTBStepKernel>().createCheckpoint(*context, stream);
+}
+
+void QTBIntegrator::loadCheckpoint(std::istream& stream) {
+    kernel.getAs<IntegrateQTBStepKernel>().loadCheckpoint(*context, stream);
+}
+
+void QTBIntegrator::serializeParameters(SerializationNode& node) const {
+    node.setIntProperty("version", 1);
+    vector<int> particles;
+    set<int> types;
+    for (int i = 0; i < context->getSystem().getNumParticles(); i++) {
+        if (particleType.find(i) == particleType.end())
+            particles.push_back(i);
+        else if (types.find(particleType.at(i)) == types.end()) {
+            particles.push_back(i);
+            types.insert(particleType.at(i));
+        }
+    }
+    vector<double> friction;
+    for (int i : particles) {
+        SerializationNode& particle = node.createChildNode("Friction").setIntProperty("particle", i);
+        getAdaptedFriction(i, friction);
+        for (double f : friction)
+            particle.createChildNode("F").setDoubleProperty("v", f);
+    }
+}
+
+void QTBIntegrator::deserializeParameters(const SerializationNode& node) {
+    if (node.getIntProperty("version") != 1)
+        throw OpenMMException("Unsupported version number");
+    for (const SerializationNode& child : node.getChildren()) {
+        int particle = child.getIntProperty("particle");
+        vector<double> friction;
+        for (const SerializationNode& f : child.getChildren())
+            friction.push_back(f.getDoubleProperty("v"));
+        kernel.getAs<IntegrateQTBStepKernel>().setAdaptedFriction(*context, particle, friction);
+    }
+}

openmmapi/src/QTBIntegratorUtilities.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/QTBIntegratorUtilities.h"
+#include "SimTKOpenMMRealType.h"
+#include <map>
+
+using namespace OpenMM;
+using namespace std;
+
+void QTBIntegratorUtilities::findTypes(const System& system, const QTBIntegrator& integrator, vector<int>& particleType,
+        vector<std::vector<int> >& typeParticles, vector<double>& typeMass, vector<double>& typeAdaptationRate) {
+    // Record information about groups defined by particle types.
+
+    particleType.resize(system.getNumParticles());
+    map<int, int> typeIndex;
+    map<int, double> massTable;
+    const auto& types = integrator.getParticleTypes();
+    double defaultAdaptationRate = integrator.getDefaultAdaptationRate();
+    for (auto particle : types) {
+        int type = particle.second;
+        double mass = system.getParticleMass(particle.first);
+        if (typeIndex.find(type) == typeIndex.end()) {
+            int index = typeIndex.size();
+            typeIndex[type] = index;
+            double rate = defaultAdaptationRate;
+            const auto& typeRates = integrator.getTypeAdaptationRates();
+            if (typeRates.find(type) != typeRates.end())
+                rate = typeRates.at(type);
+            typeAdaptationRate.push_back(rate);
+            typeParticles.push_back(vector<int>());
+            typeMass.push_back(mass);
+            massTable[type] = mass;
+        }
+        if (mass != massTable[type])
+            throw OpenMMException("QTBIntegrator: All particles of the same type must have the same mass");
+        particleType[particle.first] = typeIndex[type];
+        typeParticles[typeIndex[type]].push_back(particle.first);
+    }
+    for (int i = 0; i < system.getNumParticles(); i++)
+        if (types.find(i) == types.end()) {
+            // This particle's type isn't set, so define a new type for it.
+            particleType[i] = typeParticles.size();
+            typeAdaptationRate.push_back(defaultAdaptationRate);
+            typeParticles.push_back({i});
+            typeMass.push_back(system.getParticleMass(i));
+        }
+}
+
+void QTBIntegratorUtilities::calculateSpectrum(double temperature, double friction, double dt, int numFreq, vector<double>& theta, vector<double>& thetad, ThreadPool& threads) {
+    // Compute the standard spectrum.
+
+    double hbar = 1.054571628e-34*AVOGADRO/(1000*1e-12);
+    double kT = BOLTZ*temperature;
+    theta.resize(numFreq);
+    theta[0] = kT;
+    for (int i = 1; i < numFreq; i++) {
+        double w = M_PI*i/(numFreq*dt);
+        theta[i] = hbar*w*(0.5+1/(exp(hbar*w/kT)-1));
+    }
+
+    // Compute the deconvolved version.  The algorithm is described in the supplementary
+    // information in https://doi.org/10.1021/acs.jpclett.1c01722.
+
+    auto C = [&](double w0, double w) {
+        double t = w*w-w0*w0;
+        return (friction/M_PI)*w0*w0/(t*t+friction*friction*w*w);
+    };
+    double dw = M_PI/(numFreq*dt);
+
+    // Normalize the kernel to reduce error at low frequencies.
+
+    vector<double> sum(numFreq, 0.0), scale(numFreq);
+    for (int i = 0; i < numFreq; i++) {
+        double wi = M_PI*(i+0.5)/(numFreq*dt);
+        for (int j = 0; j < numFreq; j++) {
+            double wj = M_PI*(j+0.5)/(numFreq*dt);
+            sum[i] += C(wi, wj)*dw;
+        }
+    }
+    for (int i = 0; i < numFreq; i++)
+        scale[i] = 0.5/sum[i];
+
+    // Compute intermediate quantities.
+
+    vector<vector<double> > D(numFreq, vector<double>(numFreq));
+    vector<double> h(numFreq);
+    vector<double> fcurrent(numFreq), fnext(numFreq);
+    for (int i = 0; i < numFreq; i++)
+        fcurrent[i] = 0.5*theta[i];
+    threads.execute([&] (ThreadPool& threads, int threadIndex) {
+        for (int i = threadIndex; i < numFreq; i += threads.getNumThreads()) {
+            double wi = M_PI*(i+0.5)/(numFreq*dt);
+            h[i] = 0.0;
+            for (int j = 0; j < numFreq; j++) {
+                double wj = M_PI*(j+0.5)/(numFreq*dt);
+                h[i] += dw*C(wj, wi)*fcurrent[j]*scale[j];
+                D[i][j] = 0.0;
+                for (int k = 0; k < numFreq; k++) {
+                    double wk = M_PI*(k+0.5)/(numFreq*dt);
+                    D[i][j] += dw*C(wk, wi)*C(wk, wj)*scale[k]*scale[k];
+                }
+            }
+        }
+    });
+    threads.waitForThreads();
+
+    // Perform the iteration.
+
+    for (int iteration = 0; iteration < 20; iteration++) {
+        for (int i = 0; i < numFreq; i++) {
+            double denom = 0.0;
+            for (int j = 0; j < numFreq; j++)
+                denom += dw*D[i][j]*fcurrent[j];
+            fnext[i] = fcurrent[i]*h[i]/denom;
+        }
+        fcurrent = fnext;
+    }
+    thetad = fnext;
+}

platforms/common/include/openmm/common/CommonIntegrateQTBStepKernel.h

+#ifndef OPENMM_COMMONINTEGRATEQTBSTEPKERNEL_H_
+#define OPENMM_COMMONINTEGRATEQTBSTEPKERNEL_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:                                                              *
+ *                                                                            *
+ * This program is free software: you can redistribute it and/or modify       *
+ * it under the terms of the GNU Lesser General Public License as published   *
+ * by the Free Software Foundation, either version 3 of the License, or       *
+ * (at your option) any later version.                                        *
+ *                                                                            *
+ * This program is distributed in the hope that it will be useful,            *
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of             *
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the              *
+ * GNU Lesser General Public License for more details.                        *
+ *                                                                            *
+ * You should have received a copy of the GNU Lesser General Public License   *
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.      *
+ * -------------------------------------------------------------------------- */
+
+#include "openmm/common/ComputeArray.h"
+#include "openmm/common/ComputeContext.h"
+#include "openmm/Platform.h"
+#include "openmm/kernels.h"
+
+namespace OpenMM {
+
+/**
+ * This kernel is invoked by QTBIntegrator to take one time step.
+ */
+class CommonIntegrateQTBStepKernel : public IntegrateQTBStepKernel {
+public:
+    CommonIntegrateQTBStepKernel(std::string name, const Platform& platform, ComputeContext& cc) : IntegrateQTBStepKernel(name, platform), cc(cc),
+        hasInitializedKernels(false) {
+    }
+    /**
+     * 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:
+    std::string createFFT(int size, int inputIndex, int& outputIndex, bool forward);
+    ComputeContext& cc;
+    int segmentLength, stepIndex, numFreq, numTypes;
+    double prevTemp, dt, friction;
+    std::vector<int> particleTypeVec;
+    ComputeArray oldDelta, noise, randomForce, segmentVelocity, thetad, cutoffFunction, workspace;
+    ComputeArray particleType, typeParticleCount, typeAdaptationRate, adaptedFriction, dfdt;
+    ComputeKernel kernel1, kernel2, kernel3, noiseKernel, forceKernel, adapt1Kernel, adapt2Kernel;
+    bool hasInitializedKernels;
+};
+
+} // namespace OpenMM
+
+#endif /*OPENMM_COMMONINTEGRATEQTBSTEPKERNEL_H_*/

platforms/common/src/IntegrationUtilities.cpp

@@ -917,6 +917,8 @@ void IntegrationUtilities::createCheckpoint(ostream& stream) {
     if (!random.isInitialized())
         return;
     stream.write((char*) &randomPos, sizeof(int));
+    int numRandom = random.getSize();
+    stream.write((char*) &numRandom, sizeof(int));
     vector<mm_float4> randomVec;
     random.download(randomVec);
     stream.write((char*) &randomVec[0], sizeof(mm_float4)*random.getSize());
@@ -929,6 +931,12 @@ void IntegrationUtilities::loadCheckpoint(istream& stream) {
     if (!random.isInitialized())
         return;
     stream.read((char*) &randomPos, sizeof(int));
+    int numRandom;
+    stream.read((char*) &numRandom, sizeof(int));
+    if (numRandom != random.getSize()) {
+        random.resize(numRandom);
+        randomKernel->setArg(0, numRandom);
+    }
     vector<mm_float4> randomVec(random.getSize());
     stream.read((char*) &randomVec[0], sizeof(mm_float4)*random.getSize());
     random.upload(randomVec);

platforms/common/src/kernels/qtb.cc

+/**
+ * Perform the first part of integration: velocity step.
+ */
+KERNEL void integrateQTBPart1(int numAtoms, int paddedNumAtoms, mixed dt, int stepIndex, GLOBAL mixed4* RESTRICT velm,
+        GLOBAL const mm_long* RESTRICT force, GLOBAL mixed* RESTRICT segmentVelocity, GLOBAL const int* RESTRICT atomOrder) {
+    mixed fscale = dt/(mixed) 0x100000000;
+    for (int atom = GLOBAL_ID; atom < numAtoms; atom += GLOBAL_SIZE) {
+        int atomIndex = atomOrder[atom];
+        mixed4 velocity = velm[atom];
+        segmentVelocity[3*numAtoms*stepIndex + atomIndex] = velocity.x;
+        segmentVelocity[3*numAtoms*stepIndex + numAtoms + atomIndex] = velocity.y;
+        segmentVelocity[3*numAtoms*stepIndex + 2*numAtoms + atomIndex] = velocity.z;
+        if (velocity.w != 0.0) {
+            velocity.x += fscale*velocity.w*force[atom];
+            velocity.y += fscale*velocity.w*force[atom+paddedNumAtoms];
+            velocity.z += fscale*velocity.w*force[atom+paddedNumAtoms*2];
+            velm[atom] = velocity;
+        }
+    }
+}
+
+/**
+ * Perform the second part of integration: position half step, then interact with heat bath,
+ * then another position half step.
+ */
+KERNEL void integrateQTBPart2(int numAtoms, mixed dt, mixed friction, int stepIndex, GLOBAL mixed4* RESTRICT velm,
+        GLOBAL mixed4* RESTRICT posDelta, GLOBAL mixed4* RESTRICT oldDelta, GLOBAL const mixed* RESTRICT randomForce,
+        GLOBAL const int* RESTRICT atomOrder) {
+    mixed halfdt = 0.5f*dt;
+    mixed vscale = EXP(-dt*friction);
+    for (int atom = GLOBAL_ID; atom < numAtoms; atom += GLOBAL_SIZE) {
+        int atomIndex = atomOrder[atom];
+        mixed4 velocity = velm[atom];
+        if (velocity.w != 0.0) {
+            mixed4 delta = make_mixed4(halfdt*velocity.x, halfdt*velocity.y, halfdt*velocity.z, 0);
+            mixed fscale = dt*velocity.w;
+            velocity.x = vscale*velocity.x + fscale*randomForce[3*numAtoms*stepIndex + atomIndex];
+            velocity.y = vscale*velocity.y + fscale*randomForce[3*numAtoms*stepIndex + numAtoms + atomIndex];
+            velocity.z = vscale*velocity.z + fscale*randomForce[3*numAtoms*stepIndex + 2*numAtoms + atomIndex];
+            velm[atom] = velocity;
+            delta += make_mixed4(halfdt*velocity.x, halfdt*velocity.y, halfdt*velocity.z, 0);
+            posDelta[atom] = delta;
+            oldDelta[atom] = delta;
+        }
+    }
+}
+
+/**
+ * Perform the third part of integration: apply constraint forces to velocities, then record
+ * the constrained positions.
+ */
+KERNEL void integrateQTBPart3(int numAtoms, mixed dt, GLOBAL real4* RESTRICT posq, GLOBAL mixed4* RESTRICT velm,
+         GLOBAL const mixed4* RESTRICT posDelta, GLOBAL const mixed4* RESTRICT oldDelta
+#ifdef USE_MIXED_PRECISION
+        , GLOBAL real4* RESTRICT posqCorrection
+#endif
+        ) {
+    mixed invDt = 1/dt;
+    for (int index = GLOBAL_ID; index < numAtoms; index += GLOBAL_SIZE) {
+        mixed4 velocity = velm[index];
+        if (velocity.w != 0.0) {
+            mixed4 delta = posDelta[index];
+            velocity.x += (delta.x-oldDelta[index].x)*invDt;
+            velocity.y += (delta.y-oldDelta[index].y)*invDt;
+            velocity.z += (delta.z-oldDelta[index].z)*invDt;
+            velm[index] = velocity;
+#ifdef USE_MIXED_PRECISION
+            real4 pos1 = posq[index];
+            real4 pos2 = posqCorrection[index];
+            mixed4 pos = make_mixed4(pos1.x+(mixed)pos2.x, pos1.y+(mixed)pos2.y, pos1.z+(mixed)pos2.z, pos1.w);
+#else
+            real4 pos = posq[index];
+#endif
+            pos.x += delta.x;
+            pos.y += delta.y;
+            pos.z += delta.z;
+#ifdef USE_MIXED_PRECISION
+            posq[index] = make_real4((real) pos.x, (real) pos.y, (real) pos.z, (real) pos.w);
+            posqCorrection[index] = make_real4(pos.x-(real) pos.x, pos.y-(real) pos.y, pos.z-(real) pos.z, 0);
+#else
+            posq[index] = pos;
+#endif
+        }
+    }
+}
+
+/**
+ * Update the buffer of white noise.
+ */
+KERNEL void generateNoise(int numAtoms, int segmentLength, GLOBAL float* RESTRICT noise, GLOBAL const float* RESTRICT random, unsigned int randomIndex) {
+    randomIndex = 4*randomIndex; // Interpret it as float instead of float4
+    int fftLength = 3*segmentLength;
+    for (int i = GROUP_ID; i < 3*numAtoms; i += NUM_GROUPS) {
+        // Copy segment 2 over to segment 1
+
+        for (int j = LOCAL_ID; j < segmentLength; j += LOCAL_SIZE)
+            noise[i*fftLength+j] = noise[i*fftLength+j+segmentLength];
+        SYNC_THREADS
+
+        // Copy segment 3 over to segment 2
+
+        for (int j = LOCAL_ID; j < segmentLength; j += LOCAL_SIZE)
+            noise[i*fftLength+j+segmentLength] = noise[i*fftLength+j+2*segmentLength];
+        SYNC_THREADS
+ 
+        // Fill segment 3 with new random values.
+
+        for (int j = LOCAL_ID; j < segmentLength; j += LOCAL_SIZE)
+            noise[i*fftLength+j+2*segmentLength] = random[randomIndex+i*segmentLength+j];
+        SYNC_THREADS
+    }
+}
+
+DEVICE mixed2 multiplyComplex(mixed2 c1, mixed2 c2) {
+    return make_mixed2(c1.x*c2.x-c1.y*c2.y, c1.x*c2.y+c1.y*c2.x);
+}
+
+DEVICE mixed2 multiplyComplexConj(mixed2 c1, mixed2 c2) {
+    return make_mixed2(c1.x*c2.x+c1.y*c2.y, c1.x*c2.y-c1.y*c2.x);
+}
+
+/**
+ * Generate the random force for the next segment.
+ */
+KERNEL void generateRandomForce(int numAtoms, int segmentLength, mixed dt, mixed friction, GLOBAL float* RESTRICT noise,
+        GLOBAL mixed* RESTRICT randomForce, GLOBAL mixed4* RESTRICT velm, GLOBAL mixed* RESTRICT thetad,
+        GLOBAL mixed* RESTRICT cutoffFunction, GLOBAL int* RESTRICT particleType, GLOBAL mixed* RESTRICT adaptedFriction,
+        GLOBAL mixed2* RESTRICT workspace) {
+    const int fftLength = 3*segmentLength;
+    const int numFreq = (fftLength+1)/2;
+    GLOBAL mixed2* data0 = &workspace[GROUP_ID*3*fftLength];
+    GLOBAL mixed2* data1 = &data0[fftLength];
+    GLOBAL mixed2* w = &data1[fftLength];
+    for (int i = LOCAL_ID; i < fftLength; i += LOCAL_SIZE)
+        w[i] = make_mixed2(cos(-i*2*M_PI/fftLength), sin(-i*2*M_PI/fftLength));
+    for (int i = GROUP_ID; i < 3*numAtoms; i += NUM_GROUPS) {
+        int atom = i/3;
+        int axis = i%3;
+        mixed invMass = velm[atom].w;
+        int type = particleType[atom];
+        if (invMass != 0) {
+            for (int j = LOCAL_ID; j < fftLength; j += LOCAL_SIZE)
+                data0[j] = make_mixed2(noise[i*fftLength+j], 0);
+            SYNC_THREADS
+            FFT_FORWARD
+            for (int j = LOCAL_ID; j < numFreq; j += LOCAL_SIZE) {
+                mixed f = M_PI*j/(numFreq*dt);
+                mixed gamma = adaptedFriction[type*numFreq+j];
+                mixed cw = (1 - 2*EXP(-dt*friction)*COS(f*dt) + EXP(-2*friction*dt)) / ((friction*friction+f*f)*dt*dt);
+                mixed2 d = RECIP_DATA[j] * SQRT(cutoffFunction[j]*thetad[j]*cw*gamma/friction);
+                RECIP_DATA[j] = d;
+                if (j != 0 && j*2 != fftLength)
+                    RECIP_DATA[fftLength-j] = make_mixed2(d.x, -d.y);
+            }
+            SYNC_THREADS
+            FFT_BACKWARD
+            const mixed scale = SQRT(2*friction/(dt*invMass))/fftLength;
+            for (int j = LOCAL_ID; j < segmentLength; j += LOCAL_SIZE)
+                randomForce[numAtoms*(3*j+axis)+atom] = scale*data0[segmentLength+j].x;
+            SYNC_THREADS
+        }
+    }
+}
+
+/**
+ * Update the friction rates used for generating noise, part 1: compute the error in the
+ * fluctuation dissipation theorem.
+ */
+KERNEL void adaptFrictionPart1(int numAtoms, int segmentLength, GLOBAL const mixed4* RESTRICT velm, GLOBAL const int* RESTRICT particleType,
+        GLOBAL const mixed* RESTRICT randomForce, GLOBAL const mixed* RESTRICT segmentVelocity, GLOBAL const mixed* RESTRICT adaptedFriction,
+        GLOBAL mm_ulong* RESTRICT dfdt, GLOBAL mixed2* RESTRICT workspace) {
+    const int numFreq = (segmentLength+1)/2;
+    GLOBAL mixed2* data0 = &workspace[GROUP_ID*3*segmentLength];
+    GLOBAL mixed2* data1 = &data0[segmentLength];
+    GLOBAL mixed2* w = &data1[segmentLength];
+    for (int i = LOCAL_ID; i < segmentLength; i += LOCAL_SIZE)
+        w[i] = make_mixed2(cos(-i*2*M_PI/segmentLength), sin(-i*2*M_PI/segmentLength));
+    for (int i = GROUP_ID; i < 3*numAtoms; i += NUM_GROUPS) {
+        int atom = i/3;
+        int axis = i%3;
+        int type = particleType[atom];
+        mixed invMass = velm[atom].w;
+        if (invMass != 0) {
+            // Pack the velocities and forces together so we can transform both at once.
+            for (int j = LOCAL_ID; j < segmentLength; j += LOCAL_SIZE)
+                data0[j] = make_mixed2(segmentVelocity[numAtoms*(3*j+axis)+atom], randomForce[numAtoms*(3*j+axis)+atom]);
+            SYNC_THREADS
+            ADAPTATION_FFT
+            for (int j = LOCAL_ID; j < numFreq; j += LOCAL_SIZE) {
+                mixed2 d1 = ADAPTATION_RECIP[j];
+                mixed2 d2 = (j == 0 ? ADAPTATION_RECIP[0] : ADAPTATION_RECIP[segmentLength-j]);
+                mixed2 v = 0.5f*make_mixed2(d1.x+d2.x, d1.y-d2.y);
+                mixed2 f = 0.5f*make_mixed2(d1.y+d2.y, -d1.x+d2.x);
+                mixed cvv = v.x*v.x + v.y*v.y;
+                mixed2 cvf = multiplyComplexConj(f, v);
+                mixed dfdtinc = adaptedFriction[type*numFreq+j]*cvv/invMass - cvf.x;
+                ATOMIC_ADD(&dfdt[type*numFreq+j], (mm_ulong) realToFixedPoint(dfdtinc));
+            }
+        }
+    }
+}
+
+/**
+ * Update the friction rates used for generating noise, part 2: update the friction based
+ * on the error.
+ */
+KERNEL void adaptFrictionPart2(int numTypes, int segmentLength, mixed dt, GLOBAL const int* RESTRICT typeParticleCount,
+        GLOBAL const float* RESTRICT typeAdaptationRate, GLOBAL mixed* RESTRICT adaptedFriction, GLOBAL const mm_long* RESTRICT dfdt) {
+    int numFreq = (3*segmentLength+1)/2;
+    for (int type = GROUP_ID; type < numTypes; type += NUM_GROUPS) {
+        mixed scale = dt*typeAdaptationRate[type]/(3*typeParticleCount[type]*segmentLength)/(mixed) 0x100000000;
+        for (int i = LOCAL_ID; i < numFreq; i += LOCAL_SIZE) {
+            mixed delta = -scale*dfdt[type*numFreq+i];
+            adaptedFriction[type*numFreq+i] = max((mixed) 0, adaptedFriction[type*numFreq+i]+delta);
+        }
+    }
+}
\ No newline at end of file

platforms/cuda/src/CudaKernelFactory.cpp

@@ -36,6 +36,7 @@
 #include "openmm/common/CommonCalcCustomNonbondedForceKernel.h"
 #include "openmm/common/CommonIntegrateCustomStepKernel.h"
 #include "openmm/common/CommonIntegrateNoseHooverStepKernel.h"
+#include "openmm/common/CommonIntegrateQTBStepKernel.h"
 #include "openmm/internal/ContextImpl.h"
 #include "openmm/OpenMMException.h"
 
@@ -142,6 +143,8 @@ KernelImpl* CudaKernelFactory::createKernelImpl(std::string name, const Platform
         return new CommonIntegrateCustomStepKernel(name, platform, cu);
     if (name == IntegrateDPDStepKernel::Name())
         return new CommonIntegrateDPDStepKernel(name, platform, cu);
+    if (name == IntegrateQTBStepKernel::Name())
+        return new CommonIntegrateQTBStepKernel(name, platform, cu);
     if (name == ApplyAndersenThermostatKernel::Name())
         return new CommonApplyAndersenThermostatKernel(name, platform, cu);
     if (name == IntegrateNoseHooverStepKernel::Name())

platforms/cuda/src/CudaPlatform.cpp

@@ -105,6 +105,7 @@ CudaPlatform::CudaPlatform() {
     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/cuda/tests/TestCudaQTBIntegrator.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 "CudaTests.h"
+#include "TestQTBIntegrator.h"
+
+void runPlatformTests() {
+}

platforms/hip/src/HipKernelFactory.cpp

@@ -37,6 +37,7 @@
 #include "openmm/common/CommonCalcCustomNonbondedForceKernel.h"
 #include "openmm/common/CommonIntegrateCustomStepKernel.h"
 #include "openmm/common/CommonIntegrateNoseHooverStepKernel.h"
+#include "openmm/common/CommonIntegrateQTBStepKernel.h"
 #include "openmm/internal/ContextImpl.h"
 #include "openmm/OpenMMException.h"
 
@@ -141,6 +142,8 @@ KernelImpl* HipKernelFactory::createKernelImpl(std::string name, const Platform&
         return new CommonIntegrateCustomStepKernel(name, platform, cu);
     if (name == IntegrateDPDStepKernel::Name())
         return new CommonIntegrateDPDStepKernel(name, platform, cu);
+    if (name == IntegrateQTBStepKernel::Name())
+        return new CommonIntegrateQTBStepKernel(name, platform, cu);
     if (name == ApplyAndersenThermostatKernel::Name())
         return new CommonApplyAndersenThermostatKernel(name, platform, cu);
     if (name == IntegrateNoseHooverStepKernel::Name())

platforms/hip/src/HipPlatform.cpp

@@ -105,6 +105,7 @@ HipPlatform::HipPlatform() {
     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/hip/tests/TestHipQTBIntegrator.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 "HipTests.h"
+#include "TestQTBIntegrator.h"
+
+void runPlatformTests() {
+}

platforms/opencl/src/OpenCLKernelFactory.cpp

@@ -34,6 +34,7 @@
 #include "openmm/common/CommonCalcCustomNonbondedForceKernel.h"
 #include "openmm/common/CommonIntegrateCustomStepKernel.h"
 #include "openmm/common/CommonIntegrateNoseHooverStepKernel.h"
+#include "openmm/common/CommonIntegrateQTBStepKernel.h"
 #include "openmm/internal/ContextImpl.h"
 #include "openmm/OpenMMException.h"
 
@@ -140,6 +141,8 @@ KernelImpl* OpenCLKernelFactory::createKernelImpl(std::string name, const Platfo
         return new CommonIntegrateCustomStepKernel(name, platform, cl);
     if (name == IntegrateDPDStepKernel::Name())
         return new CommonIntegrateDPDStepKernel(name, platform, cl);
+    if (name == IntegrateQTBStepKernel::Name())
+        return new CommonIntegrateQTBStepKernel(name, platform, cl);
     if (name == ApplyAndersenThermostatKernel::Name())
         return new CommonApplyAndersenThermostatKernel(name, platform, cl);
     if (name == IntegrateNoseHooverStepKernel::Name())