Add lambda state to dhdl.xvg with AWH fep
[alexxy/gromacs.git] / src / gromacs / modularsimulator / energydata.cpp
index 3cda75f3cea59a49e5a6d861b9d8558285072894..8999abfc39920ba6f6549697776a673732d31157 100644 (file)
@@ -1,7 +1,7 @@
 /*
  * This file is part of the GROMACS molecular simulation package.
  *
- * Copyright (c) 2019,2020, by the GROMACS development team, led by
+ * Copyright (c) 2019,2020,2021, by the GROMACS development team, led by
  * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
  * and including many others, as listed in the AUTHORS file in the
  * top-level source directory and at http://www.gromacs.org.
 #include "gromacs/gmxlib/network.h"
 #include "gromacs/math/vec.h"
 #include "gromacs/mdlib/compute_io.h"
-#include "gromacs/mdlib/coupling.h"
 #include "gromacs/mdlib/enerdata_utils.h"
 #include "gromacs/mdlib/energyoutput.h"
 #include "gromacs/mdlib/mdatoms.h"
 #include "gromacs/mdlib/mdoutf.h"
 #include "gromacs/mdlib/stat.h"
+#include "gromacs/mdlib/tgroup.h"
 #include "gromacs/mdlib/update.h"
-#include "gromacs/mdrunutility/handlerestart.h"
-#include "gromacs/mdtypes/checkpointdata.h"
 #include "gromacs/mdtypes/commrec.h"
 #include "gromacs/mdtypes/enerdata.h"
 #include "gromacs/mdtypes/energyhistory.h"
+#include "gromacs/mdtypes/group.h"
 #include "gromacs/mdtypes/inputrec.h"
 #include "gromacs/mdtypes/mdatom.h"
 #include "gromacs/mdtypes/observableshistory.h"
 
 #include "freeenergyperturbationdata.h"
 #include "modularsimulator.h"
-#include "parrinellorahmanbarostat.h"
 #include "simulatoralgorithm.h"
 #include "statepropagatordata.h"
-#include "velocityscalingtemperaturecoupling.h"
 
 struct pull_t;
 class t_state;
@@ -80,7 +77,7 @@ class Awh;
 
 EnergyData::EnergyData(StatePropagatorData*        statePropagatorData,
                        FreeEnergyPerturbationData* freeEnergyPerturbationData,
-                       const gmx_mtop_t*           globalTopology,
+                       const gmx_mtop_t&           globalTopology,
                        const t_inputrec*           inputrec,
                        const MDAtoms*              mdAtoms,
                        gmx_enerdata_t*             enerd,
@@ -88,11 +85,12 @@ EnergyData::EnergyData(StatePropagatorData*        statePropagatorData,
                        const Constraints*          constr,
                        FILE*                       fplog,
                        t_fcdata*                   fcd,
-                       const MdModulesNotifier&    mdModulesNotifier,
+                       const MDModulesNotifiers&   mdModulesNotifiers,
                        bool                        isMasterRank,
                        ObservablesHistory*         observablesHistory,
                        StartingBehavior            startingBehavior,
-                       bool                        simulationsShareState) :
+                       bool                        simulationsShareState,
+                       pull_t*                     pullWork) :
     element_(std::make_unique<Element>(this, isMasterRank)),
     isMasterRank_(isMasterRank),
     forceVirialStep_(-1),
@@ -104,8 +102,6 @@ EnergyData::EnergyData(StatePropagatorData*        statePropagatorData,
     startingBehavior_(startingBehavior),
     statePropagatorData_(statePropagatorData),
     freeEnergyPerturbationData_(freeEnergyPerturbationData),
-    velocityScalingTemperatureCoupling_(nullptr),
-    parrinelloRahmanBarostat_(nullptr),
     inputrec_(inputrec),
     top_global_(globalTopology),
     mdAtoms_(mdAtoms),
@@ -114,10 +110,11 @@ EnergyData::EnergyData(StatePropagatorData*        statePropagatorData,
     constr_(constr),
     fplog_(fplog),
     fcd_(fcd),
-    mdModulesNotifier_(mdModulesNotifier),
-    groups_(&globalTopology->groups),
+    mdModulesNotifiers_(mdModulesNotifiers),
+    groups_(&globalTopology.groups),
     observablesHistory_(observablesHistory),
-    simulationsShareState_(simulationsShareState)
+    simulationsShareState_(simulationsShareState),
+    pullWork_(pullWork)
 {
     clear_mat(forceVirial_);
     clear_mat(shakeVirial_);
@@ -140,8 +137,8 @@ void EnergyData::Element::scheduleTask(Step step, Time time, const RegisterRunFu
     auto isFreeEnergyCalculationStep = freeEnergyCalculationStep_ == step;
     if (isEnergyCalculationStep || writeEnergy)
     {
-        registerRunFunction([this, time, isEnergyCalculationStep, isFreeEnergyCalculationStep]() {
-            energyData_->doStep(time, isEnergyCalculationStep, isFreeEnergyCalculationStep);
+        registerRunFunction([this, step, time, isEnergyCalculationStep, isFreeEnergyCalculationStep]() {
+            energyData_->doStep(step, time, isEnergyCalculationStep, isFreeEnergyCalculationStep);
         });
     }
     else
@@ -166,16 +163,15 @@ void EnergyData::Element::trajectoryWriterSetup(gmx_mdoutf* outf)
 
 void EnergyData::setup(gmx_mdoutf* outf)
 {
-    pull_t* pull_work = nullptr;
-    energyOutput_     = std::make_unique<EnergyOutput>(mdoutf_get_fp_ene(outf),
+    energyOutput_ = std::make_unique<EnergyOutput>(mdoutf_get_fp_ene(outf),
                                                    top_global_,
-                                                   inputrec_,
-                                                   pull_work,
+                                                   *inputrec_,
+                                                   pullWork_,
                                                    mdoutf_get_fp_dhdl(outf),
                                                    false,
                                                    startingBehavior_,
                                                    simulationsShareState_,
-                                                   mdModulesNotifier_);
+                                                   mdModulesNotifiers_);
 
     if (!isMasterRank_)
     {
@@ -187,7 +183,7 @@ void EnergyData::setup(gmx_mdoutf* outf)
     // TODO: This probably doesn't really belong here...
     //       but we have all we need in this element,
     //       so we'll leave it here for now!
-    double io = compute_io(inputrec_, top_global_->natoms, *groups_, energyOutput_->numEnergyTerms(), 1);
+    double io = compute_io(inputrec_, top_global_.natoms, *groups_, energyOutput_->numEnergyTerms(), 1);
     if ((io > 2000) && isMasterRank_)
     {
         fprintf(stderr, "\nWARNING: This run will generate roughly %.0f Mb of data\n\n", io);
@@ -195,7 +191,7 @@ void EnergyData::setup(gmx_mdoutf* outf)
     if (!inputrec_->bContinuation)
     {
         real temp = enerd_->term[F_TEMP];
-        if (inputrec_->eI != eiVV)
+        if (inputrec_->eI != IntegrationAlgorithm::VV)
         {
             /* Result of Ekin averaged over velocities of -half
              * and +half step, while we only have -half step here.
@@ -239,7 +235,7 @@ std::optional<SignallerCallback> EnergyData::Element::registerEnergyCallback(Ene
     return std::nullopt;
 }
 
-void EnergyData::doStep(Time time, bool isEnergyCalculationStep, bool isFreeEnergyCalculationStep)
+void EnergyData::doStep(Step step, Time time, bool isEnergyCalculationStep, bool isFreeEnergyCalculationStep)
 {
     enerd_->term[F_ETOT] = enerd_->term[F_EPOT] + enerd_->term[F_EKIN];
     if (freeEnergyPerturbationData_)
@@ -249,12 +245,11 @@ void EnergyData::doStep(Time time, bool isEnergyCalculationStep, bool isFreeEner
     }
     if (integratorHasConservedEnergyQuantity(inputrec_))
     {
-        enerd_->term[F_ECONSERVED] =
-                enerd_->term[F_ETOT]
-                + (velocityScalingTemperatureCoupling_
-                           ? velocityScalingTemperatureCoupling_->conservedEnergyContribution()
-                           : 0)
-                + (parrinelloRahmanBarostat_ ? parrinelloRahmanBarostat_->conservedEnergyContribution() : 0);
+        enerd_->term[F_ECONSERVED] = enerd_->term[F_ETOT];
+        for (const auto& energyContibution : conservedEnergyContributions_)
+        {
+            enerd_->term[F_ECONSERVED] += energyContibution(step, time);
+        }
     }
     matrix nullMatrix = {};
     energyOutput_->addDataAtEnergyStep(
@@ -263,17 +258,14 @@ void EnergyData::doStep(Time time, bool isEnergyCalculationStep, bool isFreeEner
             time,
             mdAtoms_->mdatoms()->tmass,
             enerd_,
-            inputrec_->fepvals,
-            inputrec_->expandedvals,
+            inputrec_->fepvals.get(),
             statePropagatorData_->constPreviousBox(),
-            PTCouplingArrays({ parrinelloRahmanBarostat_ ? parrinelloRahmanBarostat_->boxVelocities() : nullMatrix,
+            PTCouplingArrays({ parrinelloRahmanBoxVelocities_ ? parrinelloRahmanBoxVelocities_() : nullMatrix,
                                {},
                                {},
                                {},
                                {} }),
             freeEnergyPerturbationData_ ? freeEnergyPerturbationData_->currentFEPState() : 0,
-            shakeVirial_,
-            forceVirial_,
             totalVirial_,
             pressure_,
             ekind_,
@@ -374,6 +366,11 @@ gmx_enerdata_t* EnergyData::enerdata()
     return enerd_;
 }
 
+const gmx_enerdata_t* EnergyData::enerdata() const
+{
+    return enerd_;
+}
+
 gmx_ekindata_t* EnergyData::ekindata()
 {
     return ekind_;
@@ -505,14 +502,34 @@ void EnergyData::initializeEnergyHistory(StartingBehavior    startingBehavior,
     energyOutput->fillEnergyHistory(observablesHistory->energyHistory.get());
 }
 
-void EnergyData::setVelocityScalingTemperatureCoupling(const VelocityScalingTemperatureCoupling* velocityScalingTemperatureCoupling)
+void EnergyData::addConservedEnergyContribution(EnergyContribution&& energyContribution)
 {
-    velocityScalingTemperatureCoupling_ = velocityScalingTemperatureCoupling;
+    conservedEnergyContributions_.emplace_back(std::move(energyContribution));
 }
 
-void EnergyData::setParrinelloRahamnBarostat(const gmx::ParrinelloRahmanBarostat* parrinelloRahmanBarostat)
+void EnergyData::setParrinelloRahmanBoxVelocities(std::function<const rvec*()>&& parrinelloRahmanBoxVelocities)
 {
-    parrinelloRahmanBarostat_ = parrinelloRahmanBarostat;
+    GMX_RELEASE_ASSERT(!parrinelloRahmanBoxVelocities_,
+                       "Received a second callback to the Parrinello-Rahman velocities");
+    parrinelloRahmanBoxVelocities_ = parrinelloRahmanBoxVelocities;
+}
+
+void EnergyData::updateKineticEnergy()
+{
+    // The legacy sum_ekin function does not offer named types, so define variables for readability
+    // dEkin/dlambda is not handled here
+    real* dEkinDLambda = nullptr;
+    // Whether we use the full step kinetic energy (vs the average of half step KEs)
+    const bool useFullStepKineticEnergy = (inputrec_->eI == IntegrationAlgorithm::VV);
+    /* Whether we're ignoring the NHC scaling factor, only used if useFullStepKineticEnergy
+     * is true. (This parameter is confusing, as it is named `bScaleEkin`, but prompts the
+     * function to ignore scaling. There is no use case within modular simulator to ignore
+     * these, so we set this to false.) */
+    const bool ignoreScalingFactor = false;
+
+    enerd_->term[F_TEMP] = sum_ekin(
+            &(inputrec_->opts), ekind_, dEkinDLambda, useFullStepKineticEnergy, ignoreScalingFactor);
+    enerd_->term[F_EKIN] = trace(ekind_->ekin);
 }
 
 EnergyData::Element* EnergyData::element()
@@ -535,7 +552,8 @@ ISimulatorElement* EnergyData::Element::getElementPointerImpl(
         StatePropagatorData gmx_unused* statePropagatorData,
         EnergyData*                     energyData,
         FreeEnergyPerturbationData gmx_unused* freeEnergyPerturbationData,
-        GlobalCommunicationHelper gmx_unused* globalCommunicationHelper)
+        GlobalCommunicationHelper gmx_unused* globalCommunicationHelper,
+        ObservablesReducer* /*observablesReducer*/)
 {
     return energyData->element();
 }