95bca480dd83976d537c03d26d9e0167ae63a3c5
[alexxy/gromacs.git] / src / gromacs / awh / awh.cpp
1 /*
2  * This file is part of the GROMACS molecular simulation package.
3  *
4  * Copyright (c) 2015,2016,2017,2018, 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  * \ingroup module_awh
43  */
44
45 #include "gmxpre.h"
46
47 #include "awh.h"
48
49 #include <assert.h>
50
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/mdtypes/awh-history.h"
62 #include "gromacs/mdtypes/awh-params.h"
63 #include "gromacs/mdtypes/commrec.h"
64 #include "gromacs/mdtypes/forceoutput.h"
65 #include "gromacs/mdtypes/inputrec.h"
66 #include "gromacs/mdtypes/pull-params.h"
67 #include "gromacs/mdtypes/state.h"
68 #include "gromacs/pbcutil/pbc.h"
69 #include "gromacs/pulling/pull.h"
70 #include "gromacs/timing/wallcycle.h"
71 #include "gromacs/trajectory/energyframe.h"
72 #include "gromacs/utility/exceptions.h"
73 #include "gromacs/utility/gmxassert.h"
74 #include "gromacs/utility/pleasecite.h"
75
76 #include "bias.h"
77 #include "biassharing.h"
78 #include "correlationgrid.h"
79 #include "pointstate.h"
80
81 namespace gmx
82 {
83
84 /*! \internal
85  * \brief A bias and its coupling to the system.
86  *
87  * This struct is used to separate the bias machinery in the Bias class,
88  * which should be independent from the reaction coordinate, from the
89  * obtaining of the reaction coordinate values and passing the computed forces.
90  * Currently the AWH method couples to the system by mapping each
91  * AWH bias to a pull coordinate. This can easily be generalized here.
92  */
93 struct BiasCoupledToSystem
94 {
95     /*! \brief Constructor, couple a bias to a set of pull coordinates.
96      *
97      * \param[in] bias            The bias.
98      * \param[in] pullCoordIndex  The pull coordinate indices.
99      */
100     BiasCoupledToSystem(Bias                    bias,
101                         const std::vector<int> &pullCoordIndex);
102
103     Bias                   bias;           /**< The bias. */
104     const std::vector<int> pullCoordIndex; /**< The pull coordinates this bias acts on. */
105
106     /* Here AWH can be extended to work on other coordinates than pull. */
107 };
108
109 BiasCoupledToSystem::BiasCoupledToSystem(Bias                    bias,
110                                          const std::vector<int> &pullCoordIndex) :
111     bias(std::move(bias)),
112     pullCoordIndex(pullCoordIndex)
113 {
114     /* We already checked for this in grompp, but check again here. */
115     GMX_RELEASE_ASSERT(static_cast<size_t>(bias.ndim()) == pullCoordIndex.size(), "The bias dimensionality should match the number of pull coordinates.");
116 }
117
118 Awh::Awh(FILE                 *fplog,
119          const t_inputrec     &inputRecord,
120          const t_commrec      *commRecord,
121          const gmx_multisim_t *multiSimRecord,
122          const AwhParams      &awhParams,
123          const std::string    &biasInitFilename,
124          pull_t               *pull_work) :
125     seed_(awhParams.seed),
126     nstout_(awhParams.nstOut),
127     commRecord_(commRecord),
128     multiSimRecord_(multiSimRecord),
129     pull_(pull_work),
130     potentialOffset_(0)
131 {
132     /* We already checked for this in grompp, but check again here. */
133     GMX_RELEASE_ASSERT(inputRecord.pull != nullptr, "With AWH we should have pull parameters");
134     GMX_RELEASE_ASSERT(pull_work != nullptr, "With AWH pull should be initialized before initializing AWH");
135
136     if (fplog != nullptr)
137     {
138         please_cite(fplog, "Lindahl2014");
139     }
140
141     if (haveBiasSharingWithinSimulation(awhParams))
142     {
143         /* This has likely been checked by grompp, but throw anyhow. */
144         GMX_THROW(InvalidInputError("Biases within a simulation are shared, currently sharing of biases is only supported between simulations"));
145     }
146
147     int numSharingSimulations = 1;
148     if (awhParams.shareBiasMultisim && isMultiSim(multiSimRecord_))
149     {
150         numSharingSimulations = multiSimRecord_->nsim;
151     }
152
153     /* Initialize all the biases */
154     const double beta = 1/(BOLTZ*inputRecord.opts.ref_t[0]);
155     for (int k = 0; k < awhParams.numBias; k++)
156     {
157         const AwhBiasParams    &awhBiasParams = awhParams.awhBiasParams[k];
158
159         std::vector<int>        pullCoordIndex;
160         std::vector<DimParams>  dimParams;
161         for (int d = 0; d < awhBiasParams.ndim; d++)
162         {
163             const AwhDimParams &awhDimParams      = awhBiasParams.dimParams[d];
164             GMX_RELEASE_ASSERT(awhDimParams.eCoordProvider == eawhcoordproviderPULL, "Currently only the pull code is supported as coordinate provider");
165             const t_pull_coord &pullCoord         = inputRecord.pull->coord[awhDimParams.coordIndex];
166             double              conversionFactor  = pull_coordinate_is_angletype(&pullCoord) ? DEG2RAD : 1;
167             dimParams.emplace_back(conversionFactor, awhDimParams.forceConstant, beta);
168
169             pullCoordIndex.push_back(awhDimParams.coordIndex);
170         }
171
172         /* Construct the bias and couple it to the system. */
173         Bias::ThisRankWillDoIO thisRankWillDoIO = (MASTER(commRecord_) ? Bias::ThisRankWillDoIO::Yes : Bias::ThisRankWillDoIO::No);
174         biasCoupledToSystem_.emplace_back(Bias(k, awhParams, awhParams.awhBiasParams[k], dimParams, beta, inputRecord.delta_t, numSharingSimulations, biasInitFilename, thisRankWillDoIO),
175                                           pullCoordIndex);
176
177         biasCoupledToSystem_.back().bias.printInitializationToLog(fplog);
178     }
179
180     /* Need to register the AWH coordinates to be allowed to apply forces to the pull coordinates. */
181     registerAwhWithPull(awhParams, pull_);
182
183     if (numSharingSimulations > 1 && MASTER(commRecord_))
184     {
185         std::vector<size_t> pointSize;
186         for (auto const &biasCts : biasCoupledToSystem_)
187         {
188             pointSize.push_back(biasCts.bias.state().points().size());
189         }
190         /* Ensure that the shared biased are compatible between simulations */
191         biasesAreCompatibleForSharingBetweenSimulations(awhParams, pointSize, multiSimRecord_);
192     }
193 }
194
195 Awh::~Awh() = default;
196
197 bool Awh::isOutputStep(int64_t step) const
198 {
199     return (nstout_ > 0 && step % nstout_ == 0);
200 }
201
202 real Awh::applyBiasForcesAndUpdateBias(int                     ePBC,
203                                        const t_mdatoms        &mdatoms,
204                                        const matrix            box,
205                                        gmx::ForceWithVirial   *forceWithVirial,
206                                        double                  t,
207                                        int64_t                 step,
208                                        gmx_wallcycle          *wallcycle,
209                                        FILE                   *fplog)
210 {
211     GMX_ASSERT(forceWithVirial, "Need a valid ForceWithVirial object");
212
213     wallcycle_start(wallcycle, ewcAWH);
214
215     t_pbc  pbc;
216     set_pbc(&pbc, ePBC, box);
217
218     /* During the AWH update the potential can instantaneously jump due to either
219        an bias update or moving the umbrella. The jumps are kept track of and
220        subtracted from the potential in order to get a useful conserved energy quantity. */
221     double awhPotential = potentialOffset_;
222
223     for (auto &biasCts : biasCoupledToSystem_)
224     {
225         /* Update the AWH coordinate values with those of the corresponding
226          * pull coordinates.
227          */
228         awh_dvec coordValue = { 0, 0, 0, 0 };
229         for (int d = 0; d < biasCts.bias.ndim(); d++)
230         {
231             coordValue[d] = get_pull_coord_value(pull_, biasCts.pullCoordIndex[d], &pbc);
232         }
233
234         /* Perform an AWH biasing step: this means, at regular intervals,
235          * sampling observables based on the input pull coordinate value,
236          * setting the bias force and/or updating the AWH bias state.
237          */
238         double              biasPotential;
239         double              biasPotentialJump;
240         /* Note: In the near future this call will be split in calls
241          *       to supports bias sharing within a single simulation.
242          */
243         gmx::ArrayRef<const double> biasForce =
244             biasCts.bias.calcForceAndUpdateBias(coordValue,
245                                                 &biasPotential, &biasPotentialJump,
246                                                 commRecord_,
247                                                 multiSimRecord_,
248                                                 t, step, seed_, fplog);
249
250         awhPotential += biasPotential;
251
252         /* Keep track of the total potential shift needed to remove the potential jumps. */
253         potentialOffset_ -= biasPotentialJump;
254
255         /* Communicate the bias force to the pull struct.
256          * The bias potential is returned at the end of this function,
257          * so that it can be added externally to the correct energy data block.
258          */
259         for (int d = 0; d < biasCts.bias.ndim(); d++)
260         {
261             apply_external_pull_coord_force(pull_, biasCts.pullCoordIndex[d],
262                                             biasForce[d], &mdatoms,
263                                             forceWithVirial);
264         }
265
266         if (isOutputStep(step))
267         {
268             /* We might have skipped updates for part of the grid points.
269              * Ensure all points are updated before writing out their data.
270              */
271             biasCts.bias.doSkippedUpdatesForAllPoints();
272         }
273     }
274
275     wallcycle_stop(wallcycle, ewcAWH);
276
277     return MASTER(commRecord_) ? static_cast<real>(awhPotential) : 0;
278 }
279
280 std::shared_ptr<AwhHistory> Awh::initHistoryFromState() const
281 {
282     if (MASTER(commRecord_))
283     {
284         std::shared_ptr<AwhHistory> awhHistory(new AwhHistory);
285         awhHistory->bias.clear();
286         awhHistory->bias.resize(biasCoupledToSystem_.size());
287
288         for (size_t k = 0; k < awhHistory->bias.size(); k++)
289         {
290             biasCoupledToSystem_[k].bias.initHistoryFromState(&awhHistory->bias[k]);
291         }
292
293         return awhHistory;
294     }
295     else
296     {
297         /* Return an empty pointer */
298         return std::shared_ptr<AwhHistory>();
299     }
300 }
301
302 void Awh::restoreStateFromHistory(const AwhHistory *awhHistory)
303 {
304     /* Restore the history to the current state */
305     if (MASTER(commRecord_))
306     {
307         GMX_RELEASE_ASSERT(awhHistory != nullptr, "The master rank should have a valid awhHistory when restoring the state from history.");
308
309         if (awhHistory->bias.size() != biasCoupledToSystem_.size())
310         {
311             GMX_THROW(InvalidInputError("AWH state and history contain different numbers of biases. Likely you provided a checkpoint from a different simulation."));
312         }
313
314         potentialOffset_ = awhHistory->potentialOffset;
315     }
316     if (PAR(commRecord_))
317     {
318         gmx_bcast(sizeof(potentialOffset_), &potentialOffset_, commRecord_);
319     }
320
321     for (size_t k = 0; k < biasCoupledToSystem_.size(); k++)
322     {
323         biasCoupledToSystem_[k].bias.restoreStateFromHistory(awhHistory ? &awhHistory->bias[k] : nullptr, commRecord_);
324     }
325 }
326
327 void Awh::updateHistory(AwhHistory *awhHistory) const
328 {
329     if (!MASTER(commRecord_))
330     {
331         return;
332     }
333
334     /* This assert will also catch a non-master rank calling this function. */
335     GMX_RELEASE_ASSERT(awhHistory->bias.size() == biasCoupledToSystem_.size(), "AWH state and history bias count should match");
336
337     awhHistory->potentialOffset = potentialOffset_;
338
339     for (size_t k = 0; k < awhHistory->bias.size(); k++)
340     {
341         biasCoupledToSystem_[k].bias.updateHistory(&awhHistory->bias[k]);
342     }
343 }
344
345 const char * Awh::externalPotentialString()
346 {
347     return "AWH";
348 }
349
350 void Awh::registerAwhWithPull(const AwhParams &awhParams,
351                               pull_t          *pull_work)
352 {
353     GMX_RELEASE_ASSERT(pull_work, "Need a valid pull object");
354
355     for (int k = 0; k < awhParams.numBias; k++)
356     {
357         const AwhBiasParams &biasParams = awhParams.awhBiasParams[k];
358
359         for (int d = 0; d < biasParams.ndim; d++)
360         {
361             register_external_pull_potential(pull_work, biasParams.dimParams[d].coordIndex, Awh::externalPotentialString());
362         }
363     }
364 }
365
366 /* Fill the AWH data block of an energy frame with data (if there is any). */
367 void Awh::writeToEnergyFrame(int64_t      step,
368                              t_enxframe  *frame) const
369 {
370     GMX_ASSERT(MASTER(commRecord_), "writeToEnergyFrame should only be called on the master rank");
371     GMX_ASSERT(frame != nullptr, "Need a valid energy frame");
372
373     if (!isOutputStep(step))
374     {
375         /* This is not an AWH output step, don't write any AWH data */
376         return;
377     }
378
379     /* Get the total number of energy subblocks that AWH needs */
380     int numSubblocks  = 0;
381     for (auto &biasCoupledToSystem : biasCoupledToSystem_)
382     {
383         numSubblocks += biasCoupledToSystem.bias.numEnergySubblocksToWrite();
384     }
385     GMX_ASSERT(numSubblocks > 0, "We should always have data to write");
386
387     /* Add 1 energy block */
388     add_blocks_enxframe(frame, frame->nblock + 1);
389
390     /* Take the block that was just added and set the number of subblocks. */
391     t_enxblock *awhEnergyBlock = &(frame->block[frame->nblock - 1]);
392     add_subblocks_enxblock(awhEnergyBlock, numSubblocks);
393
394     /* Claim it as an AWH block. */
395     awhEnergyBlock->id = enxAWH;
396
397     /* Transfer AWH data blocks to energy sub blocks */
398     int energySubblockCount = 0;
399     for (auto &biasCoupledToSystem : biasCoupledToSystem_)
400     {
401         energySubblockCount += biasCoupledToSystem.bias.writeToEnergySubblocks(&(awhEnergyBlock->sub[energySubblockCount]));
402     }
403 }
404
405 std::unique_ptr<Awh>
406 prepareAwhModule(FILE                 *fplog,
407                  const t_inputrec     &inputRecord,
408                  t_state              *stateGlobal,
409                  const t_commrec      *commRecord,
410                  const gmx_multisim_t *multiSimRecord,
411                  const bool            startingFromCheckpoint,
412                  const bool            usingShellParticles,
413                  const std::string    &biasInitFilename,
414                  pull_t               *pull_work)
415 {
416     if (!inputRecord.bDoAwh)
417     {
418         return nullptr;
419     }
420     if (usingShellParticles)
421     {
422         GMX_THROW(InvalidInputError("AWH biasing does not support shell particles."));
423     }
424
425     auto awh = compat::make_unique<Awh>(fplog, inputRecord, commRecord, multiSimRecord, *inputRecord.awhParams,
426                                         biasInitFilename, pull_work);
427
428     if (startingFromCheckpoint)
429     {
430         // Restore the AWH history read from checkpoint
431         awh->restoreStateFromHistory(MASTER(commRecord) ? stateGlobal->awhHistory.get() : nullptr);
432     }
433     else if (MASTER(commRecord))
434     {
435         // Initialize the AWH history here
436         stateGlobal->awhHistory = awh->initHistoryFromState();
437     }
438     return awh;
439 }
440
441 } // namespace gmx