Allow elements to schedule pre-step and post-step functions

Pre-step and post-step functions allow elements to perform operations
which need to happen before any of the integration step functions run,
or after all integration step functions have run. Unlike the main
integration step scheduling, whose order is defined by the order of the
elements in the integration algorithm, the functions are not guaranteed
to run in a specific order, only before / after the integration step.

One example this is useful for is the pull code, which needs to print
some output once all external pull providers (currently only AWH) have
run. Having a post-step function allows for an easy way for the pull
element to run during the integration run, but defer printout until
after the step.

Refs #3417, #3435, #3436

diff --git a/src/gromacs/modularsimulator/modularsimulatorinterfaces.h b/src/gromacs/modularsimulator/modularsimulatorinterfaces.h
index 34eb13889610c6ab62a19ea55d13fa88781130b1..5358a108a4e6f2bf6e48f343fa2506d4dd561143 100644 (file)
--- a/src/gromacs/modularsimulator/modularsimulatorinterfaces.h
+++ b/src/gromacs/modularsimulator/modularsimulatorinterfaces.h
@@ -97,6 +97,8 @@ typedef std::function<void()> SimulatorRunFunction;
 
 //! The function type that allows to register run functions
 typedef std::function<void(SimulatorRunFunction)> RegisterRunFunction;
+//! The function type scheduling run functions for a step / time using a RegisterRunFunction reference
+typedef std::function<void(Step, Time, const RegisterRunFunction&)> SchedulingFunction;
 
 /*! \internal
  * \brief The general interface for elements of the modular simulator

diff --git a/src/gromacs/modularsimulator/simulatoralgorithm.cpp b/src/gromacs/modularsimulator/simulatoralgorithm.cpp
index 54667c5d86896d93fdee44e1352c59b89ce5bd23..0b9fb2a1fb0cc768347d7b52a9297413c4f7b8c6 100644 (file)
--- a/src/gromacs/modularsimulator/simulatoralgorithm.cpp
+++ b/src/gromacs/modularsimulator/simulatoralgorithm.cpp
@@ -364,11 +364,21 @@ void ModularSimulatorAlgorithm::populateTaskQueue()
 
         // register pre-step (task queue is local, so no problem with `this`)
         registerRunFunction([this, step, time, isNSStep]() { preStep(step, time, isNSStep); });
+        // register pre step functions
+        for (const auto& schedulingFunction : preStepScheduling_)
+        {
+            schedulingFunction(step_, time, registerRunFunction);
+        }
         // register elements for step
         for (auto& element : elementCallList_)
         {
             element->scheduleTask(step_, time, registerRunFunction);
         }
+        // register post step functions
+        for (const auto& schedulingFunction : postStepScheduling_)
+        {
+            schedulingFunction(step_, time, registerRunFunction);
+        }
         // register post-step (task queue is local, so no problem with `this`)
         registerRunFunction([this, step, time]() { postStep(step, time); });
 
@@ -659,6 +669,9 @@ ModularSimulatorAlgorithm ModularSimulatorAlgorithmBuilder::build()
 
     // Move setup / teardown list
     algorithm.elementSetupTeardownList_ = std::move(setupAndTeardownList_);
+    // Move pre- / post-step scheduling lists
+    algorithm.preStepScheduling_  = std::move(preStepScheduling_);
+    algorithm.postStepScheduling_ = std::move(postStepScheduling_);
 
     // Create element list
     // Checkpoint helper needs to be in the call list (as first element!) to react to last step
@@ -736,6 +749,16 @@ bool ModularSimulatorAlgorithmBuilderHelper::elementIsStored(const ISimulatorEle
     return builder_->elementExists(element);
 }
 
+[[maybe_unused]] void ModularSimulatorAlgorithmBuilderHelper::registerPreStepScheduling(SchedulingFunction schedulingFunction)
+{
+    builder_->preStepScheduling_.emplace_back(std::move(schedulingFunction));
+}
+
+[[maybe_unused]] void ModularSimulatorAlgorithmBuilderHelper::registerPostStepScheduling(SchedulingFunction schedulingFunction)
+{
+    builder_->postStepScheduling_.emplace_back(std::move(schedulingFunction));
+}
+
 std::optional<std::any> ModularSimulatorAlgorithmBuilderHelper::builderData(const std::string& key) const
 {
     const auto iter = builder_->builderData_.find(key);

diff --git a/src/gromacs/modularsimulator/simulatoralgorithm.h b/src/gromacs/modularsimulator/simulatoralgorithm.h
index dbf469172bfd07c26edc8b5fe0f6e821888ced0e..7eea8ee3b50e1c68e55fe3e1923d8afbba861fb5 100644 (file)
--- a/src/gromacs/modularsimulator/simulatoralgorithm.h
+++ b/src/gromacs/modularsimulator/simulatoralgorithm.h
@@ -193,6 +193,10 @@ private:
     std::vector<ISimulatorElement*> elementCallList_;
     //! List of schedulerElements (setup / teardown calling sequence)
     std::vector<ISimulatorElement*> elementSetupTeardownList_;
+    //! List of pre-step scheduling functions
+    std::vector<SchedulingFunction> preStepScheduling_;
+    //! List of post-step scheduling functions
+    std::vector<SchedulingFunction> postStepScheduling_;
 
     // Infrastructure elements
     //! The domain decomposition element
@@ -324,6 +328,22 @@ public:
     Element* storeElement(std::unique_ptr<Element> element);
     //! Check if an element is stored in the ModularSimulatorAlgorithmBuilder
     bool elementIsStored(const ISimulatorElement* element) const;
+    /*! \brief Register callback to schedule a pre-step run
+     *
+     * This allows elements to schedule a function call before the integration step.
+     * The function call is guaranteed to happen before any functions scheduled for
+     * the integration step. It is not guaranteed to happen in any specific order
+     * compared to other elements registering a pre-step scheduling function.
+     */
+    void registerPreStepScheduling(SchedulingFunction schedulingFunction);
+    /*! \brief Register callback to schedule a post-step run
+     *
+     * This allows elements to schedule a function call after the integration step.
+     * The function call is guaranteed to happen after all functions scheduled for
+     * the integration step. It is not guaranteed to happen in any specific order
+     * compared to other elements registering a post-step scheduling function.
+     */
+    void registerPostStepScheduling(SchedulingFunction schedulingFunction);
     /*! \brief Set arbitrary data in the ModularSimulatorAlgorithmBuilder
      *
      * Allows to store arbitrary data with lifetime equal to the builder. Functionality is used
@@ -491,6 +511,10 @@ private:
      * Elements should only appear once in this list
      */
     std::vector<ISimulatorElement*> setupAndTeardownList_;
+    //! List of pre-step scheduling functions
+    std::vector<SchedulingFunction> preStepScheduling_;
+    //! List of post-step scheduling functions
+    std::vector<SchedulingFunction> postStepScheduling_;
 
     //! Builder for the NeighborSearchSignaller
     SignallerBuilder<NeighborSearchSignaller> neighborSearchSignallerBuilder_;