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