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/math/utilities.h"
  62 #include "gromacs/mdrunutility/multisim.h"
  63 #include "gromacs/mdtypes/awh_history.h"
  64 #include "gromacs/mdtypes/awh_params.h"
  65 #include "gromacs/mdtypes/commrec.h"
  66 #include "gromacs/mdtypes/forceoutput.h"
  67 #include "gromacs/mdtypes/inputrec.h"
  68 #include "gromacs/mdtypes/pull_params.h"
  69 #include "gromacs/mdtypes/state.h"
  70 #include "gromacs/pbcutil/pbc.h"
  71 #include "gromacs/pulling/pull.h"
  72 #include "gromacs/timing/wallcycle.h"
  73 #include "gromacs/trajectory/energyframe.h"
  74 #include "gromacs/utility/arrayref.h"
  75 #include "gromacs/utility/exceptions.h"
  76 #include "gromacs/utility/gmxassert.h"
  77 #include "gromacs/utility/pleasecite.h"
  78
  79 #include "bias.h"
  80 #include "biassharing.h"
  81 #include "correlationgrid.h"
  82 #include "pointstate.h"
  83
  84 namespace gmx
  85 {
  86
  87 /*! \internal
  88  * \brief A bias and its coupling to the system.
  89  *
  90  * This struct is used to separate the bias machinery in the Bias class,
  91  * which should be independent from the reaction coordinate, from the
  92  * obtaining of the reaction coordinate values and passing the computed forces.
  93  * Currently the AWH method couples to the system by mapping each
  94  * AWH bias to a pull coordinate. This can easily be generalized here.
  95  */
  96 struct BiasCoupledToSystem
  97 {
  98     /*! \brief Constructor, couple a bias to a set of pull coordinates.
  99      *
 100      * \param[in] bias            The bias.
 101      * \param[in] pullCoordIndex  The pull coordinate indices.
 102      */
 103     BiasCoupledToSystem(Bias bias, const std::vector<int>& pullCoordIndex);
 104
 105     Bias                   bias_;           /**< The bias. */
 106     const std::vector<int> pullCoordIndex_; /**< The pull coordinates this bias acts on. */
 107
 108     /* Here AWH can be extended to work on other coordinates than pull. */
 109 };
 110
 111 /*! \brief Checks whether any dimension uses pulling as a coordinate provider.
 112  *
 113  * \param[in] awhBiasParams The bias params to check.
 114  * \returns true if any dimension of the bias is linked to pulling.
 115  */
 116 static bool anyDimUsesPull(const AwhBiasParams& awhBiasParams)
 117 {
 118     for (int d = 0; d < awhBiasParams.ndim; d++)
 119     {
 120         const AwhDimParams& awhDimParams = awhBiasParams.dimParams[d];
 121         if (awhDimParams.eCoordProvider == 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 (int k = 0; k < awhParams.numBias; k++)
 137     {
 138         const AwhBiasParams& awhBiasParams = awhParams.awhBiasParams[k];
 139         if (anyDimUsesPull(awhBiasParams))
 140         {
 141             return true;
 142         }
 143     }
 144     return false;
 145 }
 146
 147 /*! \brief Checks whether any dimension uses pulling as a coordinate provider.
 148  *
 149  * \param[in] biasCoupledToSystem The AWH biases to check.
 150  * \returns true if any dimension of the provided biases is linked to pulling.
 151  */
 152 static bool anyDimUsesPull(const ArrayRef<BiasCoupledToSystem> biasCoupledToSystem)
 153 {
 154     for (const auto& biasCts : biasCoupledToSystem)
 155     {
 156         if (!biasCts.pullCoordIndex_.empty())
 157         {
 158             return true;
 159         }
 160     }
 161     return false;
 162 }
 163
 164 BiasCoupledToSystem::BiasCoupledToSystem(Bias bias, const std::vector<int>& pullCoordIndex) :
 165     bias_(std::move(bias)),
 166     pullCoordIndex_(pullCoordIndex)
 167 {
 168     /* We already checked for this in grompp, but check again here. */
 169     GMX_RELEASE_ASSERT(
 170             static_cast<size_t>(bias_.ndim()) == pullCoordIndex_.size() + bias_.hasFepLambdaDimension() ? 1 : 0,
 171             "The bias dimensionality should match the number of pull and lambda coordinates.");
 172 }
 173
 174 Awh::Awh(FILE*                 fplog,
 175          const t_inputrec&     inputRecord,
 176          const t_commrec*      commRecord,
 177          const gmx_multisim_t* multiSimRecord,
 178          const AwhParams&      awhParams,
 179          const std::string&    biasInitFilename,
 180          pull_t*               pull_work,
 181          int                   numFepLambdaStates,
 182          int                   fepLambdaState) :
 183     seed_(awhParams.seed),
 184     nstout_(awhParams.nstOut),
 185     commRecord_(commRecord),
 186     multiSimRecord_(multiSimRecord),
 187     pull_(pull_work),
 188     potentialOffset_(0),
 189     numFepLambdaStates_(numFepLambdaStates),
 190     fepLambdaState_(fepLambdaState)
 191 {
 192     if (anyDimUsesPull(awhParams))
 193     {
 194         GMX_RELEASE_ASSERT(inputRecord.pull != nullptr, "With AWH we should have pull parameters");
 195         GMX_RELEASE_ASSERT(pull_work != nullptr,
 196                            "With AWH pull should be initialized before initializing AWH");
 197     }
 198
 199     if (fplog != nullptr)
 200     {
 201         please_cite(fplog, "Lindahl2014");
 202     }
 203
 204     if (haveBiasSharingWithinSimulation(awhParams))
 205     {
 206         /* This has likely been checked by grompp, but throw anyhow. */
 207         GMX_THROW(
 208                 InvalidInputError("Biases within a simulation are shared, currently sharing of "
 209                                   "biases is only supported between simulations"));
 210     }
 211
 212     int numSharingSimulations = 1;
 213     if (awhParams.shareBiasMultisim && isMultiSim(multiSimRecord_))
 214     {
 215         numSharingSimulations = multiSimRecord_->numSimulations_;
 216     }
 217
 218     /* Initialize all the biases */
 219     const double beta = 1 / (gmx::c_boltz * inputRecord.opts.ref_t[0]);
 220     for (int k = 0; k < awhParams.numBias; k++)
 221     {
 222         const AwhBiasParams& awhBiasParams = awhParams.awhBiasParams[k];
 223
 224         std::vector<int>       pullCoordIndex;
 225         std::vector<DimParams> dimParams;
 226         for (int d = 0; d < awhBiasParams.ndim; d++)
 227         {
 228             const AwhDimParams& awhDimParams = awhBiasParams.dimParams[d];
 229             if (awhDimParams.eCoordProvider != AwhCoordinateProviderType::Pull
 230                 && awhDimParams.eCoordProvider != AwhCoordinateProviderType::FreeEnergyLambda)
 231             {
 232                 GMX_THROW(
 233                         InvalidInputError("Currently only the pull code and lambda are supported "
 234                                           "as coordinate providers"));
 235             }
 236             if (awhDimParams.eCoordProvider == AwhCoordinateProviderType::Pull)
 237             {
 238                 const t_pull_coord& pullCoord = inputRecord.pull->coord[awhDimParams.coordIndex];
 239                 if (pullCoord.eGeom == PullGroupGeometry::DirectionPBC)
 240                 {
 241                     GMX_THROW(InvalidInputError(
 242                             "Pull geometry 'direction-periodic' is not supported by AWH"));
 243                 }
 244                 double conversionFactor = pull_coordinate_is_angletype(&pullCoord) ? gmx::c_deg2Rad : 1;
 245                 pullCoordIndex.push_back(awhDimParams.coordIndex);
 246                 dimParams.push_back(DimParams::pullDimParams(
 247                         conversionFactor, awhDimParams.forceConstant, beta));
 248             }
 249             else
 250             {
 251                 dimParams.push_back(DimParams::fepLambdaDimParams(numFepLambdaStates_, beta));
 252             }
 253         }
 254
 255         /* Construct the bias and couple it to the system. */
 256         Bias::ThisRankWillDoIO thisRankWillDoIO =
 257                 (MASTER(commRecord_) ? Bias::ThisRankWillDoIO::Yes : Bias::ThisRankWillDoIO::No);
 258         biasCoupledToSystem_.emplace_back(Bias(k,
 259                                                awhParams,
 260                                                awhParams.awhBiasParams[k],
 261                                                dimParams,
 262                                                beta,
 263                                                inputRecord.delta_t,
 264                                                numSharingSimulations,
 265                                                biasInitFilename,
 266                                                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 =
 351                 biasCts.bias_.calcForceAndUpdateBias(coordValue,
 352                                                      neighborLambdaEnergies,
 353                                                      neighborLambdaDhdl,
 354                                                      &biasPotential,
 355                                                      &biasPotentialJump,
 356                                                      commRecord_,
 357                                                      multiSimRecord_,
 358                                                      t,
 359                                                      step,
 360                                                      seed_,
 361                                                      fplog);
 362
 363         awhPotential += biasPotential;
 364
 365         /* Keep track of the total potential shift needed to remove the potential jumps. */
 366         potentialOffset_ -= biasPotentialJump;
 367
 368         /* Communicate the bias force to the pull struct.
 369          * The bias potential is returned at the end of this function,
 370          * so that it can be added externally to the correct energy data block.
 371          */
 372         numLambdaDimsCounted = 0;
 373         for (int d = 0; d < biasCts.bias_.ndim(); d++)
 374         {
 375             if (biasCts.bias_.dimParams()[d].isPullDimension())
 376             {
 377                 apply_external_pull_coord_force(pull_,
 378                                                 biasCts.pullCoordIndex_[d - numLambdaDimsCounted],
 379                                                 biasForce[d],
 380                                                 masses,
 381                                                 forceWithVirial);
 382             }
 383             else
 384             {
 385                 int umbrellaGridpointIndex = biasCts.bias_.state().coordState().umbrellaGridpoint();
 386                 fepLambdaState_ = biasCts.bias_.getGridCoordValue(umbrellaGridpointIndex)[d];
 387                 numLambdaDimsCounted += 1;
 388             }
 389         }
 390
 391         if (isOutputStep(step))
 392         {
 393             /* We might have skipped updates for part of the grid points.
 394              * Ensure all points are updated before writing out their data.
 395              */
 396             biasCts.bias_.doSkippedUpdatesForAllPoints();
 397         }
 398     }
 399
 400     wallcycle_stop(wallcycle, ewcAWH);
 401
 402     return MASTER(commRecord_) ? static_cast<real>(awhPotential) : 0;
 403 }
 404
 405 std::shared_ptr<AwhHistory> Awh::initHistoryFromState() const
 406 {
 407     if (MASTER(commRecord_))
 408     {
 409         std::shared_ptr<AwhHistory> awhHistory(new AwhHistory);
 410         awhHistory->bias.clear();
 411         awhHistory->bias.resize(biasCoupledToSystem_.size());
 412
 413         for (size_t k = 0; k < awhHistory->bias.size(); k++)
 414         {
 415             biasCoupledToSystem_[k].bias_.initHistoryFromState(&awhHistory->bias[k]);
 416         }
 417
 418         return awhHistory;
 419     }
 420     else
 421     {
 422         /* Return an empty pointer */
 423         return std::shared_ptr<AwhHistory>();
 424     }
 425 }
 426
 427 void Awh::restoreStateFromHistory(const AwhHistory* awhHistory)
 428 {
 429     /* Restore the history to the current state */
 430     if (MASTER(commRecord_))
 431     {
 432         GMX_RELEASE_ASSERT(awhHistory != nullptr,
 433                            "The master rank should have a valid awhHistory when restoring the "
 434                            "state from history.");
 435
 436         if (awhHistory->bias.size() != biasCoupledToSystem_.size())
 437         {
 438             GMX_THROW(InvalidInputError(
 439                     "AWH state and history contain different numbers of biases. Likely you "
 440                     "provided a checkpoint from a different simulation."));
 441         }
 442
 443         potentialOffset_ = awhHistory->potentialOffset;
 444     }
 445     if (PAR(commRecord_))
 446     {
 447         gmx_bcast(sizeof(potentialOffset_), &potentialOffset_, commRecord_->mpi_comm_mygroup);
 448     }
 449
 450     for (size_t k = 0; k < biasCoupledToSystem_.size(); k++)
 451     {
 452         biasCoupledToSystem_[k].bias_.restoreStateFromHistory(
 453                 awhHistory ? &awhHistory->bias[k] : nullptr, commRecord_);
 454     }
 455 }
 456
 457 void Awh::updateHistory(AwhHistory* awhHistory) const
 458 {
 459     if (!MASTER(commRecord_))
 460     {
 461         return;
 462     }
 463
 464     /* This assert will also catch a non-master rank calling this function. */
 465     GMX_RELEASE_ASSERT(awhHistory->bias.size() == biasCoupledToSystem_.size(),
 466                        "AWH state and history bias count should match");
 467
 468     awhHistory->potentialOffset = potentialOffset_;
 469
 470     for (size_t k = 0; k < awhHistory->bias.size(); k++)
 471     {
 472         biasCoupledToSystem_[k].bias_.updateHistory(&awhHistory->bias[k]);
 473     }
 474 }
 475
 476 const char* Awh::externalPotentialString()
 477 {
 478     return "AWH";
 479 }
 480
 481 void Awh::registerAwhWithPull(const AwhParams& awhParams, pull_t* pull_work)
 482 {
 483     GMX_RELEASE_ASSERT(!anyDimUsesPull(awhParams) || pull_work, "Need a valid pull object");
 484
 485     for (int k = 0; k < awhParams.numBias; k++)
 486     {
 487         const AwhBiasParams& biasParams = awhParams.awhBiasParams[k];
 488
 489         for (int d = 0; d < biasParams.ndim; d++)
 490         {
 491             if (biasParams.dimParams[d].eCoordProvider == AwhCoordinateProviderType::Pull)
 492             {
 493                 register_external_pull_potential(
 494                         pull_work, biasParams.dimParams[d].coordIndex, Awh::externalPotentialString());
 495             }
 496         }
 497     }
 498 }
 499
 500 /* Fill the AWH data block of an energy frame with data (if there is any). */
 501 void Awh::writeToEnergyFrame(int64_t step, t_enxframe* frame) const
 502 {
 503     GMX_ASSERT(MASTER(commRecord_), "writeToEnergyFrame should only be called on the master rank");
 504     GMX_ASSERT(frame != nullptr, "Need a valid energy frame");
 505
 506     if (!isOutputStep(step))
 507     {
 508         /* This is not an AWH output step, don't write any AWH data */
 509         return;
 510     }
 511
 512     /* Get the total number of energy subblocks that AWH needs */
 513     int numSubblocks = 0;
 514     for (auto& biasCoupledToSystem : biasCoupledToSystem_)
 515     {
 516         numSubblocks += biasCoupledToSystem.bias_.numEnergySubblocksToWrite();
 517     }
 518     GMX_ASSERT(numSubblocks > 0, "We should always have data to write");
 519
 520     /* Add 1 energy block */
 521     add_blocks_enxframe(frame, frame->nblock + 1);
 522
 523     /* Take the block that was just added and set the number of subblocks. */
 524     t_enxblock* awhEnergyBlock = &(frame->block[frame->nblock - 1]);
 525     add_subblocks_enxblock(awhEnergyBlock, numSubblocks);
 526
 527     /* Claim it as an AWH block. */
 528     awhEnergyBlock->id = enxAWH;
 529
 530     /* Transfer AWH data blocks to energy sub blocks */
 531     int energySubblockCount = 0;
 532     for (auto& biasCoupledToSystem : biasCoupledToSystem_)
 533     {
 534         energySubblockCount += biasCoupledToSystem.bias_.writeToEnergySubblocks(
 535                 &(awhEnergyBlock->sub[energySubblockCount]));
 536     }
 537 }
 538
 539 bool Awh::hasFepLambdaDimension() const
 540 {
 541     return std::any_of(
 542             std::begin(biasCoupledToSystem_),
 543             std::end(biasCoupledToSystem_),
 544             [](const auto& coupledBias) { return coupledBias.bias_.hasFepLambdaDimension(); });
 545 }
 546
 547 bool Awh::needForeignEnergyDifferences(const int64_t step) const
 548 {
 549     /* If there is no FEP lambda dimension at all in any bias there will be no need for foreign
 550      * energy differences */
 551     if (!hasFepLambdaDimension())
 552     {
 553         return false;
 554     }
 555     if (step == 0)
 556     {
 557         return true;
 558     }
 559     /* Check whether the bias(es) that has/have a FEP lambda dimension should sample coordinates
 560      * this step. Since the biases may have different sampleCoordStep it is necessary to check
 561      * this combination. */
 562     for (const auto& biasCts : biasCoupledToSystem_)
 563     {
 564         if (biasCts.bias_.hasFepLambdaDimension() && biasCts.bias_.isSampleCoordStep(step))
 565         {
 566             return true;
 567         }
 568     }
 569     return false;
 570 }
 571
 572 std::unique_ptr<Awh> prepareAwhModule(FILE*                 fplog,
 573                                       const t_inputrec&     inputRecord,
 574                                       t_state*              stateGlobal,
 575                                       const t_commrec*      commRecord,
 576                                       const gmx_multisim_t* multiSimRecord,
 577                                       const bool            startingFromCheckpoint,
 578                                       const bool            usingShellParticles,
 579                                       const std::string&    biasInitFilename,
 580                                       pull_t*               pull_work)
 581 {
 582     if (!inputRecord.bDoAwh)
 583     {
 584         return nullptr;
 585     }
 586     if (usingShellParticles)
 587     {
 588         GMX_THROW(InvalidInputError("AWH biasing does not support shell particles."));
 589     }
 590
 591     auto awh = std::make_unique<Awh>(fplog,
 592                                      inputRecord,
 593                                      commRecord,
 594                                      multiSimRecord,
 595                                      *inputRecord.awhParams,
 596                                      biasInitFilename,
 597                                      pull_work,
 598                                      inputRecord.fepvals->n_lambda,
 599                                      inputRecord.fepvals->init_fep_state);
 600
 601     if (startingFromCheckpoint)
 602     {
 603         // Restore the AWH history read from checkpoint
 604         awh->restoreStateFromHistory(MASTER(commRecord) ? stateGlobal->awhHistory.get() : nullptr);
 605     }
 606     else if (MASTER(commRecord))
 607     {
 608         // Initialize the AWH history here
 609         stateGlobal->awhHistory = awh->initHistoryFromState();
 610     }
 611     return awh;
 612 }
 613
 614 } // namespace gmx