src/gromacs/mdlib/update.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-2004, 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 "update.h"
  41
  42 #include <cmath>
  43 #include <cstdio>
  44
  45 #include <algorithm>
  46 #include <memory>
  47
  48 #include "gromacs/domdec/domdec_struct.h"
  49 #include "gromacs/fileio/confio.h"
  50 #include "gromacs/gmxlib/network.h"
  51 #include "gromacs/gmxlib/nrnb.h"
  52 #include "gromacs/listed_forces/disre.h"
  53 #include "gromacs/listed_forces/orires.h"
  54 #include "gromacs/math/functions.h"
  55 #include "gromacs/math/invertmatrix.h"
  56 #include "gromacs/math/paddedvector.h"
  57 #include "gromacs/math/units.h"
  58 #include "gromacs/math/vec.h"
  59 #include "gromacs/math/vecdump.h"
  60 #include "gromacs/mdlib/boxdeformation.h"
  61 #include "gromacs/mdlib/constr.h"
  62 #include "gromacs/mdlib/gmx_omp_nthreads.h"
  63 #include "gromacs/mdlib/mdatoms.h"
  64 #include "gromacs/mdlib/stat.h"
  65 #include "gromacs/mdlib/tgroup.h"
  66 #include "gromacs/mdtypes/commrec.h"
  67 #include "gromacs/mdtypes/fcdata.h"
  68 #include "gromacs/mdtypes/group.h"
  69 #include "gromacs/mdtypes/inputrec.h"
  70 #include "gromacs/mdtypes/md_enums.h"
  71 #include "gromacs/mdtypes/mdatom.h"
  72 #include "gromacs/mdtypes/state.h"
  73 #include "gromacs/pbcutil/boxutilities.h"
  74 #include "gromacs/pbcutil/pbc.h"
  75 #include "gromacs/pulling/pull.h"
  76 #include "gromacs/random/tabulatednormaldistribution.h"
  77 #include "gromacs/random/threefry.h"
  78 #include "gromacs/simd/simd.h"
  79 #include "gromacs/timing/wallcycle.h"
  80 #include "gromacs/topology/atoms.h"
  81 #include "gromacs/utility/exceptions.h"
  82 #include "gromacs/utility/fatalerror.h"
  83 #include "gromacs/utility/futil.h"
  84 #include "gromacs/utility/gmxassert.h"
  85 #include "gromacs/utility/gmxomp.h"
  86 #include "gromacs/utility/smalloc.h"
  87
  88 using namespace gmx; // TODO: Remove when this file is moved into gmx namespace
  89
  90 struct gmx_sd_const_t
  91 {
  92     double em = 0;
  93 };
  94
  95 struct gmx_sd_sigma_t
  96 {
  97     real V = 0;
  98 };
  99
 100 struct gmx_stochd_t
 101 {
 102     /* BD stuff */
 103     std::vector<real> bd_rf;
 104     /* SD stuff */
 105     std::vector<gmx_sd_const_t> sdc;
 106     std::vector<gmx_sd_sigma_t> sdsig;
 107     /* andersen temperature control stuff */
 108     std::vector<bool> randomize_group;
 109     std::vector<real> boltzfac;
 110
 111     explicit gmx_stochd_t(const t_inputrec& inputRecord);
 112 };
 113
 114 //! pImpled implementation for Update
 115 class Update::Impl
 116 {
 117 public:
 118     //! Constructor
 119     Impl(const t_inputrec& inputRecord, BoxDeformation* boxDeformation);
 120     //! Destructor
 121     ~Impl() = default;
 122
 123     void update_coords(const t_inputrec&                                inputRecord,
 124                        int64_t                                          step,
 125                        int                                              homenr,
 126                        bool                                             havePartiallyFrozenAtoms,
 127                        gmx::ArrayRef<const ParticleType>                ptype,
 128                        gmx::ArrayRef<const real>                        invMass,
 129                        gmx::ArrayRef<rvec>                              invMassPerDim,
 130                        t_state*                                         state,
 131                        const gmx::ArrayRefWithPadding<const gmx::RVec>& f,
 132                        const t_fcdata&                                  fcdata,
 133                        const gmx_ekindata_t*                            ekind,
 134                        const matrix                                     M,
 135                        int                                              UpdatePart,
 136                        const t_commrec*                                 cr,
 137                        bool                                             haveConstraints);
 138
 139     void finish_update(const t_inputrec& inputRecord,
 140                        const t_mdatoms*  md,
 141                        t_state*          state,
 142                        gmx_wallcycle*    wcycle,
 143                        bool              haveConstraints);
 144
 145     void update_sd_second_half(const t_inputrec&                 inputRecord,
 146                                int64_t                           step,
 147                                real*                             dvdlambda,
 148                                int                               homenr,
 149                                gmx::ArrayRef<const ParticleType> ptype,
 150                                gmx::ArrayRef<const real>         invMass,
 151                                t_state*                          state,
 152                                const t_commrec*                  cr,
 153                                t_nrnb*                           nrnb,
 154                                gmx_wallcycle*                    wcycle,
 155                                gmx::Constraints*                 constr,
 156                                bool                              do_log,
 157                                bool                              do_ene);
 158
 159     void update_for_constraint_virial(const t_inputrec&   inputRecord,
 160                                       int                 homenr,
 161                                       bool                havePartiallyFrozenAtoms,
 162                                       gmx::ArrayRef<real> invmass,
 163                                       gmx::ArrayRef<rvec> invMassPerDim,
 164                                       const t_state&      state,
 165                                       const gmx::ArrayRefWithPadding<const gmx::RVec>& f,
 166                                       const gmx_ekindata_t&                            ekind);
 167
 168     void update_temperature_constants(const t_inputrec& inputRecord);
 169
 170     const std::vector<bool>& getAndersenRandomizeGroup() const { return sd_.randomize_group; }
 171
 172     const std::vector<real>& getBoltzmanFactor() const { return sd_.boltzfac; }
 173
 174     PaddedVector<RVec>* xp() { return &xp_; }
 175
 176     BoxDeformation* deform() const { return deform_; }
 177
 178     //! Group index for freezing
 179     gmx::ArrayRef<const unsigned short> cFREEZE_;
 180     //! Group index for temperature coupling
 181     gmx::ArrayRef<const unsigned short> cTC_;
 182
 183 private:
 184     //! stochastic dynamics struct
 185     gmx_stochd_t sd_;
 186     //! xprime for constraint algorithms
 187     PaddedVector<RVec> xp_;
 188     //! Box deformation handler (or nullptr if inactive).
 189     BoxDeformation* deform_ = nullptr;
 190 };
 191
 192 Update::Update(const t_inputrec& inputRecord, BoxDeformation* boxDeformation) :
 193     impl_(new Impl(inputRecord, boxDeformation)){};
 194
 195 Update::~Update(){};
 196
 197 const std::vector<bool>& Update::getAndersenRandomizeGroup() const
 198 {
 199     return impl_->getAndersenRandomizeGroup();
 200 }
 201
 202 const std::vector<real>& Update::getBoltzmanFactor() const
 203 {
 204     return impl_->getBoltzmanFactor();
 205 }
 206
 207 PaddedVector<RVec>* Update::xp()
 208 {
 209     return impl_->xp();
 210 }
 211
 212 BoxDeformation* Update::deform() const
 213 {
 214     return impl_->deform();
 215 }
 216
 217 void Update::update_coords(const t_inputrec&                                inputRecord,
 218                            int64_t                                          step,
 219                            const t_mdatoms*                                 md,
 220                            t_state*                                         state,
 221                            const gmx::ArrayRefWithPadding<const gmx::RVec>& f,
 222                            const t_fcdata&                                  fcdata,
 223                            const gmx_ekindata_t*                            ekind,
 224                            const matrix                                     M,
 225                            int                                              updatePart,
 226                            const t_commrec*                                 cr,
 227                            const bool                                       haveConstraints)
 228 {
 229     return impl_->update_coords(inputRecord,
 230                                 step,
 231                                 md->homenr,
 232                                 md->havePartiallyFrozenAtoms,
 233                                 gmx::arrayRefFromArray(md->ptype, md->nr),
 234                                 gmx::arrayRefFromArray(md->invmass, md->nr),
 235                                 gmx::arrayRefFromArray(md->invMassPerDim, md->nr),
 236                                 state,
 237                                 f,
 238                                 fcdata,
 239                                 ekind,
 240                                 M,
 241                                 updatePart,
 242                                 cr,
 243                                 haveConstraints);
 244 }
 245
 246 void Update::finish_update(const t_inputrec& inputRecord,
 247                            const t_mdatoms*  md,
 248                            t_state*          state,
 249                            gmx_wallcycle*    wcycle,
 250                            const bool        haveConstraints)
 251 {
 252     return impl_->finish_update(inputRecord, md, state, wcycle, haveConstraints);
 253 }
 254
 255 void Update::update_sd_second_half(const t_inputrec& inputRecord,
 256                                    int64_t           step,
 257                                    real*             dvdlambda,
 258                                    const t_mdatoms*  md,
 259                                    t_state*          state,
 260                                    const t_commrec*  cr,
 261                                    t_nrnb*           nrnb,
 262                                    gmx_wallcycle*    wcycle,
 263                                    gmx::Constraints* constr,
 264                                    bool              do_log,
 265                                    bool              do_ene)
 266 {
 267     return impl_->update_sd_second_half(inputRecord,
 268                                         step,
 269                                         dvdlambda,
 270                                         md->homenr,
 271                                         gmx::arrayRefFromArray(md->ptype, md->nr),
 272                                         gmx::arrayRefFromArray(md->invmass, md->nr),
 273                                         state,
 274                                         cr,
 275                                         nrnb,
 276                                         wcycle,
 277                                         constr,
 278                                         do_log,
 279                                         do_ene);
 280 }
 281
 282 void Update::update_for_constraint_virial(const t_inputrec& inputRecord,
 283                                           const t_mdatoms&  md,
 284                                           const t_state&    state,
 285                                           const gmx::ArrayRefWithPadding<const gmx::RVec>& f,
 286                                           const gmx_ekindata_t&                            ekind)
 287 {
 288     return impl_->update_for_constraint_virial(inputRecord,
 289                                                md.homenr,
 290                                                md.havePartiallyFrozenAtoms,
 291                                                gmx::arrayRefFromArray(md.invmass, md.nr),
 292                                                gmx::arrayRefFromArray(md.invMassPerDim, md.nr),
 293                                                state,
 294                                                f,
 295                                                ekind);
 296 }
 297
 298 void Update::update_temperature_constants(const t_inputrec& inputRecord)
 299 {
 300     return impl_->update_temperature_constants(inputRecord);
 301 }
 302
 303 /*! \brief Sets whether we store the updated velocities */
 304 enum class StoreUpdatedVelocities
 305 {
 306     yes, //!< Store the updated velocities
 307     no   //!< Do not store the updated velocities
 308 };
 309
 310 /*! \brief Sets the number of different temperature coupling values */
 311 enum class NumTempScaleValues
 312 {
 313     single,  //!< Single T-scaling value (either one group or all values =1)
 314     multiple //!< Multiple T-scaling values, need to use T-group indices
 315 };
 316
 317 /*! \brief Sets if to apply no or diagonal Parrinello-Rahman pressure scaling
 318  *
 319  * Note that this enum is only used in updateMDLeapfrogSimple(), which does
 320  * not handle fully anistropic Parrinello-Rahman scaling, so we only have
 321  * options \p no and \p diagonal here and no anistropic option.
 322  */
 323 enum class ApplyParrinelloRahmanVScaling
 324 {
 325     no,      //!< Do not apply velocity scaling (not a PR-coupling run or step)
 326     diagonal //!< Apply velocity scaling using a diagonal matrix
 327 };
 328
 329 /*! \brief Integrate using leap-frog with T-scaling and optionally diagonal Parrinello-Rahman p-coupling
 330  *
 331  * \tparam       storeUpdatedVelocities Tells whether we should store the updated velocities
 332  * \tparam       numTempScaleValues     The number of different T-couple values
 333  * \tparam       applyPRVScaling        Apply Parrinello-Rahman velocity scaling
 334  * \param[in]    start                  Index of first atom to update
 335  * \param[in]    nrend                  Last atom to update: \p nrend - 1
 336  * \param[in]    dt                     The time step
 337  * \param[in]    dtPressureCouple       Time step for pressure coupling
 338  * \param[in]    invMassPerDim          1/mass per atom and dimension
 339  * \param[in]    tcstat                 Temperature coupling information
 340  * \param[in]    cTC                    T-coupling group index per atom
 341  * \param[in]    pRVScaleMatrixDiagonal Parrinello-Rahman v-scale matrix diagonal
 342  * \param[in]    x                      Input coordinates
 343  * \param[out]   xprime                 Updated coordinates
 344  * \param[inout] v                      Velocities, type either rvec* or const rvec*
 345  * \param[in]    f                      Forces
 346  *
 347  * We expect this template to get good SIMD acceleration by most compilers,
 348  * unlike the more complex general template.
 349  * Note that we might get even better SIMD acceleration when we introduce
 350  * aligned (and padded) memory, possibly with some hints for the compilers.
 351  */
 352 template<StoreUpdatedVelocities storeUpdatedVelocities, NumTempScaleValues numTempScaleValues, ApplyParrinelloRahmanVScaling applyPRVScaling, typename VelocityType>
 353 static std::enable_if_t<std::is_same<VelocityType, rvec>::value || std::is_same<VelocityType, const rvec>::value, void>
 354 updateMDLeapfrogSimple(int                                 start,
 355                        int                                 nrend,
 356                        real                                dt,
 357                        real                                dtPressureCouple,
 358                        gmx::ArrayRef<const rvec>           invMassPerDim,
 359                        gmx::ArrayRef<const t_grp_tcstat>   tcstat,
 360                        gmx::ArrayRef<const unsigned short> cTC,
 361                        const rvec                          pRVScaleMatrixDiagonal,
 362                        const rvec* gmx_restrict x,
 363                        rvec* gmx_restrict xprime,
 364                        VelocityType* gmx_restrict v,
 365                        const rvec* gmx_restrict f)
 366 {
 367     real lambdaGroup;
 368
 369     if (numTempScaleValues == NumTempScaleValues::single)
 370     {
 371         lambdaGroup = tcstat[0].lambda;
 372     }
 373
 374     for (int a = start; a < nrend; a++)
 375     {
 376         if (numTempScaleValues == NumTempScaleValues::multiple)
 377         {
 378             lambdaGroup = tcstat[cTC[a]].lambda;
 379         }
 380
 381         for (int d = 0; d < DIM; d++)
 382         {
 383             /* Note that using rvec invMassPerDim results in more efficient
 384              * SIMD code, but this increases the cache pressure.
 385              * For large systems with PME on the CPU this slows down the
 386              * (then already slow) update by 20%. If all data remains in cache,
 387              * using rvec is much faster.
 388              */
 389             real vNew = lambdaGroup * v[a][d] + f[a][d] * invMassPerDim[a][d] * dt;
 390
 391             if (applyPRVScaling == ApplyParrinelloRahmanVScaling::diagonal)
 392             {
 393                 vNew -= dtPressureCouple * pRVScaleMatrixDiagonal[d] * v[a][d];
 394             }
 395             // TODO: Remove NOLINTs once clang-tidy is updated to v11, it should be able to handle constexpr.
 396             if constexpr (storeUpdatedVelocities == StoreUpdatedVelocities::yes) // NOLINT // NOLINTNEXTLINE
 397             {
 398                 v[a][d] = vNew;
 399             }
 400             xprime[a][d] = x[a][d] + vNew * dt; // NOLINT(readability-misleading-indentation)
 401         }
 402     }
 403 }
 404
 405 #if GMX_SIMD && GMX_SIMD_HAVE_REAL
 406 #    define GMX_HAVE_SIMD_UPDATE 1
 407 #else
 408 #    define GMX_HAVE_SIMD_UPDATE 0
 409 #endif
 410
 411 #if GMX_HAVE_SIMD_UPDATE
 412
 413 /*! \brief Load (aligned) the contents of GMX_SIMD_REAL_WIDTH rvec elements sequentially into 3 SIMD registers
 414  *
 415  * The loaded output is:
 416  * \p r0: { r[index][XX], r[index][YY], ... }
 417  * \p r1: { ... }
 418  * \p r2: { ..., r[index+GMX_SIMD_REAL_WIDTH-1][YY], r[index+GMX_SIMD_REAL_WIDTH-1][ZZ] }
 419  *
 420  * \param[in]  r      Real to an rvec array, has to be aligned to SIMD register width
 421  * \param[in]  index  Index of the first rvec triplet of reals to load
 422  * \param[out] r0     Pointer to first SIMD register
 423  * \param[out] r1     Pointer to second SIMD register
 424  * \param[out] r2     Pointer to third SIMD register
 425  */
 426 static inline void simdLoadRvecs(const rvec* r, int index, SimdReal* r0, SimdReal* r1, SimdReal* r2)
 427 {
 428     const real* realPtr = r[index];
 429
 430     GMX_ASSERT(isSimdAligned(realPtr), "Pointer should be SIMD aligned");
 431
 432     *r0 = simdLoad(realPtr + 0 * GMX_SIMD_REAL_WIDTH);
 433     *r1 = simdLoad(realPtr + 1 * GMX_SIMD_REAL_WIDTH);
 434     *r2 = simdLoad(realPtr + 2 * GMX_SIMD_REAL_WIDTH);
 435 }
 436
 437 /*! \brief Store (aligned) 3 SIMD registers sequentially to GMX_SIMD_REAL_WIDTH rvec elements
 438  *
 439  * The stored output is:
 440  * \p r[index] = { { r0[0], r0[1], ... }
 441  * ...
 442  * \p r[index+GMX_SIMD_REAL_WIDTH-1] =  { ... , r2[GMX_SIMD_REAL_WIDTH-2], r2[GMX_SIMD_REAL_WIDTH-1] }
 443  *
 444  * \param[out] r      Pointer to an rvec array, has to be aligned to SIMD register width
 445  * \param[in]  index  Index of the first rvec triplet of reals to store to
 446  * \param[in]  r0     First SIMD register
 447  * \param[in]  r1     Second SIMD register
 448  * \param[in]  r2     Third SIMD register
 449  */
 450 static inline void simdStoreRvecs(rvec* r, int index, SimdReal r0, SimdReal r1, SimdReal r2)
 451 {
 452     real* realPtr = r[index];
 453
 454     GMX_ASSERT(isSimdAligned(realPtr), "Pointer should be SIMD aligned");
 455
 456     store(realPtr + 0 * GMX_SIMD_REAL_WIDTH, r0);
 457     store(realPtr + 1 * GMX_SIMD_REAL_WIDTH, r1);
 458     store(realPtr + 2 * GMX_SIMD_REAL_WIDTH, r2);
 459 }
 460
 461 /*! \brief Integrate using leap-frog with single group T-scaling and SIMD
 462  *
 463  * \tparam       storeUpdatedVelocities Tells whether we should store the updated velocities
 464  * \param[in]    start                  Index of first atom to update
 465  * \param[in]    nrend                  Last atom to update: \p nrend - 1
 466  * \param[in]    dt                     The time step
 467  * \param[in]    invMass                1/mass per atom
 468  * \param[in]    tcstat                 Temperature coupling information
 469  * \param[in]    x                      Input coordinates
 470  * \param[out]   xprime                 Updated coordinates
 471  * \param[inout] v                      Velocities, type either rvec* or const rvec*
 472  * \param[in]    f                      Forces
 473  */
 474 template<StoreUpdatedVelocities storeUpdatedVelocities, typename VelocityType>
 475 static std::enable_if_t<std::is_same<VelocityType, rvec>::value || std::is_same<VelocityType, const rvec>::value, void>
 476 updateMDLeapfrogSimpleSimd(int                               start,
 477                            int                               nrend,
 478                            real                              dt,
 479                            gmx::ArrayRef<const real>         invMass,
 480                            gmx::ArrayRef<const t_grp_tcstat> tcstat,
 481                            const rvec* gmx_restrict x,
 482                            rvec* gmx_restrict xprime,
 483                            VelocityType* gmx_restrict v,
 484                            const rvec* gmx_restrict f)
 485 {
 486     SimdReal timestep(dt);
 487     SimdReal lambdaSystem(tcstat[0].lambda);
 488
 489     /* We declare variables here, since code is often slower when declaring them inside the loop */
 490
 491     /* Note: We should implement a proper PaddedVector, so we don't need this check */
 492     GMX_ASSERT(isSimdAligned(invMass.data()), "invMass should be aligned");
 493
 494     for (int a = start; a < nrend; a += GMX_SIMD_REAL_WIDTH)
 495     {
 496         SimdReal invMass0, invMass1, invMass2;
 497         expandScalarsToTriplets(simdLoad(invMass.data() + a), &invMass0, &invMass1, &invMass2);
 498
 499         SimdReal v0, v1, v2;
 500         SimdReal f0, f1, f2;
 501         simdLoadRvecs(v, a, &v0, &v1, &v2);
 502         simdLoadRvecs(f, a, &f0, &f1, &f2);
 503
 504         v0 = fma(f0 * invMass0, timestep, lambdaSystem * v0);
 505         v1 = fma(f1 * invMass1, timestep, lambdaSystem * v1);
 506         v2 = fma(f2 * invMass2, timestep, lambdaSystem * v2);
 507
 508         // TODO: Remove NOLINTs once clang-tidy is updated to v11, it should be able to handle constexpr.
 509         if constexpr (storeUpdatedVelocities == StoreUpdatedVelocities::yes) // NOLINT // NOLINTNEXTLINE
 510         {
 511             simdStoreRvecs(v, a, v0, v1, v2);
 512         }
 513
 514         SimdReal x0, x1, x2; // NOLINT(readability-misleading-indentation)
 515         simdLoadRvecs(x, a, &x0, &x1, &x2);
 516
 517         SimdReal xprime0 = fma(v0, timestep, x0);
 518         SimdReal xprime1 = fma(v1, timestep, x1);
 519         SimdReal xprime2 = fma(v2, timestep, x2);
 520
 521         simdStoreRvecs(xprime, a, xprime0, xprime1, xprime2);
 522     }
 523 }
 524
 525 #endif // GMX_HAVE_SIMD_UPDATE
 526
 527 /*! \brief Sets the NEMD acceleration type */
 528 enum class AccelerationType
 529 {
 530     none,
 531     cosine
 532 };
 533
 534 /*! \brief Integrate using leap-frog with support for everything.
 535  *
 536  * \tparam        accelerationType  Type of NEMD acceleration.
 537  * \param[in]     start             Index of first atom to update.
 538  * \param[in]     nrend             Last atom to update: \p nrend - 1.
 539  * \param[in]     doNoseHoover      If to apply Nose-Hoover T-coupling.
 540  * \param[in]     dt                The time step.
 541  * \param[in]     dtPressureCouple  Time step for pressure coupling, is 0 when no pressure
 542  *                                  coupling should be applied at this step.
 543  * \param[in]     cTC               Temperature coupling group indices
 544  * \param[in]     invMassPerDim     Inverse mass per dimension
 545  * \param[in]     ekind             Kinetic energy data.
 546  * \param[in]     box               The box dimensions.
 547  * \param[in]     x                 Input coordinates.
 548  * \param[out]    xprime            Updated coordinates.
 549  * \param[inout]  v                 Velocities.
 550  * \param[in]     f                 Forces.
 551  * \param[in]     nh_vxi            Nose-Hoover velocity scaling factors.
 552  * \param[in]     nsttcouple        Frequency of the temperature coupling steps.
 553  * \param[in]     M                 Parrinello-Rahman scaling matrix.
 554  */
 555 template<AccelerationType accelerationType>
 556 static void updateMDLeapfrogGeneral(int                                 start,
 557                                     int                                 nrend,
 558                                     bool                                doNoseHoover,
 559                                     real                                dt,
 560                                     real                                dtPressureCouple,
 561                                     gmx::ArrayRef<const unsigned short> cTC,
 562                                     gmx::ArrayRef<const rvec>           invMassPerDim,
 563                                     const gmx_ekindata_t*               ekind,
 564                                     const matrix                        box,
 565                                     const rvec* gmx_restrict x,
 566                                     rvec* gmx_restrict xprime,
 567                                     rvec* gmx_restrict v,
 568                                     const rvec* gmx_restrict f,
 569                                     const double* gmx_restrict nh_vxi,
 570                                     const int                  nsttcouple,
 571                                     const matrix               M)
 572 {
 573     /* This is a version of the leap-frog integrator that supports
 574      * all combinations of T-coupling, P-coupling and NEMD.
 575      * Nose-Hoover uses the reversible leap-frog integrator from
 576      * Holian et al. Phys Rev E 52(3) : 2338, 1995
 577      */
 578
 579     gmx::ArrayRef<const t_grp_tcstat> tcstat = ekind->tcstat;
 580
 581     /* Initialize group values, changed later when multiple groups are used */
 582     int gt = 0;
 583
 584     real omega_Z = 2 * static_cast<real>(M_PI) / box[ZZ][ZZ];
 585
 586     for (int n = start; n < nrend; n++)
 587     {
 588         if (!cTC.empty())
 589         {
 590             gt = cTC[n];
 591         }
 592         real lg = tcstat[gt].lambda;
 593
 594         rvec vRel;
 595         real cosineZ, vCosine;
 596         switch (accelerationType)
 597         {
 598             case AccelerationType::none: copy_rvec(v[n], vRel); break;
 599             case AccelerationType::cosine:
 600                 cosineZ = std::cos(x[n][ZZ] * omega_Z);
 601                 vCosine = cosineZ * ekind->cosacc.vcos;
 602                 /* Avoid scaling the cosine profile velocity */
 603                 copy_rvec(v[n], vRel);
 604                 vRel[XX] -= vCosine;
 605                 break;
 606         }
 607
 608         real factorNH = 0.0;
 609         if (doNoseHoover)
 610         {
 611             /* Here we account for multiple time stepping, by increasing
 612              * the Nose-Hoover correction by nsttcouple
 613              * TODO: This can be pre-computed.
 614              */
 615             factorNH = 0.5 * nsttcouple * dt * nh_vxi[gt];
 616         }
 617
 618         for (int d = 0; d < DIM; d++)
 619         {
 620             real vNew = (lg * vRel[d]
 621                          + (f[n][d] * invMassPerDim[n][d] * dt - factorNH * vRel[d]
 622                             - dtPressureCouple * iprod(M[d], vRel)))
 623                         / (1 + factorNH);
 624             switch (accelerationType)
 625             {
 626                 case AccelerationType::none: break;
 627                 case AccelerationType::cosine:
 628                     if (d == XX)
 629                     {
 630                         /* Add back the mean velocity and apply acceleration */
 631                         vNew += vCosine + cosineZ * ekind->cosacc.cos_accel * dt;
 632                     }
 633                     break;
 634             }
 635             v[n][d]      = vNew;
 636             xprime[n][d] = x[n][d] + vNew * dt;
 637         }
 638     }
 639 }
 640
 641 /*! \brief Handles the Leap-frog MD x and v integration */
 642 static void do_update_md(int         start,
 643                          int         nrend,
 644                          real        dt,
 645                          int64_t     step,
 646                          const rvec* gmx_restrict x,
 647                          rvec* gmx_restrict xprime,
 648                          rvec* gmx_restrict v,
 649                          const rvec* gmx_restrict            f,
 650                          const TemperatureCoupling           etc,
 651                          const PressureCoupling              epc,
 652                          const int                           nsttcouple,
 653                          const int                           nstpcouple,
 654                          gmx::ArrayRef<const unsigned short> cTC,
 655                          gmx::ArrayRef<const real> gmx_unused invmass,
 656                          gmx::ArrayRef<const rvec>            invMassPerDim,
 657                          const gmx_ekindata_t*                ekind,
 658                          const matrix                         box,
 659                          const double* gmx_restrict nh_vxi,
 660                          const matrix               M,
 661                          bool gmx_unused havePartiallyFrozenAtoms)
 662 {
 663     GMX_ASSERT(nrend == start || xprime != x,
 664                "For SIMD optimization certain compilers need to have xprime != x");
 665
 666     /* Note: Berendsen pressure scaling is handled after do_update_md() */
 667     bool doTempCouple =
 668             (etc != TemperatureCoupling::No && do_per_step(step + nsttcouple - 1, nsttcouple));
 669     bool doNoseHoover       = (etc == TemperatureCoupling::NoseHoover && doTempCouple);
 670     bool doParrinelloRahman = (epc == PressureCoupling::ParrinelloRahman
 671                                && do_per_step(step + nstpcouple - 1, nstpcouple));
 672     bool doPROffDiagonal = (doParrinelloRahman && (M[YY][XX] != 0 || M[ZZ][XX] != 0 || M[ZZ][YY] != 0));
 673
 674     real dtPressureCouple = (doParrinelloRahman ? nstpcouple * dt : 0);
 675
 676     /* NEMD cosine acceleration is applied in updateMDLeapFrogGeneral */
 677     bool doAcceleration = (ekind->cosacc.cos_accel != 0);
 678
 679     if (doNoseHoover || doPROffDiagonal || doAcceleration)
 680     {
 681         matrix stepM;
 682         if (!doParrinelloRahman)
 683         {
 684             /* We should not apply PR scaling at this step */
 685             clear_mat(stepM);
 686         }
 687         else
 688         {
 689             copy_mat(M, stepM);
 690         }
 691
 692         if (!doAcceleration)
 693         {
 694             updateMDLeapfrogGeneral<AccelerationType::none>(start,
 695                                                             nrend,
 696                                                             doNoseHoover,
 697                                                             dt,
 698                                                             dtPressureCouple,
 699                                                             cTC,
 700                                                             invMassPerDim,
 701                                                             ekind,
 702                                                             box,
 703                                                             x,
 704                                                             xprime,
 705                                                             v,
 706                                                             f,
 707                                                             nh_vxi,
 708                                                             nsttcouple,
 709                                                             stepM);
 710         }
 711         else
 712         {
 713             updateMDLeapfrogGeneral<AccelerationType::cosine>(start,
 714                                                               nrend,
 715                                                               doNoseHoover,
 716                                                               dt,
 717                                                               dtPressureCouple,
 718                                                               cTC,
 719                                                               invMassPerDim,
 720                                                               ekind,
 721                                                               box,
 722                                                               x,
 723                                                               xprime,
 724                                                               v,
 725                                                               f,
 726                                                               nh_vxi,
 727                                                               nsttcouple,
 728                                                               stepM);
 729         }
 730     }
 731     else
 732     {
 733         /* We determine if we have a single T-coupling lambda value for all
 734          * atoms. That allows for better SIMD acceleration in the template.
 735          * If we do not do temperature coupling (in the run or this step),
 736          * all scaling values are 1, so we effectively have a single value.
 737          */
 738         bool haveSingleTempScaleValue = (!doTempCouple || ekind->ngtc == 1);
 739
 740         /* Extract some pointers needed by all cases */
 741         gmx::ArrayRef<const t_grp_tcstat> tcstat = ekind->tcstat;
 742
 743         if (doParrinelloRahman)
 744         {
 745             GMX_ASSERT(!doPROffDiagonal,
 746                        "updateMDLeapfrogSimple only support diagonal Parrinello-Rahman scaling "
 747                        "matrices");
 748
 749             rvec diagM;
 750             for (int d = 0; d < DIM; d++)
 751             {
 752                 diagM[d] = M[d][d];
 753             }
 754
 755             if (haveSingleTempScaleValue)
 756             {
 757                 updateMDLeapfrogSimple<StoreUpdatedVelocities::yes, NumTempScaleValues::single, ApplyParrinelloRahmanVScaling::diagonal>(
 758                         start, nrend, dt, dtPressureCouple, invMassPerDim, tcstat, cTC, diagM, x, xprime, v, f);
 759             }
 760             else
 761             {
 762                 updateMDLeapfrogSimple<StoreUpdatedVelocities::yes, NumTempScaleValues::multiple, ApplyParrinelloRahmanVScaling::diagonal>(
 763                         start, nrend, dt, dtPressureCouple, invMassPerDim, tcstat, cTC, diagM, x, xprime, v, f);
 764             }
 765         }
 766         else
 767         {
 768             if (haveSingleTempScaleValue)
 769             {
 770                 /* Note that modern compilers are pretty good at vectorizing
 771                  * updateMDLeapfrogSimple(). But the SIMD version will still
 772                  * be faster because invMass lowers the cache pressure
 773                  * compared to invMassPerDim.
 774                  */
 775 #if GMX_HAVE_SIMD_UPDATE
 776                 /* Check if we can use invmass instead of invMassPerDim */
 777                 if (!havePartiallyFrozenAtoms)
 778                 {
 779                     updateMDLeapfrogSimpleSimd<StoreUpdatedVelocities::yes>(
 780                             start, nrend, dt, invmass, tcstat, x, xprime, v, f);
 781                 }
 782                 else
 783 #endif
 784                 {
 785                     updateMDLeapfrogSimple<StoreUpdatedVelocities::yes, NumTempScaleValues::single, ApplyParrinelloRahmanVScaling::no>(
 786                             start, nrend, dt, dtPressureCouple, invMassPerDim, tcstat, cTC, nullptr, x, xprime, v, f);
 787                 }
 788             }
 789             else
 790             {
 791                 updateMDLeapfrogSimple<StoreUpdatedVelocities::yes, NumTempScaleValues::multiple, ApplyParrinelloRahmanVScaling::no>(
 792                         start, nrend, dt, dtPressureCouple, invMassPerDim, tcstat, cTC, nullptr, x, xprime, v, f);
 793             }
 794         }
 795     }
 796 }
 797 /*! \brief Handles the Leap-frog MD x and v integration */
 798 static void doUpdateMDDoNotUpdateVelocities(int         start,
 799                                             int         nrend,
 800                                             real        dt,
 801                                             const rvec* gmx_restrict x,
 802                                             rvec* gmx_restrict xprime,
 803                                             const rvec* gmx_restrict v,
 804                                             const rvec* gmx_restrict f,
 805                                             bool gmx_unused     havePartiallyFrozenAtoms,
 806                                             gmx::ArrayRef<real> gmx_unused invmass,
 807                                             gmx::ArrayRef<rvec>            invMassPerDim,
 808                                             const gmx_ekindata_t&          ekind)
 809 {
 810     GMX_ASSERT(nrend == start || xprime != x,
 811                "For SIMD optimization certain compilers need to have xprime != x");
 812
 813     gmx::ArrayRef<const t_grp_tcstat> tcstat = ekind.tcstat;
 814
 815     /* Check if we can use invmass instead of invMassPerDim */
 816 #if GMX_HAVE_SIMD_UPDATE
 817     if (!havePartiallyFrozenAtoms)
 818     {
 819         updateMDLeapfrogSimpleSimd<StoreUpdatedVelocities::no>(
 820                 start, nrend, dt, invmass, tcstat, x, xprime, v, f);
 821     }
 822     else
 823 #endif
 824     {
 825         updateMDLeapfrogSimple<StoreUpdatedVelocities::no, NumTempScaleValues::single, ApplyParrinelloRahmanVScaling::no>(
 826                 start, nrend, dt, dt, invMassPerDim, tcstat, gmx::ArrayRef<const unsigned short>(), nullptr, x, xprime, v, f);
 827     }
 828 }
 829
 830 static void do_update_vv_vel(int                                 start,
 831                              int                                 nrend,
 832                              real                                dt,
 833                              gmx::ArrayRef<const ivec>           nFreeze,
 834                              gmx::ArrayRef<const real>           invmass,
 835                              gmx::ArrayRef<const ParticleType>   ptype,
 836                              gmx::ArrayRef<const unsigned short> cFREEZE,
 837                              rvec                                v[],
 838                              const rvec                          f[],
 839                              gmx_bool                            bExtended,
 840                              real                                veta,
 841                              real                                alpha)
 842 {
 843     int  gf = 0;
 844     int  n, d;
 845     real g, mv1, mv2;
 846
 847     if (bExtended)
 848     {
 849         g   = 0.25 * dt * veta * alpha;
 850         mv1 = std::exp(-g);
 851         mv2 = gmx::series_sinhx(g);
 852     }
 853     else
 854     {
 855         mv1 = 1.0;
 856         mv2 = 1.0;
 857     }
 858     for (n = start; n < nrend; n++)
 859     {
 860         real w_dt = invmass[n] * dt;
 861         if (!cFREEZE.empty())
 862         {
 863             gf = cFREEZE[n];
 864         }
 865
 866         for (d = 0; d < DIM; d++)
 867         {
 868             if ((ptype[n] != ParticleType::Shell) && !nFreeze[gf][d])
 869             {
 870                 v[n][d] = mv1 * (mv1 * v[n][d] + 0.5 * (w_dt * mv2 * f[n][d]));
 871             }
 872             else
 873             {
 874                 v[n][d] = 0.0;
 875             }
 876         }
 877     }
 878 } /* do_update_vv_vel */
 879
 880 static void do_update_vv_pos(int                                 start,
 881                              int                                 nrend,
 882                              real                                dt,
 883                              gmx::ArrayRef<const ivec>           nFreeze,
 884                              gmx::ArrayRef<const ParticleType>   ptype,
 885                              gmx::ArrayRef<const unsigned short> cFREEZE,
 886                              const rvec                          x[],
 887                              rvec                                xprime[],
 888                              const rvec                          v[],
 889                              gmx_bool                            bExtended,
 890                              real                                veta)
 891 {
 892     int  gf = 0;
 893     int  n, d;
 894     real g, mr1, mr2;
 895
 896     /* Would it make more sense if Parrinello-Rahman was put here? */
 897     if (bExtended)
 898     {
 899         g   = 0.5 * dt * veta;
 900         mr1 = std::exp(g);
 901         mr2 = gmx::series_sinhx(g);
 902     }
 903     else
 904     {
 905         mr1 = 1.0;
 906         mr2 = 1.0;
 907     }
 908
 909     for (n = start; n < nrend; n++)
 910     {
 911
 912         if (!cFREEZE.empty())
 913         {
 914             gf = cFREEZE[n];
 915         }
 916
 917         for (d = 0; d < DIM; d++)
 918         {
 919             if ((ptype[n] != ParticleType::Shell) && !nFreeze[gf][d])
 920             {
 921                 xprime[n][d] = mr1 * (mr1 * x[n][d] + mr2 * dt * v[n][d]);
 922             }
 923             else
 924             {
 925                 xprime[n][d] = x[n][d];
 926             }
 927         }
 928     }
 929 } /* do_update_vv_pos */
 930
 931 gmx_stochd_t::gmx_stochd_t(const t_inputrec& inputRecord)
 932 {
 933     const t_grpopts* opts = &inputRecord.opts;
 934     const int        ngtc = opts->ngtc;
 935
 936     if (inputRecord.eI == IntegrationAlgorithm::BD)
 937     {
 938         bd_rf.resize(ngtc);
 939     }
 940     else if (EI_SD(inputRecord.eI))
 941     {
 942         sdc.resize(ngtc);
 943         sdsig.resize(ngtc);
 944
 945         for (int gt = 0; gt < ngtc; gt++)
 946         {
 947             if (opts->tau_t[gt] > 0)
 948             {
 949                 sdc[gt].em = std::exp(-inputRecord.delta_t / opts->tau_t[gt]);
 950             }
 951             else
 952             {
 953                 /* No friction and noise on this group */
 954                 sdc[gt].em = 1;
 955             }
 956         }
 957     }
 958     else if (ETC_ANDERSEN(inputRecord.etc))
 959     {
 960         randomize_group.resize(ngtc);
 961         boltzfac.resize(ngtc);
 962
 963         /* for now, assume that all groups, if randomized, are randomized at the same rate, i.e. tau_t is the same. */
 964         /* since constraint groups don't necessarily match up with temperature groups! This is checked in readir.c */
 965
 966         for (int gt = 0; gt < ngtc; gt++)
 967         {
 968             real reft = std::max<real>(0, opts->ref_t[gt]);
 969             if ((opts->tau_t[gt] > 0)
 970                 && (reft > 0)) /* tau_t or ref_t = 0 means that no randomization is done */
 971             {
 972                 randomize_group[gt] = true;
 973                 boltzfac[gt]        = gmx::c_boltz * opts->ref_t[gt];
 974             }
 975             else
 976             {
 977                 randomize_group[gt] = false;
 978             }
 979         }
 980     }
 981 }
 982
 983 void Update::Impl::update_temperature_constants(const t_inputrec& inputRecord)
 984 {
 985     if (inputRecord.eI == IntegrationAlgorithm::BD)
 986     {
 987         if (inputRecord.bd_fric != 0)
 988         {
 989             for (int gt = 0; gt < inputRecord.opts.ngtc; gt++)
 990             {
 991                 sd_.bd_rf[gt] = std::sqrt(2.0 * gmx::c_boltz * inputRecord.opts.ref_t[gt]
 992                                           / (inputRecord.bd_fric * inputRecord.delta_t));
 993             }
 994         }
 995         else
 996         {
 997             for (int gt = 0; gt < inputRecord.opts.ngtc; gt++)
 998             {
 999                 sd_.bd_rf[gt] = std::sqrt(2.0 * gmx::c_boltz * inputRecord.opts.ref_t[gt]);
1000             }
1001         }
1002     }
1003     if (inputRecord.eI == IntegrationAlgorithm::SD1)
1004     {
1005         for (int gt = 0; gt < inputRecord.opts.ngtc; gt++)
1006         {
1007             real kT = gmx::c_boltz * inputRecord.opts.ref_t[gt];
1008             /* The mass is accounted for later, since this differs per atom */
1009             sd_.sdsig[gt].V = std::sqrt(kT * (1 - sd_.sdc[gt].em * sd_.sdc[gt].em));
1010         }
1011     }
1012 }
1013
1014 Update::Impl::Impl(const t_inputrec& inputRecord, BoxDeformation* boxDeformation) :
1015     sd_(inputRecord),
1016     deform_(boxDeformation)
1017 {
1018     update_temperature_constants(inputRecord);
1019     xp_.resizeWithPadding(0);
1020 }
1021
1022 void Update::updateAfterPartition(int                                 numAtoms,
1023                                   gmx::ArrayRef<const unsigned short> cFREEZE,
1024                                   gmx::ArrayRef<const unsigned short> cTC)
1025 {
1026
1027     impl_->xp()->resizeWithPadding(numAtoms);
1028     impl_->cFREEZE_ = cFREEZE;
1029     impl_->cTC_     = cTC;
1030 }
1031
1032 /*! \brief Sets the SD update type */
1033 enum class SDUpdate : int
1034 {
1035     ForcesOnly,
1036     FrictionAndNoiseOnly,
1037     Combined
1038 };
1039
1040 /*! \brief SD integrator update
1041  *
1042  * Two phases are required in the general case of a constrained
1043  * update, the first phase from the contribution of forces, before
1044  * applying constraints, and then a second phase applying the friction
1045  * and noise, and then further constraining. For details, see
1046  * Goga2012.
1047  *
1048  * Without constraints, the two phases can be combined, for
1049  * efficiency.
1050  *
1051  * Thus three instantiations of this templated function will be made,
1052  * two with only one contribution, and one with both contributions. */
1053 template<SDUpdate updateType>
1054 static void doSDUpdateGeneral(const gmx_stochd_t&                 sd,
1055                               int                                 start,
1056                               int                                 nrend,
1057                               real                                dt,
1058                               gmx::ArrayRef<const ivec>           nFreeze,
1059                               gmx::ArrayRef<const real>           invmass,
1060                               gmx::ArrayRef<const ParticleType>   ptype,
1061                               gmx::ArrayRef<const unsigned short> cFREEZE,
1062                               gmx::ArrayRef<const unsigned short> cTC,
1063                               const rvec                          x[],
1064                               rvec                                xprime[],
1065                               rvec                                v[],
1066                               const rvec                          f[],
1067                               int64_t                             step,
1068                               int                                 seed,
1069                               const int*                          gatindex)
1070 {
1071     // cTC and cFREEZE can be nullptr any time, but various
1072     // instantiations do not make sense with particular pointer
1073     // values.
1074     if (updateType == SDUpdate::ForcesOnly)
1075     {
1076         GMX_ASSERT(f != nullptr, "SD update with only forces requires forces");
1077         GMX_ASSERT(cTC.empty(), "SD update with only forces cannot handle temperature groups");
1078     }
1079     if (updateType == SDUpdate::FrictionAndNoiseOnly)
1080     {
1081         GMX_ASSERT(f == nullptr, "SD update with only noise cannot handle forces");
1082     }
1083     if (updateType == SDUpdate::Combined)
1084     {
1085         GMX_ASSERT(f != nullptr, "SD update with forces and noise requires forces");
1086     }
1087
1088     // Even 0 bits internal counter gives 2x64 ints (more than enough for three table lookups)
1089     gmx::ThreeFry2x64<0>                       rng(seed, gmx::RandomDomain::UpdateCoordinates);
1090     gmx::TabulatedNormalDistribution<real, 14> dist;
1091
1092     for (int n = start; n < nrend; n++)
1093     {
1094         int globalAtomIndex = gatindex ? gatindex[n] : n;
1095         rng.restart(step, globalAtomIndex);
1096         dist.reset();
1097
1098         real inverseMass = invmass[n];
1099         real invsqrtMass = std::sqrt(inverseMass);
1100
1101         int freezeGroup      = !cFREEZE.empty() ? cFREEZE[n] : 0;
1102         int temperatureGroup = !cTC.empty() ? cTC[n] : 0;
1103
1104         for (int d = 0; d < DIM; d++)
1105         {
1106             if ((ptype[n] != ParticleType::Shell) && !nFreeze[freezeGroup][d])
1107             {
1108                 if (updateType == SDUpdate::ForcesOnly)
1109                 {
1110                     real vn = v[n][d] + inverseMass * f[n][d] * dt;
1111                     v[n][d] = vn;
1112                     // Simple position update.
1113                     xprime[n][d] = x[n][d] + v[n][d] * dt;
1114                 }
1115                 else if (updateType == SDUpdate::FrictionAndNoiseOnly)
1116                 {
1117                     real vn = v[n][d];
1118                     v[n][d] = (vn * sd.sdc[temperatureGroup].em
1119                                + invsqrtMass * sd.sdsig[temperatureGroup].V * dist(rng));
1120                     // The previous phase already updated the
1121                     // positions with a full v*dt term that must
1122                     // now be half removed.
1123                     xprime[n][d] = xprime[n][d] + 0.5 * (v[n][d] - vn) * dt;
1124                 }
1125                 else
1126                 {
1127                     real vn = v[n][d] + inverseMass * f[n][d] * dt;
1128                     v[n][d] = (vn * sd.sdc[temperatureGroup].em
1129                                + invsqrtMass * sd.sdsig[temperatureGroup].V * dist(rng));
1130                     // Here we include half of the friction+noise
1131                     // update of v into the position update.
1132                     xprime[n][d] = x[n][d] + 0.5 * (vn + v[n][d]) * dt;
1133                 }
1134             }
1135             else
1136             {
1137                 // When using constraints, the update is split into
1138                 // two phases, but we only need to zero the update of
1139                 // virtual, shell or frozen particles in at most one
1140                 // of the phases.
1141                 if (updateType != SDUpdate::FrictionAndNoiseOnly)
1142                 {
1143                     v[n][d]      = 0.0;
1144                     xprime[n][d] = x[n][d];
1145                 }
1146             }
1147         }
1148     }
1149 }
1150
1151 static void do_update_sd(int         start,
1152                          int         nrend,
1153                          real        dt,
1154                          int64_t     step,
1155                          const rvec* gmx_restrict x,
1156                          rvec* gmx_restrict xprime,
1157                          rvec* gmx_restrict v,
1158                          const rvec* gmx_restrict            f,
1159                          gmx::ArrayRef<const ivec>           nFreeze,
1160                          gmx::ArrayRef<const real>           invmass,
1161                          gmx::ArrayRef<const ParticleType>   ptype,
1162                          gmx::ArrayRef<const unsigned short> cFREEZE,
1163                          gmx::ArrayRef<const unsigned short> cTC,
1164                          int                                 seed,
1165                          const t_commrec*                    cr,
1166                          const gmx_stochd_t&                 sd,
1167                          bool                                haveConstraints)
1168 {
1169     if (haveConstraints)
1170     {
1171         // With constraints, the SD update is done in 2 parts
1172         doSDUpdateGeneral<SDUpdate::ForcesOnly>(sd,
1173                                                 start,
1174                                                 nrend,
1175                                                 dt,
1176                                                 nFreeze,
1177                                                 invmass,
1178                                                 ptype,
1179                                                 cFREEZE,
1180                                                 gmx::ArrayRef<const unsigned short>(),
1181                                                 x,
1182                                                 xprime,
1183                                                 v,
1184                                                 f,
1185                                                 step,
1186                                                 seed,
1187                                                 nullptr);
1188     }
1189     else
1190     {
1191         doSDUpdateGeneral<SDUpdate::Combined>(sd,
1192                                               start,
1193                                               nrend,
1194                                               dt,
1195                                               nFreeze,
1196                                               invmass,
1197                                               ptype,
1198                                               cFREEZE,
1199                                               cTC,
1200                                               x,
1201                                               xprime,
1202                                               v,
1203                                               f,
1204                                               step,
1205                                               seed,
1206                                               DOMAINDECOMP(cr) ? cr->dd->globalAtomIndices.data() : nullptr);
1207     }
1208 }
1209
1210 static void do_update_bd(int         start,
1211                          int         nrend,
1212                          real        dt,
1213                          int64_t     step,
1214                          const rvec* gmx_restrict x,
1215                          rvec* gmx_restrict xprime,
1216                          rvec* gmx_restrict v,
1217                          const rvec* gmx_restrict            f,
1218                          gmx::ArrayRef<const ivec>           nFreeze,
1219                          gmx::ArrayRef<const real>           invmass,
1220                          gmx::ArrayRef<const ParticleType>   ptype,
1221                          gmx::ArrayRef<const unsigned short> cFREEZE,
1222                          gmx::ArrayRef<const unsigned short> cTC,
1223                          real                                friction_coefficient,
1224                          const real*                         rf,
1225                          int                                 seed,
1226                          const int*                          gatindex)
1227 {
1228     /* note -- these appear to be full step velocities . . .  */
1229     int  gf = 0, gt = 0;
1230     real vn;
1231     real invfr = 0;
1232     int  n, d;
1233     // Use 1 bit of internal counters to give us 2*2 64-bits values per stream
1234     // Each 64-bit value is enough for 4 normal distribution table numbers.
1235     gmx::ThreeFry2x64<0>                       rng(seed, gmx::RandomDomain::UpdateCoordinates);
1236     gmx::TabulatedNormalDistribution<real, 14> dist;
1237
1238     if (friction_coefficient != 0)
1239     {
1240         invfr = 1.0 / friction_coefficient;
1241     }
1242
1243     for (n = start; (n < nrend); n++)
1244     {
1245         int ng = gatindex ? gatindex[n] : n;
1246
1247         rng.restart(step, ng);
1248         dist.reset();
1249
1250         if (!cFREEZE.empty())
1251         {
1252             gf = cFREEZE[n];
1253         }
1254         if (!cTC.empty())
1255         {
1256             gt = cTC[n];
1257         }
1258         for (d = 0; (d < DIM); d++)
1259         {
1260             if ((ptype[n] != ParticleType::Shell) && !nFreeze[gf][d])
1261             {
1262                 if (friction_coefficient != 0)
1263                 {
1264                     vn = invfr * f[n][d] + rf[gt] * dist(rng);
1265                 }
1266                 else
1267                 {
1268                     /* NOTE: invmass = 2/(mass*friction_constant*dt) */
1269                     vn = 0.5 * invmass[n] * f[n][d] * dt
1270                          + std::sqrt(0.5 * invmass[n]) * rf[gt] * dist(rng);
1271                 }
1272
1273                 v[n][d]      = vn;
1274                 xprime[n][d] = x[n][d] + vn * dt;
1275             }
1276             else
1277             {
1278                 v[n][d]      = 0.0;
1279                 xprime[n][d] = x[n][d];
1280             }
1281         }
1282     }
1283 }
1284
1285 extern void init_ekinstate(ekinstate_t* ekinstate, const t_inputrec* ir)
1286 {
1287     ekinstate->ekin_n = ir->opts.ngtc;
1288     snew(ekinstate->ekinh, ekinstate->ekin_n);
1289     snew(ekinstate->ekinf, ekinstate->ekin_n);
1290     snew(ekinstate->ekinh_old, ekinstate->ekin_n);
1291     ekinstate->ekinscalef_nhc.resize(ekinstate->ekin_n);
1292     ekinstate->ekinscaleh_nhc.resize(ekinstate->ekin_n);
1293     ekinstate->vscale_nhc.resize(ekinstate->ekin_n);
1294     ekinstate->dekindl          = 0;
1295     ekinstate->mvcos            = 0;
1296     ekinstate->hasReadEkinState = false;
1297 }
1298
1299 void update_ekinstate(ekinstate_t* ekinstate, const gmx_ekindata_t* ekind)
1300 {
1301     int i;
1302
1303     for (i = 0; i < ekinstate->ekin_n; i++)
1304     {
1305         copy_mat(ekind->tcstat[i].ekinh, ekinstate->ekinh[i]);
1306         copy_mat(ekind->tcstat[i].ekinf, ekinstate->ekinf[i]);
1307         copy_mat(ekind->tcstat[i].ekinh_old, ekinstate->ekinh_old[i]);
1308         ekinstate->ekinscalef_nhc[i] = ekind->tcstat[i].ekinscalef_nhc;
1309         ekinstate->ekinscaleh_nhc[i] = ekind->tcstat[i].ekinscaleh_nhc;
1310         ekinstate->vscale_nhc[i]     = ekind->tcstat[i].vscale_nhc;
1311     }
1312
1313     copy_mat(ekind->ekin, ekinstate->ekin_total);
1314     ekinstate->dekindl = ekind->dekindl;
1315     ekinstate->mvcos   = ekind->cosacc.mvcos;
1316 }
1317
1318 void restore_ekinstate_from_state(const t_commrec* cr, gmx_ekindata_t* ekind, const ekinstate_t* ekinstate)
1319 {
1320     int i, n;
1321
1322     if (MASTER(cr))
1323     {
1324         for (i = 0; i < ekinstate->ekin_n; i++)
1325         {
1326             copy_mat(ekinstate->ekinh[i], ekind->tcstat[i].ekinh);
1327             copy_mat(ekinstate->ekinf[i], ekind->tcstat[i].ekinf);
1328             copy_mat(ekinstate->ekinh_old[i], ekind->tcstat[i].ekinh_old);
1329             ekind->tcstat[i].ekinscalef_nhc = ekinstate->ekinscalef_nhc[i];
1330             ekind->tcstat[i].ekinscaleh_nhc = ekinstate->ekinscaleh_nhc[i];
1331             ekind->tcstat[i].vscale_nhc     = ekinstate->vscale_nhc[i];
1332         }
1333
1334         copy_mat(ekinstate->ekin_total, ekind->ekin);
1335
1336         ekind->dekindl      = ekinstate->dekindl;
1337         ekind->cosacc.mvcos = ekinstate->mvcos;
1338         n                   = ekinstate->ekin_n;
1339     }
1340
1341     if (PAR(cr))
1342     {
1343         gmx_bcast(sizeof(n), &n, cr->mpi_comm_mygroup);
1344         for (i = 0; i < n; i++)
1345         {
1346             gmx_bcast(DIM * DIM * sizeof(ekind->tcstat[i].ekinh[0][0]),
1347                       ekind->tcstat[i].ekinh[0],
1348                       cr->mpi_comm_mygroup);
1349             gmx_bcast(DIM * DIM * sizeof(ekind->tcstat[i].ekinf[0][0]),
1350                       ekind->tcstat[i].ekinf[0],
1351                       cr->mpi_comm_mygroup);
1352             gmx_bcast(DIM * DIM * sizeof(ekind->tcstat[i].ekinh_old[0][0]),
1353                       ekind->tcstat[i].ekinh_old[0],
1354                       cr->mpi_comm_mygroup);
1355
1356             gmx_bcast(sizeof(ekind->tcstat[i].ekinscalef_nhc),
1357                       &(ekind->tcstat[i].ekinscalef_nhc),
1358                       cr->mpi_comm_mygroup);
1359             gmx_bcast(sizeof(ekind->tcstat[i].ekinscaleh_nhc),
1360                       &(ekind->tcstat[i].ekinscaleh_nhc),
1361                       cr->mpi_comm_mygroup);
1362             gmx_bcast(sizeof(ekind->tcstat[i].vscale_nhc), &(ekind->tcstat[i].vscale_nhc), cr->mpi_comm_mygroup);
1363         }
1364         gmx_bcast(DIM * DIM * sizeof(ekind->ekin[0][0]), ekind->ekin[0], cr->mpi_comm_mygroup);
1365
1366         gmx_bcast(sizeof(ekind->dekindl), &ekind->dekindl, cr->mpi_comm_mygroup);
1367         gmx_bcast(sizeof(ekind->cosacc.mvcos), &ekind->cosacc.mvcos, cr->mpi_comm_mygroup);
1368     }
1369 }
1370
1371 void getThreadAtomRange(int numThreads, int threadIndex, int numAtoms, int* startAtom, int* endAtom)
1372 {
1373 #if GMX_HAVE_SIMD_UPDATE
1374     constexpr int blockSize = GMX_SIMD_REAL_WIDTH;
1375 #else
1376     constexpr int blockSize = 1;
1377 #endif
1378     int numBlocks = (numAtoms + blockSize - 1) / blockSize;
1379
1380     *startAtom = ((numBlocks * threadIndex) / numThreads) * blockSize;
1381     *endAtom   = ((numBlocks * (threadIndex + 1)) / numThreads) * blockSize;
1382     if (threadIndex == numThreads - 1)
1383     {
1384         *endAtom = numAtoms;
1385     }
1386 }
1387
1388 void Update::Impl::update_sd_second_half(const t_inputrec&                 inputRecord,
1389                                          int64_t                           step,
1390                                          real*                             dvdlambda,
1391                                          int                               homenr,
1392                                          gmx::ArrayRef<const ParticleType> ptype,
1393                                          gmx::ArrayRef<const real>         invMass,
1394                                          t_state*                          state,
1395                                          const t_commrec*                  cr,
1396                                          t_nrnb*                           nrnb,
1397                                          gmx_wallcycle*                    wcycle,
1398                                          gmx::Constraints*                 constr,
1399                                          bool                              do_log,
1400                                          bool                              do_ene)
1401 {
1402     if (!constr)
1403     {
1404         return;
1405     }
1406     if (inputRecord.eI == IntegrationAlgorithm::SD1)
1407     {
1408
1409         /* Cast delta_t from double to real to make the integrators faster.
1410          * The only reason for having delta_t double is to get accurate values
1411          * for t=delta_t*step when step is larger than float precision.
1412          * For integration dt the accuracy of real suffices, since with
1413          * integral += dt*integrand the increment is nearly always (much) smaller
1414          * than the integral (and the integrand has real precision).
1415          */
1416         real dt = inputRecord.delta_t;
1417
1418         wallcycle_start(wcycle, WallCycleCounter::Update);
1419
1420         int nth = gmx_omp_nthreads_get(ModuleMultiThread::Update);
1421
1422 #pragma omp parallel for num_threads(nth) schedule(static)
1423         for (int th = 0; th < nth; th++)
1424         {
1425             try
1426             {
1427                 int start_th, end_th;
1428                 getThreadAtomRange(nth, th, homenr, &start_th, &end_th);
1429
1430                 doSDUpdateGeneral<SDUpdate::FrictionAndNoiseOnly>(
1431                         sd_,
1432                         start_th,
1433                         end_th,
1434                         dt,
1435                         gmx::arrayRefFromArray(inputRecord.opts.nFreeze, inputRecord.opts.ngfrz),
1436                         invMass,
1437                         ptype,
1438                         cFREEZE_,
1439                         cTC_,
1440                         state->x.rvec_array(),
1441                         xp_.rvec_array(),
1442                         state->v.rvec_array(),
1443                         nullptr,
1444                         step,
1445                         inputRecord.ld_seed,
1446                         DOMAINDECOMP(cr) ? cr->dd->globalAtomIndices.data() : nullptr);
1447             }
1448             GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR
1449         }
1450         inc_nrnb(nrnb, eNR_UPDATE, homenr);
1451         wallcycle_stop(wcycle, WallCycleCounter::Update);
1452
1453         /* Constrain the coordinates upd->xp for half a time step */
1454         bool computeVirial = false;
1455         constr->apply(do_log,
1456                       do_ene,
1457                       step,
1458                       1,
1459                       0.5,
1460                       state->x.arrayRefWithPadding(),
1461                       xp_.arrayRefWithPadding(),
1462                       ArrayRef<RVec>(),
1463                       state->box,
1464                       state->lambda[static_cast<int>(FreeEnergyPerturbationCouplingType::Bonded)],
1465                       dvdlambda,
1466                       state->v.arrayRefWithPadding(),
1467                       computeVirial,
1468                       nullptr,
1469                       ConstraintVariable::Positions);
1470     }
1471 }
1472
1473 void Update::Impl::finish_update(const t_inputrec& inputRecord,
1474                                  const t_mdatoms*  md,
1475                                  t_state*          state,
1476                                  gmx_wallcycle*    wcycle,
1477                                  const bool        haveConstraints)
1478 {
1479     /* NOTE: Currently we always integrate to a temporary buffer and
1480      * then copy the results back here.
1481      */
1482
1483     wallcycle_start_nocount(wcycle, WallCycleCounter::Update);
1484
1485     const int homenr = md->homenr;
1486     auto      xp     = makeConstArrayRef(xp_).subArray(0, homenr);
1487     auto      x      = makeArrayRef(state->x).subArray(0, homenr);
1488
1489     if (md->havePartiallyFrozenAtoms && haveConstraints)
1490     {
1491         /* We have atoms that are frozen along some, but not all dimensions,
1492          * then constraints will have moved them also along the frozen dimensions.
1493          * To freeze such degrees of freedom we do not copy them back here.
1494          */
1495         const ivec* nFreeze = inputRecord.opts.nFreeze;
1496
1497         for (int i = 0; i < homenr; i++)
1498         {
1499             const int g = md->cFREEZE[i];
1500
1501             for (int d = 0; d < DIM; d++)
1502             {
1503                 if (nFreeze[g][d] == 0)
1504                 {
1505                     x[i][d] = xp[i][d];
1506                 }
1507             }
1508         }
1509     }
1510     else
1511     {
1512         /* We have no frozen atoms or fully frozen atoms which have not
1513          * been moved by the update, so we can simply copy all coordinates.
1514          */
1515         int gmx_unused nth = gmx_omp_nthreads_get(ModuleMultiThread::Update);
1516 #pragma omp parallel for num_threads(nth) schedule(static)
1517         for (int i = 0; i < homenr; i++)
1518         {
1519             // Trivial statement, does not throw
1520             x[i] = xp[i];
1521         }
1522     }
1523
1524     wallcycle_stop(wcycle, WallCycleCounter::Update);
1525 }
1526
1527 void Update::Impl::update_coords(const t_inputrec&                 inputRecord,
1528                                  int64_t                           step,
1529                                  int                               homenr,
1530                                  bool                              havePartiallyFrozenAtoms,
1531                                  gmx::ArrayRef<const ParticleType> ptype,
1532                                  gmx::ArrayRef<const real>         invMass,
1533                                  gmx::ArrayRef<rvec>               invMassPerDim,
1534                                  t_state*                          state,
1535                                  const gmx::ArrayRefWithPadding<const gmx::RVec>& f,
1536                                  const t_fcdata&                                  fcdata,
1537                                  const gmx_ekindata_t*                            ekind,
1538                                  const matrix                                     M,
1539                                  int                                              updatePart,
1540                                  const t_commrec*                                 cr,
1541                                  const bool                                       haveConstraints)
1542 {
1543     /* Running the velocity half does nothing except for velocity verlet */
1544     if ((updatePart == etrtVELOCITY1 || updatePart == etrtVELOCITY2) && !EI_VV(inputRecord.eI))
1545     {
1546         gmx_incons("update_coords called for velocity without VV integrator");
1547     }
1548
1549     /* Cast to real for faster code, no loss in precision (see comment above) */
1550     real dt = inputRecord.delta_t;
1551
1552     /* We need to update the NMR restraint history when time averaging is used */
1553     if (state->flags & enumValueToBitMask(StateEntry::DisreRm3Tav))
1554     {
1555         update_disres_history(*fcdata.disres, &state->hist);
1556     }
1557     if (state->flags & enumValueToBitMask(StateEntry::OrireDtav))
1558     {
1559         update_orires_history(*fcdata.orires, &state->hist);
1560     }
1561
1562     /* ############# START The update of velocities and positions ######### */
1563     int nth = gmx_omp_nthreads_get(ModuleMultiThread::Update);
1564
1565 #pragma omp parallel for num_threads(nth) schedule(static)
1566     for (int th = 0; th < nth; th++)
1567     {
1568         try
1569         {
1570             int start_th, end_th;
1571             getThreadAtomRange(nth, th, homenr, &start_th, &end_th);
1572
1573             const rvec* x_rvec  = state->x.rvec_array();
1574             rvec*       xp_rvec = xp_.rvec_array();
1575             rvec*       v_rvec  = state->v.rvec_array();
1576             const rvec* f_rvec  = as_rvec_array(f.unpaddedConstArrayRef().data());
1577
1578             switch (inputRecord.eI)
1579             {
1580                 case (IntegrationAlgorithm::MD):
1581                     do_update_md(start_th,
1582                                  end_th,
1583                                  dt,
1584                                  step,
1585                                  x_rvec,
1586                                  xp_rvec,
1587                                  v_rvec,
1588                                  f_rvec,
1589                                  inputRecord.etc,
1590                                  inputRecord.epc,
1591                                  inputRecord.nsttcouple,
1592                                  inputRecord.nstpcouple,
1593                                  cTC_,
1594                                  invMass,
1595                                  invMassPerDim,
1596                                  ekind,
1597                                  state->box,
1598                                  state->nosehoover_vxi.data(),
1599                                  M,
1600                                  havePartiallyFrozenAtoms);
1601                     break;
1602                 case (IntegrationAlgorithm::SD1):
1603                     do_update_sd(start_th,
1604                                  end_th,
1605                                  dt,
1606                                  step,
1607                                  x_rvec,
1608                                  xp_rvec,
1609                                  v_rvec,
1610                                  f_rvec,
1611                                  gmx::arrayRefFromArray(inputRecord.opts.nFreeze, inputRecord.opts.ngfrz),
1612                                  invMass,
1613                                  ptype,
1614                                  cFREEZE_,
1615                                  cTC_,
1616                                  inputRecord.ld_seed,
1617                                  cr,
1618                                  sd_,
1619                                  haveConstraints);
1620                     break;
1621                 case (IntegrationAlgorithm::BD):
1622                     do_update_bd(start_th,
1623                                  end_th,
1624                                  dt,
1625                                  step,
1626                                  x_rvec,
1627                                  xp_rvec,
1628                                  v_rvec,
1629                                  f_rvec,
1630                                  gmx::arrayRefFromArray(inputRecord.opts.nFreeze, inputRecord.opts.ngfrz),
1631                                  invMass,
1632                                  ptype,
1633                                  cFREEZE_,
1634                                  cTC_,
1635                                  inputRecord.bd_fric,
1636                                  sd_.bd_rf.data(),
1637                                  inputRecord.ld_seed,
1638                                  DOMAINDECOMP(cr) ? cr->dd->globalAtomIndices.data() : nullptr);
1639                     break;
1640                 case (IntegrationAlgorithm::VV):
1641                 case (IntegrationAlgorithm::VVAK):
1642                 {
1643                     gmx_bool bExtended = (inputRecord.etc == TemperatureCoupling::NoseHoover
1644                                           || inputRecord.epc == PressureCoupling::ParrinelloRahman
1645                                           || inputRecord.epc == PressureCoupling::Mttk);
1646
1647                     /* assuming barostat coupled to group 0 */
1648                     real alpha = 1.0 + DIM / static_cast<real>(inputRecord.opts.nrdf[0]);
1649                     switch (updatePart)
1650                     {
1651                         case etrtVELOCITY1:
1652                         case etrtVELOCITY2:
1653                             do_update_vv_vel(start_th,
1654                                              end_th,
1655                                              dt,
1656                                              gmx::arrayRefFromArray(inputRecord.opts.nFreeze,
1657                                                                     inputRecord.opts.ngfrz),
1658                                              invMass,
1659                                              ptype,
1660                                              cFREEZE_,
1661                                              v_rvec,
1662                                              f_rvec,
1663                                              bExtended,
1664                                              state->veta,
1665                                              alpha);
1666                             break;
1667                         case etrtPOSITION:
1668                             do_update_vv_pos(start_th,
1669                                              end_th,
1670                                              dt,
1671                                              gmx::arrayRefFromArray(inputRecord.opts.nFreeze,
1672                                                                     inputRecord.opts.ngfrz),
1673                                              ptype,
1674                                              cFREEZE_,
1675                                              x_rvec,
1676                                              xp_rvec,
1677                                              v_rvec,
1678                                              bExtended,
1679                                              state->veta);
1680                             break;
1681                     }
1682                     break;
1683                 }
1684                 default: gmx_fatal(FARGS, "Don't know how to update coordinates");
1685             }
1686         }
1687         GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR
1688     }
1689 }
1690
1691 void Update::Impl::update_for_constraint_virial(const t_inputrec&   inputRecord,
1692                                                 int                 homenr,
1693                                                 bool                havePartiallyFrozenAtoms,
1694                                                 gmx::ArrayRef<real> invmass,
1695                                                 gmx::ArrayRef<rvec> invMassPerDim,
1696                                                 const t_state&      state,
1697                                                 const gmx::ArrayRefWithPadding<const gmx::RVec>& f,
1698                                                 const gmx_ekindata_t& ekind)
1699 {
1700     GMX_ASSERT(inputRecord.eI == IntegrationAlgorithm::MD || inputRecord.eI == IntegrationAlgorithm::SD1,
1701                "Only leap-frog is supported here");
1702
1703     // Cast to real for faster code, no loss in precision
1704     const real dt = inputRecord.delta_t;
1705
1706     const int nth = gmx_omp_nthreads_get(ModuleMultiThread::Update);
1707
1708 #pragma omp parallel for num_threads(nth) schedule(static)
1709     for (int th = 0; th < nth; th++)
1710     {
1711         try
1712         {
1713             int start_th, end_th;
1714             getThreadAtomRange(nth, th, homenr, &start_th, &end_th);
1715
1716             const rvec* x_rvec  = state.x.rvec_array();
1717             rvec*       xp_rvec = xp_.rvec_array();
1718             const rvec* v_rvec  = state.v.rvec_array();
1719             const rvec* f_rvec  = as_rvec_array(f.unpaddedConstArrayRef().data());
1720
1721             doUpdateMDDoNotUpdateVelocities(
1722                     start_th, end_th, dt, x_rvec, xp_rvec, v_rvec, f_rvec, havePartiallyFrozenAtoms, invmass, invMassPerDim, ekind);
1723         }
1724         GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR
1725     }
1726 }