Make AWH bias sharing more flexible
[alexxy/gromacs.git] / src / gromacs / applied_forces / awh / bias.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 Bias class.
39  *
40  * \author Viveca Lindahl
41  * \author Berk Hess <hess@kth.se>
42  * \ingroup module_awh
43  */
44
45 #include "gmxpre.h"
46
47 #include "bias.h"
48
49 #include <cassert>
50 #include <cinttypes>
51 #include <cmath>
52 #include <cstdio>
53 #include <cstdlib>
54 #include <cstring>
55
56 #include <algorithm>
57 #include <memory>
58
59 #include "gromacs/fileio/gmxfio.h"
60 #include "gromacs/gmxlib/network.h"
61 #include "gromacs/math/functions.h"
62 #include "gromacs/math/utilities.h"
63 #include "gromacs/mdtypes/awh_history.h"
64 #include "gromacs/mdtypes/awh_params.h"
65 #include "gromacs/mdtypes/commrec.h"
66 #include "gromacs/utility/exceptions.h"
67 #include "gromacs/utility/gmxassert.h"
68 #include "gromacs/utility/stringutil.h"
69
70 #include "biassharing.h"
71 #include "correlationgrid.h"
72 #include "correlationhistory.h"
73 #include "pointstate.h"
74
75 namespace gmx
76 {
77
78 void Bias::warnForHistogramAnomalies(double t, int64_t step, FILE* fplog)
79 {
80     const int maxNumWarningsInCheck = 1;  /* The maximum number of warnings to print per check */
81     const int maxNumWarningsInRun   = 10; /* The maximum number of warnings to print in a run */
82
83     if (fplog == nullptr || numWarningsIssued_ >= maxNumWarningsInRun || state_.inInitialStage()
84         || !params_.isCheckHistogramForAnomaliesStep(step))
85     {
86         return;
87     }
88
89     numWarningsIssued_ +=
90             state_.warnForHistogramAnomalies(grid_, biasIndex(), t, fplog, maxNumWarningsInCheck);
91
92     if (numWarningsIssued_ >= maxNumWarningsInRun)
93     {
94         fprintf(fplog, "\nawh%d: suppressing future AWH warnings.\n", biasIndex() + 1);
95     }
96 }
97
98 void Bias::doSkippedUpdatesForAllPoints()
99 {
100     if (params_.skipUpdates())
101     {
102         state_.doSkippedUpdatesForAllPoints(params_);
103     }
104 }
105
106 gmx::ArrayRef<const double> Bias::calcForceAndUpdateBias(const awh_dvec         coordValue,
107                                                          ArrayRef<const double> neighborLambdaEnergies,
108                                                          ArrayRef<const double> neighborLambdaDhdl,
109                                                          double*                awhPotential,
110                                                          double*                potentialJump,
111                                                          double                 t,
112                                                          int64_t                step,
113                                                          int64_t                seed,
114                                                          FILE*                  fplog)
115 {
116     if (step < 0)
117     {
118         GMX_THROW(InvalidInputError(
119                 "The step number is negative which is not supported by the AWH code."));
120     }
121
122     GMX_RELEASE_ASSERT(!(params_.convolveForce && grid_.hasLambdaAxis()),
123                        "When using AWH to sample an FEP lambda dimension the AWH potential cannot "
124                        "be convolved.");
125
126     state_.setCoordValue(grid_, coordValue);
127
128     std::vector<double, AlignedAllocator<double>>& probWeightNeighbor = alignedTempWorkSpace_;
129
130     /* If the convolved force is needed or this is a sampling step,
131      * the bias in the current neighborhood needs to be up-to-date
132      * and the probablity weights need to be calculated.
133      */
134     const bool        isSampleCoordStep = params_.isSampleCoordStep(step);
135     const bool        moveUmbrella      = (isSampleCoordStep || step == 0);
136     double            convolvedBias     = 0;
137     const CoordState& coordState        = state_.coordState();
138
139     if (params_.convolveForce || moveUmbrella || isSampleCoordStep)
140     {
141         if (params_.skipUpdates())
142         {
143             state_.doSkippedUpdatesInNeighborhood(params_, grid_);
144         }
145         convolvedBias = state_.updateProbabilityWeightsAndConvolvedBias(
146                 dimParams_, grid_, moveUmbrella ? neighborLambdaEnergies : ArrayRef<const double>{}, &probWeightNeighbor);
147
148         if (isSampleCoordStep)
149         {
150             updateForceCorrelationGrid(probWeightNeighbor, neighborLambdaDhdl, t);
151
152             state_.sampleCoordAndPmf(dimParams_, grid_, probWeightNeighbor, convolvedBias);
153         }
154     }
155
156     /* Set the bias force and get the potential contribution from this bias.
157      * The potential jump occurs at different times depending on how
158      * the force is applied (and how the potential is normalized).
159      * For the convolved force it happens when the bias is updated,
160      * for the umbrella when the umbrella is moved.
161      */
162     *potentialJump = 0;
163     double potential;
164     if (params_.convolveForce)
165     {
166         state_.calcConvolvedForce(dimParams_,
167                                   grid_,
168                                   probWeightNeighbor,
169                                   moveUmbrella ? neighborLambdaDhdl : ArrayRef<const double>{},
170                                   tempForce_,
171                                   biasForce_);
172
173         potential = -convolvedBias * params_.invBeta;
174     }
175     else
176     {
177         /* Umbrella force */
178         GMX_RELEASE_ASSERT(state_.points()[coordState.umbrellaGridpoint()].inTargetRegion(),
179                            "AWH bias grid point for the umbrella reference value is outside of the "
180                            "target region.");
181         potential = state_.calcUmbrellaForceAndPotential(
182                 dimParams_,
183                 grid_,
184                 coordState.umbrellaGridpoint(),
185                 moveUmbrella ? neighborLambdaDhdl : ArrayRef<const double>{},
186                 biasForce_);
187
188         /* Moving the umbrella results in a force correction and
189          * a new potential. The umbrella center is sampled as often as
190          * the coordinate so we know the probability weights needed
191          * for moving the umbrella are up-to-date.
192          */
193         if (moveUmbrella)
194         {
195             const bool onlySampleUmbrellaGridpoint = false;
196             double     newPotential                = state_.moveUmbrella(dimParams_,
197                                                       grid_,
198                                                       probWeightNeighbor,
199                                                       neighborLambdaDhdl,
200                                                       biasForce_,
201                                                       step,
202                                                       seed,
203                                                       params_.biasIndex,
204                                                       onlySampleUmbrellaGridpoint);
205             *potentialJump                         = newPotential - potential;
206         }
207     }
208
209     /* Update the free energy estimates and bias and other history dependent method parameters */
210     if (params_.isUpdateFreeEnergyStep(step))
211     {
212         state_.updateFreeEnergyAndAddSamplesToHistogram(
213                 dimParams_, grid_, params_, t, step, fplog, &updateList_);
214
215         if (params_.convolveForce)
216         {
217             /* The update results in a potential jump, so we need the new convolved potential. */
218             double newPotential = -calcConvolvedBias(coordState.coordValue()) * params_.invBeta;
219             *potentialJump      = newPotential - potential;
220         }
221     }
222     /* If there is a lambda axis it is still controlled using an umbrella even if the force
223      * is convolved in the other dimensions. */
224     if (moveUmbrella && params_.convolveForce && grid_.hasLambdaAxis())
225     {
226         const bool onlySampleUmbrellaGridpoint = true;
227         state_.moveUmbrella(dimParams_,
228                             grid_,
229                             probWeightNeighbor,
230                             neighborLambdaDhdl,
231                             biasForce_,
232                             step,
233                             seed,
234                             params_.biasIndex,
235                             onlySampleUmbrellaGridpoint);
236     }
237
238     /* Return the potential. */
239     *awhPotential = potential;
240
241     /* Check the sampled histograms and potentially warn user if something is suspicious */
242     warnForHistogramAnomalies(t, step, fplog);
243
244     return biasForce_;
245 }
246
247 /*! \brief
248  * Count the total number of samples / sample weight over all grid points.
249  *
250  * \param[in] pointState  The state of the points in a bias.
251  * \returns the total sample count.
252  */
253 static int64_t countSamples(ArrayRef<const PointState> pointState)
254 {
255     double numSamples = 0;
256     for (const PointState& point : pointState)
257     {
258         numSamples += point.weightSumTot();
259     }
260
261     return gmx::roundToInt64(numSamples);
262 }
263
264 /*! \brief
265  * Check if the state (loaded from checkpoint) and the run are consistent.
266  *
267  * When the state and the run setup are inconsistent, an exception is thrown.
268  *
269  * \param[in] params  The parameters of the bias.
270  * \param[in] state   The state of the bias.
271  */
272 static void ensureStateAndRunConsistency(const BiasParams& params, const BiasState& state)
273 {
274     int64_t numSamples = countSamples(state.points());
275     int64_t numUpdatesFromSamples =
276             numSamples / (params.numSamplesUpdateFreeEnergy_ * params.numSharedUpdate);
277     int64_t numUpdatesExpected = state.histogramSize().numUpdates();
278     if (numUpdatesFromSamples != numUpdatesExpected)
279     {
280         std::string mesg = gmx::formatString(
281                 "The number of AWH updates in the checkpoint file (%" PRId64
282                 ") does not match the total number of AWH samples divided by the number of samples "
283                 "per update for %d sharing AWH bias(es) (%" PRId64 "/%d=%" PRId64 ")",
284                 numUpdatesExpected,
285                 params.numSharedUpdate,
286                 numSamples,
287                 params.numSamplesUpdateFreeEnergy_ * params.numSharedUpdate,
288                 numUpdatesFromSamples);
289         mesg += " Maybe you changed AWH parameters.";
290         /* Unfortunately we currently do not store the number of simulations
291          * sharing the bias or the state to checkpoint. But we can hint at
292          * a mismatch in the number of sharing simulations.
293          */
294         if (numUpdatesFromSamples % state.histogramSize().numUpdates() == 0)
295         {
296             mesg += gmx::formatString(
297                     " Or the run you continued from used %" PRId64
298                     " sharing simulations, whereas you now specified %d sharing simulations.",
299                     numUpdatesFromSamples / state.histogramSize().numUpdates(),
300                     params.numSharedUpdate);
301         }
302         GMX_THROW(InvalidInputError(mesg));
303     }
304 }
305
306 void Bias::restoreStateFromHistory(const AwhBiasHistory* biasHistory, const t_commrec* cr)
307 {
308     GMX_RELEASE_ASSERT(thisRankDoesIO_ == MASTER(cr),
309                        "The master rank should do I/O, the other ranks should not");
310
311     if (MASTER(cr))
312     {
313         GMX_RELEASE_ASSERT(biasHistory != nullptr,
314                            "On the master rank we need a valid history object to restore from");
315         state_.restoreFromHistory(*biasHistory, grid_);
316
317         /* Ensure that the state is consistent with our current run setup,
318          * since the user can have changed AWH parameters or the number
319          * of simulations sharing the bias.
320          */
321         ensureStateAndRunConsistency(params_, state_);
322
323         if (forceCorrelationGrid_ != nullptr)
324         {
325             forceCorrelationGrid_->restoreStateFromHistory(biasHistory->forceCorrelationGrid);
326         }
327     }
328
329     if (PAR(cr))
330     {
331         state_.broadcast(cr);
332     }
333 }
334
335 void Bias::initHistoryFromState(AwhBiasHistory* biasHistory) const
336 {
337     GMX_RELEASE_ASSERT(biasHistory != nullptr, "Need a valid biasHistory");
338
339     state_.initHistoryFromState(biasHistory);
340
341     if (forceCorrelationGrid_ != nullptr)
342     {
343         biasHistory->forceCorrelationGrid = initCorrelationGridHistoryFromState(forceCorrelationGrid());
344     }
345 }
346
347 void Bias::updateHistory(AwhBiasHistory* biasHistory) const
348 {
349     GMX_RELEASE_ASSERT(biasHistory != nullptr, "Need a valid biasHistory");
350
351     state_.updateHistory(biasHistory, grid_);
352
353     if (forceCorrelationGrid_ != nullptr)
354     {
355         updateCorrelationGridHistory(&biasHistory->forceCorrelationGrid, forceCorrelationGrid());
356     }
357 }
358
359 Bias::Bias(int                            biasIndexInCollection,
360            const AwhParams&               awhParams,
361            const AwhBiasParams&           awhBiasParams,
362            ArrayRef<const DimParams>      dimParamsInit,
363            double                         beta,
364            double                         mdTimeStep,
365            const BiasSharing*             biasSharing,
366            const std::string&             biasInitFilename,
367            ThisRankWillDoIO               thisRankWillDoIO,
368            BiasParams::DisableUpdateSkips disableUpdateSkips) :
369     dimParams_(dimParamsInit.begin(), dimParamsInit.end()),
370     grid_(dimParamsInit, awhBiasParams.dimParams()),
371     params_(awhParams,
372             awhBiasParams,
373             dimParams_,
374             beta,
375             mdTimeStep,
376             disableUpdateSkips,
377             biasSharing ? biasSharing->numSharingSimulations(biasIndexInCollection) : 1,
378             grid_.axis(),
379             biasIndexInCollection),
380     state_(awhBiasParams, params_.initialHistogramSize, dimParams_, grid_, biasSharing),
381     thisRankDoesIO_(thisRankWillDoIO == ThisRankWillDoIO::Yes),
382     biasForce_(ndim()),
383     tempForce_(ndim()),
384     numWarningsIssued_(0)
385 {
386     /* For a global update updateList covers all points, so reserve that */
387     updateList_.reserve(grid_.numPoints());
388
389     state_.initGridPointState(
390             awhBiasParams, dimParams_, grid_, params_, biasInitFilename, awhParams.numBias());
391
392     if (thisRankDoesIO_)
393     {
394         /* Set up the force correlation object. */
395
396         /* We let the correlation init function set its parameters
397          * to something useful for now.
398          */
399         double blockLength = 0;
400         /* Construct the force correlation object. */
401         forceCorrelationGrid_ =
402                 std::make_unique<CorrelationGrid>(state_.points().size(),
403                                                   ndim(),
404                                                   blockLength,
405                                                   CorrelationGrid::BlockLengthMeasure::Time,
406                                                   awhParams.nstSampleCoord() * mdTimeStep);
407
408         writer_ = std::make_unique<BiasWriter>(*this);
409     }
410 }
411
412 void Bias::printInitializationToLog(FILE* fplog) const
413 {
414     if (fplog != nullptr && forceCorrelationGrid_ != nullptr)
415     {
416         std::string prefix = gmx::formatString("\nawh%d:", params_.biasIndex + 1);
417
418         fprintf(fplog,
419                 "%s initial force correlation block length = %g %s"
420                 "%s force correlation number of blocks = %d",
421                 prefix.c_str(),
422                 forceCorrelationGrid().getBlockLength(),
423                 forceCorrelationGrid().blockLengthMeasure == CorrelationGrid::BlockLengthMeasure::Weight
424                         ? ""
425                         : "ps",
426                 prefix.c_str(),
427                 forceCorrelationGrid().getNumBlocks());
428     }
429 }
430
431 void Bias::updateForceCorrelationGrid(gmx::ArrayRef<const double> probWeightNeighbor,
432                                       ArrayRef<const double>      neighborLambdaDhdl,
433                                       double                      t)
434 {
435     if (forceCorrelationGrid_ == nullptr)
436     {
437         return;
438     }
439
440     const std::vector<int>& neighbor = grid_.point(state_.coordState().gridpointIndex()).neighbor;
441
442     gmx::ArrayRef<double> forceFromNeighbor = tempForce_;
443     for (size_t n = 0; n < neighbor.size(); n++)
444     {
445         double weightNeighbor = probWeightNeighbor[n];
446         int    indexNeighbor  = neighbor[n];
447
448         /* Add the force data of this neighbor point. Note: the sum of these forces is the convolved force.
449
450            We actually add the force normalized by beta which has the units of 1/length. This means that the
451            resulting correlation time integral is directly in units of friction time/length^2 which is really what
452            we're interested in. */
453         state_.calcUmbrellaForceAndPotential(
454                 dimParams_, grid_, indexNeighbor, neighborLambdaDhdl, forceFromNeighbor);
455
456         /* Note: we might want to give a whole list of data to add instead and have this loop in the data adding function */
457         forceCorrelationGrid_->addData(indexNeighbor, weightNeighbor, forceFromNeighbor, t);
458     }
459 }
460
461 /* Return the number of data blocks that have been prepared for writing. */
462 int Bias::numEnergySubblocksToWrite() const
463 {
464     GMX_RELEASE_ASSERT(writer_ != nullptr, "Should only request data from an initialized writer");
465
466     return writer_->numBlocks();
467 }
468
469 /* Write bias data blocks to energy subblocks. */
470 int Bias::writeToEnergySubblocks(t_enxsubblock* subblock) const
471 {
472     GMX_RELEASE_ASSERT(writer_ != nullptr, "Should only request data from an initialized writer");
473
474     return writer_->writeToEnergySubblocks(*this, subblock);
475 }
476
477 bool Bias::isSampleCoordStep(const int64_t step) const
478 {
479     return params_.isSampleCoordStep(step);
480 }
481
482
483 } // namespace gmx