Fix final configuration writing for modular simulator

Modular simulator was not writing final configuration (.gro) file.
This commit fixes this, and removes an outdated sentence in the
documentation of ITrajectoryWriterClient.
I02299d34 introduces a test to avoid similar errors in the future.

Note that the complexity required to introduce this feature implies that
trajectory writing in the modular simulator needs some work for
GROMACS 2021.

Change-Id: Ie7f03f5beea2d78411c7a0cdfa85f71deeb622fa

diff --git a/src/gromacs/mdlib/trajectory_writing.cpp b/src/gromacs/mdlib/trajectory_writing.cpp
index dc2cc747fc4424b2675122a1b8405813a8605599..e06b416a0a1ad0837fab92262a5e8a66619b8a0c 100644 (file)
--- a/src/gromacs/mdlib/trajectory_writing.cpp
+++ b/src/gromacs/mdlib/trajectory_writing.cpp
@@ -159,6 +159,9 @@ void do_md_trajectory_writing(FILE*                    fplog,
                 energyOutput.fillEnergyHistory(observablesHistory->energyHistory.get());
             }
         }
+        // Note that part of the following code is duplicated in StatePropagatorData::trajectoryWriterTeardown.
+        // This duplication is needed while both legacy and modular code paths are in use.
+        // TODO: Remove duplication asap, make sure to keep in sync in the meantime.
         mdoutf_write_to_trajectory_files(fplog, cr, outf, mdof_flags, top_global->natoms, step, t,
                                          state, state_global, observablesHistory, f);
         if (bLastStep && step_rel == ir->nsteps && bDoConfOut && MASTER(cr) && !bRerunMD)

diff --git a/src/gromacs/modularsimulator/modularsimulator.cpp b/src/gromacs/modularsimulator/modularsimulator.cpp
index 72c6dfa76b8cf546f2ab6e7ad99e8f9c12d8dd78..c785f5d3dddf88ab0be79ab9679b576bdc143712 100644 (file)
--- a/src/gromacs/modularsimulator/modularsimulator.cpp
+++ b/src/gromacs/modularsimulator/modularsimulator.cpp
@@ -355,6 +355,9 @@ void ModularSimulator::constructElementsAndSignallers()
     /*
      * Build data structures
      */
+    topologyHolder_ =
+            std::make_unique<TopologyHolder>(*top_global, cr, inputrec, fr, mdAtoms, constr, vsite);
+
     std::unique_ptr<FreeEnergyPerturbationElement> freeEnergyPerturbationElement    = nullptr;
     FreeEnergyPerturbationElement*                 freeEnergyPerturbationElementPtr = nullptr;
     if (inputrec->efep != efepNO)
@@ -367,7 +370,8 @@ void ModularSimulator::constructElementsAndSignallers()
     auto statePropagatorData = std::make_unique<StatePropagatorData>(
             top_global->natoms, fplog, cr, state_global, inputrec->nstxout, inputrec->nstvout,
             inputrec->nstfout, inputrec->nstxout_compressed, fr->nbv->useGpu(),
-            freeEnergyPerturbationElementPtr, inputrec, mdAtoms->mdatoms());
+            freeEnergyPerturbationElementPtr, topologyHolder_.get(), fr->bMolPBC,
+            mdrunOptions.writeConfout, opt2fn("-c", nfile, fnm), inputrec, mdAtoms->mdatoms());
     auto statePropagatorDataPtr = compat::make_not_null(statePropagatorData.get());
 
     auto energyElement = std::make_unique<EnergyElement>(
@@ -376,9 +380,6 @@ void ModularSimulator::constructElementsAndSignallers()
             startingBehavior);
     auto energyElementPtr = compat::make_not_null(energyElement.get());
 
-    topologyHolder_ =
-            std::make_unique<TopologyHolder>(*top_global, cr, inputrec, fr, mdAtoms, constr, vsite);
-
     /*
      * Build stop handler
      */
@@ -403,6 +404,7 @@ void ModularSimulator::constructElementsAndSignallers()
      */
     trajectoryElementBuilder.registerWriterClient(statePropagatorDataPtr);
     trajectoryElementBuilder.registerSignallerClient(statePropagatorDataPtr);
