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)
120 awhBiasParams.dimParams().begin(), awhBiasParams.dimParams().end(), [](const auto& awhDimParam) {
121 return awhDimParam.coordinateProvider() == AwhCoordinateProviderType::Pull;
125 /*! \brief Checks whether any dimension uses pulling as a coordinate provider.
127 * \param[in] awhParams The AWH params to check.
128 * \returns true if any dimension of awh is linked to pulling.
130 static bool anyDimUsesPull(const AwhParams& awhParams)
132 return std::any_of(awhParams.awhBiasParams().begin(),
133 awhParams.awhBiasParams().end(),
134 [](const auto& awhBiasParam) { return anyDimUsesPull(awhBiasParam); });
137 /*! \brief Checks whether any dimension uses pulling as a coordinate provider.
139 * \param[in] biasCoupledToSystem The AWH biases to check.
140 * \returns true if any dimension of the provided biases is linked to pulling.
142 static bool anyDimUsesPull(const ArrayRef<BiasCoupledToSystem> biasCoupledToSystem)
144 return std::any_of(biasCoupledToSystem.begin(), biasCoupledToSystem.end(), [](const auto& biasCts) {
145 return !biasCts.pullCoordIndex_.empty();
149 BiasCoupledToSystem::BiasCoupledToSystem(Bias bias, const std::vector<int>& pullCoordIndex) :
150 bias_(std::move(bias)), pullCoordIndex_(pullCoordIndex)
152 /* We already checked for this in grompp, but check again here. */
154 static_cast<size_t>(bias_.ndim()) == pullCoordIndex_.size() + bias_.hasFepLambdaDimension() ? 1 : 0,
155 "The bias dimensionality should match the number of pull and lambda coordinates.");
158 Awh::Awh(FILE* fplog,
159 const t_inputrec& inputRecord,
160 const t_commrec* commRecord,
161 const gmx_multisim_t* multiSimRecord,
162 const AwhParams& awhParams,
163 const std::string& biasInitFilename,
165 int numFepLambdaStates,
166 int fepLambdaState) :
167 seed_(awhParams.seed()),
168 nstout_(awhParams.nstout()),
169 commRecord_(commRecord),
170 multiSimRecord_(multiSimRecord),
173 numFepLambdaStates_(numFepLambdaStates),
174 fepLambdaState_(fepLambdaState)
176 if (anyDimUsesPull(awhParams))
178 GMX_RELEASE_ASSERT(inputRecord.pull != nullptr, "With AWH we should have pull parameters");
179 GMX_RELEASE_ASSERT(pull_work != nullptr,
180 "With AWH pull should be initialized before initializing AWH");
183 if (fplog != nullptr)
185 please_cite(fplog, "Lindahl2014");
188 if (haveBiasSharingWithinSimulation(awhParams))
190 /* This has likely been checked by grompp, but throw anyhow. */
192 InvalidInputError("Biases within a simulation are shared, currently sharing of "
193 "biases is only supported between simulations"));
196 int numSharingSimulations = 1;
197 if (awhParams.shareBiasMultisim() && isMultiSim(multiSimRecord_))
199 numSharingSimulations = multiSimRecord_->numSimulations_;
202 /* Initialize all the biases */
203 const double beta = 1 / (gmx::c_boltz * inputRecord.opts.ref_t[0]);
204 const auto& awhBiasParams = awhParams.awhBiasParams();
205 for (int k = 0; k < gmx::ssize(awhBiasParams); k++)
207 std::vector<int> pullCoordIndex;
208 std::vector<DimParams> dimParams;
209 for (const auto& awhDimParam : awhBiasParams[k].dimParams())
211 if (awhDimParam.coordinateProvider() != AwhCoordinateProviderType::Pull
212 && awhDimParam.coordinateProvider() != AwhCoordinateProviderType::FreeEnergyLambda)
215 InvalidInputError("Currently only the pull code and lambda are supported "
216 "as coordinate providers"));
218 if (awhDimParam.coordinateProvider() == AwhCoordinateProviderType::Pull)
220 const t_pull_coord& pullCoord = inputRecord.pull->coord[awhDimParam.coordinateIndex()];
221 if (pullCoord.eGeom == PullGroupGeometry::DirectionPBC)
223 GMX_THROW(InvalidInputError(
224 "Pull geometry 'direction-periodic' is not supported by AWH"));
226 double conversionFactor = pull_conversion_factor_userinput2internal(pullCoord);
227 pullCoordIndex.push_back(awhDimParam.coordinateIndex());
228 dimParams.emplace_back(DimParams::pullDimParams(
229 conversionFactor, awhDimParam.forceConstant(), beta));
233 dimParams.push_back(DimParams::fepLambdaDimParams(numFepLambdaStates_, beta));
237 /* Construct the bias and couple it to the system. */
238 Bias::ThisRankWillDoIO thisRankWillDoIO =
239 (MASTER(commRecord_) ? Bias::ThisRankWillDoIO::Yes : Bias::ThisRankWillDoIO::No);
240 biasCoupledToSystem_.emplace_back(Bias(k,
246 numSharingSimulations,
251 biasCoupledToSystem_.back().bias_.printInitializationToLog(fplog);
254 /* Need to register the AWH coordinates to be allowed to apply forces to the pull coordinates. */
255 registerAwhWithPull(awhParams, pull_);
257 if (numSharingSimulations > 1 && MASTER(commRecord_))
259 std::vector<size_t> pointSize;
260 for (auto const& biasCts : biasCoupledToSystem_)
262 pointSize.push_back(biasCts.bias_.state().points().size());
264 /* Ensure that the shared biased are compatible between simulations */
265 biasesAreCompatibleForSharingBetweenSimulations(awhParams, pointSize, multiSimRecord_);
269 Awh::~Awh() = default;
271 bool Awh::isOutputStep(int64_t step) const
273 return (nstout_ > 0 && step % nstout_ == 0);
276 real Awh::applyBiasForcesAndUpdateBias(PbcType pbcType,
277 ArrayRef<const real> masses,
278 ArrayRef<const double> neighborLambdaEnergies,
279 ArrayRef<const double> neighborLambdaDhdl,
281 gmx::ForceWithVirial* forceWithVirial,
284 gmx_wallcycle* wallcycle,
287 if (anyDimUsesPull(biasCoupledToSystem_))
289 GMX_ASSERT(forceWithVirial, "Need a valid ForceWithVirial object");
292 wallcycle_start(wallcycle, WallCycleCounter::Awh);
295 set_pbc(&pbc, pbcType, box);
297 /* During the AWH update the potential can instantaneously jump due to either
298 an bias update or moving the umbrella. The jumps are kept track of and
299 subtracted from the potential in order to get a useful conserved energy quantity. */
300 double awhPotential = potentialOffset_;
302 for (auto& biasCts : biasCoupledToSystem_)
304 /* Update the AWH coordinate values with those of the corresponding
307 awh_dvec coordValue = { 0, 0, 0, 0 };
308 int numLambdaDimsCounted = 0;
309 for (int d = 0; d < biasCts.bias_.ndim(); d++)
311 if (biasCts.bias_.dimParams()[d].isPullDimension())
313 coordValue[d] = get_pull_coord_value(
314 pull_, biasCts.pullCoordIndex_[d - numLambdaDimsCounted], pbc);
318 coordValue[d] = fepLambdaState_;
319 numLambdaDimsCounted += 1;
323 /* Perform an AWH biasing step: this means, at regular intervals,
324 * sampling observables based on the input pull coordinate value,
325 * setting the bias force and/or updating the AWH bias state.
327 double biasPotential;
328 double biasPotentialJump;
329 /* Note: In the near future this call will be split in calls
330 * to supports bias sharing within a single simulation.
332 gmx::ArrayRef<const double> biasForce =
333 biasCts.bias_.calcForceAndUpdateBias(coordValue,
334 neighborLambdaEnergies,
345 awhPotential += biasPotential;
347 /* Keep track of the total potential shift needed to remove the potential jumps. */
348 potentialOffset_ -= biasPotentialJump;
350 /* Communicate the bias force to the pull struct.
351 * The bias potential is returned at the end of this function,
352 * so that it can be added externally to the correct energy data block.
354 numLambdaDimsCounted = 0;
355 for (int d = 0; d < biasCts.bias_.ndim(); d++)
357 if (biasCts.bias_.dimParams()[d].isPullDimension())
359 apply_external_pull_coord_force(pull_,
360 biasCts.pullCoordIndex_[d - numLambdaDimsCounted],
367 int umbrellaGridpointIndex = biasCts.bias_.state().coordState().umbrellaGridpoint();
368 fepLambdaState_ = biasCts.bias_.getGridCoordValue(umbrellaGridpointIndex)[d];
369 numLambdaDimsCounted += 1;
373 if (isOutputStep(step))
375 /* We might have skipped updates for part of the grid points.
376 * Ensure all points are updated before writing out their data.
378 biasCts.bias_.doSkippedUpdatesForAllPoints();
382 wallcycle_stop(wallcycle, WallCycleCounter::Awh);
384 return MASTER(commRecord_) ? static_cast<real>(awhPotential) : 0;
387 std::shared_ptr<AwhHistory> Awh::initHistoryFromState() const
389 if (MASTER(commRecord_))
391 std::shared_ptr<AwhHistory> awhHistory(new AwhHistory);
392 awhHistory->bias.clear();
393 awhHistory->bias.resize(biasCoupledToSystem_.size());
395 for (size_t k = 0; k < awhHistory->bias.size(); k++)
397 biasCoupledToSystem_[k].bias_.initHistoryFromState(&awhHistory->bias[k]);
404 /* Return an empty pointer */
405 return std::shared_ptr<AwhHistory>();
409 void Awh::restoreStateFromHistory(const AwhHistory* awhHistory)
411 /* Restore the history to the current state */
412 if (MASTER(commRecord_))
414 GMX_RELEASE_ASSERT(awhHistory != nullptr,
415 "The master rank should have a valid awhHistory when restoring the "
416 "state from history.");
418 if (awhHistory->bias.size() != biasCoupledToSystem_.size())
420 GMX_THROW(InvalidInputError(
421 "AWH state and history contain different numbers of biases. Likely you "
422 "provided a checkpoint from a different simulation."));
425 potentialOffset_ = awhHistory->potentialOffset;
427 if (PAR(commRecord_))
429 gmx_bcast(sizeof(potentialOffset_), &potentialOffset_, commRecord_->mpi_comm_mygroup);
432 for (size_t k = 0; k < biasCoupledToSystem_.size(); k++)
434 biasCoupledToSystem_[k].bias_.restoreStateFromHistory(
435 awhHistory ? &awhHistory->bias[k] : nullptr, commRecord_);
439 void Awh::updateHistory(AwhHistory* awhHistory) const
441 if (!MASTER(commRecord_))
446 /* This assert will also catch a non-master rank calling this function. */
447 GMX_RELEASE_ASSERT(awhHistory->bias.size() == biasCoupledToSystem_.size(),
448 "AWH state and history bias count should match");
450 awhHistory->potentialOffset = potentialOffset_;
452 for (size_t k = 0; k < awhHistory->bias.size(); k++)
454 biasCoupledToSystem_[k].bias_.updateHistory(&awhHistory->bias[k]);
458 const char* Awh::externalPotentialString()
463 void Awh::registerAwhWithPull(const AwhParams& awhParams, pull_t* pull_work)
465 GMX_RELEASE_ASSERT(!anyDimUsesPull(awhParams) || pull_work, "Need a valid pull object");
467 for (const auto& biasParam : awhParams.awhBiasParams())
469 for (const auto& dimParam : biasParam.dimParams())
471 if (dimParam.coordinateProvider() == AwhCoordinateProviderType::Pull)
473 register_external_pull_potential(
474 pull_work, dimParam.coordinateIndex(), Awh::externalPotentialString());
480 /* Fill the AWH data block of an energy frame with data (if there is any). */
481 void Awh::writeToEnergyFrame(int64_t step, t_enxframe* frame) const
483 GMX_ASSERT(MASTER(commRecord_), "writeToEnergyFrame should only be called on the master rank");
484 GMX_ASSERT(frame != nullptr, "Need a valid energy frame");
486 if (!isOutputStep(step))
488 /* This is not an AWH output step, don't write any AWH data */
492 /* Get the total number of energy subblocks that AWH needs */
493 int numSubblocks = 0;
494 for (const auto& biasCoupledToSystem : biasCoupledToSystem_)
496 numSubblocks += biasCoupledToSystem.bias_.numEnergySubblocksToWrite();
498 GMX_ASSERT(numSubblocks > 0, "We should always have data to write");
500 /* Add 1 energy block */
501 add_blocks_enxframe(frame, frame->nblock + 1);
503 /* Take the block that was just added and set the number of subblocks. */
504 t_enxblock* awhEnergyBlock = &(frame->block[frame->nblock - 1]);
505 add_subblocks_enxblock(awhEnergyBlock, numSubblocks);
507 /* Claim it as an AWH block. */
508 awhEnergyBlock->id = enxAWH;
510 /* Transfer AWH data blocks to energy sub blocks */
511 int energySubblockCount = 0;
512 for (const auto& biasCoupledToSystem : biasCoupledToSystem_)
514 energySubblockCount += biasCoupledToSystem.bias_.writeToEnergySubblocks(
515 &(awhEnergyBlock->sub[energySubblockCount]));
519 bool Awh::hasFepLambdaDimension() const
522 std::begin(biasCoupledToSystem_),
523 std::end(biasCoupledToSystem_),
524 [](const auto& coupledBias) { return coupledBias.bias_.hasFepLambdaDimension(); });
527 bool Awh::needForeignEnergyDifferences(const int64_t step) const
529 /* If there is no FEP lambda dimension at all in any bias there will be no need for foreign
530 * energy differences */
531 if (!hasFepLambdaDimension())
539 /* Check whether the bias(es) that has/have a FEP lambda dimension should sample coordinates
540 * this step. Since the biases may have different sampleCoordStep it is necessary to check
541 * this combination. */
542 return std::any_of(biasCoupledToSystem_.begin(), biasCoupledToSystem_.end(), [step](const auto& biasCts) {
543 return biasCts.bias_.hasFepLambdaDimension() && biasCts.bias_.isSampleCoordStep(step);
547 std::unique_ptr<Awh> prepareAwhModule(FILE* fplog,
548 const t_inputrec& inputRecord,
549 t_state* stateGlobal,
550 const t_commrec* commRecord,
551 const gmx_multisim_t* multiSimRecord,
552 const bool startingFromCheckpoint,
553 const bool usingShellParticles,
554 const std::string& biasInitFilename,
557 if (!inputRecord.bDoAwh)
561 if (usingShellParticles)
563 GMX_THROW(InvalidInputError("AWH biasing does not support shell particles."));
566 auto awh = std::make_unique<Awh>(fplog,
570 *inputRecord.awhParams,
573 inputRecord.fepvals->n_lambda,
574 inputRecord.fepvals->init_fep_state);
576 if (startingFromCheckpoint)
578 // Restore the AWH history read from checkpoint
579 awh->restoreStateFromHistory(MASTER(commRecord) ? stateGlobal->awhHistory.get() : nullptr);
581 else if (MASTER(commRecord))
583 // Initialize the AWH history here
584 stateGlobal->awhHistory = awh->initHistoryFromState();