2 * This file is part of the GROMACS molecular simulation package.
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.
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.
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.
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.
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.
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.
36 * \brief Defines the free energy perturbation element for the modular simulator
38 * \author Pascal Merz <pascal.merz@me.com>
39 * \ingroup module_modularsimulator
44 #include "freeenergyperturbationdata.h"
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"
57 #include "modularsimulator.h"
58 #include "simulatoralgorithm.h"
63 FreeEnergyPerturbationData::FreeEnergyPerturbationData(FILE* fplog, const t_inputrec& inputrec, MDAtoms* mdAtoms) :
64 element_(std::make_unique<Element>(this, inputrec.fepvals->delta_lambda)),
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_,
79 inputrec_.simtempvals->temperatures,
80 gmx::arrayRefFromArray(inputrec_.opts.ref_t, inputrec_.opts.ngtc),
86 void FreeEnergyPerturbationData::Element::scheduleTask(Step step,
88 const RegisterRunFunction& registerRunFunction)
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));
99 registerRunFunction([this, step]() { freeEnergyPerturbationData_->updateLambdas(step); });
101 else if (needToSetExternalState)
103 registerRunFunction([this, step]() {
104 GMX_ASSERT(step == externalFepStateSetting_->newFepStateStep,
105 "FEP state setting step mismatch");
106 freeEnergyPerturbationData_->setLambdaState(step, externalFepStateSetting_->newFepState);
111 void FreeEnergyPerturbationData::updateLambdas(Step step)
113 // at beginning of step (if lambdas change...)
114 lambda_ = currentLambdas(step, *(inputrec_.fepvals), currentFEPState_);
118 void FreeEnergyPerturbationData::setLambdaState(Step step, int newState)
120 currentFEPState_ = newState;
124 ArrayRef<real> FreeEnergyPerturbationData::lambdaView()
129 ArrayRef<const real> FreeEnergyPerturbationData::constLambdaView() const
134 int FreeEnergyPerturbationData::currentFEPState() const
136 return currentFEPState_;
139 void FreeEnergyPerturbationData::updateMDAtoms()
141 update_mdatoms(mdAtoms_->mdatoms(), lambda_[FreeEnergyPerturbationCouplingType::Mass]);
144 FepStateSetting* FreeEnergyPerturbationData::enableExternalFepStateSetting() const
146 return element_->enableExternalFepStateSetting();
152 * \brief Enum describing the contents FreeEnergyPerturbationData::Element writes to modular checkpoint
154 * When changing the checkpoint content, add a new element just above Count, and adjust the
155 * checkpoint functionality.
157 enum class CheckpointVersion
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!
163 constexpr auto c_currentVersion = CheckpointVersion(int(CheckpointVersion::Count) - 1);
166 template<CheckpointDataOperation operation>
167 void FreeEnergyPerturbationData::doCheckpointData(CheckpointData<operation>* checkpointData)
169 checkpointVersion(checkpointData, "FreeEnergyPerturbationData version", c_currentVersion);
171 checkpointData->scalar("current FEP state", ¤tFEPState_);
172 checkpointData->arrayRef("lambda vector", makeCheckpointArrayRef<operation>(lambda_));
175 template<CheckpointDataOperation operation>
176 void FreeEnergyPerturbationData::Element::doCheckpointData(CheckpointData<operation>* checkpointData)
178 CheckpointVersion fileVersion = c_currentVersion;
179 if (operation == CheckpointDataOperation::Read)
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);
186 if (fileVersion > CheckpointVersion::AddedExternalLambdaSetting)
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)
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 "
201 "while the current simulation does.");
203 if (externalFepStateSetting_.has_value()) // NOLINT(readability-misleading-indentation)
205 checkpointData->scalar("Requested new FEP state", &externalFepStateSetting_->newFepState);
206 checkpointData->scalar("Step at which new FEP state is applied",
207 &externalFepStateSetting_->newFepStateStep);
212 void FreeEnergyPerturbationData::Element::saveCheckpointState(std::optional<WriteCheckpointData> checkpointData,
217 freeEnergyPerturbationData_->doCheckpointData<CheckpointDataOperation::Write>(
218 &checkpointData.value());
219 doCheckpointData<CheckpointDataOperation::Write>(&checkpointData.value());
223 void FreeEnergyPerturbationData::Element::restoreCheckpointState(std::optional<ReadCheckpointData> checkpointData,
228 freeEnergyPerturbationData_->doCheckpointData<CheckpointDataOperation::Read>(
229 &checkpointData.value());
230 doCheckpointData<CheckpointDataOperation::Read>(&checkpointData.value());
232 if (haveDDAtomOrdering(*cr))
234 dd_bcast(cr->dd, sizeof(int), &freeEnergyPerturbationData_->currentFEPState_);
236 ssize(freeEnergyPerturbationData_->lambda_) * int(sizeof(real)),
237 freeEnergyPerturbationData_->lambda_.data());
238 if (externalFepStateSetting_.has_value())
241 sizeof(externalFepStateSetting_->newFepState),
242 &externalFepStateSetting_->newFepState);
244 sizeof(externalFepStateSetting_->newFepStateStep),
245 &externalFepStateSetting_->newFepStateStep);
250 const std::string& FreeEnergyPerturbationData::Element::clientID()
252 return FreeEnergyPerturbationData::checkpointID();
255 DomDecCallback FreeEnergyPerturbationData::Element::registerDomDecCallback()
257 return [this]() { freeEnergyPerturbationData_->updateMDAtoms(); };
260 FreeEnergyPerturbationData::Element::Element(FreeEnergyPerturbationData* freeEnergyPerturbationElement,
261 double deltaLambda) :
262 freeEnergyPerturbationData_(freeEnergyPerturbationElement), doSlowGrowth_(deltaLambda != 0)
266 void FreeEnergyPerturbationData::Element::elementSetup()
268 freeEnergyPerturbationData_->updateMDAtoms();
271 FreeEnergyPerturbationData::Element* FreeEnergyPerturbationData::element()
273 return element_.get();
276 FepStateSetting* FreeEnergyPerturbationData::Element::enableExternalFepStateSetting()
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();
287 void FepStateSetting::signalSettingStep(Step step)
289 nextFepStateSettingStep = step;
292 void FepStateSetting::setNewState(int state, Step step)
295 newFepStateStep = step;
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)
306 return freeEnergyPerturbationData->element();
309 void FreeEnergyPerturbationData::readCheckpointToTrxFrame(t_trxframe* trxFrame,
310 std::optional<ReadCheckpointData> readCheckpointData)
312 if (readCheckpointData)
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_;
321 trxFrame->lambda = 0;
322 trxFrame->fep_state = 0;
324 trxFrame->bLambda = true;
327 const std::string& FreeEnergyPerturbationData::checkpointID()
329 static const std::string identifier = "FreeEnergyPerturbationData";