src/gromacs/modularsimulator/statepropagatordata.cpp

   1 /*
   2  * This file is part of the GROMACS molecular simulation package.
   3  *
   4  * Copyright (c) 2019,2020, by the GROMACS development team, led by
   5  * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
   6  * and including many others, as listed in the AUTHORS file in the
   7  * top-level source directory and at http://www.gromacs.org.
   8  *
   9  * GROMACS is free software; you can redistribute it and/or
  10  * modify it under the terms of the GNU Lesser General Public License
  11  * as published by the Free Software Foundation; either version 2.1
  12  * of the License, or (at your option) any later version.
  13  *
  14  * GROMACS is distributed in the hope that it will be useful,
  15  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  16  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  17  * Lesser General Public License for more details.
  18  *
  19  * You should have received a copy of the GNU Lesser General Public
  20  * License along with GROMACS; if not, see
  21  * http://www.gnu.org/licenses, or write to the Free Software Foundation,
  22  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
  23  *
  24  * If you want to redistribute modifications to GROMACS, please
  25  * consider that scientific software is very special. Version
  26  * control is crucial - bugs must be traceable. We will be happy to
  27  * consider code for inclusion in the official distribution, but
  28  * derived work must not be called official GROMACS. Details are found
  29  * in the README & COPYING files - if they are missing, get the
  30  * official version at http://www.gromacs.org.
  31  *
  32  * To help us fund GROMACS development, we humbly ask that you cite
  33  * the research papers on the package. Check out http://www.gromacs.org.
  34  */
  35 /*! \internal \file
  36  * \brief Defines the state for the modular simulator
  37  *
  38  * \author Pascal Merz <pascal.merz@me.com>
  39  * \ingroup module_modularsimulator
  40  */
  41
  42 #include "gmxpre.h"
  43
  44 #include "statepropagatordata.h"
  45
  46 #include "gromacs/domdec/collect.h"
  47 #include "gromacs/domdec/domdec.h"
  48 #include "gromacs/fileio/confio.h"
  49 #include "gromacs/math/vec.h"
  50 #include "gromacs/mdlib/gmx_omp_nthreads.h"
  51 #include "gromacs/mdlib/mdoutf.h"
  52 #include "gromacs/mdlib/stat.h"
  53 #include "gromacs/mdlib/update.h"
  54 #include "gromacs/mdtypes/commrec.h"
  55 #include "gromacs/mdtypes/inputrec.h"
  56 #include "gromacs/mdtypes/mdatom.h"
  57 #include "gromacs/mdtypes/state.h"
  58 #include "gromacs/pbcutil/pbc.h"
  59 #include "gromacs/topology/atoms.h"
  60 #include "gromacs/topology/topology.h"
  61
  62 #include "freeenergyperturbationelement.h"
  63
  64 namespace gmx
  65 {
  66 StatePropagatorData::StatePropagatorData(int                            numAtoms,
  67                                          FILE*                          fplog,
  68                                          const t_commrec*               cr,
  69                                          t_state*                       globalState,
  70                                          int                            nstxout,
  71                                          int                            nstvout,
  72                                          int                            nstfout,
  73                                          int                            nstxout_compressed,
  74                                          bool                           useGPU,
  75                                          FreeEnergyPerturbationElement* freeEnergyPerturbationElement,
  76                                          const TopologyHolder*          topologyHolder,
  77                                          bool              canMoleculesBeDistributedOverPBC,
  78                                          bool              writeFinalConfiguration,
  79                                          std::string       finalConfigurationFilename,
  80                                          const t_inputrec* inputrec,
  81                                          const t_mdatoms*  mdatoms) :
  82     totalNumAtoms_(numAtoms),
  83     nstxout_(nstxout),
  84     nstvout_(nstvout),
  85     nstfout_(nstfout),
  86     nstxout_compressed_(nstxout_compressed),
  87     localNAtoms_(0),
  88     ddpCount_(0),
  89     writeOutStep_(-1),
  90     vvResetVelocities_(false),
  91     freeEnergyPerturbationElement_(freeEnergyPerturbationElement),
  92     isRegularSimulationEnd_(false),
  93     lastStep_(-1),
  94     canMoleculesBeDistributedOverPBC_(canMoleculesBeDistributedOverPBC),
  95     systemHasPeriodicMolecules_(inputrec->bPeriodicMols),
  96     pbcType_(inputrec->pbcType),
  97     topologyHolder_(topologyHolder),
  98     lastPlannedStep_(inputrec->nsteps + inputrec->init_step),
  99     writeFinalConfiguration_(writeFinalConfiguration),
 100     finalConfigurationFilename_(std::move(finalConfigurationFilename)),
 101     fplog_(fplog),
 102     cr_(cr),
 103     globalState_(globalState)
 104 {
 105     // Initialize these here, as box_{{0}} in the initialization list
 106     // is confusing uncrustify and doxygen
 107     clear_mat(box_);
 108     clear_mat(previousBox_);
 109
 110     bool stateHasVelocities;
 111     // Local state only becomes valid now.
 112     if (DOMAINDECOMP(cr))
 113     {
 114         auto localState = std::make_unique<t_state>();
 115         dd_init_local_state(cr->dd, globalState, localState.get());
 116         stateHasVelocities = ((static_cast<unsigned int>(localState->flags) & (1U << estV)) != 0u);
 117         setLocalState(std::move(localState));
 118     }
 119     else
 120     {
 121         state_change_natoms(globalState, globalState->natoms);
 122         f_.resizeWithPadding(globalState->natoms);
 123         localNAtoms_ = globalState->natoms;
 124         x_           = globalState->x;
 125         v_           = globalState->v;
 126         copy_mat(globalState->box, box_);
 127         stateHasVelocities = ((static_cast<unsigned int>(globalState->flags) & (1U << estV)) != 0u);
 128         previousX_.resizeWithPadding(localNAtoms_);
 129         ddpCount_ = globalState->ddp_count;
 130         copyPosition();
 131     }
 132     if (useGPU)
 133     {
 134         changePinningPolicy(&x_, gmx::PinningPolicy::PinnedIfSupported);
 135     }
 136
 137     if (!inputrec->bContinuation)
 138     {
 139         if (stateHasVelocities)
 140         {
 141             auto v = velocitiesView().paddedArrayRef();
 142             // Set the velocities of vsites, shells and frozen atoms to zero
 143             for (int i = 0; i < mdatoms->homenr; i++)
 144             {
 145                 if (mdatoms->ptype[i] == eptVSite || mdatoms->ptype[i] == eptShell)
 146                 {
 147                     clear_rvec(v[i]);
 148                 }
 149                 else if (mdatoms->cFREEZE)
 150                 {
 151                     for (int m = 0; m < DIM; m++)
 152                     {
 153                         if (inputrec->opts.nFreeze[mdatoms->cFREEZE[i]][m])
 154                         {
 155                             v[i][m] = 0;
 156                         }
 157                     }
 158                 }
 159             }
 160         }
 161         if (inputrec->eI == eiVV)
 162         {
 163             vvResetVelocities_ = true;
 164         }
 165     }
 166 }
 167
 168 ArrayRefWithPadding<RVec> StatePropagatorData::positionsView()
 169 {
 170     return x_.arrayRefWithPadding();
 171 }
 172
 173 ArrayRefWithPadding<const RVec> StatePropagatorData::constPositionsView() const
 174 {
 175     return x_.constArrayRefWithPadding();
 176 }
 177
 178 ArrayRefWithPadding<RVec> StatePropagatorData::previousPositionsView()
 179 {
 180     return previousX_.arrayRefWithPadding();
 181 }
 182
 183 ArrayRefWithPadding<const RVec> StatePropagatorData::constPreviousPositionsView() const
 184 {
 185     return previousX_.constArrayRefWithPadding();
 186 }
 187
 188 ArrayRefWithPadding<RVec> StatePropagatorData::velocitiesView()
 189 {
 190     return v_.arrayRefWithPadding();
 191 }
 192
 193 ArrayRefWithPadding<const RVec> StatePropagatorData::constVelocitiesView() const
 194 {
 195     return v_.constArrayRefWithPadding();
 196 }
 197
 198 ArrayRefWithPadding<RVec> StatePropagatorData::forcesView()
 199 {
 200     return f_.arrayRefWithPadding();
 201 }
 202
 203 ArrayRefWithPadding<const RVec> StatePropagatorData::constForcesView() const
 204 {
 205     return f_.constArrayRefWithPadding();
 206 }
 207
 208 rvec* StatePropagatorData::box()
 209 {
 210     return box_;
 211 }
 212
 213 const rvec* StatePropagatorData::constBox()
 214 {
 215     return box_;
 216 }
 217
 218 rvec* StatePropagatorData::previousBox()
 219 {
 220     return previousBox_;
 221 }
 222
 223 const rvec* StatePropagatorData::constPreviousBox()
 224 {
 225     return previousBox_;
 226 }
 227
 228 int StatePropagatorData::localNumAtoms()
 229 {
 230     return localNAtoms_;
 231 }
 232
 233 std::unique_ptr<t_state> StatePropagatorData::localState()
 234 {
 235     auto state   = std::make_unique<t_state>();
 236     state->flags = (1U << estX) | (1U << estV) | (1U << estBOX);
 237     state_change_natoms(state.get(), localNAtoms_);
 238     state->x = x_;
 239     state->v = v_;
 240     copy_mat(box_, state->box);
 241     state->ddp_count = ddpCount_;
 242     return state;
 243 }
 244
 245 void StatePropagatorData::setLocalState(std::unique_ptr<t_state> state)
 246 {
 247     localNAtoms_ = state->natoms;
 248     x_.resizeWithPadding(localNAtoms_);
 249     previousX_.resizeWithPadding(localNAtoms_);
 250     v_.resizeWithPadding(localNAtoms_);
 251     x_ = state->x;
 252     v_ = state->v;
 253     copy_mat(state->box, box_);
 254     copyPosition();
 255     ddpCount_ = state->ddp_count;
 256
 257     if (vvResetVelocities_)
 258     {
 259         /* DomDec runs twice early in the simulation, once at setup time, and once before the first
 260          * step. Every time DD runs, it sets a new local state here. We are saving a backup during
 261          * setup time (ok for non-DD cases), so we need to update our backup to the DD state before
 262          * the first step here to avoid resetting to an earlier DD state. This is done before any
 263          * propagation that needs to be reset, so it's not very safe but correct for now.
 264          * TODO: Get rid of this once input is assumed to be at half steps
 265          */
 266         velocityBackup_ = v_;
 267     }
 268 }
 269
 270 t_state* StatePropagatorData::globalState()
 271 {
 272     return globalState_;
 273 }
 274
 275 PaddedHostVector<RVec>* StatePropagatorData::forcePointer()
 276 {
 277     return &f_;
 278 }
 279
 280 void StatePropagatorData::copyPosition()
 281 {
 282     int nth = gmx_omp_nthreads_get(emntUpdate);
 283
 284 #pragma omp parallel for num_threads(nth) schedule(static) default(none) shared(nth)
 285     for (int th = 0; th < nth; th++)
 286     {
 287         int start_th, end_th;
 288         getThreadAtomRange(nth, th, localNAtoms_, &start_th, &end_th);
 289         copyPosition(start_th, end_th);
 290     }
 291
 292     /* Box is changed in update() when we do pressure coupling,
 293      * but we should still use the old box for energy corrections and when
 294      * writing it to the energy file, so it matches the trajectory files for
 295      * the same timestep above. Make a copy in a separate array.
 296      */
 297     copy_mat(box_, previousBox_);
 298 }
 299
 300 void StatePropagatorData::copyPosition(int start, int end)
 301 {
 302     for (int i = start; i < end; ++i)
 303     {
 304         previousX_[i] = x_[i];
 305     }
 306 }
 307
 308 void StatePropagatorData::scheduleTask(Step step, Time gmx_unused time, const RegisterRunFunctionPtr& registerRunFunction)
 309 {
 310     if (vvResetVelocities_)
 311     {
 312         vvResetVelocities_ = false;
 313         (*registerRunFunction)(std::make_unique<SimulatorRunFunction>([this]() { resetVelocities(); }));
 314     }
 315     // copy x -> previousX
 316     (*registerRunFunction)(std::make_unique<SimulatorRunFunction>([this]() { copyPosition(); }));
 317     // if it's a write out step, keep a copy for writeout
 318     if (step == writeOutStep_ || (step == lastStep_ && writeFinalConfiguration_))
 319     {
 320         (*registerRunFunction)(std::make_unique<SimulatorRunFunction>([this]() { saveState(); }));
 321     }
 322 }
 323
 324 void StatePropagatorData::saveState()
 325 {
 326     GMX_ASSERT(!localStateBackup_, "Save state called again before previous state was written.");
 327     localStateBackup_ = localState();
 328     if (freeEnergyPerturbationElement_)
 329     {
 330         localStateBackup_->fep_state = freeEnergyPerturbationElement_->currentFEPState();
 331         for (unsigned long i = 0; i < localStateBackup_->lambda.size(); ++i)
 332         {
 333             localStateBackup_->lambda[i] = freeEnergyPerturbationElement_->constLambdaView()[i];
 334         }
 335         localStateBackup_->flags |= (1U << estLAMBDA) | (1U << estFEPSTATE);
 336     }
 337 }
 338
 339 SignallerCallbackPtr StatePropagatorData::registerTrajectorySignallerCallback(TrajectoryEvent event)
 340 {
 341     if (event == TrajectoryEvent::StateWritingStep)
 342     {
 343         return std::make_unique<SignallerCallback>(
 344                 [this](Step step, Time /*unused*/) { this->writeOutStep_ = step; });
 345     }
 346     return nullptr;
 347 }
 348
 349 ITrajectoryWriterCallbackPtr StatePropagatorData::registerTrajectoryWriterCallback(TrajectoryEvent event)
 350 {
 351     if (event == TrajectoryEvent::StateWritingStep)
 352     {
 353         return std::make_unique<ITrajectoryWriterCallback>(
 354                 [this](gmx_mdoutf* outf, Step step, Time time, bool writeTrajectory, bool gmx_unused writeLog) {
 355                     if (writeTrajectory)
 356                     {
 357                         write(outf, step, time);
 358                     }
 359                 });
 360     }
 361     return nullptr;
 362 }
 363
 364 void StatePropagatorData::write(gmx_mdoutf_t outf, Step currentStep, Time currentTime)
 365 {
 366     wallcycle_start(mdoutf_get_wcycle(outf), ewcTRAJ);
 367     unsigned int mdof_flags = 0;
 368     if (do_per_step(currentStep, nstxout_))
 369     {
 370         mdof_flags |= MDOF_X;
 371     }
 372     if (do_per_step(currentStep, nstvout_))
 373     {
 374         mdof_flags |= MDOF_V;
 375     }
 376     if (do_per_step(currentStep, nstfout_))
 377     {
 378         mdof_flags |= MDOF_F;
 379     }
 380     if (do_per_step(currentStep, nstxout_compressed_))
 381     {
 382         mdof_flags |= MDOF_X_COMPRESSED;
 383     }
 384     if (do_per_step(currentStep, mdoutf_get_tng_box_output_interval(outf)))
 385     {
 386         mdof_flags |= MDOF_BOX;
 387     }
 388     if (do_per_step(currentStep, mdoutf_get_tng_lambda_output_interval(outf)))
 389     {
 390         mdof_flags |= MDOF_LAMBDA;
 391     }
 392     if (do_per_step(currentStep, mdoutf_get_tng_compressed_box_output_interval(outf)))
 393     {
 394         mdof_flags |= MDOF_BOX_COMPRESSED;
 395     }
 396     if (do_per_step(currentStep, mdoutf_get_tng_compressed_lambda_output_interval(outf)))
 397     {
 398         mdof_flags |= MDOF_LAMBDA_COMPRESSED;
 399     }
 400
 401     if (mdof_flags == 0)
 402     {
 403         wallcycle_stop(mdoutf_get_wcycle(outf), ewcTRAJ);
 404         return;
 405     }
 406     GMX_ASSERT(localStateBackup_, "Trajectory writing called, but no state saved.");
 407
 408     // TODO: This is only used for CPT - needs to be filled when we turn CPT back on
 409     ObservablesHistory* observablesHistory = nullptr;
 410
 411     mdoutf_write_to_trajectory_files(fplog_, cr_, outf, static_cast<int>(mdof_flags),
 412                                      totalNumAtoms_, currentStep, currentTime,
 413                                      localStateBackup_.get(), globalState_, observablesHistory, f_);
 414
 415     if (currentStep != lastStep_ || !isRegularSimulationEnd_)
 416     {
 417         localStateBackup_.reset();
 418     }
 419     wallcycle_stop(mdoutf_get_wcycle(outf), ewcTRAJ);
 420 }
 421
 422 void StatePropagatorData::elementSetup()
 423 {
 424     if (vvResetVelocities_)
 425     {
 426         // MD-VV does the first velocity half-step only to calculate the constraint virial,
 427         // then resets the velocities since the input is assumed to be positions and velocities
 428         // at full time step. TODO: Change this to have input at half time steps.
 429         velocityBackup_ = v_;
 430     }
 431 }
 432
 433 void StatePropagatorData::resetVelocities()
 434 {
 435     v_ = velocityBackup_;
 436 }
 437
 438 void StatePropagatorData::writeCheckpoint(t_state* localState, t_state gmx_unused* globalState)
 439 {
 440     state_change_natoms(localState, localNAtoms_);
 441     localState->x = x_;
 442     localState->v = v_;
 443     copy_mat(box_, localState->box);
 444     localState->ddp_count = ddpCount_;
 445     localState->flags |= (1U << estX) | (1U << estV) | (1U << estBOX);
 446 }
 447
 448 void StatePropagatorData::trajectoryWriterTeardown(gmx_mdoutf* gmx_unused outf)
 449 {
 450     // Note that part of this code is duplicated in do_md_trajectory_writing.
 451     // This duplication is needed while both legacy and modular code paths are in use.
 452     // TODO: Remove duplication asap, make sure to keep in sync in the meantime.
 453     if (!writeFinalConfiguration_ || !isRegularSimulationEnd_)
 454     {
 455         return;
 456     }
 457
 458     GMX_ASSERT(localStateBackup_, "Final trajectory writing called, but no state saved.");
 459
 460     wallcycle_start(mdoutf_get_wcycle(outf), ewcTRAJ);
 461     if (DOMAINDECOMP(cr_))
 462     {
 463         auto globalXRef = MASTER(cr_) ? globalState_->x : gmx::ArrayRef<gmx::RVec>();
 464         dd_collect_vec(cr_->dd, localStateBackup_.get(), localStateBackup_->x, globalXRef);
 465         auto globalVRef = MASTER(cr_) ? globalState_->v : gmx::ArrayRef<gmx::RVec>();
 466         dd_collect_vec(cr_->dd, localStateBackup_.get(), localStateBackup_->v, globalVRef);
 467     }
 468     else
 469     {
 470         // We have the whole state locally: copy the local state pointer
 471         globalState_ = localStateBackup_.get();
 472     }
 473
 474     if (MASTER(cr_))
 475     {
 476         fprintf(stderr, "\nWriting final coordinates.\n");
 477         if (canMoleculesBeDistributedOverPBC_ && !systemHasPeriodicMolecules_)
 478         {
 479             // Make molecules whole only for confout writing
 480             do_pbc_mtop(pbcType_, localStateBackup_->box, &topologyHolder_->globalTopology(),
 481                         globalState_->x.rvec_array());
 482         }
 483         write_sto_conf_mtop(finalConfigurationFilename_.c_str(), *topologyHolder_->globalTopology().name,
 484                             &topologyHolder_->globalTopology(), globalState_->x.rvec_array(),
 485                             globalState_->v.rvec_array(), pbcType_, localStateBackup_->box);
 486     }
 487     wallcycle_stop(mdoutf_get_wcycle(outf), ewcTRAJ);
 488 }
 489
 490 SignallerCallbackPtr StatePropagatorData::registerLastStepCallback()
 491 {
 492     return std::make_unique<SignallerCallback>([this](Step step, Time) {
 493         lastStep_               = step;
 494         isRegularSimulationEnd_ = (step == lastPlannedStep_);
 495     });
 496 }
 497
 498 } // namespace gmx