b40abb9b1e6e2dadf7780ded7e1c69f6fdcab835
[alexxy/gromacs.git] / src / gromacs / modularsimulator / propagator.cpp
1 /*
2  * This file is part of the GROMACS molecular simulation package.
3  *
4  * Copyright (c) 2019,2020, 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 propagator 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 "propagator.h"
45
46 #include "gromacs/utility.h"
47 #include "gromacs/math/vec.h"
48 #include "gromacs/math/vectypes.h"
49 #include "gromacs/mdlib/gmx_omp_nthreads.h"
50 #include "gromacs/mdlib/mdatoms.h"
51 #include "gromacs/mdlib/update.h"
52 #include "gromacs/mdtypes/inputrec.h"
53 #include "gromacs/mdtypes/mdatom.h"
54 #include "gromacs/timing/wallcycle.h"
55 #include "gromacs/utility/fatalerror.h"
56
57 #include "modularsimulator.h"
58 #include "simulatoralgorithm.h"
59 #include "statepropagatordata.h"
60
61 namespace gmx
62 {
63 //! Update velocities
64 template<NumVelocityScalingValues numVelocityScalingValues, ParrinelloRahmanVelocityScaling parrinelloRahmanVelocityScaling>
65 static void inline updateVelocities(int         a,
66                                     real        dt,
67                                     real        lambda,
68                                     const rvec* gmx_restrict invMassPerDim,
69                                     rvec* gmx_restrict v,
70                                     const rvec* gmx_restrict f,
71                                     const rvec               diagPR,
72                                     const matrix             matrixPR)
73 {
74     for (int d = 0; d < DIM; d++)
75     {
76         // TODO: Extract this into policy classes
77         if (numVelocityScalingValues != NumVelocityScalingValues::None
78             && parrinelloRahmanVelocityScaling == ParrinelloRahmanVelocityScaling::No)
79         {
80             v[a][d] *= lambda;
81         }
82         if (numVelocityScalingValues != NumVelocityScalingValues::None
83             && parrinelloRahmanVelocityScaling == ParrinelloRahmanVelocityScaling::Diagonal)
84         {
85             v[a][d] *= (lambda - diagPR[d]);
86         }
87         if (numVelocityScalingValues != NumVelocityScalingValues::None
88             && parrinelloRahmanVelocityScaling == ParrinelloRahmanVelocityScaling::Full)
89         {
90             v[a][d] = lambda * v[a][d] - iprod(matrixPR[d], v[a]);
91         }
92         if (numVelocityScalingValues == NumVelocityScalingValues::None
93             && parrinelloRahmanVelocityScaling == ParrinelloRahmanVelocityScaling::Diagonal)
94         {
95             v[a][d] *= (1 - diagPR[d]);
96         }
97         if (numVelocityScalingValues == NumVelocityScalingValues::None
98             && parrinelloRahmanVelocityScaling == ParrinelloRahmanVelocityScaling::Full)
99         {
100             v[a][d] -= iprod(matrixPR[d], v[a]);
101         }
102         v[a][d] += f[a][d] * invMassPerDim[a][d] * dt;
103     }
104 }
105
106 //! Update positions
107 static void inline updatePositions(int         a,
108                                    real        dt,
109                                    const rvec* gmx_restrict x,
110                                    rvec* gmx_restrict xprime,
111                                    const rvec* gmx_restrict v)
112 {
113     for (int d = 0; d < DIM; d++)
114     {
115         xprime[a][d] = x[a][d] + v[a][d] * dt;
116     }
117 }
118
119 //! Helper function diagonalizing the PR matrix if possible
120 template<ParrinelloRahmanVelocityScaling parrinelloRahmanVelocityScaling>
121 static inline bool diagonalizePRMatrix(matrix matrixPR, rvec diagPR)
122 {
123     if (parrinelloRahmanVelocityScaling != ParrinelloRahmanVelocityScaling::Full)
124     {
125         return false;
126     }
127     else
128     {
129         if (matrixPR[YY][XX] == 0 && matrixPR[ZZ][XX] == 0 && matrixPR[ZZ][YY] == 0)
130         {
131             diagPR[XX] = matrixPR[XX][XX];
132             diagPR[YY] = matrixPR[YY][YY];
133             diagPR[ZZ] = matrixPR[ZZ][ZZ];
134             return true;
135         }
136         else
137         {
138             return false;
139         }
140     }
141 }
142
143 //! Propagation (position only)
144 template<>
145 template<NumVelocityScalingValues numVelocityScalingValues, ParrinelloRahmanVelocityScaling parrinelloRahmanVelocityScaling>
146 void Propagator<IntegrationStep::PositionsOnly>::run()
147 {
148     wallcycle_start(wcycle_, ewcUPDATE);
149
150     auto xp = as_rvec_array(statePropagatorData_->positionsView().paddedArrayRef().data());
151     auto x = as_rvec_array(statePropagatorData_->constPreviousPositionsView().paddedArrayRef().data());
152     auto v = as_rvec_array(statePropagatorData_->constVelocitiesView().paddedArrayRef().data());
153
154     int nth    = gmx_omp_nthreads_get(emntUpdate);
155     int homenr = mdAtoms_->mdatoms()->homenr;
156
157 #pragma omp parallel for num_threads(nth) schedule(static) default(none) shared(nth, homenr, x, xp, v)
158     for (int th = 0; th < nth; th++)
159     {
160         try
161         {
162             int start_th, end_th;
163             getThreadAtomRange(nth, th, homenr, &start_th, &end_th);
164
165             for (int a = start_th; a < end_th; a++)
166             {
167                 updatePositions(a, timestep_, x, xp, v);
168             }
169         }
170         GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR
171     }
172     wallcycle_stop(wcycle_, ewcUPDATE);
173 }
174
175 //! Propagation (velocity only)
176 template<>
177 template<NumVelocityScalingValues numVelocityScalingValues, ParrinelloRahmanVelocityScaling parrinelloRahmanVelocityScaling>
178 void Propagator<IntegrationStep::VelocitiesOnly>::run()
179 {
180     wallcycle_start(wcycle_, ewcUPDATE);
181
182     auto v = as_rvec_array(statePropagatorData_->velocitiesView().paddedArrayRef().data());
183     auto f = as_rvec_array(statePropagatorData_->constForcesView().force().data());
184     auto invMassPerDim = mdAtoms_->mdatoms()->invMassPerDim;
185
186     const real lambda =
187             (numVelocityScalingValues == NumVelocityScalingValues::Single) ? velocityScaling_[0] : 1.0;
188
189     const bool isFullScalingMatrixDiagonal =
190             diagonalizePRMatrix<parrinelloRahmanVelocityScaling>(matrixPR_, diagPR_);
191
192     const int nth    = gmx_omp_nthreads_get(emntUpdate);
193     const int homenr = mdAtoms_->mdatoms()->homenr;
194
195 // const variables could be shared, but gcc-8 & gcc-9 don't agree how to write that...
196 // https://www.gnu.org/software/gcc/gcc-9/porting_to.html -> OpenMP data sharing
197 #pragma omp parallel for num_threads(nth) schedule(static) default(none) \
198         shared(v, f, invMassPerDim) firstprivate(nth, homenr, lambda, isFullScalingMatrixDiagonal)
199     for (int th = 0; th < nth; th++)
200     {
201         try
202         {
203             int start_th, end_th;
204             getThreadAtomRange(nth, th, homenr, &start_th, &end_th);
205
206             for (int a = start_th; a < end_th; a++)
207             {
208                 if (isFullScalingMatrixDiagonal)
209                 {
210                     updateVelocities<numVelocityScalingValues, ParrinelloRahmanVelocityScaling::Diagonal>(
211                             a, timestep_,
212                             numVelocityScalingValues == NumVelocityScalingValues::Multiple
213                                     ? velocityScaling_[mdAtoms_->mdatoms()->cTC[a]]
214                                     : lambda,
215                             invMassPerDim, v, f, diagPR_, matrixPR_);
216                 }
217                 else
218                 {
219                     updateVelocities<numVelocityScalingValues, parrinelloRahmanVelocityScaling>(
220                             a, timestep_,
221                             numVelocityScalingValues == NumVelocityScalingValues::Multiple
222                                     ? velocityScaling_[mdAtoms_->mdatoms()->cTC[a]]
223                                     : lambda,
224                             invMassPerDim, v, f, diagPR_, matrixPR_);
225                 }
226             }
227         }
228         GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR
229     }
230     wallcycle_stop(wcycle_, ewcUPDATE);
231 }
232
233 //! Propagation (leapfrog case - position and velocity)
234 template<>
235 template<NumVelocityScalingValues numVelocityScalingValues, ParrinelloRahmanVelocityScaling parrinelloRahmanVelocityScaling>
236 void Propagator<IntegrationStep::LeapFrog>::run()
237 {
238     wallcycle_start(wcycle_, ewcUPDATE);
239
240     auto xp = as_rvec_array(statePropagatorData_->positionsView().paddedArrayRef().data());
241     auto x = as_rvec_array(statePropagatorData_->constPreviousPositionsView().paddedArrayRef().data());
242     auto v = as_rvec_array(statePropagatorData_->velocitiesView().paddedArrayRef().data());
243     auto f = as_rvec_array(statePropagatorData_->constForcesView().force().data());
244     auto invMassPerDim = mdAtoms_->mdatoms()->invMassPerDim;
245
246     const real lambda =
247             (numVelocityScalingValues == NumVelocityScalingValues::Single) ? velocityScaling_[0] : 1.0;
248
249     const bool isFullScalingMatrixDiagonal =
250             diagonalizePRMatrix<parrinelloRahmanVelocityScaling>(matrixPR_, diagPR_);
251
252     const int nth    = gmx_omp_nthreads_get(emntUpdate);
253     const int homenr = mdAtoms_->mdatoms()->homenr;
254
255 // const variables could be shared, but gcc-8 & gcc-9 don't agree how to write that...
256 // https://www.gnu.org/software/gcc/gcc-9/porting_to.html -> OpenMP data sharing
257 #pragma omp parallel for num_threads(nth) schedule(static) default(none) shared( \
258         x, xp, v, f, invMassPerDim) firstprivate(nth, homenr, lambda, isFullScalingMatrixDiagonal)
259     for (int th = 0; th < nth; th++)
260     {
261         try
262         {
263             int start_th, end_th;
264             getThreadAtomRange(nth, th, homenr, &start_th, &end_th);
265
266             for (int a = start_th; a < end_th; a++)
267             {
268                 if (isFullScalingMatrixDiagonal)
269                 {
270                     updateVelocities<numVelocityScalingValues, ParrinelloRahmanVelocityScaling::Diagonal>(
271                             a, timestep_,
272                             numVelocityScalingValues == NumVelocityScalingValues::Multiple
273                                     ? velocityScaling_[mdAtoms_->mdatoms()->cTC[a]]
274                                     : lambda,
275                             invMassPerDim, v, f, diagPR_, matrixPR_);
276                 }
277                 else
278                 {
279                     updateVelocities<numVelocityScalingValues, parrinelloRahmanVelocityScaling>(
280                             a, timestep_,
281                             numVelocityScalingValues == NumVelocityScalingValues::Multiple
282                                     ? velocityScaling_[mdAtoms_->mdatoms()->cTC[a]]
283                                     : lambda,
284                             invMassPerDim, v, f, diagPR_, matrixPR_);
285                 }
286                 updatePositions(a, timestep_, x, xp, v);
287             }
288         }
289         GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR
290     }
291     wallcycle_stop(wcycle_, ewcUPDATE);
292 }
293
294 //! Propagation (velocity verlet stage 2 - velocity and position)
295 template<>
296 template<NumVelocityScalingValues numVelocityScalingValues, ParrinelloRahmanVelocityScaling parrinelloRahmanVelocityScaling>
297 void Propagator<IntegrationStep::VelocityVerletPositionsAndVelocities>::run()
298 {
299     wallcycle_start(wcycle_, ewcUPDATE);
300
301     auto xp = as_rvec_array(statePropagatorData_->positionsView().paddedArrayRef().data());
302     auto x = as_rvec_array(statePropagatorData_->constPreviousPositionsView().paddedArrayRef().data());
303     auto v = as_rvec_array(statePropagatorData_->velocitiesView().paddedArrayRef().data());
304     auto f = as_rvec_array(statePropagatorData_->constForcesView().force().data());
305     auto invMassPerDim = mdAtoms_->mdatoms()->invMassPerDim;
306
307     const real lambda =
308             (numVelocityScalingValues == NumVelocityScalingValues::Single) ? velocityScaling_[0] : 1.0;
309
310     const bool isFullScalingMatrixDiagonal =
311             diagonalizePRMatrix<parrinelloRahmanVelocityScaling>(matrixPR_, diagPR_);
312
313     const int nth    = gmx_omp_nthreads_get(emntUpdate);
314     const int homenr = mdAtoms_->mdatoms()->homenr;
315
316 // const variables could be shared, but gcc-8 & gcc-9 don't agree how to write that...
317 // https://www.gnu.org/software/gcc/gcc-9/porting_to.html -> OpenMP data sharing
318 #pragma omp parallel for num_threads(nth) schedule(static) default(none) shared( \
319         x, xp, v, f, invMassPerDim) firstprivate(nth, homenr, lambda, isFullScalingMatrixDiagonal)
320     for (int th = 0; th < nth; th++)
321     {
322         try
323         {
324             int start_th, end_th;
325             getThreadAtomRange(nth, th, homenr, &start_th, &end_th);
326
327             for (int a = start_th; a < end_th; a++)
328             {
329                 if (isFullScalingMatrixDiagonal)
330                 {
331                     updateVelocities<numVelocityScalingValues, ParrinelloRahmanVelocityScaling::Diagonal>(
332                             a, 0.5 * timestep_,
333                             numVelocityScalingValues == NumVelocityScalingValues::Multiple
334                                     ? velocityScaling_[mdAtoms_->mdatoms()->cTC[a]]
335                                     : lambda,
336                             invMassPerDim, v, f, diagPR_, matrixPR_);
337                 }
338                 else
339                 {
340                     updateVelocities<numVelocityScalingValues, parrinelloRahmanVelocityScaling>(
341                             a, 0.5 * timestep_,
342                             numVelocityScalingValues == NumVelocityScalingValues::Multiple
343                                     ? velocityScaling_[mdAtoms_->mdatoms()->cTC[a]]
344                                     : lambda,
345                             invMassPerDim, v, f, diagPR_, matrixPR_);
346                 }
347                 updatePositions(a, timestep_, x, xp, v);
348             }
349         }
350         GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR
351     }
352     wallcycle_stop(wcycle_, ewcUPDATE);
353 }
354
355 template<IntegrationStep algorithm>
356 Propagator<algorithm>::Propagator(double               timestep,
357                                   StatePropagatorData* statePropagatorData,
358                                   const MDAtoms*       mdAtoms,
359                                   gmx_wallcycle*       wcycle) :
360     timestep_(timestep),
361     statePropagatorData_(statePropagatorData),
362     doSingleVelocityScaling_(false),
363     doGroupVelocityScaling_(false),
364     scalingStepVelocity_(-1),
365     diagPR_{ 0 },
366     matrixPR_{ { 0 } },
367     scalingStepPR_(-1),
368     mdAtoms_(mdAtoms),
369     wcycle_(wcycle)
370 {
371 }
372
373 template<IntegrationStep algorithm>
374 void Propagator<algorithm>::scheduleTask(Step gmx_unused step,
375                                          Time gmx_unused            time,
376                                          const RegisterRunFunction& registerRunFunction)
377 {
378     const bool doSingleVScalingThisStep = (doSingleVelocityScaling_ && (step == scalingStepVelocity_));
379     const bool doGroupVScalingThisStep = (doGroupVelocityScaling_ && (step == scalingStepVelocity_));
380
381     const bool doParrinelloRahmanThisStep = (step == scalingStepPR_);
382
383     if (doSingleVScalingThisStep)
384     {
385         if (doParrinelloRahmanThisStep)
386         {
387             registerRunFunction([this]() {
388                 run<NumVelocityScalingValues::Single, ParrinelloRahmanVelocityScaling::Full>();
389             });
390         }
391         else
392         {
393             registerRunFunction([this]() {
394                 run<NumVelocityScalingValues::Single, ParrinelloRahmanVelocityScaling::No>();
395             });
396         }
397     }
398     else if (doGroupVScalingThisStep)
399     {
400         if (doParrinelloRahmanThisStep)
401         {
402             registerRunFunction([this]() {
403                 run<NumVelocityScalingValues::Multiple, ParrinelloRahmanVelocityScaling::Full>();
404             });
405         }
406         else
407         {
408             registerRunFunction([this]() {
409                 run<NumVelocityScalingValues::Multiple, ParrinelloRahmanVelocityScaling::No>();
410             });
411         }
412     }
413     else
414     {
415         if (doParrinelloRahmanThisStep)
416         {
417             registerRunFunction([this]() {
418                 run<NumVelocityScalingValues::None, ParrinelloRahmanVelocityScaling::Full>();
419             });
420         }
421         else
422         {
423             registerRunFunction([this]() {
424                 run<NumVelocityScalingValues::None, ParrinelloRahmanVelocityScaling::No>();
425             });
426         }
427     }
428 }
429
430 template<IntegrationStep algorithm>
431 void Propagator<algorithm>::setNumVelocityScalingVariables(int numVelocityScalingVariables)
432 {
433     if (algorithm == IntegrationStep::PositionsOnly)
434     {
435         gmx_fatal(FARGS, "Velocity scaling not implemented for IntegrationStep::PositionsOnly.");
436     }
437     GMX_ASSERT(velocityScaling_.empty(),
438                "Number of velocity scaling variables cannot be changed once set.");
439
440     velocityScaling_.resize(numVelocityScalingVariables, 1.);
441     doSingleVelocityScaling_ = numVelocityScalingVariables == 1;
442     doGroupVelocityScaling_  = numVelocityScalingVariables > 1;
443 }
444
445 template<IntegrationStep algorithm>
446 ArrayRef<real> Propagator<algorithm>::viewOnVelocityScaling()
447 {
448     if (algorithm == IntegrationStep::PositionsOnly)
449     {
450         gmx_fatal(FARGS, "Velocity scaling not implemented for IntegrationStep::PositionsOnly.");
451     }
452     GMX_ASSERT(!velocityScaling_.empty(), "Number of velocity scaling variables not set.");
453
454     return velocityScaling_;
455 }
456
457 template<IntegrationStep algorithm>
458 PropagatorCallback Propagator<algorithm>::velocityScalingCallback()
459 {
460     if (algorithm == IntegrationStep::PositionsOnly)
461     {
462         gmx_fatal(FARGS, "Velocity scaling not implemented for IntegrationStep::PositionsOnly.");
463     }
464
465     return [this](Step step) { scalingStepVelocity_ = step; };
466 }
467
468 template<IntegrationStep algorithm>
469 ArrayRef<rvec> Propagator<algorithm>::viewOnPRScalingMatrix()
470 {
471     GMX_RELEASE_ASSERT(
472             algorithm != IntegrationStep::PositionsOnly,
473             "Parrinello-Rahman scaling not implemented for IntegrationStep::PositionsOnly.");
474
475     clear_mat(matrixPR_);
476     // gcc-5 needs this to be explicit (all other tested compilers would be ok
477     // with simply returning matrixPR)
478     return ArrayRef<rvec>(matrixPR_);
479 }
480
481 template<IntegrationStep algorithm>
482 PropagatorCallback Propagator<algorithm>::prScalingCallback()
483 {
484     GMX_RELEASE_ASSERT(
485             algorithm != IntegrationStep::PositionsOnly,
486             "Parrinello-Rahman scaling not implemented for IntegrationStep::PositionsOnly.");
487
488     return [this](Step step) { scalingStepPR_ = step; };
489 }
490
491 template<IntegrationStep algorithm>
492 ISimulatorElement* Propagator<algorithm>::getElementPointerImpl(
493         LegacySimulatorData*                    legacySimulatorData,
494         ModularSimulatorAlgorithmBuilderHelper* builderHelper,
495         StatePropagatorData*                    statePropagatorData,
496         EnergyData gmx_unused*     energyData,
497         FreeEnergyPerturbationData gmx_unused* freeEnergyPerturbationData,
498         GlobalCommunicationHelper gmx_unused* globalCommunicationHelper,
499         double                                timestep,
500         RegisterWithThermostat                registerWithThermostat,
501         RegisterWithBarostat                  registerWithBarostat)
502 {
503     auto* element = builderHelper->storeElement(std::make_unique<Propagator<algorithm>>(
504             timestep, statePropagatorData, legacySimulatorData->mdAtoms, legacySimulatorData->wcycle));
505     if (registerWithThermostat == RegisterWithThermostat::True)
506     {
507         auto* propagator = static_cast<Propagator<algorithm>*>(element);
508         builderHelper->registerWithThermostat(
509                 { [propagator](int num) { propagator->setNumVelocityScalingVariables(num); },
510                   [propagator]() { return propagator->viewOnVelocityScaling(); },
511                   [propagator]() { return propagator->velocityScalingCallback(); } });
512     }
513     if (registerWithBarostat == RegisterWithBarostat::True)
514     {
515         auto* propagator = static_cast<Propagator<algorithm>*>(element);
516         builderHelper->registerWithBarostat(
517                 { [propagator]() { return propagator->viewOnPRScalingMatrix(); },
518                   [propagator]() { return propagator->prScalingCallback(); } });
519     }
520     return element;
521 }
522
523 // Explicit template initializations
524 template class Propagator<IntegrationStep::PositionsOnly>;
525 template class Propagator<IntegrationStep::VelocitiesOnly>;
526 template class Propagator<IntegrationStep::LeapFrog>;
527 template class Propagator<IntegrationStep::VelocityVerletPositionsAndVelocities>;
528
529 } // namespace gmx