src/gromacs/modularsimulator/expandedensembleelement.cpp

   1 /*
   2  * This file is part of the GROMACS molecular simulation package.
   3  *
   4  * Copyright (c) 2021, by the GROMACS development team, led by
   5  * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
   6  * and including many others, as listed in the AUTHORS file in the
   7  * top-level source directory and at http://www.gromacs.org.
   8  *
   9  * GROMACS is free software; you can redistribute it and/or
  10  * modify it under the terms of the GNU Lesser General Public License
  11  * as published by the Free Software Foundation; either version 2.1
  12  * of the License, or (at your option) any later version.
  13  *
  14  * GROMACS is distributed in the hope that it will be useful,
  15  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  16  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  17  * Lesser General Public License for more details.
  18  *
  19  * You should have received a copy of the GNU Lesser General Public
  20  * License along with GROMACS; if not, see
  21  * http://www.gnu.org/licenses, or write to the Free Software Foundation,
  22  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
  23  *
  24  * If you want to redistribute modifications to GROMACS, please
  25  * consider that scientific software is very special. Version
  26  * control is crucial - bugs must be traceable. We will be happy to
  27  * consider code for inclusion in the official distribution, but
  28  * derived work must not be called official GROMACS. Details are found
  29  * in the README & COPYING files - if they are missing, get the
  30  * official version at http://www.gromacs.org.
  31  *
  32  * To help us fund GROMACS development, we humbly ask that you cite
  33  * the research papers on the package. Check out http://www.gromacs.org.
  34  */
  35 /*! \internal \file
  36  * \brief Defines the expanded ensemble element for the modular simulator
  37  *
  38  * \author Pascal Merz <pascal.merz@me.com>
  39  * \ingroup module_modularsimulator
  40  */
  41
  42 #include "gmxpre.h"
  43
  44 #include "expandedensembleelement.h"
  45
  46 #include "gromacs/domdec/distribute.h"
  47 #include "gromacs/mdlib/expanded.h"
  48 #include "gromacs/mdlib/stat.h"
  49 #include "gromacs/mdtypes/checkpointdata.h"
  50 #include "gromacs/mdtypes/commrec.h"
  51 #include "gromacs/mdtypes/df_history.h"
  52 #include "gromacs/mdtypes/inputrec.h"
  53 #include "gromacs/utility/fatalerror.h"
  54
  55 #include "energydata.h"
  56 #include "simulatoralgorithm.h"
  57
  58 namespace gmx
  59 {
  60
  61 void ExpandedEnsembleElement::apply(Step step, bool doLambdaStep, bool doLog)
  62 {
  63     if (doLambdaStep)
  64     {
  65         const int newFepState =
  66                 expandedEnsembleUpdateLambdaState(fplog_,
  67                                                   inputrec_,
  68                                                   energyData_->enerdata(),
  69                                                   freeEnergyPerturbationData_->currentFEPState(),
  70                                                   dfhist_.get(),
  71                                                   step);
  72         // Set new state at next step
  73         fepStateSetting_->setNewState(newFepState, step + 1);
  74     }
  75     if (doLog)
  76     {
  77         /* only needed if doing expanded ensemble */
  78         PrintFreeEnergyInfoToFile(fplog_,
  79                                   inputrec_->fepvals.get(),
  80                                   inputrec_->expandedvals.get(),
  81                                   inputrec_->bSimTemp ? inputrec_->simtempvals.get() : nullptr,
  82                                   dfhist_.get(),
  83                                   freeEnergyPerturbationData_->currentFEPState(),
  84                                   inputrec_->nstlog,
  85                                   step);
  86     }
  87 }
  88
  89 void ExpandedEnsembleElement::elementSetup()
  90 {
  91     // Check nstexpanded here, because the grompp check was broken (#2714)
  92     if (inputrec_->expandedvals->nstexpanded % inputrec_->nstcalcenergy != 0)
  93     {
  94         gmx_fatal(FARGS,
  95                   "With expanded ensemble, nstexpanded should be a multiple of nstcalcenergy");
  96     }
  97     init_expanded_ensemble(restoredFromCheckpoint_, inputrec_, dfhist_.get());
  98 }
  99
 100 void ExpandedEnsembleElement::scheduleTask(Step step, Time /*unused*/, const RegisterRunFunction& registerRunFunction)
 101 {
 102     const bool isFirstStep  = (step == initialStep_);
 103     const bool doLambdaStep = (do_per_step(step, frequency_) && !isFirstStep);
 104     const bool doLog        = (isMasterRank_ && step == nextLogWritingStep_ && (fplog_ != nullptr));
 105
 106     if (doLambdaStep || doLog)
 107     {
 108         registerRunFunction([this, step, doLambdaStep, doLog]() { apply(step, doLambdaStep, doLog); });
 109     }
 110     if (doLambdaStep)
 111     {
 112         // We'll compute a new lambda state and want it applied for next step
 113         fepStateSetting_->signalSettingStep(step + 1);
 114     }
 115 }
 116
 117 namespace
 118 {
 119 /*!
 120  * \brief Enum describing the contents FreeEnergyPerturbationData::Element writes to modular checkpoint
 121  *
 122  * When changing the checkpoint content, add a new element just above Count, and adjust the
 123  * checkpoint functionality.
 124  */
 125 enum class CheckpointVersion
 126 {
 127     Base, //!< First version of modular checkpointing
 128     Count //!< Number of entries. Add new versions right above this!
 129 };
 130 constexpr auto c_currentVersion = CheckpointVersion(int(CheckpointVersion::Count) - 1);
 131 } // namespace
 132
 133 template<CheckpointDataOperation operation>
 134 void ExpandedEnsembleElement::doCheckpointData(CheckpointData<operation>* checkpointData)
 135 {
 136     checkpointVersion(checkpointData, "ExpandedEnsembleElement version", c_currentVersion);
 137
 138     dfhist_->doCheckpoint<operation>(checkpointData->subCheckpointData("dfhist"),
 139                                      inputrec_->expandedvals->elamstats);
 140 }
 141
 142 void ExpandedEnsembleElement::saveCheckpointState(std::optional<WriteCheckpointData> checkpointData,
 143                                                   const t_commrec*                   cr)
 144 {
 145     if (MASTER(cr))
 146     {
 147         doCheckpointData<CheckpointDataOperation::Write>(&checkpointData.value());
 148     }
 149 }
 150
 151 void ExpandedEnsembleElement::restoreCheckpointState(std::optional<ReadCheckpointData> checkpointData,
 152                                                      const t_commrec*                  cr)
 153 {
 154     if (MASTER(cr))
 155     {
 156         doCheckpointData<CheckpointDataOperation::Read>(&checkpointData.value());
 157     }
 158     if (DOMAINDECOMP(cr))
 159     {
 160         dd_distribute_dfhist(cr->dd, dfhist_.get());
 161     }
 162     restoredFromCheckpoint_ = true;
 163 }
 164
 165 const std::string& ExpandedEnsembleElement::clientID()
 166 {
 167     return identifier_;
 168 }
 169
 170 std::optional<SignallerCallback> ExpandedEnsembleElement::registerLoggingCallback()
 171 {
 172     if (isMasterRank_)
 173     {
 174         return [this](Step step, Time /*unused*/) { nextLogWritingStep_ = step; };
 175     }
 176     else
 177     {
 178         return std::nullopt;
 179     }
 180 }
 181
 182 ExpandedEnsembleElement::ExpandedEnsembleElement(bool                              isMasterRank,
 183                                                  Step                              initialStep,
 184                                                  int                               frequency,
 185                                                  const EnergyData*                 energyData,
 186                                                  const FreeEnergyPerturbationData* freeEnergyPerturbationData,
 187                                                  FILE*                             fplog,
 188                                                  const t_inputrec*                 inputrec) :
 189     fepStateSetting_(freeEnergyPerturbationData->enableExternalFepStateSetting()),
 190     isMasterRank_(isMasterRank),
 191     initialStep_(initialStep),
 192     frequency_(frequency),
 193     nextLogWritingStep_(-1),
 194     dfhist_(std::make_unique<df_history_t>()),
 195     restoredFromCheckpoint_(false),
 196     energyData_(energyData),
 197     freeEnergyPerturbationData_(freeEnergyPerturbationData),
 198     fplog_(fplog),
 199     inputrec_(inputrec)
 200 {
 201     init_df_history(dfhist_.get(), inputrec_->fepvals->n_lambda);
 202 }
 203
 204 ISimulatorElement* ExpandedEnsembleElement::getElementPointerImpl(
 205         LegacySimulatorData*                    legacySimulatorData,
 206         ModularSimulatorAlgorithmBuilderHelper* builderHelper,
 207         StatePropagatorData gmx_unused* statePropagatorData,
 208         EnergyData*                     energyData,
 209         FreeEnergyPerturbationData*     freeEnergyPerturbationData,
 210         GlobalCommunicationHelper gmx_unused* globalCommunicationHelper,
 211         ObservablesReducer* /*observablesReducer*/)
 212 {
 213     return builderHelper->storeElement(std::make_unique<ExpandedEnsembleElement>(
 214             MASTER(legacySimulatorData->cr),
 215             legacySimulatorData->inputrec->init_step,
 216             legacySimulatorData->inputrec->expandedvals->nstexpanded,
 217             energyData,
 218             freeEnergyPerturbationData,
 219             legacySimulatorData->fplog,
 220             legacySimulatorData->inputrec));
 221 }
 222
 223 } // namespace gmx