+    lastStepSignallerBuilder.registerSignallerClient(statePropagatorDataPtr);
 
     trajectoryElementBuilder.registerWriterClient(energyElementPtr);
     trajectoryElementBuilder.registerSignallerClient(energyElementPtr);

diff --git a/src/gromacs/modularsimulator/modularsimulatorinterfaces.h b/src/gromacs/modularsimulator/modularsimulatorinterfaces.h
index a93b2070fb344cf631c66e6fbd0b320fcdfc5421..6d4986d2dd9137f97ed21b9fd6fe8a9fc0d34a54 100644 (file)
--- a/src/gromacs/modularsimulator/modularsimulatorinterfaces.h
+++ b/src/gromacs/modularsimulator/modularsimulatorinterfaces.h
@@ -296,8 +296,6 @@ typedef std::unique_ptr<ITrajectoryWriterCallback> ITrajectoryWriterCallbackPtr;
  * callback called by the TrajectoryElement when trajectory writing happens.
  *
  * Setup and teardown methods allow clients to perform tasks with a valid output pointer.
- * These functions have a standard implementation to allow clients not to implement
- * them if no setup / teardown is needed.
  */
 class ITrajectoryWriterClient
 {

diff --git a/src/gromacs/modularsimulator/statepropagatordata.cpp b/src/gromacs/modularsimulator/statepropagatordata.cpp
index c59f6d1e8bade7404add682771e19c4ebd038b35..17b8eb47cb30ccf1cdafda2a07c38fa3ea9c9097 100644 (file)
--- a/src/gromacs/modularsimulator/statepropagatordata.cpp
+++ b/src/gromacs/modularsimulator/statepropagatordata.cpp
@@ -43,7 +43,9 @@
 
 #include "statepropagatordata.h"
 
+#include "gromacs/domdec/collect.h"
 #include "gromacs/domdec/domdec.h"
+#include "gromacs/fileio/confio.h"
 #include "gromacs/math/vec.h"
 #include "gromacs/mdlib/gmx_omp_nthreads.h"
 #include "gromacs/mdlib/mdoutf.h"
@@ -53,7 +55,9 @@
 #include "gromacs/mdtypes/inputrec.h"
 #include "gromacs/mdtypes/mdatom.h"
 #include "gromacs/mdtypes/state.h"
+#include "gromacs/pbcutil/pbc.h"
 #include "gromacs/topology/atoms.h"
+#include "gromacs/topology/topology.h"
 
 #include "freeenergyperturbationelement.h"
 
@@ -69,8 +73,12 @@ StatePropagatorData::StatePropagatorData(int                            numAtoms
                                          int                            nstxout_compressed,
                                          bool                           useGPU,
                                          FreeEnergyPerturbationElement* freeEnergyPerturbationElement,
-                                         const t_inputrec*              inputrec,
-                                         const t_mdatoms*               mdatoms) :
+                                         const TopologyHolder*          topologyHolder,
+                                         bool              canMoleculesBeDistributedOverPBC,
+                                         bool              writeFinalConfiguration,
+                                         std::string       finalConfigurationFilename,
+                                         const t_inputrec* inputrec,
+                                         const t_mdatoms*  mdatoms) :
     totalNumAtoms_(numAtoms),
     nstxout_(nstxout),
     nstvout_(nstvout),
