Add lambda state to dhdl.xvg with AWH fep

[alexxy/gromacs.git] / src / gromacs / modularsimulator / energydata.cpp

diff --git a/src/gromacs/modularsimulator/energydata.cpp b/src/gromacs/modularsimulator/energydata.cpp
index 32e3b1892226ed36fd285be1eb1b74261ad95291..8999abfc39920ba6f6549697776a673732d31157 100644 (file)
--- a/src/gromacs/modularsimulator/energydata.cpp
+++ b/src/gromacs/modularsimulator/energydata.cpp
@@ -51,10 +51,12 @@
 #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/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"
@@ -83,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),
@@ -107,10 +110,11 @@ EnergyData::EnergyData(StatePropagatorData*        statePropagatorData,
     constr_(constr),
     fplog_(fplog),
     fcd_(fcd),
-    mdModulesNotifier_(mdModulesNotifier),
+    mdModulesNotifiers_(mdModulesNotifiers),
     groups_(&globalTopology.groups),
     observablesHistory_(observablesHistory),
-    simulationsShareState_(simulationsShareState)
+    simulationsShareState_(simulationsShareState),
+    pullWork_(pullWork)
 {
     clear_mat(forceVirial_);
     clear_mat(shakeVirial_);
@@ -159,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,
+                                                   pullWork_,
                                                    mdoutf_get_fp_dhdl(outf),
                                                    false,
                                                    startingBehavior_,
                                                    simulationsShareState_,
-                                                   mdModulesNotifier_);
+                                                   mdModulesNotifiers_);
 
     if (!isMasterRank_)
     {
@@ -256,7 +259,6 @@ void EnergyData::doStep(Step step, Time time, bool isEnergyCalculationStep, bool
             mdAtoms_->mdatoms()->tmass,
             enerd_,
             inputrec_->fepvals.get(),
-            inputrec_->expandedvals.get(),
             statePropagatorData_->constPreviousBox(),
             PTCouplingArrays({ parrinelloRahmanBoxVelocities_ ? parrinelloRahmanBoxVelocities_() : nullMatrix,
                                {},
@@ -264,8 +266,6 @@ void EnergyData::doStep(Step step, Time time, bool isEnergyCalculationStep, bool
                                {},
                                {} }),
             freeEnergyPerturbationData_ ? freeEnergyPerturbationData_->currentFEPState() : 0,
-            shakeVirial_,
-            forceVirial_,
             totalVirial_,
             pressure_,
             ekind_,
@@ -366,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_;
@@ -509,6 +514,24 @@ void EnergyData::setParrinelloRahmanBoxVelocities(std::function<const rvec*()>&&
     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()
 {
     return element_.get();
@@ -529,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();
 }