src/gromacs/applied_forces/densityfitting.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
  37  * Declares data structure and utilities for density fitting
  38  *
  39  * \author Christian Blau <blau@kth.se>
  40  * \ingroup module_applied_forces
  41  */
  42 #include "gmxpre.h"
  43
  44 #include "densityfitting.h"
  45
  46 #include <memory>
  47 #include <numeric>
  48
  49 #include "gromacs/domdec/localatomsetmanager.h"
  50 #include "gromacs/fileio/checkpoint.h"
  51 #include "gromacs/fileio/mrcdensitymap.h"
  52 #include "gromacs/math/multidimarray.h"
  53 #include "gromacs/mdlib/broadcaststructs.h"
  54 #include "gromacs/mdtypes/commrec.h"
  55 #include "gromacs/mdtypes/imdmodule.h"
  56 #include "gromacs/utility/classhelpers.h"
  57 #include "gromacs/utility/exceptions.h"
  58 #include "gromacs/utility/keyvaluetreebuilder.h"
  59 #include "gromacs/utility/mdmodulenotification.h"
  60
  61 #include "densityfittingforceprovider.h"
  62 #include "densityfittingoptions.h"
  63 #include "densityfittingoutputprovider.h"
  64
  65 namespace gmx
  66 {
  67
  68 class IMdpOptionProvider;
  69 class DensityFittingForceProvider;
  70
  71 namespace
  72 {
  73
  74 /*! \internal
  75  * \brief Collect density fitting parameters only available during simulation setup.
  76  *
  77  * \todo Implement builder pattern that will not use unqiue_ptr to check if
  78  *       parameters have been set or not.
  79  *
  80  * To build the density fitting force provider during simulation setup,
  81  * the DensityFitting class needs access to parameters that become available
  82  * only during simulation setup.
  83  *
  84  * This class collects these parameters via MdModuleNotifications in the
  85  * simulation setup phase and provides a check if all necessary parameters have
  86  * been provided.
  87  */
  88 class DensityFittingSimulationParameterSetup
  89 {
  90 public:
  91     DensityFittingSimulationParameterSetup() = default;
  92
  93     /*! \brief Set the local atom set for the density fitting.
  94      * \param[in] localAtomSet of atoms to be fitted
  95      */
  96     void setLocalAtomSet(const LocalAtomSet& localAtomSet)
  97     {
  98         localAtomSet_ = std::make_unique<LocalAtomSet>(localAtomSet);
  99     }
 100
 101     /*! \brief Return local atom set for density fitting.
 102      * \throws InternalError if local atom set is not set
 103      * \returns local atom set for density fitting
 104      */
 105     const LocalAtomSet& localAtomSet() const
 106     {
 107         if (localAtomSet_ == nullptr)
 108         {
 109             GMX_THROW(
 110                     InternalError("Local atom set is not set for density "
 111                                   "guided simulation."));
 112         }
 113         return *localAtomSet_;
 114     }
 115
 116     /*! \brief Return transformation into density lattice.
 117      * \throws InternalError if transformation into density lattice is not set
 118      * \returns transformation into density lattice
 119      */
 120     const TranslateAndScale& transformationToDensityLattice() const
 121     {
 122         if (transformationToDensityLattice_ == nullptr)
 123         {
 124             GMX_THROW(InternalError(
 125                     "Transformation to reference density not set for density guided simulation."));
 126         }
 127         return *transformationToDensityLattice_;
 128     }
 129     /*! \brief Return reference density
 130      * \throws InternalError if reference density is not set
 131      * \returns the reference density
 132      */
 133     basic_mdspan<const float, dynamicExtents3D> referenceDensity() const
 134     {
 135         if (referenceDensity_ == nullptr)
 136         {
 137             GMX_THROW(InternalError("Reference density not set for density guided simulation."));
 138         }
 139         return referenceDensity_->asConstView();
 140     }
 141
 142     /*! \brief Reads the reference density from file.
 143      *
 144      * Reads and check file, then set and communicate the internal
 145      * parameters related to the reference density with the file data.
 146      *
 147      * \throws FileIOError if reading from file was not successful
 148      */
 149     void readReferenceDensityFromFile(const std::string& referenceDensityFileName)
 150     {
 151         MrcDensityMapOfFloatFromFileReader reader(referenceDensityFileName);
 152         referenceDensity_ = std::make_unique<MultiDimArray<std::vector<float>, dynamicExtents3D>>(
 153                 reader.densityDataCopy());
 154         transformationToDensityLattice_ =
 155                 std::make_unique<TranslateAndScale>(reader.transformationToDensityLattice());
 156     }
 157
 158     //! Normalize the reference density so that the sum over all voxels is unity
 159     void normalizeReferenceDensity()
 160     {
 161         if (referenceDensity_ == nullptr)
 162         {
 163             GMX_THROW(InternalError("Need to set reference density before normalizing it."));
 164         }
 165
 166         const real sumOfDensityData = std::accumulate(begin(referenceDensity_->asView()),
 167                                                       end(referenceDensity_->asView()), 0.);
 168         for (float& referenceDensityVoxel : referenceDensity_->asView())
 169         {
 170             referenceDensityVoxel /= sumOfDensityData;
 171         }
 172     }
 173     /*! \brief Set the periodic boundary condition via MdModuleNotifier.
 174      *
 175      * The pbc type is wrapped in PeriodicBoundaryConditionType to
 176      * allow the MdModuleNotifier to statically distinguish the callback
 177      * function type from other 'int' function callbacks.
 178      *
 179      * \param[in] pbcType enumerates the periodic boundary condition.
 180      */
 181     void setPeriodicBoundaryConditionType(const PbcType& pbcType)
 182     {
 183         pbcType_ = std::make_unique<PbcType>(pbcType);
 184     }
 185
 186     //! Get the periodic boundary conditions
 187     PbcType periodicBoundaryConditionType()
 188     {
 189         if (pbcType_ == nullptr)
 190         {
 191             GMX_THROW(InternalError(
 192                     "Periodic boundary condition enum not set for density guided simulation."));
 193         }
 194         return *pbcType_;
 195     }
 196
 197     //! Set the simulation time step
 198     void setSimulationTimeStep(double timeStep) { simulationTimeStep_ = timeStep; }
 199
 200     //! Return the simulation time step
 201     double simulationTimeStep() { return simulationTimeStep_; }
 202
 203 private:
 204     //! The reference density to fit to
 205     std::unique_ptr<MultiDimArray<std::vector<float>, dynamicExtents3D>> referenceDensity_;
 206     //! The coordinate transformation into the reference density
 207     std::unique_ptr<TranslateAndScale> transformationToDensityLattice_;
 208     //! The local atom set to act on
 209     std::unique_ptr<LocalAtomSet> localAtomSet_;
 210     //! The type of periodic boundary conditions in the simulation
 211     std::unique_ptr<PbcType> pbcType_;
 212     //! The simulation time step
 213     double simulationTimeStep_ = 1;
 214
 215     GMX_DISALLOW_COPY_AND_ASSIGN(DensityFittingSimulationParameterSetup);
 216 };
 217
 218 /*! \internal
 219  * \brief Density fitting
 220  *
 221  * Class that implements the density fitting forces and potential
 222  * \note the virial calculation is not yet implemented
 223  */
 224 class DensityFitting final : public IMDModule
 225 {
 226 public:
 227     /*! \brief Construct the density fitting module.
 228      *
 229      * \param[in] notifier allows the module to subscribe to notifications from MdModules.
 230      *
 231      * The density fitting code subscribes to these notifications:
 232      *   - setting atom group indices in the densityFittingOptions_ from an
 233      *     index group string by taking a parmeter const IndexGroupsAndNames &
 234      *   - storing its internal parameters in a tpr file by writing to a
 235      *     key-value-tree during pre-processing by a function taking a
 236      *     KeyValueTreeObjectBuilder as parameter
 237      *   - reading its internal parameters from a key-value-tree during
 238      *     simulation setup by taking a const KeyValueTreeObject & parameter
 239      *   - constructing local atom sets in the simulation parameter setup
 240      *     by taking a LocalAtomSetManager * as parameter
 241      *   - the type of periodic boundary conditions that are used
 242      *     by taking a PeriodicBoundaryConditionType as parameter
 243      *   - the writing of checkpoint data
 244      *     by taking a MdModulesWriteCheckpointData as parameter
 245      *   - the reading of checkpoint data
 246      *     by taking a MdModulesCheckpointReadingDataOnMaster as parameter
 247      *   - the broadcasting of checkpoint data
 248      *     by taking MdModulesCheckpointReadingBroadcast as parameter
 249      */
 250     explicit DensityFitting(MdModulesNotifier* notifier)
 251     {
 252         // Callbacks for several kinds of MdModuleNotification are created
 253         // and subscribed, and will be dispatched correctly at run time
 254         // based on the type of the parameter required by the lambda.
 255
 256         // Setting the atom group indices from index group string
 257         const auto setFitGroupIndicesFunction = [this](const IndexGroupsAndNames& indexGroupsAndNames) {
 258             densityFittingOptions_.setFitGroupIndices(indexGroupsAndNames);
 259         };
 260         notifier->preProcessingNotifications_.subscribe(setFitGroupIndicesFunction);
 261
 262         // Writing internal parameters during pre-processing
 263         const auto writeInternalParametersFunction = [this](KeyValueTreeObjectBuilder treeBuilder) {
 264             densityFittingOptions_.writeInternalParametersToKvt(treeBuilder);
 265         };
 266         notifier->preProcessingNotifications_.subscribe(writeInternalParametersFunction);
 267
 268         // Reading internal parameters during simulation setup
 269         const auto readInternalParametersFunction = [this](const KeyValueTreeObject& tree) {
 270             densityFittingOptions_.readInternalParametersFromKvt(tree);
 271         };
 272         notifier->simulationSetupNotifications_.subscribe(readInternalParametersFunction);
 273
 274         // Checking for consistency with all .mdp options
 275         const auto checkEnergyCaluclationFrequencyFunction =
 276                 [this](EnergyCalculationFrequencyErrors* energyCalculationFrequencyErrors) {
 277                     densityFittingOptions_.checkEnergyCaluclationFrequency(energyCalculationFrequencyErrors);
 278                 };
 279         notifier->preProcessingNotifications_.subscribe(checkEnergyCaluclationFrequencyFunction);
 280
 281         // constructing local atom sets during simulation setup
 282         const auto setLocalAtomSetFunction = [this](LocalAtomSetManager* localAtomSetManager) {
 283             this->constructLocalAtomSet(localAtomSetManager);
 284         };
 285         notifier->simulationSetupNotifications_.subscribe(setLocalAtomSetFunction);
 286
 287         // constructing local atom sets during simulation setup
 288         const auto setPeriodicBoundaryContionsFunction = [this](const PbcType& pbc) {
 289             this->densityFittingSimulationParameters_.setPeriodicBoundaryConditionType(pbc);
 290         };
 291         notifier->simulationSetupNotifications_.subscribe(setPeriodicBoundaryContionsFunction);
 292
 293         // setting the simulation time step
 294         const auto setSimulationTimeStepFunction = [this](const SimulationTimeStep& simulationTimeStep) {
 295             this->densityFittingSimulationParameters_.setSimulationTimeStep(simulationTimeStep.delta_t);
 296         };
 297         notifier->simulationSetupNotifications_.subscribe(setSimulationTimeStepFunction);
 298
 299         // adding output to energy file
 300         const auto requestEnergyOutput =
 301                 [this](MdModulesEnergyOutputToDensityFittingRequestChecker* energyOutputRequest) {
 302                     this->setEnergyOutputRequest(energyOutputRequest);
 303                 };
 304         notifier->simulationSetupNotifications_.subscribe(requestEnergyOutput);
 305
 306         // writing checkpoint data
 307         const auto checkpointDataWriting = [this](MdModulesWriteCheckpointData checkpointData) {
 308             this->writeCheckpointData(checkpointData);
 309         };
 310         notifier->checkpointingNotifications_.subscribe(checkpointDataWriting);
 311
 312         // reading checkpoint data
 313         const auto checkpointDataReading = [this](MdModulesCheckpointReadingDataOnMaster checkpointData) {
 314             this->readCheckpointDataOnMaster(checkpointData);
 315         };
 316         notifier->checkpointingNotifications_.subscribe(checkpointDataReading);
 317
 318         // broadcasting checkpoint data
 319         const auto checkpointDataBroadcast = [this](MdModulesCheckpointReadingBroadcast checkpointData) {
 320             this->broadcastCheckpointData(checkpointData);
 321         };
 322         notifier->checkpointingNotifications_.subscribe(checkpointDataBroadcast);
 323     }
 324
 325     //! From IMDModule
 326     IMdpOptionProvider* mdpOptionProvider() override { return &densityFittingOptions_; }
 327
 328     //! Add this module to the force providers if active
 329     void initForceProviders(ForceProviders* forceProviders) override
 330     {
 331         if (densityFittingOptions_.active())
 332         {
 333             const auto& parameters = densityFittingOptions_.buildParameters();
 334             densityFittingSimulationParameters_.readReferenceDensityFromFile(
 335                     densityFittingOptions_.referenceDensityFileName());
 336             if (parameters.normalizeDensities_)
 337             {
 338                 densityFittingSimulationParameters_.normalizeReferenceDensity();
 339             }
 340             forceProvider_ = std::make_unique<DensityFittingForceProvider>(
 341                     parameters, densityFittingSimulationParameters_.referenceDensity(),
 342                     densityFittingSimulationParameters_.transformationToDensityLattice(),
 343                     densityFittingSimulationParameters_.localAtomSet(),
 344                     densityFittingSimulationParameters_.periodicBoundaryConditionType(),
 345                     densityFittingSimulationParameters_.simulationTimeStep(), densityFittingState_);
 346             forceProviders->addForceProvider(forceProvider_.get());
 347         }
 348     }
 349
 350     //! This MDModule provides its own output
 351     IMDOutputProvider* outputProvider() override { return &densityFittingOutputProvider_; }
 352
 353     /*! \brief Set up the local atom sets that are used by this module.
 354      *
 355      * \note When density fitting is set up with MdModuleNotification in
 356      *       the constructor, this function is called back.
 357      *
 358      * \param[in] localAtomSetManager the manager to add local atom sets.
 359      */
 360     void constructLocalAtomSet(LocalAtomSetManager* localAtomSetManager)
 361     {
 362         if (densityFittingOptions_.active())
 363         {
 364             LocalAtomSet atomSet =
 365                     localAtomSetManager->add(densityFittingOptions_.buildParameters().indices_);
 366             densityFittingSimulationParameters_.setLocalAtomSet(atomSet);
 367         }
 368     }
 369
 370     /*! \brief Request energy output to energy file during simulation.
 371      */
 372     void setEnergyOutputRequest(MdModulesEnergyOutputToDensityFittingRequestChecker* energyOutputRequest)
 373     {
 374         energyOutputRequest->energyOutputToDensityFitting_ = densityFittingOptions_.active();
 375     }
 376
 377     /*! \brief Write internal density fitting data to checkpoint file.
 378      * \param[in] checkpointWriting enables writing to the Key-Value-Tree
 379      *                              that is used for storing the checkpoint
 380      *                              information
 381      *
 382      * \note The provided state to checkpoint has to change if checkpointing
 383      *       is moved before the force provider call in the MD-loop.
 384      */
 385     void writeCheckpointData(MdModulesWriteCheckpointData checkpointWriting)
 386     {
 387         if (densityFittingOptions_.active())
 388         {
 389             const DensityFittingForceProviderState& state = forceProvider_->stateToCheckpoint();
 390             checkpointWriting.builder_.addValue<std::int64_t>(
 391                     DensityFittingModuleInfo::name_ + "-stepsSinceLastCalculation",
 392                     state.stepsSinceLastCalculation_);
 393             checkpointWriting.builder_.addValue<real>(
 394                     DensityFittingModuleInfo::name_ + "-adaptiveForceConstantScale",
 395                     state.adaptiveForceConstantScale_);
 396             KeyValueTreeObjectBuilder exponentialMovingAverageKvtEntry =
 397                     checkpointWriting.builder_.addObject(DensityFittingModuleInfo::name_
 398                                                          + "-exponentialMovingAverageState");
 399             exponentialMovingAverageStateAsKeyValueTree(exponentialMovingAverageKvtEntry,
 400                                                         state.exponentialMovingAverageState_);
 401         }
 402     }
 403
 404     /*! \brief Read the internal parameters from the checkpoint file on master
 405      * \param[in] checkpointReading holding the checkpoint information
 406      */
 407     void readCheckpointDataOnMaster(MdModulesCheckpointReadingDataOnMaster checkpointReading)
 408     {
 409         if (densityFittingOptions_.active())
 410         {
 411             if (checkpointReading.checkpointedData_.keyExists(DensityFittingModuleInfo::name_
 412                                                               + "-stepsSinceLastCalculation"))
 413             {
 414                 densityFittingState_.stepsSinceLastCalculation_ =
 415                         checkpointReading
 416                                 .checkpointedData_[DensityFittingModuleInfo::name_
 417                                                    + "-stepsSinceLastCalculation"]
 418                                 .cast<std::int64_t>();
 419             }
 420             if (checkpointReading.checkpointedData_.keyExists(DensityFittingModuleInfo::name_
 421                                                               + "-adaptiveForceConstantScale"))
 422             {
 423                 densityFittingState_.adaptiveForceConstantScale_ =
 424                         checkpointReading
 425                                 .checkpointedData_[DensityFittingModuleInfo::name_
 426                                                    + "-adaptiveForceConstantScale"]
 427                                 .cast<real>();
 428             }
 429             if (checkpointReading.checkpointedData_.keyExists(DensityFittingModuleInfo::name_
 430                                                               + "-exponentialMovingAverageState"))
 431             {
 432                 densityFittingState_.exponentialMovingAverageState_ = exponentialMovingAverageStateFromKeyValueTree(
 433                         checkpointReading
 434                                 .checkpointedData_[DensityFittingModuleInfo::name_ + "-exponentialMovingAverageState"]
 435                                 .asObject());
 436             }
 437         }
 438     }
 439
 440     /*! \brief Broadcast the internal parameters.
 441      * \param[in] checkpointBroadcast containing the communication record to
 442      *                                broadcast the checkpoint information
 443      */
 444     void broadcastCheckpointData(MdModulesCheckpointReadingBroadcast checkpointBroadcast)
 445     {
 446         if (densityFittingOptions_.active())
 447         {
 448             if (PAR(&(checkpointBroadcast.cr_)))
 449             {
 450                 block_bc(&(checkpointBroadcast.cr_), densityFittingState_.stepsSinceLastCalculation_);
 451                 block_bc(&(checkpointBroadcast.cr_), densityFittingState_.adaptiveForceConstantScale_);
 452                 block_bc(&(checkpointBroadcast.cr_), densityFittingState_.exponentialMovingAverageState_);
 453             }
 454         }
 455     }
 456
 457 private:
 458     //! The output provider
 459     DensityFittingOutputProvider densityFittingOutputProvider_;
 460     //! The options provided for density fitting
 461     DensityFittingOptions densityFittingOptions_;
 462     //! Object that evaluates the forces
 463     std::unique_ptr<DensityFittingForceProvider> forceProvider_;
 464     /*! \brief Parameters for density fitting that become available at
 465      * simulation setup time.
 466      */
 467     DensityFittingSimulationParameterSetup densityFittingSimulationParameters_;
 468     //! The internal parameters of density fitting force provider
 469     DensityFittingForceProviderState densityFittingState_;
 470
 471     GMX_DISALLOW_COPY_AND_ASSIGN(DensityFitting);
 472 };
 473
 474 } // namespace
 475
 476 std::unique_ptr<IMDModule> DensityFittingModuleInfo::create(MdModulesNotifier* notifier)
 477 {
 478     return std::make_unique<DensityFitting>(notifier);
 479 }
 480
 481 const std::string DensityFittingModuleInfo::name_ = "density-guided-simulation";
 482
 483 } // namespace gmx