Introduce expanded ensemble in modular simulator

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

diff --git a/src/gromacs/modularsimulator/modularsimulator.cpp b/src/gromacs/modularsimulator/modularsimulator.cpp
index 4e91aa95d633aaf613b979739f3afc8df147f1ec..d3fce2c0b5563bf8bce6df5c8821aaf171f55e58 100644 (file)
--- a/src/gromacs/modularsimulator/modularsimulator.cpp
+++ b/src/gromacs/modularsimulator/modularsimulator.cpp
@@ -78,6 +78,7 @@
 #include "andersentemperaturecoupling.h"
 #include "computeglobalselement.h"
 #include "constraintelement.h"
+#include "expandedensembleelement.h"
 #include "firstorderpressurecoupling.h"
 #include "forceelement.h"
 #include "mttk.h"
@@ -154,7 +155,12 @@ void ModularSimulator::addIntegrationElements(ModularSimulatorAlgorithmBuilder*
             builder->add<ConstraintsElement<ConstraintVariable::Velocities>>();
         }
         builder->add<ComputeGlobalsElement<ComputeGlobalsAlgorithm::VelocityVerlet>>();
+        // Here, we have x / v / f at the full time step
         builder->add<StatePropagatorData::Element>();
+        if (legacySimulatorData_->inputrec->bExpanded)
+        {
+            builder->add<ExpandedEnsembleElement>();
+        }
         if (legacySimulatorData_->inputrec->etc == TemperatureCoupling::VRescale
             || legacySimulatorData_->inputrec->etc == TemperatureCoupling::Berendsen)
         {
@@ -246,6 +252,10 @@ void ModularSimulator::addIntegrationElements(ModularSimulatorAlgorithmBuilder*
         }
         // We have a full state at time t here
         builder->add<StatePropagatorData::Element>();
+        if (legacySimulatorData_->inputrec->bExpanded)
+        {
+            builder->add<ExpandedEnsembleElement>();
+        }
 
         // Propagate extended system variables from t to t+dt/2
         if (legacySimulatorData_->inputrec->epc == PressureCoupling::Mttk)
@@ -370,13 +380,6 @@ bool ModularSimulator::isInputCompatible(bool                             exitOn
     isInputCompatible =
             isInputCompatible
             && conditionalAssert(!doRerun, "Rerun is not supported by the modular simulator.");
-    isInputCompatible =
-            isInputCompatible
-            && conditionalAssert(inputrec->efep == FreeEnergyPerturbationType::No
-                                         || inputrec->efep == FreeEnergyPerturbationType::Yes
-                                         || inputrec->efep == FreeEnergyPerturbationType::SlowGrowth,
-                                 "Expanded ensemble free energy calculation is not "
-                                 "supported by the modular simulator.");
     isInputCompatible = isInputCompatible
                         && conditionalAssert(!inputrec->bPull,
                                              "Pulling is not supported by the modular simulator.");
@@ -447,10 +450,6 @@ bool ModularSimulator::isInputCompatible(bool                             exitOn
             isInputCompatible
             && conditionalAssert(!inputrec->bSimTemp,
                                  "Simulated tempering is not supported by the modular simulator.");
-    isInputCompatible = isInputCompatible
-                        && conditionalAssert(!inputrec->bExpanded,
-                                             "Expanded ensemble simulations are not supported by "
-                                             "the modular simulator.");
     isInputCompatible =
             isInputCompatible
             && conditionalAssert(!doEssentialDynamics,