Improved temperature reporting for Drude particles (#3486)

* DrudeLangevinIntegrator has getSystemTemperature()

* DrudeNoseHooverIntegrator has getSystemTemperature()

* StateDataReporter reports system temperature for Drude systems

* Fixed incorrect return type

olla/include/openmm/kernels.h

@@ -9,7 +9,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2008-2021 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2022 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -184,6 +184,15 @@ public:
      * @param velocities  a vector containing the particle velocities
      */
     virtual void setVelocities(ContextImpl& context, const std::vector<Vec3>& velocities) = 0;
+    /**
+     * Compute velocities, shifted in time to account for a leapfrog integrator.  The shift
+     * is based on the most recently computed forces.
+     * 
+     * @param context     the context in which to execute this kernel
+     * @param timeShift   the amount by which to shift the velocities in time
+     * @param velocities  the shifted velocities are returned in this
+     */
+    virtual void computeShiftedVelocities(ContextImpl& context, double timeShift, std::vector<Vec3>& velocities) = 0;
     /**
      * Get the current forces on all particles.
      *

openmmapi/include/openmm/internal/ContextImpl.h

@@ -9,7 +9,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2008-2021 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2022 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -217,6 +217,14 @@ public:
      * Calculate the kinetic energy of the system (in kJ/mol).
      */
     double calcKineticEnergy();
+    /**
+     * Compute velocities, shifted in time to account for a leapfrog integrator.  The shift
+     * is based on the most recently computed forces.
+     * 
+     * @param timeShift   the amount by which to shift the velocities in time
+     * @param velocities  the shifted velocities are returned in this
+     */
+    void computeShiftedVelocities(double timeShift, std::vector<Vec3>& velocities);
     /**
      * This should be called at the start of each time step.  It calls updateContextState() on each
      * ForceImpl in the system, allowing them to modify the values of state variables.

openmmapi/src/ContextImpl.cpp

@@ -6,7 +6,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2008-2016 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2022 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -323,6 +323,10 @@ double ContextImpl::calcKineticEnergy() {
     return integrator.computeKineticEnergy();
 }
 
+void ContextImpl::computeShiftedVelocities(double timeShift, std::vector<Vec3>& velocities) {
+    updateStateDataKernel.getAs<UpdateStateDataKernel>().computeShiftedVelocities(*this, timeShift, velocities);
+}
+
 bool ContextImpl::updateContextState() {
     bool forcesInvalid = false;
     for (auto force : forceImpls)

platforms/common/include/openmm/common/IntegrationUtilities.h

@@ -9,7 +9,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2009-2019 Stanford University and the Authors.      *
+ * Portions copyright (c) 2009-2022 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -130,6 +130,13 @@ public:
      * @param timeShift   the amount by which to shift the velocities in time
      */
     double computeKineticEnergy(double timeShift);
+    /**
+     * Compute the current velocities, shifting them in time to account for a leapfrog integrator.
+     * 
+     * @param timeShift   the amount by which to shift the velocities in time
+     * @param velocities  the shifted velocities are returned in this
+     */
+    void computeShiftedVelocities(double timeShift, std::vector<Vec3>& velocities);
 protected:
     virtual void applyConstraintsImpl(bool constrainVelocities, double tol) = 0;
     ComputeContext& context;

platforms/common/src/IntegrationUtilities.cpp

@@ -6,7 +6,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2009-2021 Stanford University and the Authors.      *
+ * Portions copyright (c) 2009-2022 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -861,3 +861,49 @@ double IntegrationUtilities::computeKineticEnergy(double timeShift) {
         posDelta.copyTo(context.getVelm());
     return 0.5*energy;
 }
+
+void IntegrationUtilities::computeShiftedVelocities(double timeShift, vector<Vec3>& velocities) {
+    ContextSelector selector(context);
+    int numParticles = context.getNumAtoms();
+    if (timeShift != 0) {
+        // Copy the velocities into the posDelta array while we temporarily modify them.
+
+        context.getVelm().copyTo(posDelta);
+
+        // Apply the time shift.
+
+        timeShiftKernel->setArg(0, context.getVelm());
+        timeShiftKernel->setArg(1, context.getLongForceBuffer());
+        if (context.getUseDoublePrecision())
+            timeShiftKernel->setArg(2, timeShift);
+        else
+            timeShiftKernel->setArg(2, (float) timeShift);
+        timeShiftKernel->execute(numParticles);
+        applyConstraintsImpl(true, 1e-4);
+    }
+    
+    // Retrieve the velocities.
+    
+    velocities.resize(numParticles);
+    if (context.getUseDoublePrecision() || context.getUseMixedPrecision()) {
+        auto velm = (mm_double4*)context.getPinnedBuffer();
+        context.getVelm().download(velm);
+        for (int i = 0; i < numParticles; i++) {
+            mm_double4 v = velm[i];
+            velocities[i] = Vec3(v.x, v.y, v.z);
+        }
+    }
+    else {
+        auto velm = (mm_float4*)context.getPinnedBuffer();
+        context.getVelm().download(velm);
+        for (int i = 0; i < numParticles; i++) {
+            mm_float4 v = velm[i];
+            velocities[i] = Vec3(v.x, v.y, v.z);
+        }
+    }
+    
+    // Restore the velocities.
+    
+    if (timeShift != 0)
+        posDelta.copyTo(context.getVelm());
+}

platforms/cuda/include/CudaKernels.h

@@ -9,7 +9,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2008-2021 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2022 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -170,6 +170,15 @@ public:
      * @param velocities  a vector containg the particle velocities
      */
     void setVelocities(ContextImpl& context, const std::vector<Vec3>& velocities);
+    /**
+     * Compute velocities, shifted in time to account for a leapfrog integrator.  The shift
+     * is based on the most recently computed forces.
+     * 
+     * @param context     the context in which to execute this kernel
+     * @param timeShift   the amount by which to shift the velocities in time
+     * @param velocities  the shifted velocities are returned in this
+     */
+    void computeShiftedVelocities(ContextImpl& context, double timeShift, std::vector<Vec3>& velocities);
     /**
      * Get the current forces on all particles.
      *

platforms/cuda/src/CudaKernels.cpp

@@ -280,6 +280,10 @@ void CudaUpdateStateDataKernel::setVelocities(ContextImpl& context, const vector
     }
 }
 
+void CudaUpdateStateDataKernel::computeShiftedVelocities(ContextImpl& context, double timeShift, vector<Vec3>& velocities) {
+    cu.getIntegrationUtilities().computeShiftedVelocities(timeShift, velocities);
+}
+
 void CudaUpdateStateDataKernel::getForces(ContextImpl& context, vector<Vec3>& forces) {
     ContextSelector selector(cu);
     long long* force = (long long*) cu.getPinnedBuffer();

platforms/opencl/include/OpenCLKernels.h

@@ -9,7 +9,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2008-2019 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2022 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -144,6 +144,15 @@ public:
      * @param velocities  a vector containg the particle velocities
      */
     void setVelocities(ContextImpl& context, const std::vector<Vec3>& velocities);
+    /**
+     * Compute velocities, shifted in time to account for a leapfrog integrator.  The shift
+     * is based on the most recently computed forces.
+     * 
+     * @param context     the context in which to execute this kernel
+     * @param timeShift   the amount by which to shift the velocities in time
+     * @param velocities  the shifted velocities are returned in this
+     */
+    void computeShiftedVelocities(ContextImpl& context, double timeShift, std::vector<Vec3>& velocities);
     /**
      * Get the current forces on all particles.
      *

platforms/opencl/src/OpenCLKernels.cpp

@@ -291,6 +291,10 @@ void OpenCLUpdateStateDataKernel::setVelocities(ContextImpl& context, const vect
     }
 }
 
+void OpenCLUpdateStateDataKernel::computeShiftedVelocities(ContextImpl& context, double timeShift, vector<Vec3>& velocities) {
+    cl.getIntegrationUtilities().computeShiftedVelocities(timeShift, velocities);
+}
+
 void OpenCLUpdateStateDataKernel::getForces(ContextImpl& context, vector<Vec3>& forces) {
     const vector<cl_int>& order = cl.getAtomIndex();
     int numParticles = context.getSystem().getNumParticles();

platforms/reference/include/ReferenceKernels.h

@@ -174,6 +174,15 @@ public:
      * @param velocities  a vector containg the particle velocities
      */
     void setVelocities(ContextImpl& context, const std::vector<Vec3>& velocities);
+    /**
+     * Compute velocities, shifted in time to account for a leapfrog integrator.  The shift
+     * is based on the most recently computed forces.
+     * 
+     * @param context     the context in which to execute this kernel
+     * @param timeShift   the amount by which to shift the velocities in time
+     * @param velocities  the shifted velocities are returned in this
+     */
+    void computeShiftedVelocities(ContextImpl& context, double timeShift, std::vector<Vec3>& velocities);
     /**
      * Get the current forces on all particles.
      *
@@ -216,6 +225,7 @@ public:
     void loadCheckpoint(ContextImpl& context, std::istream& stream);
 private:
     ReferencePlatform::PlatformData& data;
+    std::vector<double> masses;
 };
 
 /**

platforms/reference/src/ReferenceKernels.cpp

@@ -144,30 +144,9 @@ static void validateVariables(const Lepton::ExpressionTreeNode& node, const set<
  * for a leapfrog integrator.
  */
 static double computeShiftedKineticEnergy(ContextImpl& context, vector<double>& masses, double timeShift) {
-    vector<Vec3>& posData = extractPositions(context);
-    vector<Vec3>& velData = extractVelocities(context);
-    vector<Vec3>& forceData = extractForces(context);
     int numParticles = context.getSystem().getNumParticles();
-    
-    // Compute the shifted velocities.
-    
     vector<Vec3> shiftedVel(numParticles);
-    for (int i = 0; i < numParticles; ++i) {
-        if (masses[i] > 0)
-            shiftedVel[i] = velData[i]+forceData[i]*(timeShift/masses[i]);
-        else
-            shiftedVel[i] = velData[i];
-    }
-    
-    // Apply constraints to them.
-    
-    vector<double> inverseMasses(numParticles);
-    for (int i = 0; i < numParticles; i++)
-        inverseMasses[i] = (masses[i] == 0 ? 0 : 1/masses[i]);
-    extractConstraints(context).applyToVelocities(posData, shiftedVel, inverseMasses, 1e-4);
-    
-    // Compute the kinetic energy.
-    
+    context.computeShiftedVelocities(timeShift, shiftedVel);
     double energy = 0.0;
     for (int i = 0; i < numParticles; ++i)
         if (masses[i] > 0)
@@ -203,6 +182,10 @@ double ReferenceCalcForcesAndEnergyKernel::finishComputation(ContextImpl& contex
 }
 
 void ReferenceUpdateStateDataKernel::initialize(const System& system) {
+    int numParticles = system.getNumParticles();
+    masses.resize(numParticles);
+    for (int i = 0; i < numParticles; ++i)
+        masses[i] = system.getParticleMass(i);
 }
 
 double ReferenceUpdateStateDataKernel::getTime(const ContextImpl& context) const {
@@ -257,6 +240,32 @@ void ReferenceUpdateStateDataKernel::setVelocities(ContextImpl& context, const s
     }
 }
 
+void ReferenceUpdateStateDataKernel::computeShiftedVelocities(ContextImpl& context, double timeShift, std::vector<Vec3>& velocities) {
+    int numParticles = context.getSystem().getNumParticles();
+    vector<Vec3>& posData = extractPositions(context);
+    vector<Vec3>& velData = extractVelocities(context);
+    vector<Vec3>& forceData = extractForces(context);
+    
+    // Compute the shifted velocities.
+    
+    velocities.resize(numParticles);
+    vector<double> inverseMasses(numParticles);
+    for (int i = 0; i < numParticles; ++i) {
+        if (masses[i] == 0) {
+            velocities[i] = velData[i];
+            inverseMasses[i] = 0;
+        }
+        else {
+            velocities[i] = velData[i]+forceData[i]*(timeShift/masses[i]);
+            inverseMasses[i] = 1/masses[i];
+        }
+    }
+    
+    // Apply constraints to them.
+    
+    extractConstraints(context).applyToVelocities(posData, velocities, inverseMasses, 1e-4);
+}
+
 void ReferenceUpdateStateDataKernel::getForces(ContextImpl& context, std::vector<Vec3>& forces) {
     int numParticles = context.getSystem().getNumParticles();
     vector<Vec3>& forceData = extractForces(context);

plugins/drude/openmmapi/include/openmm/DrudeLangevinIntegrator.h

@@ -9,7 +9,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2008-2021 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2022 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -117,6 +117,15 @@ public:
      * @param steps   the number of time steps to take
      */
     void step(int steps) override;
+    /**
+     * Compute the instantaneous temperature of the System, measured in Kelvin.
+     * This is calculated based on the kinetic energy of the ordinary particles (ones
+     * not attached to a Drude particle), as well as the center of mass motion of the
+     * Drude particle pairs.  It does not include the internal motion of the pairs.
+     * On average, this should be approximately equal to the value returned by
+     * getTemperature().
+     */
+    double computeSystemTemperature();
 protected:
     /**
      * This will be called by the Context when it is created.  It informs the Integrator

plugins/drude/openmmapi/include/openmm/DrudeNoseHooverIntegrator.h

@@ -9,9 +9,9 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2019 Stanford University and the Authors.           *
+ * Portions copyright (c) 2019-2022 Stanford University and the Authors.      *
  * Authors: Andreas Krämer and Andrew C. Simmonett                            *
- * Contributors:                                                              *
+ * Contributors: Peter Eastman                                                *
  *                                                                            *
  * Permission is hereby granted, free of charge, to any person obtaining a    *
  * copy of this software and associated documentation files (the "Software"), *
@@ -93,6 +93,16 @@ public:
      * Compute the kinetic energy of all (real and drude) particles at the current time.
      */
     double computeTotalKineticEnergy();
+    /**
+     * Compute the instantaneous temperature of the System, measured in Kelvin.
+     * This is calculated based on the kinetic energy of the ordinary particles (ones
+     * not attached to a Drude particle), as well as the center of mass motion of the
+     * Drude particle pairs.  It does not include the internal motion of the pairs.
+     * On average, this should be approximately equal to the value returned by
+     * getTemperature().
+     */
+    double computeSystemTemperature();
+protected:
     /**
      * Return a list of velocities normally distributed around a target temperature, with the Drude
      * temperatures assigned according to the Drude temperature assigned to the integrator.
@@ -103,7 +113,6 @@ public:
      */
     virtual std::vector<Vec3> getVelocitiesForTemperature(const System &system, double temperature,
                                                           int randomSeed) const override;
-protected:
     double drudeTemperature;
 };
 

plugins/drude/openmmapi/src/DrudeHelpers.cpp

@@ -6,7 +6,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2008-2015 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2022 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -31,6 +31,7 @@
 
 #include "sfmt/SFMT.h"
 #include "SimTKOpenMMRealType.h"
+#include "openmm/CMMotionRemover.h"
 #include "openmm/DrudeForce.h"
 #include "openmm/DrudeLangevinIntegrator.h"
 #include "openmm/Context.h"
@@ -39,9 +40,14 @@
 
 #include <set>
 
+using namespace std;
+
 namespace OpenMM {
 
-std::vector<Vec3> assignDrudeVelocities(const System &system, double temperature, double drudeTemperature, int randomSeed) {
+/**
+ * Identify normal particles (not part of a pair) and Drude particle pairs.
+ */
+void findParticlesAndPairs(const System &system, vector<int>& normalParticles, vector<pair<int, int> >& pairParticles) {
     // Find the underlying Drude force object
     const DrudeForce* drudeForce = NULL;
     for (int i = 0; i < system.getNumForces(); i++)
@@ -55,11 +61,9 @@ std::vector<Vec3> assignDrudeVelocities(const System &system, double temperature
         throw OpenMMException("The System does not contain a DrudeForce");
 
     // Figure out which particles are individual and which are Drude pairs
-    std::set<int> particles;
-    std::vector<std::pair<int, int>> pairParticles;
-    for (int i = 0; i < system.getNumParticles(); i++) {
+    set<int> particles;
+    for (int i = 0; i < system.getNumParticles(); i++)
         particles.insert(i);
-    }
     for (int i = 0; i < drudeForce->getNumParticles(); i++) {
         int p, p1, p2, p3, p4;
         double charge, polarizability, aniso12, aniso34;
@@ -68,12 +72,18 @@ std::vector<Vec3> assignDrudeVelocities(const System &system, double temperature
         particles.erase(p1);
         pairParticles.emplace_back(p, p1);
     }
-    std::vector<int> normalParticles(particles.begin(), particles.end());
+    normalParticles.insert(normalParticles.begin(), particles.begin(), particles.end());
+}
+
+vector<Vec3> assignDrudeVelocities(const System &system, double temperature, double drudeTemperature, int randomSeed) {
+    vector<int> normalParticles;
+    vector<pair<int, int> > pairParticles;
+    findParticlesAndPairs(system, normalParticles, pairParticles);
 
     // Generate the list of Gaussian random numbers.
     OpenMM_SFMT::SFMT sfmt;
     init_gen_rand(randomSeed, sfmt);
-    std::vector<double> randoms;
+    vector<double> randoms;
     while (randoms.size() < system.getNumParticles()*3) {
         double x, y, r2;
         do {
@@ -81,13 +91,13 @@ std::vector<Vec3> assignDrudeVelocities(const System &system, double temperature
             y = 2.0*genrand_real2(sfmt)-1.0;
             r2 = x*x + y*y;
         } while (r2 >= 1.0 || r2 == 0.0);
-        double multiplier = sqrt((-2.0*std::log(r2))/r2);
+        double multiplier = sqrt((-2.0*log(r2))/r2);
         randoms.push_back(x*multiplier);
         randoms.push_back(y*multiplier);
     }
 
     // Assign the velocities.
-    std::vector<Vec3> velocities(system.getNumParticles(), Vec3());
+    vector<Vec3> velocities(system.getNumParticles(), Vec3());
     int nextRandom = 0;
     // First the indivitual atoms
     for (const auto &atom : normalParticles ) {
@@ -117,4 +127,56 @@ std::vector<Vec3> assignDrudeVelocities(const System &system, double temperature
     return velocities;
 }
 
+double computeSystemTemperatureFromVelocities(const System& system, const vector<Vec3>& velocities) {
+    vector<int> normalParticles;
+    vector<pair<int, int> > pairParticles;
+    findParticlesAndPairs(system, normalParticles, pairParticles);
+    double energy = 0.0;
+    int dof = 0;
+    
+    // Kinetic energy and degrees of freedom from normal particles.
+    
+    for (int i : normalParticles) {
+        double mass = system.getParticleMass(i);
+        if (mass > 0) {
+            energy += mass*(velocities[i].dot(velocities[i]));
+            dof += 3;
+        }
+    }
+    
+    // Kinetic energy and degrees of freedom from Drude particle pairs.
+    
+    for (auto pair : pairParticles) {
+        int p1 = pair.first;
+        int p2 = pair.second;
+        double mass1 = system.getParticleMass(p1);
+        double mass2 = system.getParticleMass(p2);
+        if (mass1 != 0 || mass2 != 0) {
+            Vec3 momentum = mass1*velocities[p1] + mass2*velocities[p2];
+            energy += momentum.dot(momentum)/(mass1+mass2);
+            dof += 3;
+        }
+    }
+    
+    // Reduce degrees of freedom for constraints.
+
+    for (int i = 0; i < system.getNumConstraints(); i++) {
+        int p1, p2;
+        double distance;
+        system.getConstraintParameters(i, p1, p2, distance);
+        if (system.getParticleMass(p1) > 0 || system.getParticleMass(p2) > 0)
+            dof--;
+    }
+
+    // Reduce degrees of freedom if there is a CMMotionRemover.
+
+    for (int i = 0; i < system.getNumForces(); i++)
+        if (dynamic_cast<const CMMotionRemover*>(&system.getForce(i)) != NULL) {
+            dof -= 3;
+            break;
+        }
+    energy *= 0.5;
+    return 2*energy/(dof*BOLTZ);
+}
+
 }

plugins/drude/openmmapi/src/DrudeLangevinIntegrator.cpp

@@ -6,7 +6,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2008-2021 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2022 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -43,6 +43,10 @@ using namespace OpenMM;
 using std::string;
 using std::vector;
 
+namespace OpenMM {
+    double computeSystemTemperatureFromVelocities(const System& system, const vector<Vec3>& velocities);
+}
+
 DrudeLangevinIntegrator::DrudeLangevinIntegrator(double temperature, double frictionCoeff, double drudeTemperature, double drudeFrictionCoeff, double stepSize) : DrudeIntegrator(stepSize) {
     setTemperature(temperature);
     setFriction(frictionCoeff);
@@ -115,3 +119,12 @@ void DrudeLangevinIntegrator::step(int steps) {
         kernel.getAs<IntegrateDrudeLangevinStepKernel>().execute(*context, *this);
     }
 }
+
+double DrudeLangevinIntegrator::computeSystemTemperature() {
+    if (context == NULL)
+        throw OpenMMException("This Integrator is not bound to a context!");  
+    context->calcForcesAndEnergy(true, false, getIntegrationForceGroups());
+    vector<Vec3> velocities;
+    context->computeShiftedVelocities(getVelocityTimeOffset(), velocities);
+    return computeSystemTemperatureFromVelocities(context->getSystem(), velocities);
+}

plugins/drude/openmmapi/src/DrudeNoseHooverIntegrator.cpp

@@ -6,9 +6,9 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2019 Stanford University and the Authors.           *
+ * Portions copyright (c) 2019-2022 Stanford University and the Authors.      *
  * Authors: Andreas Krämer and Andrew C. Simmonett                            *
- * Contributors:                                                              *
+ * Contributors: Peter Eastman                                                *
  *                                                                            *
  * Permission is hereby granted, free of charge, to any person obtaining a    *
  * copy of this software and associated documentation files (the "Software"), *
@@ -48,6 +48,7 @@ using std::vector;
 
 namespace OpenMM {
     extern std::vector<Vec3> assignDrudeVelocities(const System &system, double temperature, double drudeTemperature, int randomSeed);
+    double computeSystemTemperatureFromVelocities(const System& system, const vector<Vec3>& velocities);
 }
 
 DrudeNoseHooverIntegrator::DrudeNoseHooverIntegrator(double temperature, double collisionFrequency, 
@@ -146,6 +147,14 @@ double DrudeNoseHooverIntegrator::computeTotalKineticEnergy() {
     return kernel.getAs<IntegrateNoseHooverStepKernel>().computeKineticEnergy(*context, *this);
 }
 
+double DrudeNoseHooverIntegrator::computeSystemTemperature() {
+    if (context == NULL)
+        throw OpenMMException("This Integrator is not bound to a context!");  
+    context->calcForcesAndEnergy(true, false, getIntegrationForceGroups());
+    vector<Vec3> velocities;
+    context->computeShiftedVelocities(getVelocityTimeOffset(), velocities);
+    return computeSystemTemperatureFromVelocities(context->getSystem(), velocities);
+}
 std::vector<Vec3> DrudeNoseHooverIntegrator::getVelocitiesForTemperature(const System &system, double temperature,
                                                                          int randomSeedIn) const {
     return assignDrudeVelocities(system, temperature, drudeTemperature, randomSeedIn);

plugins/drude/platforms/reference/src/ReferenceDrudeKernels.cpp

@@ -6,7 +6,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2011-2020 Stanford University and the Authors.      *
+ * Portions copyright (c) 2011-2022 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -64,26 +64,8 @@ static ReferenceConstraints& extractConstraints(ContextImpl& context) {
 static double computeShiftedKineticEnergy(ContextImpl& context, vector<double>& inverseMasses, double timeShift) {
     const System& system = context.getSystem();
     int numParticles = system.getNumParticles();
-    vector<Vec3>& posData = extractPositions(context);
-    vector<Vec3>& velData = extractVelocities(context);
-    vector<Vec3>& forceData = extractForces(context);
-    
-    // Compute the shifted velocities.
-    
     vector<Vec3> shiftedVel(numParticles);
-    for (int i = 0; i < numParticles; ++i) {
-        if (inverseMasses[i] > 0)
-            shiftedVel[i] = velData[i]+forceData[i]*(timeShift*inverseMasses[i]);
-        else
-            shiftedVel[i] = velData[i];
-    }
-    
-    // Apply constraints to them.
-    
-    extractConstraints(context).applyToVelocities(posData, shiftedVel, inverseMasses, 1e-4);
-    
-    // Compute the kinetic energy.
-    
+    context.computeShiftedVelocities(timeShift, shiftedVel);
     double energy = 0.0;
     for (int i = 0; i < numParticles; ++i)
         if (inverseMasses[i] > 0)

plugins/drude/tests/TestDrudeLangevinIntegrator.h

@@ -158,17 +158,21 @@ void testWater() {
     // Compute the internal and center of mass temperatures.
 
     double ke = 0;
+    double systemTemp = 0;
     int numSteps = 8000;
     for (int i = 0; i < numSteps; i++) {
         integ.step(1);
         ke += context.getState(State::Energy).getKineticEnergy();
+        systemTemp += integ.computeSystemTemperature();
     }
     ke /= numSteps;
+    systemTemp /= numSteps;
     int numStandardDof = 3*3*numMolecules-system.getNumConstraints();
     int numDrudeDof = 3*numMolecules;
     int numDof = numStandardDof+numDrudeDof;
     double expectedTemp = (numStandardDof*temperature+numDrudeDof*temperatureDrude)/numDof;
     ASSERT_USUALLY_EQUAL_TOL(expectedTemp, ke/(0.5*numDof*BOLTZ), 0.03);
+    ASSERT_USUALLY_EQUAL_TOL(temperature, systemTemp, 0.03);
 }
 
 void testForceEnergyConsistency() {

plugins/drude/tests/TestDrudeNoseHoover.h

@@ -6,9 +6,9 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2019 Stanford University and the Authors.           *
+ * Portions copyright (c) 2019-2022 Stanford University and the Authors.      *
  * Authors: Andreas Krämer and Andrew C. Simmonett                            *
- * Contributors:                                                              *
+ * Contributors: Peter Eastman                                                *
  *                                                                            *
  * Permission is hereby granted, free of charge, to any person obtaining a    *
  * copy of this software and associated documentation files (the "Software"), *
@@ -151,7 +151,9 @@ void testWaterBox() {
     }
 
     // Compute the internal and center of mass temperatures.
+
     double totalKE = 0;
+    double systemTemp = 0;
     const int numSteps = 500;
     double meanTemp = 0.0;
     double meanDrudeTemp = 0.0;
@@ -171,11 +173,14 @@ void testWaterBox() {
         double conserved = PE + fullKE + heatBathEnergy;
         meanConserved = (i*meanConserved + conserved)/(i+1);
         totalKE += KE;
+        systemTemp += integ.computeSystemTemperature();
         ASSERT(fabs(meanConserved - conserved) < TOL);
     }
     totalKE /= numSteps;
+    systemTemp /= numSteps;
     ASSERT_USUALLY_EQUAL_TOL(temperature, meanTemp,  0.03);
     ASSERT_USUALLY_EQUAL_TOL(temperatureDrude, meanDrudeTemp,  0.03);
+    ASSERT_USUALLY_EQUAL_TOL(temperature, systemTemp, 0.03);
 }
 
 
@@ -192,7 +197,6 @@ double testWaterBoxWithHardWallConstraint(double hardWallConstraint){
     const int numMolecules = gridSize*gridSize*gridSize;
     int numStandardDof = 3*3*numMolecules - system.getNumConstraints();
     int numDrudeDof = 3*numMolecules;
-    int numDof = numStandardDof+numDrudeDof;
     const double temperature = 300.0;
     const double temperatureDrude = 10.0;
 
@@ -216,11 +220,7 @@ double testWaterBoxWithHardWallConstraint(double hardWallConstraint){
     integ.step(50);
 
     // Compute the internal and center of mass temperatures.
-    double totalKE = 0;
     const int numSteps = 500;
-    double meanTemp = 0.0;
-    double meanDrudeTemp = 0.0;
-    double meanConserved = 0.0;
     double maxR = 0.0;
     for (int i = 0; i < numSteps; i++) {
         integ.step(1);

wrappers/python/openmm/app/statedatareporter.py

@@ -248,7 +248,11 @@ class StateDataReporter(object):
         if self._totalEnergy:
             values.append((state.getKineticEnergy()+state.getPotentialEnergy()).value_in_unit(unit.kilojoules_per_mole))
         if self._temperature:
-            values.append((2*state.getKineticEnergy()/(self._dof*unit.MOLAR_GAS_CONSTANT_R)).value_in_unit(unit.kelvin))
+            integrator = simulation.context.getIntegrator()
+            if hasattr(integrator, 'computeSystemTemperature'):
+                values.append(integrator.computeSystemTemperature().value_in_unit(unit.kelvin))
+            else:
+                values.append((2*state.getKineticEnergy()/(self._dof*unit.MOLAR_GAS_CONSTANT_R)).value_in_unit(unit.kelvin))
         if self._volume:
             values.append(volume.value_in_unit(unit.nanometer**3))
         if self._density: