2 * This file is part of the GROMACS molecular simulation package.
4 * Copyright (c) 2015,2016,2017,2018,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.
38 * Implements the Awh class.
40 * \author Viveca Lindahl
41 * \author Berk Hess <hess@kth.se>
42 * \author Magnus Lundborg
59 #include "gromacs/fileio/enxio.h"
60 #include "gromacs/gmxlib/network.h"
61 #include "gromacs/math/units.h"
62 #include "gromacs/math/utilities.h"
63 #include "gromacs/mdrunutility/multisim.h"
64 #include "gromacs/mdtypes/awh_history.h"
65 #include "gromacs/mdtypes/awh_params.h"
66 #include "gromacs/mdtypes/commrec.h"
67 #include "gromacs/mdtypes/forceoutput.h"
68 #include "gromacs/mdtypes/inputrec.h"
69 #include "gromacs/mdtypes/pull_params.h"
70 #include "gromacs/mdtypes/state.h"
71 #include "gromacs/pbcutil/pbc.h"
72 #include "gromacs/pulling/pull.h"
73 #include "gromacs/timing/wallcycle.h"
74 #include "gromacs/trajectory/energyframe.h"
75 #include "gromacs/utility/arrayref.h"
76 #include "gromacs/utility/exceptions.h"
77 #include "gromacs/utility/gmxassert.h"
78 #include "gromacs/utility/pleasecite.h"
81 #include "biassharing.h"
82 #include "correlationgrid.h"
83 #include "pointstate.h"
89 * \brief A bias and its coupling to the system.
91 * This struct is used to separate the bias machinery in the Bias class,
92 * which should be independent from the reaction coordinate, from the
93 * obtaining of the reaction coordinate values and passing the computed forces.
94 * Currently the AWH method couples to the system by mapping each
95 * AWH bias to a pull coordinate. This can easily be generalized here.
97 struct BiasCoupledToSystem
99 /*! \brief Constructor, couple a bias to a set of pull coordinates.
101 * \param[in] bias The bias.
102 * \param[in] pullCoordIndex The pull coordinate indices.
104 BiasCoupledToSystem(Bias bias, const std::vector<int>& pullCoordIndex);
106 Bias bias_; /**< The bias. */
107 const std::vector<int> pullCoordIndex_; /**< The pull coordinates this bias acts on. */
109 /* Here AWH can be extended to work on other coordinates than pull. */
112 /*! \brief Checks whether any dimension uses pulling as a coordinate provider.
114 * \param[in] awhBiasParams The bias params to check.
115 * \returns true if any dimension of the bias is linked to pulling.
117 static bool anyDimUsesPull(const AwhBiasParams& awhBiasParams)
119 for (const auto& awhDimParam : awhBiasParams.dimParams())
121 if (awhDimParam.coordinateProvider() == AwhCoordinateProviderType::Pull)
129 /*! \brief Checks whether any dimension uses pulling as a coordinate provider.
131 * \param[in] awhParams The AWH params to check.
132 * \returns true if any dimension of awh is linked to pulling.
134 static bool anyDimUsesPull(const AwhParams& awhParams)
136 for (const auto& awhBiasParam : awhParams.awhBiasParams())
138 if (anyDimUsesPull(awhBiasParam))
146 /*! \brief Checks whether any dimension uses pulling as a coordinate provider.
148 * \param[in] biasCoupledToSystem The AWH biases to check.
149 * \returns true if any dimension of the provided biases is linked to pulling.
151 static bool anyDimUsesPull(const ArrayRef<BiasCoupledToSystem> biasCoupledToSystem)
153 for (const auto& biasCts : biasCoupledToSystem)
155 if (!biasCts.pullCoordIndex_.empty())
163 BiasCoupledToSystem::BiasCoupledToSystem(Bias bias, const std::vector<int>& pullCoordIndex) :
164 bias_(std::move(bias)),
165 pullCoordIndex_(pullCoordIndex)
167 /* We already checked for this in grompp, but check again here. */
169 static_cast<size_t>(bias_.ndim()) == pullCoordIndex_.size() + bias_.hasFepLambdaDimension() ? 1 : 0,
170 "The bias dimensionality should match the number of pull and lambda coordinates.");
173 Awh::Awh(FILE* fplog,
174 const t_inputrec& inputRecord,
175 const t_commrec* commRecord,
176 const gmx_multisim_t* multiSimRecord,
177 const AwhParams& awhParams,
178 const std::string& biasInitFilename,
180 int numFepLambdaStates,
181 int fepLambdaState) :
182 seed_(awhParams.seed()),
183 nstout_(awhParams.nstout()),
184 commRecord_(commRecord),
185 multiSimRecord_(multiSimRecord),
188 numFepLambdaStates_(numFepLambdaStates),
189 fepLambdaState_(fepLambdaState)
191 if (anyDimUsesPull(awhParams))
193 GMX_RELEASE_ASSERT(inputRecord.pull != nullptr, "With AWH we should have pull parameters");
194 GMX_RELEASE_ASSERT(pull_work != nullptr,
195 "With AWH pull should be initialized before initializing AWH");
198 if (fplog != nullptr)
200 please_cite(fplog, "Lindahl2014");
203 if (haveBiasSharingWithinSimulation(awhParams))
205 /* This has likely been checked by grompp, but throw anyhow. */
207 InvalidInputError("Biases within a simulation are shared, currently sharing of "
208 "biases is only supported between simulations"));
211 int numSharingSimulations = 1;
212 if (awhParams.shareBiasMultisim() && isMultiSim(multiSimRecord_))
214 numSharingSimulations = multiSimRecord_->numSimulations_;
217 /* Initialize all the biases */
218 const double beta = 1 / (gmx::c_boltz * inputRecord.opts.ref_t[0]);
219 const auto& awhBiasParams = awhParams.awhBiasParams();
220 for (int k = 0; k < gmx::ssize(awhBiasParams); k++)
222 std::vector<int> pullCoordIndex;
223 std::vector<DimParams> dimParams;
224 for (const auto& awhDimParam : awhBiasParams[k].dimParams())
226 if (awhDimParam.coordinateProvider() != AwhCoordinateProviderType::Pull
227 && awhDimParam.coordinateProvider() != AwhCoordinateProviderType::FreeEnergyLambda)
230 InvalidInputError("Currently only the pull code and lambda are supported "
231 "as coordinate providers"));
233 if (awhDimParam.coordinateProvider() == AwhCoordinateProviderType::Pull)
235 const t_pull_coord& pullCoord = inputRecord.pull->coord[awhDimParam.coordinateIndex()];
236 if (pullCoord.eGeom == PullGroupGeometry::DirectionPBC)
238 GMX_THROW(InvalidInputError(
239 "Pull geometry 'direction-periodic' is not supported by AWH"));
241 double conversionFactor = pull_conversion_factor_userinput2internal(pullCoord);
242 pullCoordIndex.push_back(awhDimParam.coordinateIndex());
243 dimParams.emplace_back(DimParams::pullDimParams(
244 conversionFactor, awhDimParam.forceConstant(), beta));
248 dimParams.push_back(DimParams::fepLambdaDimParams(numFepLambdaStates_, beta));
252 /* Construct the bias and couple it to the system. */
253 Bias::ThisRankWillDoIO thisRankWillDoIO =
254 (MASTER(commRecord_) ? Bias::ThisRankWillDoIO::Yes : Bias::ThisRankWillDoIO::No);
255 biasCoupledToSystem_.emplace_back(Bias(k,
261 numSharingSimulations,
266 biasCoupledToSystem_.back().bias_.printInitializationToLog(fplog);
269 /* Need to register the AWH coordinates to be allowed to apply forces to the pull coordinates. */
270 registerAwhWithPull(awhParams, pull_);
272 if (numSharingSimulations > 1 && MASTER(commRecord_))
274 std::vector<size_t> pointSize;
275 for (auto const& biasCts : biasCoupledToSystem_)
277 pointSize.push_back(biasCts.bias_.state().points().size());
279 /* Ensure that the shared biased are compatible between simulations */
280 biasesAreCompatibleForSharingBetweenSimulations(awhParams, pointSize, multiSimRecord_);
284 Awh::~Awh() = default;
286 bool Awh::isOutputStep(int64_t step) const
288 return (nstout_ > 0 && step % nstout_ == 0);
291 real Awh::applyBiasForcesAndUpdateBias(PbcType pbcType,
292 ArrayRef<const real> masses,
293 ArrayRef<const double> neighborLambdaEnergies,
294 ArrayRef<const double> neighborLambdaDhdl,
296 gmx::ForceWithVirial* forceWithVirial,
299 gmx_wallcycle* wallcycle,
302 if (anyDimUsesPull(biasCoupledToSystem_))
304 GMX_ASSERT(forceWithVirial, "Need a valid ForceWithVirial object");
307 wallcycle_start(wallcycle, WallCycleCounter::Awh);
310 set_pbc(&pbc, pbcType, box);
312 /* During the AWH update the potential can instantaneously jump due to either
313 an bias update or moving the umbrella. The jumps are kept track of and
314 subtracted from the potential in order to get a useful conserved energy quantity. */
315 double awhPotential = potentialOffset_;
317 for (auto& biasCts : biasCoupledToSystem_)
319 /* Update the AWH coordinate values with those of the corresponding
322 awh_dvec coordValue = { 0, 0, 0, 0 };
323 int numLambdaDimsCounted = 0;
324 for (int d = 0; d < biasCts.bias_.ndim(); d++)
326 if (biasCts.bias_.dimParams()[d].isPullDimension())
328 coordValue[d] = get_pull_coord_value(
329 pull_, biasCts.pullCoordIndex_[d - numLambdaDimsCounted], &pbc);
333 coordValue[d] = fepLambdaState_;
334 numLambdaDimsCounted += 1;
338 /* Perform an AWH biasing step: this means, at regular intervals,
339 * sampling observables based on the input pull coordinate value,
340 * setting the bias force and/or updating the AWH bias state.
342 double biasPotential;
343 double biasPotentialJump;
344 /* Note: In the near future this call will be split in calls
345 * to supports bias sharing within a single simulation.
347 gmx::ArrayRef<const double> biasForce =
348 biasCts.bias_.calcForceAndUpdateBias(coordValue,
349 neighborLambdaEnergies,
360 awhPotential += biasPotential;
362 /* Keep track of the total potential shift needed to remove the potential jumps. */
363 potentialOffset_ -= biasPotentialJump;
365 /* Communicate the bias force to the pull struct.
366 * The bias potential is returned at the end of this function,
367 * so that it can be added externally to the correct energy data block.
369 numLambdaDimsCounted = 0;
370 for (int d = 0; d < biasCts.bias_.ndim(); d++)
372 if (biasCts.bias_.dimParams()[d].isPullDimension())
374 apply_external_pull_coord_force(pull_,
375 biasCts.pullCoordIndex_[d - numLambdaDimsCounted],
382 int umbrellaGridpointIndex = biasCts.bias_.state().coordState().umbrellaGridpoint();
383 fepLambdaState_ = biasCts.bias_.getGridCoordValue(umbrellaGridpointIndex)[d];
384 numLambdaDimsCounted += 1;
388 if (isOutputStep(step))
390 /* We might have skipped updates for part of the grid points.
391 * Ensure all points are updated before writing out their data.
393 biasCts.bias_.doSkippedUpdatesForAllPoints();
397 wallcycle_stop(wallcycle, WallCycleCounter::Awh);
399 return MASTER(commRecord_) ? static_cast<real>(awhPotential) : 0;
402 std::shared_ptr<AwhHistory> Awh::initHistoryFromState() const
404 if (MASTER(commRecord_))
406 std::shared_ptr<AwhHistory> awhHistory(new AwhHistory);
407 awhHistory->bias.clear();
408 awhHistory->bias.resize(biasCoupledToSystem_.size());
410 for (size_t k = 0; k < awhHistory->bias.size(); k++)
412 biasCoupledToSystem_[k].bias_.initHistoryFromState(&awhHistory->bias[k]);
419 /* Return an empty pointer */
420 return std::shared_ptr<AwhHistory>();
424 void Awh::restoreStateFromHistory(const AwhHistory* awhHistory)
426 /* Restore the history to the current state */
427 if (MASTER(commRecord_))
429 GMX_RELEASE_ASSERT(awhHistory != nullptr,
430 "The master rank should have a valid awhHistory when restoring the "
431 "state from history.");
433 if (awhHistory->bias.size() != biasCoupledToSystem_.size())
435 GMX_THROW(InvalidInputError(
436 "AWH state and history contain different numbers of biases. Likely you "
437 "provided a checkpoint from a different simulation."));
440 potentialOffset_ = awhHistory->potentialOffset;
442 if (PAR(commRecord_))
444 gmx_bcast(sizeof(potentialOffset_), &potentialOffset_, commRecord_->mpi_comm_mygroup);
447 for (size_t k = 0; k < biasCoupledToSystem_.size(); k++)
449 biasCoupledToSystem_[k].bias_.restoreStateFromHistory(
450 awhHistory ? &awhHistory->bias[k] : nullptr, commRecord_);
454 void Awh::updateHistory(AwhHistory* awhHistory) const
456 if (!MASTER(commRecord_))
461 /* This assert will also catch a non-master rank calling this function. */
462 GMX_RELEASE_ASSERT(awhHistory->bias.size() == biasCoupledToSystem_.size(),
463 "AWH state and history bias count should match");
465 awhHistory->potentialOffset = potentialOffset_;
467 for (size_t k = 0; k < awhHistory->bias.size(); k++)
469 biasCoupledToSystem_[k].bias_.updateHistory(&awhHistory->bias[k]);
473 const char* Awh::externalPotentialString()
478 void Awh::registerAwhWithPull(const AwhParams& awhParams, pull_t* pull_work)
480 GMX_RELEASE_ASSERT(!anyDimUsesPull(awhParams) || pull_work, "Need a valid pull object");
482 for (const auto& biasParam : awhParams.awhBiasParams())
484 for (const auto& dimParam : biasParam.dimParams())
486 if (dimParam.coordinateProvider() == AwhCoordinateProviderType::Pull)
488 register_external_pull_potential(
489 pull_work, dimParam.coordinateIndex(), Awh::externalPotentialString());
495 /* Fill the AWH data block of an energy frame with data (if there is any). */
496 void Awh::writeToEnergyFrame(int64_t step, t_enxframe* frame) const
498 GMX_ASSERT(MASTER(commRecord_), "writeToEnergyFrame should only be called on the master rank");
499 GMX_ASSERT(frame != nullptr, "Need a valid energy frame");
501 if (!isOutputStep(step))
503 /* This is not an AWH output step, don't write any AWH data */
507 /* Get the total number of energy subblocks that AWH needs */
508 int numSubblocks = 0;
509 for (auto& biasCoupledToSystem : biasCoupledToSystem_)
511 numSubblocks += biasCoupledToSystem.bias_.numEnergySubblocksToWrite();
513 GMX_ASSERT(numSubblocks > 0, "We should always have data to write");
515 /* Add 1 energy block */
516 add_blocks_enxframe(frame, frame->nblock + 1);
518 /* Take the block that was just added and set the number of subblocks. */
519 t_enxblock* awhEnergyBlock = &(frame->block[frame->nblock - 1]);
520 add_subblocks_enxblock(awhEnergyBlock, numSubblocks);
522 /* Claim it as an AWH block. */
523 awhEnergyBlock->id = enxAWH;
525 /* Transfer AWH data blocks to energy sub blocks */
526 int energySubblockCount = 0;
527 for (auto& biasCoupledToSystem : biasCoupledToSystem_)
529 energySubblockCount += biasCoupledToSystem.bias_.writeToEnergySubblocks(
530 &(awhEnergyBlock->sub[energySubblockCount]));
534 bool Awh::hasFepLambdaDimension() const
537 std::begin(biasCoupledToSystem_),
538 std::end(biasCoupledToSystem_),
539 [](const auto& coupledBias) { return coupledBias.bias_.hasFepLambdaDimension(); });
542 bool Awh::needForeignEnergyDifferences(const int64_t step) const
544 /* If there is no FEP lambda dimension at all in any bias there will be no need for foreign
545 * energy differences */
546 if (!hasFepLambdaDimension())
554 /* Check whether the bias(es) that has/have a FEP lambda dimension should sample coordinates
555 * this step. Since the biases may have different sampleCoordStep it is necessary to check
556 * this combination. */
557 for (const auto& biasCts : biasCoupledToSystem_)
559 if (biasCts.bias_.hasFepLambdaDimension() && biasCts.bias_.isSampleCoordStep(step))
567 std::unique_ptr<Awh> prepareAwhModule(FILE* fplog,
568 const t_inputrec& inputRecord,
569 t_state* stateGlobal,
570 const t_commrec* commRecord,
571 const gmx_multisim_t* multiSimRecord,
572 const bool startingFromCheckpoint,
573 const bool usingShellParticles,
574 const std::string& biasInitFilename,
577 if (!inputRecord.bDoAwh)
581 if (usingShellParticles)
583 GMX_THROW(InvalidInputError("AWH biasing does not support shell particles."));
586 auto awh = std::make_unique<Awh>(fplog,
590 *inputRecord.awhParams,
593 inputRecord.fepvals->n_lambda,
594 inputRecord.fepvals->init_fep_state);
596 if (startingFromCheckpoint)
598 // Restore the AWH history read from checkpoint
599 awh->restoreStateFromHistory(MASTER(commRecord) ? stateGlobal->awhHistory.get() : nullptr);
601 else if (MASTER(commRecord))
603 // Initialize the AWH history here
604 stateGlobal->awhHistory = awh->initHistoryFromState();