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