@@ -81,6 +89,15 @@ StatePropagatorData::StatePropagatorData(int                            numAtoms
     writeOutStep_(-1),
     vvResetVelocities_(false),
     freeEnergyPerturbationElement_(freeEnergyPerturbationElement),
+    isRegularSimulationEnd_(false),
+    lastStep_(-1),
+    canMoleculesBeDistributedOverPBC_(canMoleculesBeDistributedOverPBC),
+    systemHasPeriodicMolecules_(inputrec->bPeriodicMols),
+    pbcType_(inputrec->ePBC),
+    topologyHolder_(topologyHolder),
+    lastPlannedStep_(inputrec->nsteps + inputrec->init_step),
+    writeFinalConfiguration_(writeFinalConfiguration),
+    finalConfigurationFilename_(std::move(finalConfigurationFilename)),
     fplog_(fplog),
     cr_(cr),
     globalState_(globalState)
@@ -298,7 +315,7 @@ void StatePropagatorData::scheduleTask(Step step, Time gmx_unused time, const Re
     // copy x -> previousX
     (*registerRunFunction)(std::make_unique<SimulatorRunFunction>([this]() { copyPosition(); }));
     // if it's a write out step, keep a copy for writeout
-    if (step == writeOutStep_)
+    if (step == writeOutStep_ || (step == lastStep_ && writeFinalConfiguration_))
     {
         (*registerRunFunction)(std::make_unique<SimulatorRunFunction>([this]() { saveState(); }));
     }
@@ -346,6 +363,7 @@ ITrajectoryWriterCallbackPtr StatePropagatorData::registerTrajectoryWriterCallba
 
 void StatePropagatorData::write(gmx_mdoutf_t outf, Step currentStep, Time currentTime)
 {
+    wallcycle_start(mdoutf_get_wcycle(outf), ewcTRAJ);
     unsigned int mdof_flags = 0;
     if (do_per_step(currentStep, nstxout_))
     {
@@ -382,6 +400,7 @@ void StatePropagatorData::write(gmx_mdoutf_t outf, Step currentStep, Time curren
 
     if (mdof_flags == 0)
     {
+        wallcycle_stop(mdoutf_get_wcycle(outf), ewcTRAJ);
         return;
     }
     GMX_ASSERT(localStateBackup_, "Trajectory writing called, but no state saved.");
@@ -393,7 +412,11 @@ void StatePropagatorData::write(gmx_mdoutf_t outf, Step currentStep, Time curren
                                      totalNumAtoms_, currentStep, currentTime,
                                      localStateBackup_.get(), globalState_, observablesHistory, f_);
 
-    localStateBackup_.reset();
+    if (currentStep != lastStep_ || !isRegularSimulationEnd_)
+    {
+        localStateBackup_.reset();
+    }
+    wallcycle_stop(mdoutf_get_wcycle(outf), ewcTRAJ);
 }
 
 void StatePropagatorData::elementSetup()
@@ -422,4 +445,54 @@ void StatePropagatorData::writeCheckpoint(t_state* localState, t_state gmx_unuse
     localState->flags |= (1U << estX) | (1U << estV) | (1U << estBOX);
 }
 
+void StatePropagatorData::trajectoryWriterTeardown(gmx_mdoutf* gmx_unused outf)
+{
+    // Note that part of this code is duplicated in do_md_trajectory_writing.
+    // This duplication is needed while both legacy and modular code paths are in use.
+    // TODO: Remove duplication asap, make sure to keep in sync in the meantime.
+    if (!writeFinalConfiguration_ || !isRegularSimulationEnd_)
+    {
+        return;
+    }
+
+    GMX_ASSERT(localStateBackup_, "Final trajectory writing called, but no state saved.");
+
+    wallcycle_start(mdoutf_get_wcycle(outf), ewcTRAJ);
+    if (DOMAINDECOMP(cr_))
+    {
+        auto globalXRef = MASTER(cr_) ? globalState_->x : gmx::ArrayRef<gmx::RVec>();
+        dd_collect_vec(cr_->dd, localStateBackup_.get(), localStateBackup_->x, globalXRef);
+        auto globalVRef = MASTER(cr_) ? globalState_->v : gmx::ArrayRef<gmx::RVec>();
+        dd_collect_vec(cr_->dd, localStateBackup_.get(), localStateBackup_->v, globalVRef);
+    }
+    else
+    {
+        // We have the whole state locally: copy the local state pointer
+        globalState_ = localStateBackup_.get();
+    }
+
+    if (MASTER(cr_))
+    {
+        fprintf(stderr, "\nWriting final coordinates.\n");
+        if (canMoleculesBeDistributedOverPBC_ && !systemHasPeriodicMolecules_)
+        {
+            // Make molecules whole only for confout writing
+            do_pbc_mtop(pbcType_, localStateBackup_->box, &topologyHolder_->globalTopology(),
+                        globalState_->x.rvec_array());
+        }
+        write_sto_conf_mtop(finalConfigurationFilename_.c_str(), *topologyHolder_->globalTopology().name,
+                            &topologyHolder_->globalTopology(), globalState_->x.rvec_array(),
+                            globalState_->v.rvec_array(), pbcType_, localStateBackup_->box);
+    }
+    wallcycle_stop(mdoutf_get_wcycle(outf), ewcTRAJ);
+}
+
+SignallerCallbackPtr StatePropagatorData::registerLastStepCallback()
+{
+    return std::make_unique<SignallerCallback>([this](Step step, Time) {
+        lastStep_               = step;
+        isRegularSimulationEnd_ = (step == lastPlannedStep_);
+    });
+}
+
 } // namespace gmx

diff --git a/src/gromacs/modularsimulator/statepropagatordata.h b/src/gromacs/modularsimulator/statepropagatordata.h
index eed831f0e7c6cfe6ee9c9082485d8edbfa2444bc..d6a7fac8d385700b3e20373b05a8d8187b1e8a77 100644 (file)
--- a/src/gromacs/modularsimulator/statepropagatordata.h
+++ b/src/gromacs/modularsimulator/statepropagatordata.h
@@ -47,6 +47,7 @@
 #include "gromacs/math/vectypes.h"
 
 #include "modularsimulatorinterfaces.h"
+#include "topologyholder.h"
 
 struct gmx_mdoutf;
 struct t_commrec;
@@ -100,7 +101,8 @@ class StatePropagatorData final :
     public ISimulatorElement,
     public ITrajectoryWriterClient,
     public ITrajectorySignallerClient,
-    public ICheckpointHelperClient
+    public ICheckpointHelperClient,
+    public ILastStepSignallerClient
 {
 public:
     //! Constructor
@@ -114,6 +116,10 @@ public:
                         int                            nstxout_compressed,
                         bool                           useGPU,
                         FreeEnergyPerturbationElement* freeEnergyPerturbationElement,
+                        const TopologyHolder*          topologyHolder,
+                        bool                           canMoleculesBeDistributedOverPBC,
+                        bool                           writeFinalConfiguration,
+                        std::string                    finalConfigurationFilename,
                         const t_inputrec*              inputrec,
                         const t_mdatoms*               mdatoms);
 
@@ -240,6 +246,9 @@ private:
     //! ICheckpointHelperClient implementation
     void writeCheckpoint(t_state* localState, t_state* globalState) override;
 
+    //! ILastStepSignallerClient implementation (used for final output only)
+    SignallerCallbackPtr registerLastStepCallback() override;
+
     //! Callback writing the state to file
     void write(gmx_mdoutf* outf, Step step, Time time);
 
@@ -253,6 +262,26 @@ private:
     //! Pointer to the free energy perturbation element (for trajectory writing only)
     FreeEnergyPerturbationElement* freeEnergyPerturbationElement_;
 
+    //! Whether planned total number of steps was reached (used for final output only)
+    bool isRegularSimulationEnd_;
+    //! The signalled last step (used for final output only)
+    Step lastStep_;
+
+    //! Whether system can have molecules distributed over PBC boundaries (used for final output only)
+    const bool canMoleculesBeDistributedOverPBC_;
+    //! Whether system has molecules self-interacting through PBC (used for final output only)
+    const bool systemHasPeriodicMolecules_;
+    //! The PBC type (used for final output only)
+    const int pbcType_;
+    //! Pointer to the topology (used for final output only)
+    const TopologyHolder* topologyHolder_;
+    //! The (planned) last step - determines whether final configuration is written (used for final output only)
+    const Step lastPlannedStep_;
+    //! Whether final configuration was chosen in mdrun options (used for final output only)
+    const bool writeFinalConfiguration_;
+    //! The filename of the final configuration file (used for final output only)
+    const std::string finalConfigurationFilename_;
+
     // Access to ISimulator data
     //! Handles logging.
     FILE* fplog_;
@@ -263,8 +292,8 @@ private:
 
     //! No trajectory writer setup needed
     void trajectoryWriterSetup(gmx_mdoutf gmx_unused* outf) override {}
-    //! No trajectory writer teardown needed
-    void trajectoryWriterTeardown(gmx_mdoutf gmx_unused* outf) override {}
+    //! Trajectory writer teardown - write final coordinates
+    void trajectoryWriterTeardown(gmx_mdoutf* outf) override;
 };
 
 } // namespace gmx