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