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 #include <vector>
  58
  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"
  79
  80 #include "bias.h"
  81 #include "biassharing.h"
  82 #include "correlationgrid.h"
  83 #include "pointstate.h"
  84
  85 namespace gmx
  86 {
  87
  88 /*! \internal
  89  * \brief A bias and its coupling to the system.
  90  *
  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.
  96  */
  97 struct BiasCoupledToSystem
  98 {
  99     /*! \brief Constructor, couple a bias to a set of pull coordinates.
 100      *
 101      * \param[in] bias            The bias.
 102      * \param[in] pullCoordIndex  The pull coordinate indices.
 103      */
 104     BiasCoupledToSystem(Bias bias, const std::vector<int>& pullCoordIndex);
 105
 106     Bias                   bias_;           /**< The bias. */
 107     const std::vector<int> pullCoordIndex_; /**< The pull coordinates this bias acts on. */
 108
 109     /* Here AWH can be extended to work on other coordinates than pull. */
 110 };
 111
 112 /*! \brief Checks whether any dimension uses pulling as a coordinate provider.
 113  *
 114  * \param[in] awhBiasParams The bias params to check.
 115  * \returns true if any dimension of the bias is linked to pulling.
 116  */
 117 static bool anyDimUsesPull(const AwhBiasParams& awhBiasParams)
 118 {
 119     return std::any_of(
 120             awhBiasParams.dimParams().begin(), awhBiasParams.dimParams().end(), [](const auto& awhDimParam) {
 121                 return awhDimParam.coordinateProvider() == AwhCoordinateProviderType::Pull;
 122             });
 123 }
 124
 125 /*! \brief Checks whether any dimension uses pulling as a coordinate provider.
 126  *
 127  * \param[in] awhParams The AWH params to check.
 128  * \returns true if any dimension of awh is linked to pulling.
 129  */
 130 static bool anyDimUsesPull(const AwhParams& awhParams)
 131 {
 132     return std::any_of(awhParams.awhBiasParams().begin(),
 133                        awhParams.awhBiasParams().end(),
 134                        [](const auto& awhBiasParam) { return anyDimUsesPull(awhBiasParam); });
 135 }
 136
 137 /*! \brief Checks whether any dimension uses pulling as a coordinate provider.
 138  *
 139  * \param[in] biasCoupledToSystem The AWH biases to check.
 140  * \returns true if any dimension of the provided biases is linked to pulling.
 141  */
 142 static bool anyDimUsesPull(const ArrayRef<BiasCoupledToSystem> biasCoupledToSystem)
 143 {
 144     return std::any_of(biasCoupledToSystem.begin(), biasCoupledToSystem.end(), [](const auto& biasCts) {
 145         return !biasCts.pullCoordIndex_.empty();
 146     });
 147 }
 148
 149 BiasCoupledToSystem::BiasCoupledToSystem(Bias bias, const std::vector<int>& pullCoordIndex) :
 150     bias_(std::move(bias)), pullCoordIndex_(pullCoordIndex)
 151 {
 152     /* We already checked for this in grompp, but check again here. */
 153     GMX_RELEASE_ASSERT(
 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.");
 156 }
 157
 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,
 164          pull_t*               pull_work,
 165          int                   numFepLambdaStates,
 166          int                   fepLambdaState) :
 167     seed_(awhParams.seed()),
 168     nstout_(awhParams.nstout()),
 169     commRecord_(commRecord),
 170     multiSimRecord_(multiSimRecord),
 171     pull_(pull_work),
 172     potentialOffset_(0),
 173     numFepLambdaStates_(numFepLambdaStates),
 174     fepLambdaState_(fepLambdaState)
 175 {
 176     if (anyDimUsesPull(awhParams))
 177     {
 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");
 181     }
 182
 183     if (fplog != nullptr)
 184     {
 185         please_cite(fplog, "Lindahl2014");
 186     }
 187
 188     if (haveBiasSharingWithinSimulation(awhParams))
 189     {
 190         /* This has likely been checked by grompp, but throw anyhow. */
 191         GMX_THROW(
 192                 InvalidInputError("Biases within a simulation are shared, currently sharing of "
 193                                   "biases is only supported between simulations"));
 194     }
 195
 196     int numSharingSimulations = 1;
 197     if (awhParams.shareBiasMultisim() && isMultiSim(multiSimRecord_))
 198     {
 199         numSharingSimulations = multiSimRecord_->numSimulations_;
 200     }
 201
 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++)
 206     {
 207         std::vector<int>       pullCoordIndex;
 208         std::vector<DimParams> dimParams;
 209         for (const auto& awhDimParam : awhBiasParams[k].dimParams())
 210         {
 211             if (awhDimParam.coordinateProvider() != AwhCoordinateProviderType::Pull
 212                 && awhDimParam.coordinateProvider() != AwhCoordinateProviderType::FreeEnergyLambda)
 213             {
 214                 GMX_THROW(
 215                         InvalidInputError("Currently only the pull code and lambda are supported "
 216                                           "as coordinate providers"));
 217             }
 218             if (awhDimParam.coordinateProvider() == AwhCoordinateProviderType::Pull)
 219             {
 220                 const t_pull_coord& pullCoord = inputRecord.pull->coord[awhDimParam.coordinateIndex()];
 221                 if (pullCoord.eGeom == PullGroupGeometry::DirectionPBC)
 222                 {
 223                     GMX_THROW(InvalidInputError(
 224                             "Pull geometry 'direction-periodic' is not supported by AWH"));
 225                 }
 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));
 230             }
 231             else
 232             {
 233                 dimParams.push_back(DimParams::fepLambdaDimParams(numFepLambdaStates_, beta));
 234             }
 235         }
 236
 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,
 241                                                awhParams,
 242                                                awhBiasParams[k],
 243                                                dimParams,
 244                                                beta,
 245                                                inputRecord.delta_t,
 246                                                numSharingSimulations,
 247                                                biasInitFilename,
 248                                                thisRankWillDoIO),
 249                                           pullCoordIndex);
 250
 251         biasCoupledToSystem_.back().bias_.printInitializationToLog(fplog);
 252     }
 253
 254     /* Need to register the AWH coordinates to be allowed to apply forces to the pull coordinates. */
 255     registerAwhWithPull(awhParams, pull_);
 256
 257     if (numSharingSimulations > 1 && MASTER(commRecord_))
 258     {
 259         std::vector<size_t> pointSize;
 260         for (auto const& biasCts : biasCoupledToSystem_)
 261         {
 262             pointSize.push_back(biasCts.bias_.state().points().size());
 263         }
 264         /* Ensure that the shared biased are compatible between simulations */
 265         biasesAreCompatibleForSharingBetweenSimulations(awhParams, pointSize, multiSimRecord_);
 266     }
 267 }
 268
 269 Awh::~Awh() = default;
 270
 271 bool Awh::isOutputStep(int64_t step) const
 272 {
 273     return (nstout_ > 0 && step % nstout_ == 0);
 274 }
 275
 276 real Awh::applyBiasForcesAndUpdateBias(PbcType                pbcType,
 277                                        ArrayRef<const real>   masses,
 278                                        ArrayRef<const double> neighborLambdaEnergies,
 279                                        ArrayRef<const double> neighborLambdaDhdl,
 280                                        const matrix           box,
 281                                        gmx::ForceWithVirial*  forceWithVirial,
 282                                        double                 t,
 283                                        int64_t                step,
 284                                        gmx_wallcycle*         wallcycle,
 285                                        FILE*                  fplog)
 286 {
 287     if (anyDimUsesPull(biasCoupledToSystem_))
 288     {
 289         GMX_ASSERT(forceWithVirial, "Need a valid ForceWithVirial object");
 290     }
 291
 292     wallcycle_start(wallcycle, WallCycleCounter::Awh);
 293
 294     t_pbc pbc;
 295     set_pbc(&pbc, pbcType, box);
 296
 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_;
 301
 302     for (auto& biasCts : biasCoupledToSystem_)
 303     {
 304         /* Update the AWH coordinate values with those of the corresponding
 305          * pull coordinates.
 306          */
 307         awh_dvec coordValue           = { 0, 0, 0, 0 };
 308         int      numLambdaDimsCounted = 0;
 309         for (int d = 0; d < biasCts.bias_.ndim(); d++)
 310         {
 311             if (biasCts.bias_.dimParams()[d].isPullDimension())
 312             {
 313                 coordValue[d] = get_pull_coord_value(
 314                         pull_, biasCts.pullCoordIndex_[d - numLambdaDimsCounted], &pbc);
 315             }
 316             else
 317             {
 318                 coordValue[d] = fepLambdaState_;
 319                 numLambdaDimsCounted += 1;
 320             }
 321         }
 322
 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.
 326          */
 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.
 331          */
 332         gmx::ArrayRef<const double> biasForce =
 333                 biasCts.bias_.calcForceAndUpdateBias(coordValue,
 334                                                      neighborLambdaEnergies,
 335                                                      neighborLambdaDhdl,
 336                                                      &biasPotential,
 337                                                      &biasPotentialJump,
 338                                                      commRecord_,
 339                                                      multiSimRecord_,
 340                                                      t,
 341                                                      step,
 342                                                      seed_,
 343                                                      fplog);
 344
 345         awhPotential += biasPotential;
 346
 347         /* Keep track of the total potential shift needed to remove the potential jumps. */
 348         potentialOffset_ -= biasPotentialJump;
 349
 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.
 353          */
 354         numLambdaDimsCounted = 0;
 355         for (int d = 0; d < biasCts.bias_.ndim(); d++)
 356         {
 357             if (biasCts.bias_.dimParams()[d].isPullDimension())
 358             {
 359                 apply_external_pull_coord_force(pull_,
 360                                                 biasCts.pullCoordIndex_[d - numLambdaDimsCounted],
 361                                                 biasForce[d],
 362                                                 masses,
 363                                                 forceWithVirial);
 364             }
 365             else
 366             {
 367                 int umbrellaGridpointIndex = biasCts.bias_.state().coordState().umbrellaGridpoint();
 368                 fepLambdaState_ = biasCts.bias_.getGridCoordValue(umbrellaGridpointIndex)[d];
 369                 numLambdaDimsCounted += 1;
 370             }
 371         }
 372
 373         if (isOutputStep(step))
 374         {
 375             /* We might have skipped updates for part of the grid points.
 376              * Ensure all points are updated before writing out their data.
 377              */
 378             biasCts.bias_.doSkippedUpdatesForAllPoints();
 379         }
 380     }
 381
 382     wallcycle_stop(wallcycle, WallCycleCounter::Awh);
 383
 384     return MASTER(commRecord_) ? static_cast<real>(awhPotential) : 0;
 385 }
 386
 387 std::shared_ptr<AwhHistory> Awh::initHistoryFromState() const
 388 {
 389     if (MASTER(commRecord_))
 390     {
 391         std::shared_ptr<AwhHistory> awhHistory(new AwhHistory);
 392         awhHistory->bias.clear();
 393         awhHistory->bias.resize(biasCoupledToSystem_.size());
 394
 395         for (size_t k = 0; k < awhHistory->bias.size(); k++)
 396         {
 397             biasCoupledToSystem_[k].bias_.initHistoryFromState(&awhHistory->bias[k]);
 398         }
 399
 400         return awhHistory;
 401     }
 402     else
 403     {
 404         /* Return an empty pointer */
 405         return std::shared_ptr<AwhHistory>();
 406     }
 407 }
 408
 409 void Awh::restoreStateFromHistory(const AwhHistory* awhHistory)
 410 {
 411     /* Restore the history to the current state */
 412     if (MASTER(commRecord_))
 413     {
 414         GMX_RELEASE_ASSERT(awhHistory != nullptr,
 415                            "The master rank should have a valid awhHistory when restoring the "
 416                            "state from history.");
 417
 418         if (awhHistory->bias.size() != biasCoupledToSystem_.size())
 419         {
 420             GMX_THROW(InvalidInputError(
 421                     "AWH state and history contain different numbers of biases. Likely you "
 422                     "provided a checkpoint from a different simulation."));
 423         }
 424
 425         potentialOffset_ = awhHistory->potentialOffset;
 426     }
 427     if (PAR(commRecord_))
 428     {
 429         gmx_bcast(sizeof(potentialOffset_), &potentialOffset_, commRecord_->mpi_comm_mygroup);
 430     }
 431
 432     for (size_t k = 0; k < biasCoupledToSystem_.size(); k++)
 433     {
 434         biasCoupledToSystem_[k].bias_.restoreStateFromHistory(
 435                 awhHistory ? &awhHistory->bias[k] : nullptr, commRecord_);
 436     }
 437 }
 438
 439 void Awh::updateHistory(AwhHistory* awhHistory) const
 440 {
 441     if (!MASTER(commRecord_))
 442     {
 443         return;
 444     }
 445
 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");
 449
 450     awhHistory->potentialOffset = potentialOffset_;
 451
 452     for (size_t k = 0; k < awhHistory->bias.size(); k++)
 453     {
 454         biasCoupledToSystem_[k].bias_.updateHistory(&awhHistory->bias[k]);
 455     }
 456 }
 457
 458 const char* Awh::externalPotentialString()
 459 {
 460     return "AWH";
 461 }
 462
 463 void Awh::registerAwhWithPull(const AwhParams& awhParams, pull_t* pull_work)
 464 {
 465     GMX_RELEASE_ASSERT(!anyDimUsesPull(awhParams) || pull_work, "Need a valid pull object");
 466
 467     for (const auto& biasParam : awhParams.awhBiasParams())
 468     {
 469         for (const auto& dimParam : biasParam.dimParams())
 470         {
 471             if (dimParam.coordinateProvider() == AwhCoordinateProviderType::Pull)
 472             {
 473                 register_external_pull_potential(
 474                         pull_work, dimParam.coordinateIndex(), Awh::externalPotentialString());
 475             }
 476         }
 477     }
 478 }
 479
 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
 482 {
 483     GMX_ASSERT(MASTER(commRecord_), "writeToEnergyFrame should only be called on the master rank");
 484     GMX_ASSERT(frame != nullptr, "Need a valid energy frame");
 485
 486     if (!isOutputStep(step))
 487     {
 488         /* This is not an AWH output step, don't write any AWH data */
 489         return;
 490     }
 491
 492     /* Get the total number of energy subblocks that AWH needs */
 493     int numSubblocks = 0;
 494     for (const auto& biasCoupledToSystem : biasCoupledToSystem_)
 495     {
 496         numSubblocks += biasCoupledToSystem.bias_.numEnergySubblocksToWrite();
 497     }
 498     GMX_ASSERT(numSubblocks > 0, "We should always have data to write");
 499
 500     /* Add 1 energy block */
 501     add_blocks_enxframe(frame, frame->nblock + 1);
 502
 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);
 506
 507     /* Claim it as an AWH block. */
 508     awhEnergyBlock->id = enxAWH;
 509
 510     /* Transfer AWH data blocks to energy sub blocks */
 511     int energySubblockCount = 0;
 512     for (const auto& biasCoupledToSystem : biasCoupledToSystem_)
 513     {
 514         energySubblockCount += biasCoupledToSystem.bias_.writeToEnergySubblocks(
 515                 &(awhEnergyBlock->sub[energySubblockCount]));
 516     }
 517 }
 518
 519 bool Awh::hasFepLambdaDimension() const
 520 {
 521     return std::any_of(
 522             std::begin(biasCoupledToSystem_),
 523             std::end(biasCoupledToSystem_),
 524             [](const auto& coupledBias) { return coupledBias.bias_.hasFepLambdaDimension(); });
 525 }
 526
 527 bool Awh::needForeignEnergyDifferences(const int64_t step) const
 528 {
 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())
 532     {
 533         return false;
 534     }
 535     if (step == 0)
 536     {
 537         return true;
 538     }
 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);
 544     });
 545 }
 546
 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,
 555                                       pull_t*               pull_work)
 556 {
 557     if (!inputRecord.bDoAwh)
 558     {
 559         return nullptr;
 560     }
 561     if (usingShellParticles)
 562     {
 563         GMX_THROW(InvalidInputError("AWH biasing does not support shell particles."));
 564     }
 565
 566     auto awh = std::make_unique<Awh>(fplog,
 567                                      inputRecord,
 568                                      commRecord,
 569                                      multiSimRecord,
 570                                      *inputRecord.awhParams,
 571                                      biasInitFilename,
 572                                      pull_work,
 573                                      inputRecord.fepvals->n_lambda,
 574                                      inputRecord.fepvals->init_fep_state);
 575
 576     if (startingFromCheckpoint)
 577     {
 578         // Restore the AWH history read from checkpoint
 579         awh->restoreStateFromHistory(MASTER(commRecord) ? stateGlobal->awhHistory.get() : nullptr);
 580     }
 581     else if (MASTER(commRecord))
 582     {
 583         // Initialize the AWH history here
 584         stateGlobal->awhHistory = awh->initHistoryFromState();
 585     }
 586     return awh;
 587 }
 588
 589 } // namespace gmx