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