src/gromacs/applied_forces/awh/awh.cpp

   1 /*
   2  * This file is part of the GROMACS molecular simulation package.
   3  *
   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.
   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
  36 /*! \internal \file
  37  * \brief
  38  * Implements the Awh class.
  39  *
  40  * \author Viveca Lindahl
  41  * \author Berk Hess <hess@kth.se>
  42  * \author Magnus Lundborg
  43  * \ingroup module_awh
  44  */
  45
  46 #include "gmxpre.h"
  47
  48 #include "awh.h"
  49
  50 #include <cassert>
  51 #include <cmath>
  52 #include <cstdio>
  53 #include <cstdlib>
  54 #include <cstring>
  55
  56 #include <algorithm>
  57
  58 #include "gromacs/fileio/enxio.h"
  59 #include "gromacs/gmxlib/network.h"
  60 #include "gromacs/math/units.h"
  61 #include "gromacs/mdrunutility/multisim.h"
  62 #include "gromacs/mdtypes/awh_history.h"
  63 #include "gromacs/mdtypes/awh_params.h"
  64 #include "gromacs/mdtypes/commrec.h"
  65 #include "gromacs/mdtypes/forceoutput.h"
  66 #include "gromacs/mdtypes/inputrec.h"
  67 #include "gromacs/mdtypes/pull_params.h"
  68 #include "gromacs/mdtypes/state.h"
  69 #include "gromacs/pbcutil/pbc.h"
  70 #include "gromacs/pulling/pull.h"
  71 #include "gromacs/timing/wallcycle.h"
  72 #include "gromacs/trajectory/energyframe.h"
  73 #include "gromacs/utility/arrayref.h"
  74 #include "gromacs/utility/exceptions.h"
  75 #include "gromacs/utility/gmxassert.h"
  76 #include "gromacs/utility/pleasecite.h"
  77
  78 #include "bias.h"
  79 #include "biassharing.h"
  80 #include "correlationgrid.h"
  81 #include "pointstate.h"
  82
  83 namespace gmx
  84 {
  85
  86 /*! \internal
  87  * \brief A bias and its coupling to the system.
  88  *
  89  * This struct is used to separate the bias machinery in the Bias class,
  90  * which should be independent from the reaction coordinate, from the
  91  * obtaining of the reaction coordinate values and passing the computed forces.
  92  * Currently the AWH method couples to the system by mapping each
  93  * AWH bias to a pull coordinate. This can easily be generalized here.
  94  */
  95 struct BiasCoupledToSystem
  96 {
  97     /*! \brief Constructor, couple a bias to a set of pull coordinates.
  98      *
  99      * \param[in] bias            The bias.
 100      * \param[in] pullCoordIndex  The pull coordinate indices.
 101      */
 102     BiasCoupledToSystem(Bias bias, const std::vector<int>& pullCoordIndex);
 103
 104     Bias                   bias_;           /**< The bias. */
 105     const std::vector<int> pullCoordIndex_; /**< The pull coordinates this bias acts on. */
 106
 107     /* Here AWH can be extended to work on other coordinates than pull. */
 108 };
 109
 110 /*! \brief Checks whether any dimension uses the given coordinate provider type.
 111  *
 112  * \param[in] awhBiasParams The bias params to check.
 113  * \param[in] awhCoordProvider The type of coordinate provider
 114  * \returns true if any dimension of the bias is linked to the given provider
 115  */
 116 static bool anyDimUsesProvider(const AwhBiasParams& awhBiasParams, const int awhCoordProvider)
 117 {
 118     for (int d = 0; d < awhBiasParams.ndim; d++)
 119     {
 120         const AwhDimParams& awhDimParams = awhBiasParams.dimParams[d];
 121         if (awhDimParams.eCoordProvider == awhCoordProvider)
 122         {
 123             return true;
 124         }
 125     }
 126     return false;
 127 }
 128
 129 /*! \brief Checks whether any dimension uses the given coordinate provider type.
 130  *
 131  * \param[in] awhParams The AWH params to check.
 132  * \param[in] awhCoordProvider The type of coordinate provider
 133  * \returns true if any dimension of awh is linked to the given provider type.
 134  */
 135 static bool anyDimUsesProvider(const AwhParams& awhParams, const int awhCoordProvider)
 136 {
 137     for (int k = 0; k < awhParams.numBias; k++)
 138     {
 139         const AwhBiasParams& awhBiasParams = awhParams.awhBiasParams[k];
 140         if (anyDimUsesProvider(awhBiasParams, awhCoordProvider))
 141         {
 142             return true;
 143         }
 144     }
 145     return false;
 146 }
 147
 148 /*! \brief Checks whether any dimension uses pulling as a coordinate provider.
 149  *
 150  * \param[in] biasCoupledToSystem The AWH biases to check.
 151  * \returns true if any dimension of the provided biases is linked to pulling.
 152  */
 153 static bool anyDimUsesPull(const ArrayRef<BiasCoupledToSystem> biasCoupledToSystem)
 154 {
 155     for (const auto& biasCts : biasCoupledToSystem)
 156     {
 157         if (!biasCts.pullCoordIndex_.empty())
 158         {
 159             return true;
 160         }
 161     }
 162     return false;
 163 }
 164
 165 BiasCoupledToSystem::BiasCoupledToSystem(Bias bias, const std::vector<int>& pullCoordIndex) :
 166     bias_(std::move(bias)),
 167     pullCoordIndex_(pullCoordIndex)
 168 {
 169     /* We already checked for this in grompp, but check again here. */
 170     GMX_RELEASE_ASSERT(
 171             static_cast<size_t>(bias_.ndim()) == pullCoordIndex_.size() + bias_.hasFepLambdaDimension() ? 1 : 0,
 172             "The bias dimensionality should match the number of pull and lambda coordinates.");
 173 }
 174
 175 Awh::Awh(FILE*                 fplog,
 176          const t_inputrec&     inputRecord,
 177          const t_commrec*      commRecord,
 178          const gmx_multisim_t* multiSimRecord,
 179          const AwhParams&      awhParams,
 180          const std::string&    biasInitFilename,
 181          pull_t*               pull_work,
 182          int                   numFepLambdaStates,
 183          int                   fepLambdaState) :
 184     seed_(awhParams.seed),
 185     nstout_(awhParams.nstOut),
 186     commRecord_(commRecord),
 187     multiSimRecord_(multiSimRecord),
 188     pull_(pull_work),
 189     potentialOffset_(0),
 190     numFepLambdaStates_(numFepLambdaStates),
 191     fepLambdaState_(fepLambdaState)
 192 {
 193     if (anyDimUsesProvider(awhParams, eawhcoordproviderPULL))
 194     {
 195         GMX_RELEASE_ASSERT(inputRecord.pull != nullptr, "With AWH we should have pull parameters");
 196         GMX_RELEASE_ASSERT(pull_work != nullptr,
 197                            "With AWH pull should be initialized before initializing AWH");
 198     }
 199
 200     if (fplog != nullptr)
 201     {
 202         please_cite(fplog, "Lindahl2014");
 203
 204         if (anyDimUsesProvider(awhParams, eawhcoordproviderFREE_ENERGY_LAMBDA))
 205         {
 206             please_cite(fplog, "Lundborg2021");
 207         }
 208     }
 209
 210     if (haveBiasSharingWithinSimulation(awhParams))
 211     {
 212         /* This has likely been checked by grompp, but throw anyhow. */
 213         GMX_THROW(
 214                 InvalidInputError("Biases within a simulation are shared, currently sharing of "
 215                                   "biases is only supported between simulations"));
 216     }
 217
 218     int numSharingSimulations = 1;
 219     if (awhParams.shareBiasMultisim && isMultiSim(multiSimRecord_))
 220     {
 221         numSharingSimulations = multiSimRecord_->numSimulations_;
 222     }
 223
 224     /* Initialize all the biases */
 225     const double beta = 1 / (BOLTZ * inputRecord.opts.ref_t[0]);
 226     for (int k = 0; k < awhParams.numBias; k++)
 227     {
 228         const AwhBiasParams& awhBiasParams = awhParams.awhBiasParams[k];
 229
 230         std::vector<int>       pullCoordIndex;
 231         std::vector<DimParams> dimParams;
 232         for (int d = 0; d < awhBiasParams.ndim; d++)
 233         {
 234             const AwhDimParams& awhDimParams = awhBiasParams.dimParams[d];
 235             if (awhDimParams.eCoordProvider != eawhcoordproviderPULL
 236                 && awhDimParams.eCoordProvider != eawhcoordproviderFREE_ENERGY_LAMBDA)
 237             {
 238                 GMX_THROW(
 239                         InvalidInputError("Currently only the pull code and lambda are supported "
 240                                           "as coordinate providers"));
 241             }
 242             if (awhDimParams.eCoordProvider == eawhcoordproviderPULL)
 243             {
 244                 const t_pull_coord& pullCoord = inputRecord.pull->coord[awhDimParams.coordIndex];
 245                 if (pullCoord.eGeom == epullgDIRPBC)
 246                 {
 247                     GMX_THROW(InvalidInputError(
 248                             "Pull geometry 'direction-periodic' is not supported by AWH"));
 249                 }
 250                 double conversionFactor = pull_coordinate_is_angletype(&pullCoord) ? DEG2RAD : 1;
 251                 pullCoordIndex.push_back(awhDimParams.coordIndex);
 252                 dimParams.push_back(DimParams::pullDimParams(conversionFactor,
 253                                                              awhDimParams.forceConstant, beta));
 254             }
 255             else
 256             {
 257                 dimParams.push_back(DimParams::fepLambdaDimParams(numFepLambdaStates_, beta));
 258             }
 259         }
 260
 261         /* Construct the bias and couple it to the system. */
 262         Bias::ThisRankWillDoIO thisRankWillDoIO =
 263                 (MASTER(commRecord_) ? Bias::ThisRankWillDoIO::Yes : Bias::ThisRankWillDoIO::No);
 264         biasCoupledToSystem_.emplace_back(
 265                 Bias(k, awhParams, awhParams.awhBiasParams[k], dimParams, beta, inputRecord.delta_t,
 266                      numSharingSimulations, biasInitFilename, thisRankWillDoIO),
 267                 pullCoordIndex);
 268
 269         biasCoupledToSystem_.back().bias_.printInitializationToLog(fplog);
 270     }
 271
 272     /* Need to register the AWH coordinates to be allowed to apply forces to the pull coordinates. */
 273     registerAwhWithPull(awhParams, pull_);
 274
 275     if (numSharingSimulations > 1 && MASTER(commRecord_))
 276     {
 277         std::vector<size_t> pointSize;
 278         for (auto const& biasCts : biasCoupledToSystem_)
 279         {
 280             pointSize.push_back(biasCts.bias_.state().points().size());
 281         }
 282         /* Ensure that the shared biased are compatible between simulations */
 283         biasesAreCompatibleForSharingBetweenSimulations(awhParams, pointSize, multiSimRecord_);
 284     }
 285 }
 286
 287 Awh::~Awh() = default;
 288
 289 bool Awh::isOutputStep(int64_t step) const
 290 {
 291     return (nstout_ > 0 && step % nstout_ == 0);
 292 }
 293
 294 real Awh::applyBiasForcesAndUpdateBias(PbcType                pbcType,
 295                                        const real*            masses,
 296                                        ArrayRef<const double> neighborLambdaEnergies,
 297                                        ArrayRef<const double> neighborLambdaDhdl,
 298                                        const matrix           box,
 299                                        gmx::ForceWithVirial*  forceWithVirial,
 300                                        double                 t,
 301                                        int64_t                step,
 302                                        gmx_wallcycle*         wallcycle,
 303                                        FILE*                  fplog)
 304 {
 305     if (anyDimUsesPull(biasCoupledToSystem_))
 306     {
 307         GMX_ASSERT(forceWithVirial, "Need a valid ForceWithVirial object");
 308     }
 309
 310     wallcycle_start(wallcycle, ewcAWH);
 311
 312     t_pbc pbc;
 313     set_pbc(&pbc, pbcType, box);
 314
 315     /* During the AWH update the potential can instantaneously jump due to either
 316        an bias update or moving the umbrella. The jumps are kept track of and
 317        subtracted from the potential in order to get a useful conserved energy quantity. */
 318     double awhPotential = potentialOffset_;
 319
 320     for (auto& biasCts : biasCoupledToSystem_)
 321     {
 322         /* Update the AWH coordinate values with those of the corresponding
 323          * pull coordinates.
 324          */
 325         awh_dvec coordValue           = { 0, 0, 0, 0 };
 326         int      numLambdaDimsCounted = 0;
 327         for (int d = 0; d < biasCts.bias_.ndim(); d++)
 328         {
 329             if (biasCts.bias_.dimParams()[d].isPullDimension())
 330             {
 331                 coordValue[d] = get_pull_coord_value(
 332                         pull_, biasCts.pullCoordIndex_[d - numLambdaDimsCounted], &pbc);
 333             }
 334             else
 335             {
 336                 coordValue[d] = fepLambdaState_;
 337                 numLambdaDimsCounted += 1;
 338             }
 339         }
 340
 341         /* Perform an AWH biasing step: this means, at regular intervals,
 342          * sampling observables based on the input pull coordinate value,
 343          * setting the bias force and/or updating the AWH bias state.
 344          */
 345         double biasPotential;
 346         double biasPotentialJump;
 347         /* Note: In the near future this call will be split in calls
 348          *       to supports bias sharing within a single simulation.
 349          */
 350         gmx::ArrayRef<const double> biasForce = biasCts.bias_.calcForceAndUpdateBias(
 351                 coordValue, neighborLambdaEnergies, neighborLambdaDhdl, &biasPotential,
 352                 &biasPotentialJump, commRecord_, multiSimRecord_, t, step, seed_, fplog);
 353
 354         awhPotential += biasPotential;
 355
 356         /* Keep track of the total potential shift needed to remove the potential jumps. */
 357         potentialOffset_ -= biasPotentialJump;
 358
 359         /* Communicate the bias force to the pull struct.
 360          * The bias potential is returned at the end of this function,
 361          * so that it can be added externally to the correct energy data block.
 362          */
 363         numLambdaDimsCounted = 0;
 364         for (int d = 0; d < biasCts.bias_.ndim(); d++)
 365         {
 366             if (biasCts.bias_.dimParams()[d].isPullDimension())
 367             {
 368                 apply_external_pull_coord_force(pull_, biasCts.pullCoordIndex_[d - numLambdaDimsCounted],
 369                                                 biasForce[d], masses, forceWithVirial);
 370             }
 371             else
 372             {
 373                 int umbrellaGridpointIndex = biasCts.bias_.state().coordState().umbrellaGridpoint();
 374                 fepLambdaState_ = biasCts.bias_.getGridCoordValue(umbrellaGridpointIndex)[d];
 375                 numLambdaDimsCounted += 1;
 376             }
 377         }
 378
 379         if (isOutputStep(step))
 380         {
 381             /* We might have skipped updates for part of the grid points.
 382              * Ensure all points are updated before writing out their data.
 383              */
 384             biasCts.bias_.doSkippedUpdatesForAllPoints();
 385         }
 386     }
 387
 388     wallcycle_stop(wallcycle, ewcAWH);
 389
 390     return MASTER(commRecord_) ? static_cast<real>(awhPotential) : 0;
 391 }
 392
 393 std::shared_ptr<AwhHistory> Awh::initHistoryFromState() const
 394 {
 395     if (MASTER(commRecord_))
 396     {
 397         std::shared_ptr<AwhHistory> awhHistory(new AwhHistory);
 398         awhHistory->bias.clear();
 399         awhHistory->bias.resize(biasCoupledToSystem_.size());
 400
 401         for (size_t k = 0; k < awhHistory->bias.size(); k++)
 402         {
 403             biasCoupledToSystem_[k].bias_.initHistoryFromState(&awhHistory->bias[k]);
 404         }
 405
 406         return awhHistory;
 407     }
 408     else
 409     {
 410         /* Return an empty pointer */
 411         return std::shared_ptr<AwhHistory>();
 412     }
 413 }
 414
 415 void Awh::restoreStateFromHistory(const AwhHistory* awhHistory)
 416 {
 417     /* Restore the history to the current state */
 418     if (MASTER(commRecord_))
 419     {
 420         GMX_RELEASE_ASSERT(awhHistory != nullptr,
 421                            "The master rank should have a valid awhHistory when restoring the "
 422                            "state from history.");
 423
 424         if (awhHistory->bias.size() != biasCoupledToSystem_.size())
 425         {
 426             GMX_THROW(InvalidInputError(
 427                     "AWH state and history contain different numbers of biases. Likely you "
 428                     "provided a checkpoint from a different simulation."));
 429         }
 430
 431         potentialOffset_ = awhHistory->potentialOffset;
 432     }
 433     if (PAR(commRecord_))
 434     {
 435         gmx_bcast(sizeof(potentialOffset_), &potentialOffset_, commRecord_->mpi_comm_mygroup);
 436     }
 437
 438     for (size_t k = 0; k < biasCoupledToSystem_.size(); k++)
 439     {
 440         biasCoupledToSystem_[k].bias_.restoreStateFromHistory(
 441                 awhHistory ? &awhHistory->bias[k] : nullptr, commRecord_);
 442     }
 443 }
 444
 445 void Awh::updateHistory(AwhHistory* awhHistory) const
 446 {
 447     if (!MASTER(commRecord_))
 448     {
 449         return;
 450     }
 451
 452     /* This assert will also catch a non-master rank calling this function. */
 453     GMX_RELEASE_ASSERT(awhHistory->bias.size() == biasCoupledToSystem_.size(),
 454                        "AWH state and history bias count should match");
 455
 456     awhHistory->potentialOffset = potentialOffset_;
 457
 458     for (size_t k = 0; k < awhHistory->bias.size(); k++)
 459     {
 460         biasCoupledToSystem_[k].bias_.updateHistory(&awhHistory->bias[k]);
 461     }
 462 }
 463
 464 const char* Awh::externalPotentialString()
 465 {
 466     return "AWH";
 467 }
 468
 469 void Awh::registerAwhWithPull(const AwhParams& awhParams, pull_t* pull_work)
 470 {
 471     GMX_RELEASE_ASSERT(!anyDimUsesProvider(awhParams, eawhcoordproviderPULL) || pull_work,
 472                        "Need a valid pull object");
 473
 474     for (int k = 0; k < awhParams.numBias; k++)
 475     {
 476         const AwhBiasParams& biasParams = awhParams.awhBiasParams[k];
 477
 478         for (int d = 0; d < biasParams.ndim; d++)
 479         {
 480             if (biasParams.dimParams[d].eCoordProvider == eawhcoordproviderPULL)
 481             {
 482                 register_external_pull_potential(pull_work, biasParams.dimParams[d].coordIndex,
 483                                                  Awh::externalPotentialString());
 484             }
 485         }
 486     }
 487 }
 488
 489 /* Fill the AWH data block of an energy frame with data (if there is any). */
 490 void Awh::writeToEnergyFrame(int64_t step, t_enxframe* frame) const
 491 {
 492     GMX_ASSERT(MASTER(commRecord_), "writeToEnergyFrame should only be called on the master rank");
 493     GMX_ASSERT(frame != nullptr, "Need a valid energy frame");
 494
 495     if (!isOutputStep(step))
 496     {
 497         /* This is not an AWH output step, don't write any AWH data */
 498         return;
 499     }
 500
 501     /* Get the total number of energy subblocks that AWH needs */
 502     int numSubblocks = 0;
 503     for (auto& biasCoupledToSystem : biasCoupledToSystem_)
 504     {
 505         numSubblocks += biasCoupledToSystem.bias_.numEnergySubblocksToWrite();
 506     }
 507     GMX_ASSERT(numSubblocks > 0, "We should always have data to write");
 508
 509     /* Add 1 energy block */
 510     add_blocks_enxframe(frame, frame->nblock + 1);
 511
 512     /* Take the block that was just added and set the number of subblocks. */
 513     t_enxblock* awhEnergyBlock = &(frame->block[frame->nblock - 1]);
 514     add_subblocks_enxblock(awhEnergyBlock, numSubblocks);
 515
 516     /* Claim it as an AWH block. */
 517     awhEnergyBlock->id = enxAWH;
 518
 519     /* Transfer AWH data blocks to energy sub blocks */
 520     int energySubblockCount = 0;
 521     for (auto& biasCoupledToSystem : biasCoupledToSystem_)
 522     {
 523         energySubblockCount += biasCoupledToSystem.bias_.writeToEnergySubblocks(
 524                 &(awhEnergyBlock->sub[energySubblockCount]));
 525     }
 526 }
 527
 528 bool Awh::hasFepLambdaDimension() const
 529 {
 530     return std::any_of(
 531             std::begin(biasCoupledToSystem_), std::end(biasCoupledToSystem_),
 532             [](const auto& coupledBias) { return coupledBias.bias_.hasFepLambdaDimension(); });
 533 }
 534
 535 bool Awh::needForeignEnergyDifferences(const int64_t step) const
 536 {
 537     /* If there is no FEP lambda dimension at all in any bias there will be no need for foreign
 538      * energy differences */
 539     if (!hasFepLambdaDimension())
 540     {
 541         return false;
 542     }
 543     if (step == 0)
 544     {
 545         return true;
 546     }
 547     /* Check whether the bias(es) that has/have a FEP lambda dimension should sample coordinates
 548      * this step. Since the biases may have different sampleCoordStep it is necessary to check
 549      * this combination. */
 550     for (const auto& biasCts : biasCoupledToSystem_)
 551     {
 552         if (biasCts.bias_.hasFepLambdaDimension() && biasCts.bias_.isSampleCoordStep(step))
 553         {
 554             return true;
 555         }
 556     }
 557     return false;
 558 }
 559
 560 std::unique_ptr<Awh> prepareAwhModule(FILE*                 fplog,
 561                                       const t_inputrec&     inputRecord,
 562                                       t_state*              stateGlobal,
 563                                       const t_commrec*      commRecord,
 564                                       const gmx_multisim_t* multiSimRecord,
 565                                       const bool            startingFromCheckpoint,
 566                                       const bool            usingShellParticles,
 567                                       const std::string&    biasInitFilename,
 568                                       pull_t*               pull_work)
 569 {
 570     if (!inputRecord.bDoAwh)
 571     {
 572         return nullptr;
 573     }
 574     if (usingShellParticles)
 575     {
 576         GMX_THROW(InvalidInputError("AWH biasing does not support shell particles."));
 577     }
 578
 579     auto awh = std::make_unique<Awh>(
 580             fplog, inputRecord, commRecord, multiSimRecord, *inputRecord.awhParams, biasInitFilename,
 581             pull_work, inputRecord.fepvals->n_lambda, inputRecord.fepvals->init_fep_state);
 582
 583     if (startingFromCheckpoint)
 584     {
 585         // Restore the AWH history read from checkpoint
 586         awh->restoreStateFromHistory(MASTER(commRecord) ? stateGlobal->awhHistory.get() : nullptr);
 587     }
 588     else if (MASTER(commRecord))
 589     {
 590         // Initialize the AWH history here
 591         stateGlobal->awhHistory = awh->initHistoryFromState();
 592     }
 593     return awh;
 594 }
 595
 596 } // namespace gmx