Reimplement constant acceleration groups
[alexxy/gromacs.git] / src / gromacs / mdrun / shellfc.cpp
1 /*
2  * This file is part of the GROMACS molecular simulation package.
3  *
4  * Copyright (c) 1991-2000, University of Groningen, The Netherlands.
5  * Copyright (c) 2001-2008, The GROMACS development team.
6  * Copyright (c) 2013,2014,2015,2016,2017 by the GROMACS development team.
7  * Copyright (c) 2018,2019,2020,2021, by the GROMACS development team, led by
8  * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
9  * and including many others, as listed in the AUTHORS file in the
10  * top-level source directory and at http://www.gromacs.org.
11  *
12  * GROMACS is free software; you can redistribute it and/or
13  * modify it under the terms of the GNU Lesser General Public License
14  * as published by the Free Software Foundation; either version 2.1
15  * of the License, or (at your option) any later version.
16  *
17  * GROMACS is distributed in the hope that it will be useful,
18  * but WITHOUT ANY WARRANTY; without even the implied warranty of
19  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
20  * Lesser General Public License for more details.
21  *
22  * You should have received a copy of the GNU Lesser General Public
23  * License along with GROMACS; if not, see
24  * http://www.gnu.org/licenses, or write to the Free Software Foundation,
25  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
26  *
27  * If you want to redistribute modifications to GROMACS, please
28  * consider that scientific software is very special. Version
29  * control is crucial - bugs must be traceable. We will be happy to
30  * consider code for inclusion in the official distribution, but
31  * derived work must not be called official GROMACS. Details are found
32  * in the README & COPYING files - if they are missing, get the
33  * official version at http://www.gromacs.org.
34  *
35  * To help us fund GROMACS development, we humbly ask that you cite
36  * the research papers on the package. Check out http://www.gromacs.org.
37  */
38 #include "gmxpre.h"
39
40 #include "shellfc.h"
41
42 #include <cmath>
43 #include <cstdint>
44 #include <cstdlib>
45 #include <cstring>
46
47 #include <algorithm>
48 #include <array>
49
50 #include "gromacs/domdec/dlbtiming.h"
51 #include "gromacs/domdec/domdec.h"
52 #include "gromacs/domdec/domdec_struct.h"
53 #include "gromacs/domdec/mdsetup.h"
54 #include "gromacs/gmxlib/network.h"
55 #include "gromacs/math/functions.h"
56 #include "gromacs/math/units.h"
57 #include "gromacs/math/utilities.h"
58 #include "gromacs/math/vec.h"
59 #include "gromacs/math/vecdump.h"
60 #include "gromacs/mdlib/constr.h"
61 #include "gromacs/mdlib/enerdata_utils.h"
62 #include "gromacs/mdlib/force.h"
63 #include "gromacs/mdlib/force_flags.h"
64 #include "gromacs/mdlib/gmx_omp_nthreads.h"
65 #include "gromacs/mdlib/vsite.h"
66 #include "gromacs/mdtypes/commrec.h"
67 #include "gromacs/mdtypes/enerdata.h"
68 #include "gromacs/mdtypes/forcebuffers.h"
69 #include "gromacs/mdtypes/forcerec.h"
70 #include "gromacs/mdtypes/inputrec.h"
71 #include "gromacs/mdtypes/md_enums.h"
72 #include "gromacs/mdtypes/mdatom.h"
73 #include "gromacs/mdtypes/state.h"
74 #include "gromacs/pbcutil/pbc.h"
75 #include "gromacs/topology/ifunc.h"
76 #include "gromacs/topology/mtop_lookup.h"
77 #include "gromacs/topology/mtop_util.h"
78 #include "gromacs/utility/arrayref.h"
79 #include "gromacs/utility/arraysize.h"
80 #include "gromacs/utility/cstringutil.h"
81 #include "gromacs/utility/fatalerror.h"
82 #include "gromacs/utility/gmxassert.h"
83
84 using gmx::ArrayRef;
85 using gmx::ArrayRefWithPadding;
86 using gmx::RVec;
87
88 struct t_shell
89 {
90     int  nnucl      = 0;  /* The number of nuclei */
91     int  shellIndex = -1; /* The shell index */
92     int  nucl1      = -1; /* The first nuclei connected to the shell    */
93     int  nucl2      = -1; /* The second nuclei connected to the shell   */
94     int  nucl3      = -1; /* The third nuclei connected to the shell    */
95     real k          = 0;  /* force constant                     */
96     real k_1        = 0;  /* 1 over force constant              */
97     rvec xold;            /* The old shell coordinates */
98     rvec fold;            /* The old force on the shell */
99     rvec step;            /* Step size for steepest descents */
100 };
101
102 struct gmx_shellfc_t
103 {
104     /* Shell counts, indices, parameters and working data */
105     std::vector<t_shell> shell_gl;              /* All the shells (for DD only)              */
106     std::vector<int>     shell_index_gl;        /* Global shell index (for DD only)          */
107     gmx_bool             bInterCG;              /* Are there inter charge-group shells?      */
108     std::vector<t_shell> shells;                /* The local shells                          */
109     bool                 predictShells = false; /* Predict shell positions                   */
110     bool                 requireInit   = false; /* Require initialization of shell positions */
111     int                  nflexcon      = 0;     /* The number of flexible constraints        */
112
113     std::array<PaddedHostVector<RVec>, 2> x; /* Coordinate buffers for iterative minimization */
114     std::array<PaddedHostVector<RVec>, 2> f; /* Force buffers for iterative minimization */
115
116     /* Flexible constraint working data */
117     std::vector<RVec>       acc_dir;                /* Acceleration direction for flexcon        */
118     gmx::PaddedVector<RVec> x_old;                  /* Old coordinates for flexcon               */
119     gmx::PaddedVector<RVec> adir_xnold;             /* Work space for init_adir                  */
120     gmx::PaddedVector<RVec> adir_xnew;              /* Work space for init_adir                  */
121     std::int64_t            numForceEvaluations;    /* Total number of force evaluations         */
122     int                     numConvergedIterations; /* Total number of iterations that converged */
123 };
124
125
126 static void pr_shell(FILE* fplog, ArrayRef<const t_shell> shells)
127 {
128     fprintf(fplog, "SHELL DATA\n");
129     fprintf(fplog, "%5s  %8s  %5s  %5s  %5s\n", "Shell", "Force k", "Nucl1", "Nucl2", "Nucl3");
130     for (const t_shell& shell : shells)
131     {
132         fprintf(fplog, "%5d  %8.3f  %5d", shell.shellIndex, 1.0 / shell.k_1, shell.nucl1);
133         if (shell.nnucl == 2)
134         {
135             fprintf(fplog, "  %5d\n", shell.nucl2);
136         }
137         else if (shell.nnucl == 3)
138         {
139             fprintf(fplog, "  %5d  %5d\n", shell.nucl2, shell.nucl3);
140         }
141         else
142         {
143             fprintf(fplog, "\n");
144         }
145     }
146 }
147
148 /* TODO The remain call of this function passes non-NULL mass and NULL
149  * mtop, so this routine can be simplified.
150  *
151  * The other code path supported doing prediction before the MD loop
152  * started, but even when called, the prediction was always
153  * over-written by a subsequent call in the MD loop, so has been
154  * removed. */
155 static void predict_shells(FILE*                     fplog,
156                            ArrayRef<RVec>            x,
157                            ArrayRef<RVec>            v,
158                            real                      dt,
159                            ArrayRef<const t_shell>   shells,
160                            gmx::ArrayRef<const real> mass,
161                            gmx_bool                  bInit)
162 {
163     int  m, n1, n2, n3;
164     real dt_1, fudge, tm, m1, m2, m3;
165
166     /* We introduce a fudge factor for performance reasons: with this choice
167      * the initial force on the shells is about a factor of two lower than
168      * without
169      */
170     fudge = 1.0;
171
172     ArrayRef<RVec> xOrV;
173     if (bInit)
174     {
175         if (fplog)
176         {
177             fprintf(fplog, "RELAX: Using prediction for initial shell placement\n");
178         }
179         xOrV = x;
180         dt_1 = 1;
181     }
182     else
183     {
184         xOrV = v;
185         dt_1 = fudge * dt;
186     }
187
188     for (const t_shell& shell : shells)
189     {
190         const int s1 = shell.shellIndex;
191         if (bInit)
192         {
193             clear_rvec(x[s1]);
194         }
195         switch (shell.nnucl)
196         {
197             case 1:
198                 n1 = shell.nucl1;
199                 for (m = 0; (m < DIM); m++)
200                 {
201                     x[s1][m] += xOrV[n1][m] * dt_1;
202                 }
203                 break;
204             case 2:
205                 n1 = shell.nucl1;
206                 n2 = shell.nucl2;
207                 m1 = mass[n1];
208                 m2 = mass[n2];
209                 tm = dt_1 / (m1 + m2);
210                 for (m = 0; (m < DIM); m++)
211                 {
212                     x[s1][m] += (m1 * xOrV[n1][m] + m2 * xOrV[n2][m]) * tm;
213                 }
214                 break;
215             case 3:
216                 n1 = shell.nucl1;
217                 n2 = shell.nucl2;
218                 n3 = shell.nucl3;
219                 m1 = mass[n1];
220                 m2 = mass[n2];
221                 m3 = mass[n3];
222                 tm = dt_1 / (m1 + m2 + m3);
223                 for (m = 0; (m < DIM); m++)
224                 {
225                     x[s1][m] += (m1 * xOrV[n1][m] + m2 * xOrV[n2][m] + m3 * xOrV[n3][m]) * tm;
226                 }
227                 break;
228             default: gmx_fatal(FARGS, "Shell %d has %d nuclei!", s1, shell.nnucl);
229         }
230     }
231 }
232
233 gmx_shellfc_t* init_shell_flexcon(FILE*             fplog,
234                                   const gmx_mtop_t& mtop,
235                                   int               nflexcon,
236                                   int               nstcalcenergy,
237                                   bool              usingDomainDecomposition,
238                                   bool              usingPmeOnGpu)
239 {
240     gmx_shellfc_t* shfc;
241
242     int  ns, nshell, nsi;
243     int  i, j, type, a_offset, mol, ftype, nra;
244     real qS, alpha;
245     int  aS, aN = 0; /* Shell and nucleus */
246     int bondtypes[] = { F_BONDS, F_HARMONIC, F_CUBICBONDS, F_POLARIZATION, F_ANHARM_POL, F_WATER_POL };
247 #define NBT asize(bondtypes)
248     const gmx_ffparams_t* ffparams;
249
250     const gmx::EnumerationArray<ParticleType, int> numParticles = gmx_mtop_particletype_count(mtop);
251     if (fplog)
252     {
253         /* Print the number of each particle type */
254         for (const auto entry : gmx::keysOf(numParticles))
255         {
256             const int number = numParticles[entry];
257             if (number != 0)
258             {
259                 fprintf(fplog, "There are: %d %ss\n", number, enumValueToString(entry));
260             }
261         }
262     }
263
264     nshell = numParticles[ParticleType::Shell];
265
266     if (nshell == 0 && nflexcon == 0)
267     {
268         /* We're not doing shells or flexible constraints */
269         return nullptr;
270     }
271
272     shfc           = new gmx_shellfc_t;
273     shfc->nflexcon = nflexcon;
274
275     if (nshell == 0)
276     {
277         /* Only flexible constraints, no shells.
278          * Note that make_local_shells() does not need to be called.
279          */
280         return shfc;
281     }
282
283     if (nstcalcenergy != 1)
284     {
285         gmx_fatal(FARGS,
286                   "You have nstcalcenergy set to a value (%d) that is different from 1.\nThis is "
287                   "not supported in combination with shell particles.\nPlease make a new tpr file.",
288                   nstcalcenergy);
289     }
290     if (nshell > 0 && usingDomainDecomposition)
291     {
292         gmx_fatal(
293                 FARGS,
294                 "Shell particles are not implemented with domain decomposition, use a single rank");
295     }
296
297     /* We have shells: fill the shell data structure */
298
299     /* Global system sized array, this should be avoided */
300     std::vector<int> shell_index(mtop.natoms);
301
302     nshell = 0;
303     for (const AtomProxy atomP : AtomRange(mtop))
304     {
305         const t_atom& local = atomP.atom();
306         int           i     = atomP.globalAtomNumber();
307         if (local.ptype == ParticleType::Shell)
308         {
309             shell_index[i] = nshell++;
310         }
311     }
312
313     std::vector<t_shell> shell(nshell);
314
315     ffparams = &mtop.ffparams;
316
317     /* Now fill the structures */
318     /* TODO: See if we can use update groups that cover shell constructions */
319     shfc->bInterCG = FALSE;
320     ns             = 0;
321     a_offset       = 0;
322     for (size_t mb = 0; mb < mtop.molblock.size(); mb++)
323     {
324         const gmx_molblock_t& molb = mtop.molblock[mb];
325         const gmx_moltype_t&  molt = mtop.moltype[molb.type];
326
327         const t_atom* atom = molt.atoms.atom;
328         for (mol = 0; mol < molb.nmol; mol++)
329         {
330             for (j = 0; (j < NBT); j++)
331             {
332                 const int* ia = molt.ilist[bondtypes[j]].iatoms.data();
333                 for (i = 0; (i < molt.ilist[bondtypes[j]].size());)
334                 {
335                     type  = ia[0];
336                     ftype = ffparams->functype[type];
337                     nra   = interaction_function[ftype].nratoms;
338
339                     /* Check whether we have a bond with a shell */
340                     aS = -1;
341
342                     switch (bondtypes[j])
343                     {
344                         case F_BONDS:
345                         case F_HARMONIC:
346                         case F_CUBICBONDS:
347                         case F_POLARIZATION:
348                         case F_ANHARM_POL:
349                             if (atom[ia[1]].ptype == ParticleType::Shell)
350                             {
351                                 aS = ia[1];
352                                 aN = ia[2];
353                             }
354                             else if (atom[ia[2]].ptype == ParticleType::Shell)
355                             {
356                                 aS = ia[2];
357                                 aN = ia[1];
358                             }
359                             break;
360                         case F_WATER_POL:
361                             aN = ia[4]; /* Dummy */
362                             aS = ia[5]; /* Shell */
363                             break;
364                         default: gmx_fatal(FARGS, "Death Horror: %s, %d", __FILE__, __LINE__);
365                     }
366
367                     if (aS != -1)
368                     {
369                         qS = atom[aS].q;
370
371                         /* Check whether one of the particles is a shell... */
372                         nsi = shell_index[a_offset + aS];
373                         if ((nsi < 0) || (nsi >= nshell))
374                         {
375                             gmx_fatal(FARGS, "nsi is %d should be within 0 - %d. aS = %d", nsi, nshell, aS);
376                         }
377                         if (shell[nsi].shellIndex == -1)
378                         {
379                             shell[nsi].shellIndex = a_offset + aS;
380                             ns++;
381                         }
382                         else if (shell[nsi].shellIndex != a_offset + aS)
383                         {
384                             gmx_fatal(FARGS, "Weird stuff in %s, %d", __FILE__, __LINE__);
385                         }
386
387                         if (shell[nsi].nucl1 == -1)
388                         {
389                             shell[nsi].nucl1 = a_offset + aN;
390                         }
391                         else if (shell[nsi].nucl2 == -1)
392                         {
393                             shell[nsi].nucl2 = a_offset + aN;
394                         }
395                         else if (shell[nsi].nucl3 == -1)
396                         {
397                             shell[nsi].nucl3 = a_offset + aN;
398                         }
399                         else
400                         {
401                             if (fplog)
402                             {
403                                 pr_shell(fplog, shell);
404                             }
405                             gmx_fatal(FARGS, "Can not handle more than three bonds per shell\n");
406                         }
407                         if (aS != aN)
408                         {
409                             /* shell[nsi].bInterCG = TRUE; */
410                             shfc->bInterCG = TRUE;
411                         }
412
413                         switch (bondtypes[j])
414                         {
415                             case F_BONDS:
416                             case F_HARMONIC:
417                                 shell[nsi].k += ffparams->iparams[type].harmonic.krA;
418                                 break;
419                             case F_CUBICBONDS:
420                                 shell[nsi].k += ffparams->iparams[type].cubic.kb;
421                                 break;
422                             case F_POLARIZATION:
423                             case F_ANHARM_POL:
424                                 if (!gmx_within_tol(qS, atom[aS].qB, GMX_REAL_EPS * 10))
425                                 {
426                                     gmx_fatal(FARGS,
427                                               "polarize can not be used with qA(%e) != qB(%e) for "
428                                               "atom %d of molecule block %zu",
429                                               qS,
430                                               atom[aS].qB,
431                                               aS + 1,
432                                               mb + 1);
433                                 }
434                                 shell[nsi].k += gmx::square(qS) * gmx::c_one4PiEps0
435                                                 / ffparams->iparams[type].polarize.alpha;
436                                 break;
437                             case F_WATER_POL:
438                                 if (!gmx_within_tol(qS, atom[aS].qB, GMX_REAL_EPS * 10))
439                                 {
440                                     gmx_fatal(FARGS,
441                                               "water_pol can not be used with qA(%e) != qB(%e) for "
442                                               "atom %d of molecule block %zu",
443                                               qS,
444                                               atom[aS].qB,
445                                               aS + 1,
446                                               mb + 1);
447                                 }
448                                 alpha = (ffparams->iparams[type].wpol.al_x
449                                          + ffparams->iparams[type].wpol.al_y
450                                          + ffparams->iparams[type].wpol.al_z)
451                                         / 3.0;
452                                 shell[nsi].k += gmx::square(qS) * gmx::c_one4PiEps0 / alpha;
453                                 break;
454                             default: gmx_fatal(FARGS, "Death Horror: %s, %d", __FILE__, __LINE__);
455                         }
456                         shell[nsi].nnucl++;
457                     }
458                     ia += nra + 1;
459                     i += nra + 1;
460                 }
461             }
462             a_offset += molt.atoms.nr;
463         }
464         /* Done with this molecule type */
465     }
466
467     /* Verify whether it's all correct */
468     if (ns != nshell)
469     {
470         gmx_fatal(FARGS, "Something weird with shells. They may not be bonded to something");
471     }
472
473     for (i = 0; (i < ns); i++)
474     {
475         shell[i].k_1 = 1.0 / shell[i].k;
476     }
477
478     if (debug)
479     {
480         pr_shell(debug, shell);
481     }
482
483
484     shfc->shell_gl       = shell;
485     shfc->shell_index_gl = shell_index;
486
487     shfc->predictShells = (getenv("GMX_NOPREDICT") == nullptr);
488     shfc->requireInit   = false;
489     if (!shfc->predictShells)
490     {
491         if (fplog)
492         {
493             fprintf(fplog, "\nWill never predict shell positions\n");
494         }
495     }
496     else
497     {
498         shfc->requireInit = (getenv("GMX_REQUIRE_SHELL_INIT") != nullptr);
499         if (shfc->requireInit && fplog)
500         {
501             fprintf(fplog, "\nWill always initiate shell positions\n");
502         }
503     }
504
505     if (shfc->predictShells)
506     {
507         if (shfc->bInterCG)
508         {
509             if (fplog)
510             {
511                 fprintf(fplog,
512                         "\nNOTE: in the current version shell prediction during the crun is "
513                         "disabled\n\n");
514             }
515             /* Prediction improves performance, so we should implement either:
516              * 1. communication for the atoms needed for prediction
517              * 2. prediction using the velocities of shells; currently the
518              *    shell velocities are zeroed, it's a bit tricky to keep
519              *    track of the shell displacements and thus the velocity.
520              */
521             shfc->predictShells = false;
522         }
523     }
524
525     /* shfc->x is used as a coordinate buffer for the sim_util's `do_force` function, and
526      * when using PME it must be pinned. */
527     if (usingPmeOnGpu)
528     {
529         for (i = 0; i < 2; i++)
530         {
531             changePinningPolicy(&shfc->x[i], gmx::PinningPolicy::PinnedIfSupported);
532         }
533     }
534
535     return shfc;
536 }
537
538 void gmx::make_local_shells(const t_commrec* cr, const t_mdatoms& md, gmx_shellfc_t* shfc)
539 {
540     int           a0, a1;
541     gmx_domdec_t* dd = nullptr;
542
543     if (haveDDAtomOrdering(*cr))
544     {
545         dd = cr->dd;
546         a0 = 0;
547         a1 = dd_numHomeAtoms(*dd);
548     }
549     else
550     {
551         /* Single node: we need all shells, copy them */
552         shfc->shells = shfc->shell_gl;
553
554         return;
555     }
556
557     ArrayRef<const int> ind = shfc->shell_index_gl;
558
559     std::vector<t_shell>& shells = shfc->shells;
560     shells.clear();
561     auto* ptype = md.ptype;
562     for (int i = a0; i < a1; i++)
563     {
564         if (ptype[i] == ParticleType::Shell)
565         {
566             if (dd)
567             {
568                 shells.push_back(shfc->shell_gl[ind[dd->globalAtomIndices[i]]]);
569             }
570             else
571             {
572                 shells.push_back(shfc->shell_gl[ind[i]]);
573             }
574             t_shell& shell = shells.back();
575
576             /* With inter-cg shells we can no do shell prediction,
577              * so we do not need the nuclei numbers.
578              */
579             if (!shfc->bInterCG)
580             {
581                 shell.nucl1 = i + shell.nucl1 - shell.shellIndex;
582                 if (shell.nnucl > 1)
583                 {
584                     shell.nucl2 = i + shell.nucl2 - shell.shellIndex;
585                 }
586                 if (shell.nnucl > 2)
587                 {
588                     shell.nucl3 = i + shell.nucl3 - shell.shellIndex;
589                 }
590             }
591             shell.shellIndex = i;
592         }
593     }
594 }
595
596 static void do_1pos(rvec xnew, const rvec xold, const rvec f, real step)
597 {
598     real xo, yo, zo;
599     real dx, dy, dz;
600
601     xo = xold[XX];
602     yo = xold[YY];
603     zo = xold[ZZ];
604
605     dx = f[XX] * step;
606     dy = f[YY] * step;
607     dz = f[ZZ] * step;
608
609     xnew[XX] = xo + dx;
610     xnew[YY] = yo + dy;
611     xnew[ZZ] = zo + dz;
612 }
613
614 static void do_1pos3(rvec xnew, const rvec xold, const rvec f, const rvec step)
615 {
616     real xo, yo, zo;
617     real dx, dy, dz;
618
619     xo = xold[XX];
620     yo = xold[YY];
621     zo = xold[ZZ];
622
623     dx = f[XX] * step[XX];
624     dy = f[YY] * step[YY];
625     dz = f[ZZ] * step[ZZ];
626
627     xnew[XX] = xo + dx;
628     xnew[YY] = yo + dy;
629     xnew[ZZ] = zo + dz;
630 }
631
632 static void directional_sd(ArrayRef<const RVec> xold,
633                            ArrayRef<RVec>       xnew,
634                            ArrayRef<const RVec> acc_dir,
635                            int                  homenr,
636                            real                 step)
637 {
638     const rvec* xo = as_rvec_array(xold.data());
639     rvec*       xn = as_rvec_array(xnew.data());
640
641     for (int i = 0; i < homenr; i++)
642     {
643         do_1pos(xn[i], xo[i], acc_dir[i], step);
644     }
645 }
646
647 static void shell_pos_sd(ArrayRef<const RVec> xcur,
648                          ArrayRef<RVec>       xnew,
649                          ArrayRef<const RVec> f,
650                          ArrayRef<t_shell>    shells,
651                          int                  count)
652 {
653     const real step_scale_min = 0.8, step_scale_increment = 0.2, step_scale_max = 1.2,
654                step_scale_multiple = (step_scale_max - step_scale_min) / step_scale_increment;
655     int        d;
656     real       dx, df, k_est;
657     const real zero = 0;
658 #ifdef PRINT_STEP
659     real step_min, step_max;
660
661     step_min = 1e30;
662     step_max = 0;
663 #endif
664     for (t_shell& shell : shells)
665     {
666         const int ind = shell.shellIndex;
667         if (count == 1)
668         {
669             for (d = 0; d < DIM; d++)
670             {
671                 shell.step[d] = shell.k_1;
672 #ifdef PRINT_STEP
673                 step_min = std::min(step_min, shell.step[d]);
674                 step_max = std::max(step_max, shell.step[d]);
675 #endif
676             }
677         }
678         else
679         {
680             for (d = 0; d < DIM; d++)
681             {
682                 dx = xcur[ind][d] - shell.xold[d];
683                 df = f[ind][d] - shell.fold[d];
684                 /* -dx/df gets used to generate an interpolated value, but would
685                  * cause a NaN if df were binary-equal to zero. Values close to
686                  * zero won't cause problems (because of the min() and max()), so
687                  * just testing for binary inequality is OK. */
688                 if (zero != df)
689                 {
690                     k_est = -dx / df;
691                     /* Scale the step size by a factor interpolated from
692                      * step_scale_min to step_scale_max, as k_est goes from 0 to
693                      * step_scale_multiple * shell.step[d] */
694                     shell.step[d] = step_scale_min * shell.step[d]
695                                     + step_scale_increment
696                                               * std::min(step_scale_multiple * shell.step[d],
697                                                          std::max(k_est, zero));
698                 }
699                 else
700                 {
701                     /* Here 0 == df */
702                     if (gmx_numzero(dx)) /* 0 == dx */
703                     {
704                         /* Likely this will never happen, but if it does just
705                          * don't scale the step. */
706                     }
707                     else /* 0 != dx */
708                     {
709                         shell.step[d] *= step_scale_max;
710                     }
711                 }
712 #ifdef PRINT_STEP
713                 step_min = std::min(step_min, shell.step[d]);
714                 step_max = std::max(step_max, shell.step[d]);
715 #endif
716             }
717         }
718         copy_rvec(xcur[ind], shell.xold);
719         copy_rvec(f[ind], shell.fold);
720
721         do_1pos3(xnew[ind], xcur[ind], f[ind], shell.step);
722
723         if (gmx_debug_at)
724         {
725             fprintf(debug, "shell = %d\n", ind);
726             pr_rvec(debug, 0, "fshell", f[ind], DIM, TRUE);
727             pr_rvec(debug, 0, "xold", xcur[ind], DIM, TRUE);
728             pr_rvec(debug, 0, "step", shell.step, DIM, TRUE);
729             pr_rvec(debug, 0, "xnew", xnew[ind], DIM, TRUE);
730         }
731     }
732 #ifdef PRINT_STEP
733     printf("step %.3e %.3e\n", step_min, step_max);
734 #endif
735 }
736
737 static void decrease_step_size(ArrayRef<t_shell> shells)
738 {
739     for (t_shell& shell : shells)
740     {
741         svmul(0.8, shell.step, shell.step);
742     }
743 }
744
745 static void print_epot(FILE* fp, int64_t mdstep, int count, real epot, real df, int ndir, real sf_dir)
746 {
747     char buf[22];
748
749     fprintf(fp, "MDStep=%5s/%2d EPot: %12.8e, rmsF: %6.2e", gmx_step_str(mdstep, buf), count, epot, df);
750     if (ndir)
751     {
752         fprintf(fp, ", dir. rmsF: %6.2e\n", std::sqrt(sf_dir / ndir));
753     }
754     else
755     {
756         fprintf(fp, "\n");
757     }
758 }
759
760
761 static real rms_force(const t_commrec*        cr,
762                       ArrayRef<const RVec>    force,
763                       ArrayRef<const t_shell> shells,
764                       int                     ndir,
765                       real*                   sf_dir,
766                       real*                   Epot)
767 {
768     double      buf[4];
769     const rvec* f = as_rvec_array(force.data());
770
771     buf[0] = *sf_dir;
772     for (const t_shell& shell : shells)
773     {
774         buf[0] += norm2(f[shell.shellIndex]);
775     }
776     int ntot = shells.ssize();
777
778     if (PAR(cr))
779     {
780         buf[1] = ntot;
781         buf[2] = *sf_dir;
782         buf[3] = *Epot;
783         gmx_sumd(4, buf, cr);
784         ntot    = gmx::roundToInt(buf[1]);
785         *sf_dir = buf[2];
786         *Epot   = buf[3];
787     }
788     ntot += ndir;
789
790     return (ntot ? std::sqrt(buf[0] / ntot) : 0);
791 }
792
793 static void dump_shells(FILE* fp, ArrayRef<RVec> f, real ftol, ArrayRef<const t_shell> shells)
794 {
795     real ft2, ff2;
796
797     ft2 = gmx::square(ftol);
798
799     for (const t_shell& shell : shells)
800     {
801         const int ind = shell.shellIndex;
802         ff2           = iprod(f[ind], f[ind]);
803         if (ff2 > ft2)
804         {
805             fprintf(fp,
806                     "SHELL %5d, force %10.5f  %10.5f  %10.5f, |f| %10.5f\n",
807                     ind,
808                     f[ind][XX],
809                     f[ind][YY],
810                     f[ind][ZZ],
811                     std::sqrt(ff2));
812         }
813     }
814 }
815
816 static void init_adir(gmx_shellfc_t*            shfc,
817                       gmx::Constraints*         constr,
818                       const t_inputrec*         ir,
819                       const t_commrec*          cr,
820                       int                       dd_ac1,
821                       int64_t                   step,
822                       const t_mdatoms&          md,
823                       int                       end,
824                       ArrayRefWithPadding<RVec> xOld,
825                       ArrayRef<RVec>            x_init,
826                       ArrayRefWithPadding<RVec> xCurrent,
827                       ArrayRef<RVec>            f,
828                       ArrayRef<RVec>            acc_dir,
829                       const matrix              box,
830                       ArrayRef<const real>      lambda,
831                       real*                     dvdlambda)
832 {
833     double dt, w_dt;
834     int    n, d;
835
836     if (haveDDAtomOrdering(*cr))
837     {
838         n = dd_ac1;
839     }
840     else
841     {
842         n = end;
843     }
844     shfc->adir_xnold.resizeWithPadding(n);
845     shfc->adir_xnew.resizeWithPadding(n);
846     rvec* xnold = as_rvec_array(shfc->adir_xnold.data());
847     rvec* xnew  = as_rvec_array(shfc->adir_xnew.data());
848     rvec* x_old = as_rvec_array(xOld.paddedArrayRef().data());
849     rvec* x     = as_rvec_array(xCurrent.paddedArrayRef().data());
850
851     auto* ptype   = md.ptype;
852     auto  invmass = gmx::arrayRefFromArray(md.invmass, md.nr);
853     dt            = ir->delta_t;
854
855     /* Does NOT work with freeze or acceleration groups (yet) */
856     for (n = 0; n < end; n++)
857     {
858         w_dt = invmass[n] * dt;
859
860         for (d = 0; d < DIM; d++)
861         {
862             if ((ptype[n] != ParticleType::VSite) && (ptype[n] != ParticleType::Shell))
863             {
864                 xnold[n][d] = x[n][d] - (x_init[n][d] - x_old[n][d]);
865                 xnew[n][d]  = 2 * x[n][d] - x_old[n][d] + f[n][d] * w_dt * dt;
866             }
867             else
868             {
869                 xnold[n][d] = x[n][d];
870                 xnew[n][d]  = x[n][d];
871             }
872         }
873     }
874     bool needsLogging  = false;
875     bool computeEnergy = false;
876     bool computeVirial = false;
877     constr->apply(needsLogging,
878                   computeEnergy,
879                   step,
880                   0,
881                   1.0,
882                   xCurrent,
883                   shfc->adir_xnold.arrayRefWithPadding(),
884                   {},
885                   box,
886                   lambda[static_cast<int>(FreeEnergyPerturbationCouplingType::Bonded)],
887                   &(dvdlambda[static_cast<int>(FreeEnergyPerturbationCouplingType::Bonded)]),
888                   {},
889                   computeVirial,
890                   nullptr,
891                   gmx::ConstraintVariable::Positions);
892     constr->apply(needsLogging,
893                   computeEnergy,
894                   step,
895                   0,
896                   1.0,
897                   xCurrent,
898                   shfc->adir_xnew.arrayRefWithPadding(),
899                   {},
900                   box,
901                   lambda[static_cast<int>(FreeEnergyPerturbationCouplingType::Bonded)],
902                   &(dvdlambda[static_cast<int>(FreeEnergyPerturbationCouplingType::Bonded)]),
903                   {},
904                   computeVirial,
905                   nullptr,
906                   gmx::ConstraintVariable::Positions);
907
908     for (n = 0; n < end; n++)
909     {
910         for (d = 0; d < DIM; d++)
911         {
912             xnew[n][d] = -(2 * x[n][d] - xnold[n][d] - xnew[n][d]) / gmx::square(dt)
913                          - f[n][d] * invmass[n];
914         }
915         clear_rvec(acc_dir[n]);
916     }
917
918     /* Project the acceleration on the old bond directions */
919     constr->apply(needsLogging,
920                   computeEnergy,
921                   step,
922                   0,
923                   1.0,
924                   xOld,
925                   shfc->adir_xnew.arrayRefWithPadding(),
926                   acc_dir,
927                   box,
928                   lambda[static_cast<int>(FreeEnergyPerturbationCouplingType::Bonded)],
929                   &(dvdlambda[static_cast<int>(FreeEnergyPerturbationCouplingType::Bonded)]),
930                   {},
931                   computeVirial,
932                   nullptr,
933                   gmx::ConstraintVariable::Deriv_FlexCon);
934 }
935
936 void relax_shell_flexcon(FILE*                         fplog,
937                          const t_commrec*              cr,
938                          const gmx_multisim_t*         ms,
939                          gmx_bool                      bVerbose,
940                          gmx_enfrot*                   enforcedRotation,
941                          int64_t                       mdstep,
942                          const t_inputrec*             inputrec,
943                          gmx::ImdSession*              imdSession,
944                          pull_t*                       pull_work,
945                          gmx_bool                      bDoNS,
946                          int                           force_flags,
947                          const gmx_localtop_t*         top,
948                          gmx::Constraints*             constr,
949                          gmx_enerdata_t*               enerd,
950                          int                           natoms,
951                          ArrayRefWithPadding<RVec>     xPadded,
952                          ArrayRefWithPadding<RVec>     vPadded,
953                          const matrix                  box,
954                          ArrayRef<real>                lambda,
955                          const history_t*              hist,
956                          gmx::ForceBuffersView*        f,
957                          tensor                        force_vir,
958                          const t_mdatoms&              md,
959                          CpuPpLongRangeNonbondeds*     longRangeNonbondeds,
960                          t_nrnb*                       nrnb,
961                          gmx_wallcycle*                wcycle,
962                          gmx_shellfc_t*                shfc,
963                          t_forcerec*                   fr,
964                          gmx::MdrunScheduleWorkload*   runScheduleWork,
965                          double                        t,
966                          rvec                          mu_tot,
967                          gmx::VirtualSitesHandler*     vsite,
968                          const DDBalanceRegionHandler& ddBalanceRegionHandler)
969 {
970     real Epot[2], df[2];
971     real sf_dir, invdt;
972     real dum = 0;
973     char sbuf[22];
974     int  nat, dd_ac0, dd_ac1 = 0, i;
975     int  homenr = md.homenr, end = homenr;
976     int  d, Min = 0, count = 0;
977 #define Try (1 - Min) /* At start Try = 1 */
978
979     const bool        bCont        = (mdstep == inputrec->init_step) && inputrec->bContinuation;
980     const bool        bInit        = (mdstep == inputrec->init_step) || shfc->requireInit;
981     const real        ftol         = inputrec->em_tol;
982     const int         number_steps = inputrec->niter;
983     ArrayRef<t_shell> shells       = shfc->shells;
984     const int         nflexcon     = shfc->nflexcon;
985
986     if (haveDDAtomOrdering(*cr))
987     {
988         nat = dd_natoms_vsite(*cr->dd);
989         if (nflexcon > 0)
990         {
991             dd_get_constraint_range(*cr->dd, &dd_ac0, &dd_ac1);
992             nat = std::max(nat, dd_ac1);
993         }
994     }
995     else
996     {
997         nat = natoms;
998     }
999
1000     for (i = 0; (i < 2); i++)
1001     {
1002         shfc->x[i].resizeWithPadding(nat);
1003         shfc->f[i].resizeWithPadding(nat);
1004     }
1005
1006     /* Create views that we can swap for trail and minimum for positions and forces */
1007     ArrayRefWithPadding<RVec> posWithPadding[2];
1008     ArrayRefWithPadding<RVec> forceWithPadding[2];
1009     ArrayRef<RVec>            pos[2];
1010     ArrayRef<RVec>            force[2];
1011     for (i = 0; (i < 2); i++)
1012     {
1013         posWithPadding[i]   = shfc->x[i].arrayRefWithPadding();
1014         pos[i]              = posWithPadding[i].paddedArrayRef();
1015         forceWithPadding[i] = shfc->f[i].arrayRefWithPadding();
1016         force[i]            = forceWithPadding[i].paddedArrayRef();
1017     }
1018
1019     ArrayRef<RVec> x = xPadded.unpaddedArrayRef();
1020     ArrayRef<RVec> v = vPadded.unpaddedArrayRef();
1021
1022     if (bDoNS && inputrec->pbcType != PbcType::No && !haveDDAtomOrdering(*cr))
1023     {
1024         /* This is the only time where the coordinates are used
1025          * before do_force is called, which normally puts all
1026          * charge groups in the box.
1027          */
1028         put_atoms_in_box_omp(
1029                 fr->pbcType, box, x.subArray(0, md.homenr), gmx_omp_nthreads_get(ModuleMultiThread::Default));
1030     }
1031
1032     if (nflexcon)
1033     {
1034         shfc->acc_dir.resize(nat);
1035         shfc->x_old.resizeWithPadding(nat);
1036         ArrayRef<RVec> x_old = shfc->x_old.arrayRefWithPadding().unpaddedArrayRef();
1037         for (i = 0; i < homenr; i++)
1038         {
1039             for (d = 0; d < DIM; d++)
1040             {
1041                 x_old[i][d] = x[i][d] - v[i][d] * inputrec->delta_t;
1042             }
1043         }
1044     }
1045
1046     auto massT = gmx::arrayRefFromArray(md.massT, md.nr);
1047     /* Do a prediction of the shell positions, when appropriate.
1048      * Without velocities (EM, NM, BD) we only do initial prediction.
1049      */
1050     if (shfc->predictShells && !bCont && (EI_STATE_VELOCITY(inputrec->eI) || bInit))
1051     {
1052         predict_shells(fplog, x, v, inputrec->delta_t, shells, massT, bInit);
1053     }
1054
1055     /* Calculate the forces first time around */
1056     if (gmx_debug_at)
1057     {
1058         pr_rvecs(debug, 0, "x b4 do_force", as_rvec_array(x.data()), homenr);
1059     }
1060     int                   shellfc_flags = force_flags | (bVerbose ? GMX_FORCE_ENERGY : 0);
1061     gmx::ForceBuffersView forceViewInit = gmx::ForceBuffersView(forceWithPadding[Min], {}, false);
1062     do_force(fplog,
1063              cr,
1064              ms,
1065              *inputrec,
1066              nullptr,
1067              enforcedRotation,
1068              imdSession,
1069              pull_work,
1070              mdstep,
1071              nrnb,
1072              wcycle,
1073              top,
1074              box,
1075              xPadded,
1076              hist,
1077              &forceViewInit,
1078              force_vir,
1079              &md,
1080              enerd,
1081              lambda,
1082              fr,
1083              runScheduleWork,
1084              vsite,
1085              mu_tot,
1086              t,
1087              nullptr,
1088              longRangeNonbondeds,
1089              (bDoNS ? GMX_FORCE_NS : 0) | shellfc_flags,
1090              ddBalanceRegionHandler);
1091
1092     sf_dir = 0;
1093     if (nflexcon)
1094     {
1095         init_adir(shfc,
1096                   constr,
1097                   inputrec,
1098                   cr,
1099                   dd_ac1,
1100                   mdstep,
1101                   md,
1102                   end,
1103                   shfc->x_old.arrayRefWithPadding(),
1104                   x,
1105                   xPadded,
1106                   force[Min],
1107                   shfc->acc_dir,
1108                   box,
1109                   lambda,
1110                   &dum);
1111
1112         for (i = 0; i < end; i++)
1113         {
1114             sf_dir += massT[i] * norm2(shfc->acc_dir[i]);
1115         }
1116     }
1117     accumulatePotentialEnergies(enerd, lambda, inputrec->fepvals.get());
1118     Epot[Min] = enerd->term[F_EPOT];
1119
1120     df[Min] = rms_force(cr, forceWithPadding[Min].paddedArrayRef(), shells, nflexcon, &sf_dir, &Epot[Min]);
1121     df[Try] = 0;
1122     if (debug)
1123     {
1124         fprintf(debug, "df = %g  %g\n", df[Min], df[Try]);
1125     }
1126
1127     if (gmx_debug_at)
1128     {
1129         pr_rvecs(debug, 0, "force0", as_rvec_array(force[Min].data()), md.nr);
1130     }
1131
1132     if (!shells.empty() || nflexcon > 0)
1133     {
1134         /* Copy x to pos[Min] & pos[Try]: during minimization only the
1135          * shell positions are updated, therefore the other particles must
1136          * be set here, in advance.
1137          */
1138         std::copy(xPadded.paddedArrayRef().begin(),
1139                   xPadded.paddedArrayRef().end(),
1140                   posWithPadding[Min].paddedArrayRef().begin());
1141         std::copy(xPadded.paddedArrayRef().begin(),
1142                   xPadded.paddedArrayRef().end(),
1143                   posWithPadding[Try].paddedArrayRef().begin());
1144     }
1145
1146     if (bVerbose && MASTER(cr))
1147     {
1148         print_epot(stdout, mdstep, 0, Epot[Min], df[Min], nflexcon, sf_dir);
1149     }
1150
1151     if (debug)
1152     {
1153         fprintf(debug, "%17s: %14.10e\n", interaction_function[F_EKIN].longname, enerd->term[F_EKIN]);
1154         fprintf(debug, "%17s: %14.10e\n", interaction_function[F_EPOT].longname, enerd->term[F_EPOT]);
1155         fprintf(debug, "%17s: %14.10e\n", interaction_function[F_ETOT].longname, enerd->term[F_ETOT]);
1156         fprintf(debug, "SHELLSTEP %s\n", gmx_step_str(mdstep, sbuf));
1157     }
1158
1159     /* First check whether we should do shells, or whether the force is
1160      * low enough even without minimization.
1161      */
1162     bool bConverged = (df[Min] < ftol);
1163
1164     for (count = 1; (!(bConverged) && (count < number_steps)); count++)
1165     {
1166         if (vsite)
1167         {
1168             vsite->construct(pos[Min], v, box, gmx::VSiteOperation::PositionsAndVelocities);
1169         }
1170
1171         if (nflexcon)
1172         {
1173             init_adir(shfc,
1174                       constr,
1175                       inputrec,
1176                       cr,
1177                       dd_ac1,
1178                       mdstep,
1179                       md,
1180                       end,
1181                       shfc->x_old.arrayRefWithPadding(),
1182                       x,
1183                       posWithPadding[Min],
1184                       force[Min],
1185                       shfc->acc_dir,
1186                       box,
1187                       lambda,
1188                       &dum);
1189
1190             directional_sd(pos[Min], pos[Try], shfc->acc_dir, end, fr->fc_stepsize);
1191         }
1192
1193         /* New positions, Steepest descent */
1194         shell_pos_sd(pos[Min], pos[Try], force[Min], shells, count);
1195
1196         if (gmx_debug_at)
1197         {
1198             pr_rvecs(debug, 0, "RELAX: pos[Min]  ", as_rvec_array(pos[Min].data()), homenr);
1199             pr_rvecs(debug, 0, "RELAX: pos[Try]  ", as_rvec_array(pos[Try].data()), homenr);
1200         }
1201         /* Try the new positions */
1202         gmx::ForceBuffersView forceViewTry = gmx::ForceBuffersView(forceWithPadding[Try], {}, false);
1203         do_force(fplog,
1204                  cr,
1205                  ms,
1206                  *inputrec,
1207                  nullptr,
1208                  enforcedRotation,
1209                  imdSession,
1210                  pull_work,
1211                  1,
1212                  nrnb,
1213                  wcycle,
1214                  top,
1215                  box,
1216                  posWithPadding[Try],
1217                  hist,
1218                  &forceViewTry,
1219                  force_vir,
1220                  &md,
1221                  enerd,
1222                  lambda,
1223                  fr,
1224                  runScheduleWork,
1225                  vsite,
1226                  mu_tot,
1227                  t,
1228                  nullptr,
1229                  longRangeNonbondeds,
1230                  shellfc_flags,
1231                  ddBalanceRegionHandler);
1232         accumulatePotentialEnergies(enerd, lambda, inputrec->fepvals.get());
1233         if (gmx_debug_at)
1234         {
1235             pr_rvecs(debug, 0, "RELAX: force[Min]", as_rvec_array(force[Min].data()), homenr);
1236             pr_rvecs(debug, 0, "RELAX: force[Try]", as_rvec_array(force[Try].data()), homenr);
1237         }
1238         sf_dir = 0;
1239         if (nflexcon)
1240         {
1241             init_adir(shfc,
1242                       constr,
1243                       inputrec,
1244                       cr,
1245                       dd_ac1,
1246                       mdstep,
1247                       md,
1248                       end,
1249                       shfc->x_old.arrayRefWithPadding(),
1250                       x,
1251                       posWithPadding[Try],
1252                       force[Try],
1253                       shfc->acc_dir,
1254                       box,
1255                       lambda,
1256                       &dum);
1257
1258             ArrayRef<const RVec> acc_dir = shfc->acc_dir;
1259             for (i = 0; i < end; i++)
1260             {
1261                 sf_dir += massT[i] * norm2(acc_dir[i]);
1262             }
1263         }
1264
1265         Epot[Try] = enerd->term[F_EPOT];
1266
1267         df[Try] = rms_force(cr, force[Try], shells, nflexcon, &sf_dir, &Epot[Try]);
1268
1269         if (debug)
1270         {
1271             fprintf(debug, "df = %g  %g\n", df[Min], df[Try]);
1272         }
1273
1274         if (debug)
1275         {
1276             if (gmx_debug_at)
1277             {
1278                 pr_rvecs(debug, 0, "F na do_force", as_rvec_array(force[Try].data()), homenr);
1279             }
1280             if (gmx_debug_at)
1281             {
1282                 fprintf(debug, "SHELL ITER %d\n", count);
1283                 dump_shells(debug, force[Try], ftol, shells);
1284             }
1285         }
1286
1287         if (bVerbose && MASTER(cr))
1288         {
1289             print_epot(stdout, mdstep, count, Epot[Try], df[Try], nflexcon, sf_dir);
1290         }
1291
1292         bConverged = (df[Try] < ftol);
1293
1294         if ((df[Try] < df[Min]))
1295         {
1296             if (debug)
1297             {
1298                 fprintf(debug, "Swapping Min and Try\n");
1299             }
1300             if (nflexcon)
1301             {
1302                 /* Correct the velocities for the flexible constraints */
1303                 invdt = 1 / inputrec->delta_t;
1304                 for (i = 0; i < end; i++)
1305                 {
1306                     for (d = 0; d < DIM; d++)
1307                     {
1308                         v[i][d] += (pos[Try][i][d] - pos[Min][i][d]) * invdt;
1309                     }
1310                 }
1311             }
1312             Min = Try;
1313         }
1314         else
1315         {
1316             decrease_step_size(shells);
1317         }
1318     }
1319     shfc->numForceEvaluations += count;
1320     if (bConverged)
1321     {
1322         shfc->numConvergedIterations++;
1323     }
1324     if (MASTER(cr) && !(bConverged))
1325     {
1326         /* Note that the energies and virial are incorrect when not converged */
1327         if (fplog)
1328         {
1329             fprintf(fplog,
1330                     "step %s: EM did not converge in %d iterations, RMS force %6.2e\n",
1331                     gmx_step_str(mdstep, sbuf),
1332                     number_steps,
1333                     df[Min]);
1334         }
1335         fprintf(stderr,
1336                 "step %s: EM did not converge in %d iterations, RMS force %6.2e\n",
1337                 gmx_step_str(mdstep, sbuf),
1338                 number_steps,
1339                 df[Min]);
1340     }
1341
1342     /* Copy back the coordinates and the forces */
1343     std::copy(pos[Min].begin(), pos[Min].end(), x.data());
1344     std::copy(force[Min].begin(), force[Min].end(), f->force().begin());
1345 }
1346
1347 void done_shellfc(FILE* fplog, gmx_shellfc_t* shfc, int64_t numSteps)
1348 {
1349     if (shfc && fplog && numSteps > 0)
1350     {
1351         double numStepsAsDouble = static_cast<double>(numSteps);
1352         fprintf(fplog,
1353                 "Fraction of iterations that converged:           %.2f %%\n",
1354                 (shfc->numConvergedIterations * 100.0) / numStepsAsDouble);
1355         fprintf(fplog,
1356                 "Average number of force evaluations per MD step: %.2f\n\n",
1357                 shfc->numForceEvaluations / numStepsAsDouble);
1358     }
1359
1360     delete shfc;
1361 }