src/gromacs/modularsimulator/freeenergyperturbationdata.cpp

   1 /*
   2  * This file is part of the GROMACS molecular simulation package.
   3  *
   4  * Copyright (c) 2019,2020,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 free energy perturbation 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 "freeenergyperturbationdata.h"
  45
  46 #include "gromacs/domdec/domdec_network.h"
  47 #include "gromacs/mdlib/freeenergyparameters.h"
  48 #include "gromacs/mdlib/md_support.h"
  49 #include "gromacs/mdlib/mdatoms.h"
  50 #include "gromacs/mdtypes/checkpointdata.h"
  51 #include "gromacs/mdtypes/commrec.h"
  52 #include "gromacs/mdtypes/inputrec.h"
  53 #include "gromacs/mdtypes/mdatom.h"
  54 #include "gromacs/mdtypes/state.h"
  55 #include "gromacs/trajectory/trajectoryframe.h"
  56
  57 #include "modularsimulator.h"
  58 #include "simulatoralgorithm.h"
  59
  60 namespace gmx
  61 {
  62
  63 FreeEnergyPerturbationData::FreeEnergyPerturbationData(FILE* fplog, const t_inputrec& inputrec, MDAtoms* mdAtoms) :
  64     element_(std::make_unique<Element>(this, inputrec.fepvals->delta_lambda)),
  65     lambda_(),
  66     currentFEPState_(0),
  67     fplog_(fplog),
  68     inputrec_(inputrec),
  69     mdAtoms_(mdAtoms)
  70 {
  71     std::fill(lambda_.begin(), lambda_.end(), 0);
  72     // The legacy implementation only filled the lambda vector in state_global, which is only
  73     // available on master. We have the lambda vector available everywhere, so we pass a `true`
  74     // for isMaster on all ranks. See #3647.
  75     initialize_lambdas(fplog_,
  76                        inputrec_.efep,
  77                        inputrec_.bSimTemp,
  78                        *inputrec_.fepvals,
  79                        inputrec_.simtempvals->temperatures,
  80                        gmx::arrayRefFromArray(inputrec_.opts.ref_t, inputrec_.opts.ngtc),
  81                        true,
  82                        &currentFEPState_,
  83                        lambda_);
  84 }
  85
  86 void FreeEnergyPerturbationData::Element::scheduleTask(Step                       step,
  87                                                        Time gmx_unused            time,
  88                                                        const RegisterRunFunction& registerRunFunction)
  89 {
  90     // If we do slow growth, we update lambda every step
  91     // If it's set externally, we get notified, so we only update when necessary (at nextLambdaSettingStep_)
  92     // However, if we reload from checkpoint, it might be that checkpointing happened right between the
  93     // external caller setting the state and us applying it, so we also check newFepStateStep_.
  94     const bool needToSetExternalState = externalFepStateSetting_
  95                                         && ((step == externalFepStateSetting_->nextFepStateSettingStep)
  96                                             || (step == externalFepStateSetting_->newFepStateStep));
  97     if (doSlowGrowth_)
  98     {
  99         registerRunFunction([this, step]() { freeEnergyPerturbationData_->updateLambdas(step); });
 100     }
 101     else if (needToSetExternalState)
 102     {
 103         registerRunFunction([this, step]() {
 104             GMX_ASSERT(step == externalFepStateSetting_->newFepStateStep,
 105                        "FEP state setting step mismatch");
 106             freeEnergyPerturbationData_->setLambdaState(step, externalFepStateSetting_->newFepState);
 107         });
 108     }
 109 }
 110
 111 void FreeEnergyPerturbationData::updateLambdas(Step step)
 112 {
 113     // at beginning of step (if lambdas change...)
 114     lambda_ = currentLambdas(step, *(inputrec_.fepvals), currentFEPState_);
 115     updateMDAtoms();
 116 }
 117
 118 void FreeEnergyPerturbationData::setLambdaState(Step step, int newState)
 119 {
 120     currentFEPState_ = newState;
 121     updateLambdas(step);
 122 }
 123
 124 ArrayRef<real> FreeEnergyPerturbationData::lambdaView()
 125 {
 126     return lambda_;
 127 }
 128
 129 ArrayRef<const real> FreeEnergyPerturbationData::constLambdaView() const
 130 {
 131     return lambda_;
 132 }
 133
 134 int FreeEnergyPerturbationData::currentFEPState() const
 135 {
 136     return currentFEPState_;
 137 }
 138
 139 void FreeEnergyPerturbationData::updateMDAtoms()
 140 {
 141     update_mdatoms(mdAtoms_->mdatoms(), lambda_[FreeEnergyPerturbationCouplingType::Mass]);
 142 }
 143
 144 FepStateSetting* FreeEnergyPerturbationData::enableExternalFepStateSetting() const
 145 {
 146     return element_->enableExternalFepStateSetting();
 147 }
 148
 149 namespace
 150 {
 151 /*!
 152  * \brief Enum describing the contents FreeEnergyPerturbationData::Element writes to modular checkpoint
 153  *
 154  * When changing the checkpoint content, add a new element just above Count, and adjust the
 155  * checkpoint functionality.
 156  */
 157 enum class CheckpointVersion
 158 {
 159     Base,                       //!< First version of modular checkpointing
 160     AddedExternalLambdaSetting, //!< Additional values to ensure no state setting info is lost
 161     Count                       //!< Number of entries. Add new versions right above this!
 162 };
 163 constexpr auto c_currentVersion = CheckpointVersion(int(CheckpointVersion::Count) - 1);
 164 } // namespace
 165
 166 template<CheckpointDataOperation operation>
 167 void FreeEnergyPerturbationData::doCheckpointData(CheckpointData<operation>* checkpointData)
 168 {
 169     checkpointVersion(checkpointData, "FreeEnergyPerturbationData version", c_currentVersion);
 170
 171     checkpointData->scalar("current FEP state", &currentFEPState_);
 172     checkpointData->arrayRef("lambda vector", makeCheckpointArrayRef<operation>(lambda_));
 173 }
 174
 175 template<CheckpointDataOperation operation>
 176 void FreeEnergyPerturbationData::Element::doCheckpointData(CheckpointData<operation>* checkpointData)
 177 {
 178     CheckpointVersion fileVersion = c_currentVersion;
 179     if (operation == CheckpointDataOperation::Read)
 180     {
 181         // We can read the same key as above - we can only write it once, though!
 182         fileVersion = checkpointVersion(
 183                 checkpointData, "FreeEnergyPerturbationData version", c_currentVersion);
 184     }
 185
 186     if (fileVersion > CheckpointVersion::AddedExternalLambdaSetting)
 187     {
 188         // If checkpointing happens between receiving the request and actually setting the new
 189         // lambda state, we need to preserve this information.
 190         bool externalLambdaSetting = externalFepStateSetting_.has_value();
 191         checkpointData->scalar("External lambda setting", &externalLambdaSetting);
 192         if constexpr (operation == CheckpointDataOperation::Read)
 193         {
 194             GMX_RELEASE_ASSERT(
 195                     !(externalLambdaSetting && !externalFepStateSetting_.has_value()),
 196                     "Checkpoint mismatch: Checkpointed simulation used external lambda setting, "
 197                     "while the current simulation does not.");
 198             GMX_RELEASE_ASSERT(!(!externalLambdaSetting && externalFepStateSetting_.has_value()),
 199                                "Checkpoint mismatch: Checkpointed simulation dit not use external "
 200                                "lambda setting, "
 201                                "while the current simulation does.");
 202         }
 203         if (externalFepStateSetting_.has_value()) // NOLINT(readability-misleading-indentation)
 204         {
 205             checkpointData->scalar("Requested new FEP state", &externalFepStateSetting_->newFepState);
 206             checkpointData->scalar("Step at which new FEP state is applied",
 207                                    &externalFepStateSetting_->newFepStateStep);
 208         }
 209     }
 210 }
 211
 212 void FreeEnergyPerturbationData::Element::saveCheckpointState(std::optional<WriteCheckpointData> checkpointData,
 213                                                               const t_commrec*                   cr)
 214 {
 215     if (MASTER(cr))
 216     {
 217         freeEnergyPerturbationData_->doCheckpointData<CheckpointDataOperation::Write>(
 218                 &checkpointData.value());
 219         doCheckpointData<CheckpointDataOperation::Write>(&checkpointData.value());
 220     }
 221 }
 222
 223 void FreeEnergyPerturbationData::Element::restoreCheckpointState(std::optional<ReadCheckpointData> checkpointData,
 224                                                                  const t_commrec* cr)
 225 {
 226     if (MASTER(cr))
 227     {
 228         freeEnergyPerturbationData_->doCheckpointData<CheckpointDataOperation::Read>(
 229                 &checkpointData.value());
 230         doCheckpointData<CheckpointDataOperation::Read>(&checkpointData.value());
 231     }
 232     if (DOMAINDECOMP(cr))
 233     {
 234         dd_bcast(cr->dd, sizeof(int), &freeEnergyPerturbationData_->currentFEPState_);
 235         dd_bcast(cr->dd,
 236                  ssize(freeEnergyPerturbationData_->lambda_) * int(sizeof(real)),
 237                  freeEnergyPerturbationData_->lambda_.data());
 238         if (externalFepStateSetting_.has_value())
 239         {
 240             dd_bcast(cr->dd,
 241                      sizeof(externalFepStateSetting_->newFepState),
 242                      &externalFepStateSetting_->newFepState);
 243             dd_bcast(cr->dd,
 244                      sizeof(externalFepStateSetting_->newFepStateStep),
 245                      &externalFepStateSetting_->newFepStateStep);
 246         }
 247     }
 248 }
 249
 250 const std::string& FreeEnergyPerturbationData::Element::clientID()
 251 {
 252     return FreeEnergyPerturbationData::checkpointID();
 253 }
 254
 255 DomDecCallback FreeEnergyPerturbationData::Element::registerDomDecCallback()
 256 {
 257     return [this]() { freeEnergyPerturbationData_->updateMDAtoms(); };
 258 }
 259
 260 FreeEnergyPerturbationData::Element::Element(FreeEnergyPerturbationData* freeEnergyPerturbationElement,
 261                                              double                      deltaLambda) :
 262     freeEnergyPerturbationData_(freeEnergyPerturbationElement), doSlowGrowth_(deltaLambda != 0)
 263 {
 264 }
 265
 266 void FreeEnergyPerturbationData::Element::elementSetup()
 267 {
 268     freeEnergyPerturbationData_->updateMDAtoms();
 269 }
 270
 271 FreeEnergyPerturbationData::Element* FreeEnergyPerturbationData::element()
 272 {
 273     return element_.get();
 274 }
 275
 276 FepStateSetting* FreeEnergyPerturbationData::Element::enableExternalFepStateSetting()
 277 {
 278     GMX_RELEASE_ASSERT(!doSlowGrowth_,
 279                        "External FEP state setting is incompatible with slow growth.");
 280     // This could be implemented with some sanity checks, but there's no use case right now
 281     GMX_RELEASE_ASSERT(!externalFepStateSetting_.has_value(),
 282                        "External FEP state setting by more than one element not supported.");
 283     externalFepStateSetting_ = FepStateSetting();
 284     return &externalFepStateSetting_.value();
 285 }
 286
 287 void FepStateSetting::signalSettingStep(Step step)
 288 {
 289     nextFepStateSettingStep = step;
 290 }
 291
 292 void FepStateSetting::setNewState(int state, Step step)
 293 {
 294     newFepState     = state;
 295     newFepStateStep = step;
 296 }
 297
 298 ISimulatorElement* FreeEnergyPerturbationData::Element::getElementPointerImpl(
 299         LegacySimulatorData gmx_unused*        legacySimulatorData,
 300         ModularSimulatorAlgorithmBuilderHelper gmx_unused* builderHelper,
 301         StatePropagatorData gmx_unused* statePropagatorData,
 302         EnergyData gmx_unused*      energyData,
 303         FreeEnergyPerturbationData* freeEnergyPerturbationData,
 304         GlobalCommunicationHelper gmx_unused* globalCommunicationHelper)
 305 {
 306     return freeEnergyPerturbationData->element();
 307 }
 308
 309 void FreeEnergyPerturbationData::readCheckpointToTrxFrame(t_trxframe* trxFrame,
 310                                                           std::optional<ReadCheckpointData> readCheckpointData)
 311 {
 312     if (readCheckpointData)
 313     {
 314         FreeEnergyPerturbationData freeEnergyPerturbationData(nullptr, t_inputrec(), nullptr);
 315         freeEnergyPerturbationData.doCheckpointData(&readCheckpointData.value());
 316         trxFrame->lambda = freeEnergyPerturbationData.lambda_[FreeEnergyPerturbationCouplingType::Fep];
 317         trxFrame->fep_state = freeEnergyPerturbationData.currentFEPState_;
 318     }
 319     else
 320     {
 321         trxFrame->lambda    = 0;
 322         trxFrame->fep_state = 0;
 323     }
 324     trxFrame->bLambda = true;
 325 }
 326
 327 const std::string& FreeEnergyPerturbationData::checkpointID()
 328 {
 329     static const std::string identifier = "FreeEnergyPerturbationData";
 330     return identifier;
 331 }
 332
 333 } // namespace gmx