Use ObservablesReducer for LINCS RMSD computation
[alexxy/gromacs.git] / src / gromacs / modularsimulator / computeglobalselement.cpp
1 /*
2  * This file is part of the GROMACS molecular simulation package.
3  *
4  * Copyright (c) 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 /*! \internal \file
36  * \brief Defines the global reduction element for the modular simulator
37  *
38  * \author Pascal Merz <pascal.merz@me.com>
39  * \ingroup module_modularsimulator
40  */
41
42 #include "gmxpre.h"
43
44 #include "computeglobalselement.h"
45
46 #include "gromacs/domdec/domdec.h"
47 #include "gromacs/gmxlib/network.h"
48 #include "gromacs/gmxlib/nrnb.h"
49 #include "gromacs/math/vec.h"
50 #include "gromacs/mdlib/constr.h"
51 #include "gromacs/mdlib/md_support.h"
52 #include "gromacs/mdlib/mdatoms.h"
53 #include "gromacs/mdlib/stat.h"
54 #include "gromacs/mdlib/update.h"
55 #include "gromacs/mdlib/vcm.h"
56 #include "gromacs/mdtypes/commrec.h"
57 #include "gromacs/mdtypes/group.h"
58 #include "gromacs/mdtypes/inputrec.h"
59 #include "gromacs/mdtypes/mdatom.h"
60 #include "gromacs/topology/topology.h"
61
62 #include "freeenergyperturbationdata.h"
63 #include "modularsimulator.h"
64 #include "simulatoralgorithm.h"
65
66 namespace gmx
67 {
68 template<ComputeGlobalsAlgorithm algorithm>
69 ComputeGlobalsElement<algorithm>::ComputeGlobalsElement(StatePropagatorData* statePropagatorData,
70                                                         EnergyData*          energyData,
71                                                         FreeEnergyPerturbationData* freeEnergyPerturbationData,
72                                                         SimulationSignals*  signals,
73                                                         int                 nstglobalcomm,
74                                                         FILE*               fplog,
75                                                         const MDLogger&     mdlog,
76                                                         t_commrec*          cr,
77                                                         const t_inputrec*   inputrec,
78                                                         const MDAtoms*      mdAtoms,
79                                                         t_nrnb*             nrnb,
80                                                         gmx_wallcycle*      wcycle,
81                                                         t_forcerec*         fr,
82                                                         const gmx_mtop_t&   global_top,
83                                                         Constraints*        constr,
84                                                         ObservablesReducer* observablesReducer) :
85     energyReductionStep_(-1),
86     virialReductionStep_(-1),
87     vvSchedulingStep_(-1),
88     doStopCM_(inputrec->comm_mode != ComRemovalAlgorithm::No),
89     nstcomm_(inputrec->nstcomm),
90     nstglobalcomm_(nstglobalcomm),
91     lastStep_(inputrec->nsteps + inputrec->init_step),
92     initStep_(inputrec->init_step),
93     nullSignaller_(std::make_unique<SimulationSignaller>(nullptr, nullptr, nullptr, false, false)),
94     statePropagatorData_(statePropagatorData),
95     energyData_(energyData),
96     freeEnergyPerturbationData_(freeEnergyPerturbationData),
97     vcm_(global_top.groups, *inputrec),
98     signals_(signals),
99     fplog_(fplog),
100     mdlog_(mdlog),
101     cr_(cr),
102     inputrec_(inputrec),
103     top_global_(global_top),
104     mdAtoms_(mdAtoms),
105     constr_(constr),
106     nrnb_(nrnb),
107     wcycle_(wcycle),
108     fr_(fr),
109     observablesReducer_(observablesReducer)
110 {
111     reportComRemovalInfo(fplog, vcm_);
112     gstat_ = global_stat_init(inputrec_);
113 }
114
115 template<ComputeGlobalsAlgorithm algorithm>
116 ComputeGlobalsElement<algorithm>::~ComputeGlobalsElement()
117 {
118     global_stat_destroy(gstat_);
119 }
120
121 template<ComputeGlobalsAlgorithm algorithm>
122 void ComputeGlobalsElement<algorithm>::elementSetup()
123 {
124     if (doStopCM_ && !inputrec_->bContinuation)
125     {
126         // To minimize communication, compute_globals computes the COM velocity
127         // and the kinetic energy for the velocities without COM motion removed.
128         // Thus to get the kinetic energy without the COM contribution, we need
129         // to call compute_globals twice.
130
131         compute(-1, CGLO_GSTAT | CGLO_STOPCM, nullSignaller_.get(), false, true);
132         // Clean up after pre-step use of compute()
133         observablesReducer_->markAsReadyToReduce();
134
135         auto v = statePropagatorData_->velocitiesView();
136         // At initialization, do not pass x with acceleration-correction mode
137         // to avoid (incorrect) correction of the initial coordinates.
138         auto x = vcm_.mode == ComRemovalAlgorithm::LinearAccelerationCorrection
139                          ? ArrayRefWithPadding<RVec>()
140                          : statePropagatorData_->positionsView();
141         process_and_stopcm_grp(
142                 fplog_, &vcm_, *mdAtoms_->mdatoms(), x.unpaddedArrayRef(), v.unpaddedArrayRef());
143         inc_nrnb(nrnb_, eNR_STOPCM, mdAtoms_->mdatoms()->homenr);
144     }
145
146     unsigned int cglo_flags = (CGLO_TEMPERATURE | CGLO_GSTAT
147                                | (energyData_->hasReadEkinFromCheckpoint() ? CGLO_READEKIN : 0));
148
149     if (algorithm == ComputeGlobalsAlgorithm::VelocityVerlet)
150     {
151         cglo_flags |= CGLO_PRESSURE | CGLO_CONSTRAINT;
152     }
153
154     compute(-1, cglo_flags, nullSignaller_.get(), false, true);
155
156     // Calculate the initial half step temperature, and save the ekinh_old
157     for (int i = 0; (i < inputrec_->opts.ngtc); i++)
158     {
159         copy_mat(energyData_->ekindata()->tcstat[i].ekinh, energyData_->ekindata()->tcstat[i].ekinh_old);
160     }
161
162     // Clean up after pre-step use of compute()
163     observablesReducer_->markAsReadyToReduce();
164 }
165
166 template<ComputeGlobalsAlgorithm algorithm>
167 void ComputeGlobalsElement<algorithm>::scheduleTask(Step                       step,
168                                                     Time gmx_unused            time,
169                                                     const RegisterRunFunction& registerRunFunction)
170 {
171     const bool needComReduction    = doStopCM_ && do_per_step(step, nstcomm_);
172     const bool needGlobalReduction = step == energyReductionStep_ || step == virialReductionStep_
173                                      || needComReduction || do_per_step(step, nstglobalcomm_)
174                                      || (EI_VV(inputrec_->eI) && inputrecNvtTrotter(inputrec_)
175                                          && do_per_step(step - 1, nstglobalcomm_));
176
177     // TODO: CGLO_GSTAT is only used for needToSumEkinhOld_, i.e. to signal that we do or do not
178     //       sum the previous kinetic energy. We should simplify / clarify this.
179
180     if (algorithm == ComputeGlobalsAlgorithm::LeapFrog)
181     {
182         // With Leap-Frog we can skip compute_globals at
183         // non-communication steps, but we need to calculate
184         // the kinetic energy one step before communication.
185
186         // With leap-frog we also need to compute the half-step
187         // kinetic energy at the step before we need to write
188         // the full-step kinetic energy
189         const bool needEkinAtNextStep = (do_per_step(step + 1, nstglobalcomm_) || step + 1 == lastStep_);
190
191         if (!needGlobalReduction && !needEkinAtNextStep)
192         {
193             return;
194         }
195
196         const bool doEnergy = step == energyReductionStep_;
197         int        flags    = (needGlobalReduction ? CGLO_GSTAT : 0) | (doEnergy ? CGLO_ENERGY : 0)
198                     | (needComReduction ? CGLO_STOPCM : 0) | CGLO_TEMPERATURE | CGLO_PRESSURE
199                     | CGLO_CONSTRAINT;
200
201         // Since we're already communicating at this step, we
202         // can propagate intra-simulation signals. Note that
203         // check_nstglobalcomm has the responsibility for
204         // choosing the value of nstglobalcomm which satisfies
205         // the need of the different signallers.
206         const bool doIntraSimSignal = true;
207         // Disable functionality
208         const bool doInterSimSignal = false;
209
210         // Make signaller to signal stop / reset / checkpointing signals
211         auto signaller = std::make_shared<SimulationSignaller>(
212                 signals_, cr_, nullptr, doInterSimSignal, doIntraSimSignal);
213
214         registerRunFunction([this, step, flags, signaller = std::move(signaller)]() {
215             compute(step, flags, signaller.get(), true);
216         });
217     }
218     else if (algorithm == ComputeGlobalsAlgorithm::VelocityVerlet)
219     {
220         // For VV, we schedule two calls to compute globals per step.
221         if (step != vvSchedulingStep_)
222         {
223             // This is the first scheduling call for this step (positions & velocities at full time
224             // step) Set this as the current scheduling step
225             vvSchedulingStep_ = step;
226
227             // For vv, the state at the beginning of the step is positions at time t, velocities at time t - dt/2
228             // The first velocity propagation (+dt/2) therefore actually corresponds to the previous step.
229             // So we need information from the last step in the first half of the integration
230             if (!needGlobalReduction && !do_per_step(step - 1, nstglobalcomm_))
231             {
232                 return;
233             }
234
235             const bool doTemperature = step != initStep_ || inputrec_->bContinuation;
236             const bool doEnergy      = step == energyReductionStep_;
237
238             int flags = (needGlobalReduction ? CGLO_GSTAT : 0) | (doEnergy ? CGLO_ENERGY : 0)
239                         | (doTemperature ? CGLO_TEMPERATURE : 0) | CGLO_PRESSURE | CGLO_CONSTRAINT
240                         | (needComReduction ? CGLO_STOPCM : 0) | CGLO_SCALEEKIN;
241
242             registerRunFunction(
243                     [this, step, flags]() { compute(step, flags, nullSignaller_.get(), false); });
244         }
245         else
246         {
247             // second call to compute_globals for this step
248             // Reset the scheduling step to avoid confusion if scheduling needs
249             // to be repeated (in case of unexpected simulation termination)
250             vvSchedulingStep_ = -1;
251
252             if (!needGlobalReduction)
253             {
254                 return;
255             }
256             int flags = CGLO_GSTAT | CGLO_CONSTRAINT;
257
258             // Since we're already communicating at this step, we
259             // can propagate intra-simulation signals. Note that
260             // check_nstglobalcomm has the responsibility for
261             // choosing the value of nstglobalcomm which satisfies
262             // the need of the different signallers.
263             const bool doIntraSimSignal = true;
264             // Disable functionality
265             const bool doInterSimSignal = false;
266
267             auto signaller = std::make_shared<SimulationSignaller>(
268                     signals_, cr_, nullptr, doInterSimSignal, doIntraSimSignal);
269
270             registerRunFunction([this, step, flags, signaller = std::move(signaller)]() {
271                 compute(step, flags, signaller.get(), true);
272             });
273         }
274     }
275 }
276
277 template<ComputeGlobalsAlgorithm algorithm>
278 void ComputeGlobalsElement<algorithm>::compute(gmx::Step            step,
279                                                unsigned int         flags,
280                                                SimulationSignaller* signaller,
281                                                bool                 useLastBox,
282                                                bool                 isInit)
283 {
284     auto        x       = statePropagatorData_->positionsView().unpaddedArrayRef();
285     auto        v       = statePropagatorData_->velocitiesView().unpaddedArrayRef();
286     const auto* box     = statePropagatorData_->constBox();
287     const auto* lastbox = useLastBox ? statePropagatorData_->constPreviousBox()
288                                      : statePropagatorData_->constBox();
289
290     compute_globals(gstat_,
291                     cr_,
292                     inputrec_,
293                     fr_,
294                     energyData_->ekindata(),
295                     x,
296                     v,
297                     box,
298                     mdAtoms_->mdatoms(),
299                     nrnb_,
300                     &vcm_,
301                     step != -1 ? wcycle_ : nullptr,
302                     energyData_->enerdata(),
303                     energyData_->forceVirial(step),
304                     energyData_->constraintVirial(step),
305                     energyData_->totalVirial(step),
306                     energyData_->pressure(step),
307                     signaller,
308                     lastbox,
309                     energyData_->needToSumEkinhOld(),
310                     flags,
311                     step,
312                     observablesReducer_);
313     if (flags & CGLO_STOPCM && !isInit)
314     {
315         process_and_stopcm_grp(fplog_, &vcm_, *mdAtoms_->mdatoms(), x, v);
316         inc_nrnb(nrnb_, eNR_STOPCM, mdAtoms_->mdatoms()->homenr);
317     }
318 }
319
320 template<ComputeGlobalsAlgorithm algorithm>
321 std::optional<SignallerCallback> ComputeGlobalsElement<algorithm>::registerEnergyCallback(EnergySignallerEvent event)
322 {
323     if (event == EnergySignallerEvent::EnergyCalculationStep)
324     {
325         return [this](Step step, Time /*unused*/) { energyReductionStep_ = step; };
326     }
327     if (event == EnergySignallerEvent::VirialCalculationStep)
328     {
329         return [this](Step step, Time /*unused*/) { virialReductionStep_ = step; };
330     }
331     return std::nullopt;
332 }
333
334 template<ComputeGlobalsAlgorithm algorithm>
335 std::optional<SignallerCallback>
336 ComputeGlobalsElement<algorithm>::registerTrajectorySignallerCallback(TrajectoryEvent event)
337 {
338     if (event == TrajectoryEvent::EnergyWritingStep)
339     {
340         return [this](Step step, Time /*unused*/) { energyReductionStep_ = step; };
341     }
342     return std::nullopt;
343 }
344
345 namespace
346 {
347
348 /*! \brief Schedule a function for actions that must happen at the end of each step
349  *
350  * After reduction, an ObservablesReducer is marked as unavailable for
351  * further reduction this step. This needs to be reset in order to be
352  * used on the next step.
353  *
354  * \param[in]  observablesReducer The ObservablesReducer to mark as ready for use
355  */
356 SchedulingFunction registerPostStepSchedulingFunction(ObservablesReducer* observablesReducer)
357 {
358     SchedulingFunction postStepSchedulingFunction =
359             [observablesReducer](
360                     Step /*step*/, Time /*time*/, const RegisterRunFunction& registerRunFunction) {
361                 SimulatorRunFunction completeObservablesReducerStep = [&observablesReducer]() {
362                     observablesReducer->markAsReadyToReduce();
363                 };
364                 registerRunFunction(completeObservablesReducerStep);
365             };
366     return postStepSchedulingFunction;
367 }
368
369 } // namespace
370
371 //! Explicit template instantiation
372 //! \{
373 template class ComputeGlobalsElement<ComputeGlobalsAlgorithm::LeapFrog>;
374 template class ComputeGlobalsElement<ComputeGlobalsAlgorithm::VelocityVerlet>;
375 //! \}
376
377 template<>
378 ISimulatorElement* ComputeGlobalsElement<ComputeGlobalsAlgorithm::LeapFrog>::getElementPointerImpl(
379         LegacySimulatorData*                    legacySimulatorData,
380         ModularSimulatorAlgorithmBuilderHelper* builderHelper,
381         StatePropagatorData*                    statePropagatorData,
382         EnergyData*                             energyData,
383         FreeEnergyPerturbationData*             freeEnergyPerturbationData,
384         GlobalCommunicationHelper*              globalCommunicationHelper,
385         ObservablesReducer*                     observablesReducer)
386 {
387     ComputeGlobalsElement* element = builderHelper->storeElement(
388             std::make_unique<ComputeGlobalsElement<ComputeGlobalsAlgorithm::LeapFrog>>(
389                     statePropagatorData,
390                     energyData,
391                     freeEnergyPerturbationData,
392                     globalCommunicationHelper->simulationSignals(),
393                     globalCommunicationHelper->nstglobalcomm(),
394                     legacySimulatorData->fplog,
395                     legacySimulatorData->mdlog,
396                     legacySimulatorData->cr,
397                     legacySimulatorData->inputrec,
398                     legacySimulatorData->mdAtoms,
399                     legacySimulatorData->nrnb,
400                     legacySimulatorData->wcycle,
401                     legacySimulatorData->fr,
402                     legacySimulatorData->top_global,
403                     legacySimulatorData->constr,
404                     observablesReducer));
405     builderHelper->registerPostStepScheduling(
406             registerPostStepSchedulingFunction(element->observablesReducer_));
407
408     return element;
409 }
410
411 template<>
412 ISimulatorElement* ComputeGlobalsElement<ComputeGlobalsAlgorithm::VelocityVerlet>::getElementPointerImpl(
413         LegacySimulatorData*                    simulator,
414         ModularSimulatorAlgorithmBuilderHelper* builderHelper,
415         StatePropagatorData*                    statePropagatorData,
416         EnergyData*                             energyData,
417         FreeEnergyPerturbationData*             freeEnergyPerturbationData,
418         GlobalCommunicationHelper*              globalCommunicationHelper,
419         ObservablesReducer*                     observablesReducer)
420 {
421     // We allow this element to be added multiple times to the call list, but we only want one
422     // actual element built
423     static const std::string key("vvComputeGlobalsElement");
424
425     const std::optional<std::any> cachedValue = builderHelper->builderData(key);
426
427     if (cachedValue)
428     {
429         return std::any_cast<ComputeGlobalsElement*>(cachedValue.value());
430     }
431     else
432     {
433         ComputeGlobalsElement* vvComputeGlobalsElement = builderHelper->storeElement(
434                 std::make_unique<ComputeGlobalsElement<ComputeGlobalsAlgorithm::VelocityVerlet>>(
435                         statePropagatorData,
436                         energyData,
437                         freeEnergyPerturbationData,
438                         globalCommunicationHelper->simulationSignals(),
439                         globalCommunicationHelper->nstglobalcomm(),
440                         simulator->fplog,
441                         simulator->mdlog,
442                         simulator->cr,
443                         simulator->inputrec,
444                         simulator->mdAtoms,
445                         simulator->nrnb,
446                         simulator->wcycle,
447                         simulator->fr,
448                         simulator->top_global,
449                         simulator->constr,
450                         observablesReducer));
451         builderHelper->storeBuilderData(key, vvComputeGlobalsElement);
452         builderHelper->registerPostStepScheduling(
453                 registerPostStepSchedulingFunction(vvComputeGlobalsElement->observablesReducer_));
454         return vvComputeGlobalsElement;
455     }
456 }
457 } // namespace gmx