Refactor v-rescale thermostat to allow for other algorithms
authorPascal Merz <pascal.merz@me.com>
Tue, 8 Sep 2020 02:40:22 +0000 (20:40 -0600)
committerMark Abraham <mark.j.abraham@gmail.com>
Fri, 18 Sep 2020 15:26:03 +0000 (15:26 +0000)
This is pure refactoring, splitting the element class (interfaces with
the simulator and other objects) from the actual implementation of the
thermostat. This will allow to implement other thermostat algorithms
using the same infrastructure.

Refs #3423

src/gromacs/modularsimulator/velocityscalingtemperaturecoupling.cpp
src/gromacs/modularsimulator/velocityscalingtemperaturecoupling.h

index 54b3a104b78db3ef5605401e43e314e234bed2fc..8a9bf53ca990baba0bfc81c8bc07b9f39396c160 100644 (file)
 namespace gmx
 {
 
+/*! \internal
+ * \brief Data used by the concrete temperature coupling implementations
+ */
+struct TemperatureCouplingData
+{
+    //! The coupling time step - simulation time step x nstcouple_
+    const double couplingTimeStep;
+    //! Coupling temperature per group
+    ArrayRef<const real> referenceTemperature;
+    //! Coupling time per group
+    ArrayRef<const real> couplingTime;
+    //! Number of degrees of freedom per group
+    ArrayRef<const real> numDegreesOfFreedom;
+    //! Work exerted by thermostat per group
+    ArrayRef<const double> temperatureCouplingIntegral;
+};
+
+/*! \internal
+ * \brief Interface for temperature coupling implementations
+ */
+class ITemperatureCouplingImpl
+{
+public:
+    //! Allow access to the scaling vectors
+    virtual void connectWithPropagator(const PropagatorThermostatConnection& connectionData,
+                                       int numTemperatureGroups) = 0;
+
+    /*! \brief Make a temperature control step
+     *
+     * \param step                     The current step
+     * \param temperatureGroup         The current temperature group
+     * \param currentKineticEnergy     The kinetic energy of the temperature group
+     * \param currentTemperature       The temperature of the temperature group
+     * \param temperatureCouplingData  Access to general temperature coupling data
+     *
+     * \return  The temperature coupling integral for the current temperature group
+     */
+    [[nodiscard]] virtual real apply(Step                           step,
+                                     int                            temperatureGroup,
+                                     real                           currentKineticEnergy,
+                                     real                           currentTemperature,
+                                     const TemperatureCouplingData& temperatureCouplingData) = 0;
+
+    //! Write private data to checkpoint
+    virtual void writeCheckpoint(WriteCheckpointData checkpointData, const t_commrec* cr) = 0;
+    //! Read private data from checkpoint
+    virtual void readCheckpoint(ReadCheckpointData checkpointData, const t_commrec* cr) = 0;
+
+    //! Standard virtual destructor
+    virtual ~ITemperatureCouplingImpl() = default;
+};
+
+/*! \internal
+ * \brief Implements v-rescale temperature coupling
+ */
+class VRescaleTemperatureCoupling final : public ITemperatureCouplingImpl
+{
+public:
+    //! Apply the v-rescale temperature control
+    real apply(Step step,
+               int  temperatureGroup,
+               real currentKineticEnergy,
+               real gmx_unused                currentTemperature,
+               const TemperatureCouplingData& temperatureCouplingData) override
+    {
+        if (!(temperatureCouplingData.couplingTime[temperatureGroup] >= 0
+              && temperatureCouplingData.numDegreesOfFreedom[temperatureGroup] > 0
+              && currentKineticEnergy > 0))
+        {
+            lambdaStartVelocities_[temperatureGroup] = 1.0;
+            return temperatureCouplingData.temperatureCouplingIntegral[temperatureGroup];
+        }
+
+        const real referenceKineticEnergy =
+                0.5 * temperatureCouplingData.referenceTemperature[temperatureGroup] * BOLTZ
+                * temperatureCouplingData.numDegreesOfFreedom[temperatureGroup];
+
+        const real newKineticEnergy =
+                vrescale_resamplekin(currentKineticEnergy, referenceKineticEnergy,
+                                     temperatureCouplingData.numDegreesOfFreedom[temperatureGroup],
+                                     temperatureCouplingData.couplingTime[temperatureGroup]
+                                             / temperatureCouplingData.couplingTimeStep,
+                                     step, seed_);
+
+        // Analytically newKineticEnergy >= 0, but we check for rounding errors
+        if (newKineticEnergy <= 0)
+        {
+            lambdaStartVelocities_[temperatureGroup] = 0.0;
+        }
+        else
+        {
+            lambdaStartVelocities_[temperatureGroup] = std::sqrt(newKineticEnergy / currentKineticEnergy);
+        }
+
+        if (debug)
+        {
+            fprintf(debug, "TC: group %d: Ekr %g, Ek %g, Ek_new %g, Lambda: %g\n", temperatureGroup,
+                    referenceKineticEnergy, currentKineticEnergy, newKineticEnergy,
+                    lambdaStartVelocities_[temperatureGroup]);
+        }
+
+        return temperatureCouplingData.temperatureCouplingIntegral[temperatureGroup]
+               - (newKineticEnergy - currentKineticEnergy);
+    }
+
+    //! Connect with propagator - v-rescale only scales start step velocities
+    void connectWithPropagator(const PropagatorThermostatConnection& connectionData,
+                               int                                   numTemperatureGroups) override
+    {
+        connectionData.setNumVelocityScalingVariables(numTemperatureGroups);
+        lambdaStartVelocities_ = connectionData.getViewOnVelocityScaling();
+    }
+
+    //! No data to write to checkpoint
+    void writeCheckpoint(WriteCheckpointData gmx_unused checkpointData, const t_commrec gmx_unused* cr) override
+    {
+    }
+    //! No data to read from checkpoints
+    void readCheckpoint(ReadCheckpointData gmx_unused checkpointData, const t_commrec gmx_unused* cr) override
+    {
+    }
+
+    //! Constructor
+    VRescaleTemperatureCoupling(int64_t seed) : seed_(seed) {}
+
+private:
+    //! The random seed
+    const int64_t seed_;
+
+    //! View on the scaling factor of the propagator (pre-step velocities)
+    ArrayRef<real> lambdaStartVelocities_;
+};
+
 VelocityScalingTemperatureCoupling::VelocityScalingTemperatureCoupling(
         int                               nstcouple,
         int                               offset,
@@ -74,37 +207,45 @@ VelocityScalingTemperatureCoupling::VelocityScalingTemperatureCoupling(
         const real*                       referenceTemperature,
         const real*                       couplingTime,
         const real*                       numDegreesOfFreedom,
-        EnergyData*                       energyData) :
+        EnergyData*                       energyData,
+        TemperatureCouplingType           couplingType) :
     nstcouple_(nstcouple),
     offset_(offset),
     useFullStepKE_(useFullStepKE),
     reportPreviousConservedEnergy_(reportPreviousConservedEnergy),
-    seed_(seed),
     numTemperatureGroups_(numTemperatureGroups),
     couplingTimeStep_(couplingTimeStep),
     referenceTemperature_(referenceTemperature, referenceTemperature + numTemperatureGroups),
     couplingTime_(couplingTime, couplingTime + numTemperatureGroups),
     numDegreesOfFreedom_(numDegreesOfFreedom, numDegreesOfFreedom + numTemperatureGroups),
-    thermostatIntegral_(numTemperatureGroups, 0.0),
+    temperatureCouplingIntegral_(numTemperatureGroups, 0.0),
     energyData_(energyData)
 {
     if (reportPreviousConservedEnergy_ == ReportPreviousStepConservedEnergy::Yes)
     {
-        thermostatIntegralPreviousStep_ = thermostatIntegral_;
+        temperatureCouplingIntegralPreviousStep_ = temperatureCouplingIntegral_;
     }
     energyData->setVelocityScalingTemperatureCoupling(this);
+    if (couplingType == etcVRESCALE)
+    {
+        temperatureCouplingImpl_ = std::make_unique<VRescaleTemperatureCoupling>(seed);
+    }
+    else
+    {
+        throw NotImplementedError("Temperature coupling " + std::string(ETCOUPLTYPE(couplingType))
+                                  + " is not implemented for modular simulator.");
+    }
 }
 
 void VelocityScalingTemperatureCoupling::connectWithPropagator(const PropagatorThermostatConnection& connectionData)
 {
-    connectionData.setNumVelocityScalingVariables(numTemperatureGroups_);
-    lambda_             = connectionData.getViewOnVelocityScaling();
+    temperatureCouplingImpl_->connectWithPropagator(connectionData, numTemperatureGroups_);
     propagatorCallback_ = connectionData.getVelocityScalingCallback();
 }
 
 void VelocityScalingTemperatureCoupling::elementSetup()
 {
-    if (!propagatorCallback_ || lambda_.empty())
+    if (!propagatorCallback_)
     {
         throw MissingElementConnectionError(
                 "Velocity scaling temperature coupling was not connected to a propagator.\n"
@@ -141,54 +282,24 @@ void VelocityScalingTemperatureCoupling::setLambda(Step step)
     // if we report the previous energy, calculate before the step
     if (reportPreviousConservedEnergy_ == ReportPreviousStepConservedEnergy::Yes)
     {
-        thermostatIntegralPreviousStep_ = thermostatIntegral_;
+        temperatureCouplingIntegralPreviousStep_ = temperatureCouplingIntegral_;
     }
 
-    real currentKineticEnergy, referenceKineticEnergy, newKineticEnergy;
+    const auto*             ekind          = energyData_->ekindata();
+    TemperatureCouplingData thermostatData = { couplingTimeStep_, referenceTemperature_, couplingTime_,
+                                               numDegreesOfFreedom_, temperatureCouplingIntegral_ };
 
-    const auto* ekind = energyData_->ekindata();
-
-    for (int i = 0; (i < numTemperatureGroups_); i++)
+    for (int temperatureGroup = 0; (temperatureGroup < numTemperatureGroups_); temperatureGroup++)
     {
-        if (useFullStepKE_ == UseFullStepKE::Yes)
-        {
-            currentKineticEnergy = trace(ekind->tcstat[i].ekinf);
-        }
-        else
-        {
-            currentKineticEnergy = trace(ekind->tcstat[i].ekinh);
-        }
+        const real currentKineticEnergy = useFullStepKE_ == UseFullStepKE::Yes
+                                                  ? trace(ekind->tcstat[temperatureGroup].ekinf)
+                                                  : trace(ekind->tcstat[temperatureGroup].ekinh);
+        const real currentTemperature = useFullStepKE_ == UseFullStepKE::Yes
+                                                ? ekind->tcstat[temperatureGroup].T
+                                                : ekind->tcstat[temperatureGroup].Th;
 
-        if (couplingTime_[i] >= 0 && numDegreesOfFreedom_[i] > 0 && currentKineticEnergy > 0)
-        {
-            referenceKineticEnergy = 0.5 * referenceTemperature_[i] * BOLTZ * numDegreesOfFreedom_[i];
-
-            newKineticEnergy = vrescale_resamplekin(currentKineticEnergy, referenceKineticEnergy,
-                                                    numDegreesOfFreedom_[i],
-                                                    couplingTime_[i] / couplingTimeStep_, step, seed_);
-
-            // Analytically newKineticEnergy >= 0, but we check for rounding errors
-            if (newKineticEnergy <= 0)
-            {
-                lambda_[i] = 0.0;
-            }
-            else
-            {
-                lambda_[i] = std::sqrt(newKineticEnergy / currentKineticEnergy);
-            }
-
-            thermostatIntegral_[i] -= newKineticEnergy - currentKineticEnergy;
-
-            if (debug)
-            {
-                fprintf(debug, "TC: group %d: Ekr %g, Ek %g, Ek_new %g, Lambda: %g\n", i,
-                        referenceKineticEnergy, currentKineticEnergy, newKineticEnergy, lambda_[i]);
-            }
-        }
-        else
-        {
-            lambda_[i] = 1.0;
-        }
+        temperatureCouplingIntegral_[temperatureGroup] = temperatureCouplingImpl_->apply(
+                step, temperatureGroup, currentKineticEnergy, currentTemperature, thermostatData);
     }
 }
 
@@ -217,11 +328,12 @@ void VelocityScalingTemperatureCoupling::doCheckpointData(CheckpointData<operati
         checkpointVersion(checkpointData, "VRescaleThermostat version", c_currentVersion);
 
         checkpointData->arrayRef("thermostat integral",
-                                 makeCheckpointArrayRef<operation>(thermostatIntegral_));
+                                 makeCheckpointArrayRef<operation>(temperatureCouplingIntegral_));
     }
     if (operation == CheckpointDataOperation::Read && DOMAINDECOMP(cr))
     {
-        dd_bcast(cr->dd, thermostatIntegral_.size() * sizeof(double), thermostatIntegral_.data());
+        dd_bcast(cr->dd, temperatureCouplingIntegral_.size() * sizeof(double),
+                 temperatureCouplingIntegral_.data());
     }
 }
 
@@ -229,12 +341,14 @@ void VelocityScalingTemperatureCoupling::writeCheckpoint(WriteCheckpointData che
                                                          const t_commrec*    cr)
 {
     doCheckpointData<CheckpointDataOperation::Write>(&checkpointData, cr);
+    temperatureCouplingImpl_->writeCheckpoint(checkpointData.subCheckpointData("thermostat impl"), cr);
 }
 
 void VelocityScalingTemperatureCoupling::readCheckpoint(ReadCheckpointData checkpointData,
                                                         const t_commrec*   cr)
 {
     doCheckpointData<CheckpointDataOperation::Read>(&checkpointData, cr);
+    temperatureCouplingImpl_->readCheckpoint(checkpointData.subCheckpointData("thermostat impl"), cr);
 }
 
 const std::string& VelocityScalingTemperatureCoupling::clientID()
@@ -245,9 +359,10 @@ const std::string& VelocityScalingTemperatureCoupling::clientID()
 real VelocityScalingTemperatureCoupling::conservedEnergyContribution() const
 {
     return (reportPreviousConservedEnergy_ == ReportPreviousStepConservedEnergy::Yes)
-                   ? std::accumulate(thermostatIntegralPreviousStep_.begin(),
-                                     thermostatIntegralPreviousStep_.end(), 0.0)
-                   : std::accumulate(thermostatIntegral_.begin(), thermostatIntegral_.end(), 0.0);
+                   ? std::accumulate(temperatureCouplingIntegralPreviousStep_.begin(),
+                                     temperatureCouplingIntegralPreviousStep_.end(), 0.0)
+                   : std::accumulate(temperatureCouplingIntegral_.begin(),
+                                     temperatureCouplingIntegral_.end(), 0.0);
 }
 
 ISimulatorElement* VelocityScalingTemperatureCoupling::getElementPointerImpl(
@@ -266,7 +381,7 @@ ISimulatorElement* VelocityScalingTemperatureCoupling::getElementPointerImpl(
             legacySimulatorData->inputrec->ld_seed, legacySimulatorData->inputrec->opts.ngtc,
             legacySimulatorData->inputrec->delta_t * legacySimulatorData->inputrec->nsttcouple,
             legacySimulatorData->inputrec->opts.ref_t, legacySimulatorData->inputrec->opts.tau_t,
-            legacySimulatorData->inputrec->opts.nrdf, energyData));
+            legacySimulatorData->inputrec->opts.nrdf, energyData, legacySimulatorData->inputrec->etc));
     auto* thermostat = static_cast<VelocityScalingTemperatureCoupling*>(element);
     builderHelper->registerThermostat([thermostat](const PropagatorThermostatConnection& connection) {
         thermostat->connectWithPropagator(connection);
index 388bd4c4a913369686fcf7533879cf24024ce26c..e68ab6411a14c683a1f4d5a123318cb2d2059340 100644 (file)
@@ -55,7 +55,9 @@ struct t_commrec;
 
 namespace gmx
 {
+class ITemperatureCouplingImpl;
 class LegacySimulatorData;
+struct TemperatureCouplingData;
 
 //! Enum describing whether the thermostat is using full or half step kinetic energy
 enum class UseFullStepKE
@@ -73,12 +75,18 @@ enum class ReportPreviousStepConservedEnergy
     Count
 };
 
+using TemperatureCouplingType = int;
+
 /*! \internal
  * \ingroup module_modularsimulator
- * \brief Element implementing the v-rescale thermostat
+ * \brief Element implementing the a velocity-scaling thermostat
  *
  * This element takes a callback to the propagator and updates the velocity
- * scaling factor according to the v-rescale thermostat.
+ * scaling factor according to the internal temperature coupling implementation.
+ *
+ * Note that the concrete implementation is handled by the concrete
+ * implementations of the ITemperatureCouplingImpl interface, while the element
+ * handles the scheduling and interfacing with other elements.
  */
 class VelocityScalingTemperatureCoupling final : public ISimulatorElement, public ICheckpointHelperClient
 {
@@ -94,7 +102,8 @@ public:
                                        const real*                       referenceTemperature,
                                        const real*                       couplingTime,
                                        const real*                       numDegreesOfFreedom,
-                                       EnergyData*                       energyData);
+                                       EnergyData*                       energyData,
+                                       TemperatureCouplingType           couplingType);
 
     /*! \brief Register run function for step / time
      *
@@ -156,8 +165,6 @@ private:
     const UseFullStepKE useFullStepKE_;
     //! Whether we are reporting the conserved energy from the previous step
     const ReportPreviousStepConservedEnergy reportPreviousConservedEnergy_;
-    //! The random seed
-    const int64_t seed_;
 
     //! The number of temperature groups
     const int numTemperatureGroups_;
@@ -170,22 +177,23 @@ private:
     //! Number of degrees of freedom per group
     const std::vector<real> numDegreesOfFreedom_;
     //! Work exerted by thermostat per group
-    std::vector<double> thermostatIntegral_;
+    std::vector<double> temperatureCouplingIntegral_;
     //! Work exerted by thermostat per group (backup from previous step)
-    std::vector<double> thermostatIntegralPreviousStep_;
+    std::vector<double> temperatureCouplingIntegralPreviousStep_;
 
     // TODO: Clarify relationship to data objects and find a more robust alternative to raw pointers (#3583)
     //! Pointer to the energy data (for ekindata)
     EnergyData* energyData_;
 
-    //! View on the scaling factor of the propagator
-    ArrayRef<real> lambda_;
     //! Callback to let propagator know that we updated lambda
     PropagatorCallback propagatorCallback_;
 
     //! Set new lambda value (at T-coupling steps)
     void setLambda(Step step);
 
+    //! The temperature coupling implementation
+    std::unique_ptr<ITemperatureCouplingImpl> temperatureCouplingImpl_;
+
     //! CheckpointHelper identifier
     const std::string identifier_ = "VelocityScalingTemperatureCoupling";
     //! Helper function to read from / write to CheckpointData