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_, *inputrec_, true, &currentFEPState_, lambda_);
  76 }
  77
  78 void FreeEnergyPerturbationData::Element::scheduleTask(Step step,
  79                                                        Time gmx_unused            time,
  80                                                        const RegisterRunFunction& registerRunFunction)
  81 {
  82     if (lambdasChange_)
  83     {
  84         registerRunFunction([this, step]() { freeEnergyPerturbationData_->updateLambdas(step); });
  85     }
  86 }
  87
  88 void FreeEnergyPerturbationData::updateLambdas(Step step)
  89 {
  90     // at beginning of step (if lambdas change...)
  91     lambda_ = currentLambdas(step, *(inputrec_->fepvals), currentFEPState_);
  92     updateMDAtoms();
  93 }
  94
  95 ArrayRef<real> FreeEnergyPerturbationData::lambdaView()
  96 {
  97     return lambda_;
  98 }
  99
 100 ArrayRef<const real> FreeEnergyPerturbationData::constLambdaView()
 101 {
 102     return lambda_;
 103 }
 104
 105 int FreeEnergyPerturbationData::currentFEPState()
 106 {
 107     return currentFEPState_;
 108 }
 109
 110 void FreeEnergyPerturbationData::updateMDAtoms()
 111 {
 112     update_mdatoms(mdAtoms_->mdatoms(), lambda_[FreeEnergyPerturbationCouplingType::Mass]);
 113 }
 114
 115 namespace
 116 {
 117 /*!
 118  * \brief Enum describing the contents FreeEnergyPerturbationData::Element writes to modular checkpoint
 119  *
 120  * When changing the checkpoint content, add a new element just above Count, and adjust the
 121  * checkpoint functionality.
 122  */
 123 enum class CheckpointVersion
 124 {
 125     Base, //!< First version of modular checkpointing
 126     Count //!< Number of entries. Add new versions right above this!
 127 };
 128 constexpr auto c_currentVersion = CheckpointVersion(int(CheckpointVersion::Count) - 1);
 129 } // namespace
 130
 131 template<CheckpointDataOperation operation>
 132 void FreeEnergyPerturbationData::doCheckpointData(CheckpointData<operation>* checkpointData)
 133 {
 134     checkpointVersion(checkpointData, "FreeEnergyPerturbationData version", c_currentVersion);
 135
 136     checkpointData->scalar("current FEP state", &currentFEPState_);
 137     checkpointData->arrayRef("lambda vector", makeCheckpointArrayRef<operation>(lambda_));
 138 }
 139
 140 void FreeEnergyPerturbationData::Element::saveCheckpointState(std::optional<WriteCheckpointData> checkpointData,
 141                                                               const t_commrec*                   cr)
 142 {
 143     if (MASTER(cr))
 144     {
 145         freeEnergyPerturbationData_->doCheckpointData<CheckpointDataOperation::Write>(
 146                 &checkpointData.value());
 147     }
 148 }
 149
 150 void FreeEnergyPerturbationData::Element::restoreCheckpointState(std::optional<ReadCheckpointData> checkpointData,
 151                                                                  const t_commrec* cr)
 152 {
 153     if (MASTER(cr))
 154     {
 155         freeEnergyPerturbationData_->doCheckpointData<CheckpointDataOperation::Read>(
 156                 &checkpointData.value());
 157     }
 158     if (DOMAINDECOMP(cr))
 159     {
 160         dd_bcast(cr->dd, sizeof(int), &freeEnergyPerturbationData_->currentFEPState_);
 161         dd_bcast(cr->dd,
 162                  ssize(freeEnergyPerturbationData_->lambda_) * int(sizeof(real)),
 163                  freeEnergyPerturbationData_->lambda_.data());
 164     }
 165 }
 166
 167 const std::string& FreeEnergyPerturbationData::Element::clientID()
 168 {
 169     return FreeEnergyPerturbationData::checkpointID();
 170 }
 171
 172 FreeEnergyPerturbationData::Element::Element(FreeEnergyPerturbationData* freeEnergyPerturbationElement,
 173                                              double                      deltaLambda) :
 174     freeEnergyPerturbationData_(freeEnergyPerturbationElement),
 175     lambdasChange_(deltaLambda != 0)
 176 {
 177 }
 178
 179 void FreeEnergyPerturbationData::Element::elementSetup()
 180 {
 181     freeEnergyPerturbationData_->updateMDAtoms();
 182 }
 183
 184 FreeEnergyPerturbationData::Element* FreeEnergyPerturbationData::element()
 185 {
 186     return element_.get();
 187 }
 188
 189 ISimulatorElement* FreeEnergyPerturbationData::Element::getElementPointerImpl(
 190         LegacySimulatorData gmx_unused*        legacySimulatorData,
 191         ModularSimulatorAlgorithmBuilderHelper gmx_unused* builderHelper,
 192         StatePropagatorData gmx_unused* statePropagatorData,
 193         EnergyData gmx_unused*      energyData,
 194         FreeEnergyPerturbationData* freeEnergyPerturbationData,
 195         GlobalCommunicationHelper gmx_unused* globalCommunicationHelper)
 196 {
 197     return freeEnergyPerturbationData->element();
 198 }
 199
 200 void FreeEnergyPerturbationData::readCheckpointToTrxFrame(t_trxframe* trxFrame,
 201                                                           std::optional<ReadCheckpointData> readCheckpointData)
 202 {
 203     if (readCheckpointData)
 204     {
 205         FreeEnergyPerturbationData freeEnergyPerturbationData;
 206         freeEnergyPerturbationData.doCheckpointData(&readCheckpointData.value());
 207         trxFrame->lambda = freeEnergyPerturbationData.lambda_[FreeEnergyPerturbationCouplingType::Fep];
 208         trxFrame->fep_state = freeEnergyPerturbationData.currentFEPState_;
 209     }
 210     else
 211     {
 212         trxFrame->lambda    = 0;
 213         trxFrame->fep_state = 0;
 214     }
 215     trxFrame->bLambda = true;
 216 }
 217
 218 const std::string& FreeEnergyPerturbationData::checkpointID()
 219 {
 220     static const std::string identifier = "FreeEnergyPerturbationData";
 221     return identifier;
 222 }
 223
 224 } // namespace gmx