src/gromacs/applied_forces/densityfitting/densityfittingforceprovider.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  * Implements force provider for density fitting
  38  *
  39  * \author Christian Blau <blau@kth.se>
  40  * \ingroup module_applied_forces
  41  */
  42 #include "gmxpre.h"
  43
  44 #include "densityfittingforceprovider.h"
  45
  46 #include <numeric>
  47 #include <optional>
  48
  49 #include "gromacs/domdec/localatomset.h"
  50 #include "gromacs/gmxlib/network.h"
  51 #include "gromacs/math/coordinatetransformation.h"
  52 #include "gromacs/math/densityfit.h"
  53 #include "gromacs/math/densityfittingforce.h"
  54 #include "gromacs/math/gausstransform.h"
  55 #include "gromacs/mdlib/broadcaststructs.h"
  56 #include "gromacs/mdtypes/commrec.h"
  57 #include "gromacs/mdtypes/enerdata.h"
  58 #include "gromacs/mdtypes/forceoutput.h"
  59 #include "gromacs/pbcutil/pbc.h"
  60 #include "gromacs/utility/strconvert.h"
  61
  62 #include "densityfittingamplitudelookup.h"
  63 #include "densityfittingparameters.h"
  64
  65 namespace gmx
  66 {
  67
  68 namespace
  69 {
  70
  71 /*! \internal \brief Generate the spread kernel from Gaussian parameters.
  72  *
  73  * \param[in] sigma the width of the Gaussian to be spread
  74  * \param[in] nSigma the range of the Gaussian in multiples of sigma
  75  * \param[in] scaleToLattice the coordinate transformation into the spreading lattice
  76  * \returns A Gauss-transform kernel shape
  77  */
  78 GaussianSpreadKernelParameters::Shape makeSpreadKernel(real sigma, real nSigma, const ScaleCoordinates& scaleToLattice)
  79 {
  80     RVec sigmaInLatticeCoordinates{ sigma, sigma, sigma };
  81     scaleToLattice(&sigmaInLatticeCoordinates);
  82     return { DVec{ sigmaInLatticeCoordinates[XX], sigmaInLatticeCoordinates[YY],
  83                    sigmaInLatticeCoordinates[ZZ] },
  84              nSigma };
  85 }
  86
  87 } // namespace
  88
  89 /********************************************************************
  90  * DensityFittingForceProviderState
  91  */
  92
  93 const std::string DensityFittingForceProviderState::adaptiveForceConstantScaleName_ =
  94         "adaptiveForceConstantScale";
  95
  96 const std::string DensityFittingForceProviderState::exponentialMovingAverageStateName_ =
  97         "exponentialMovingAverageState";
  98
  99 const std::string DensityFittingForceProviderState::stepsSinceLastCalculationName_ =
 100         "stepsSinceLastCalculation";
 101
 102 void DensityFittingForceProviderState::writeState(KeyValueTreeObjectBuilder kvtBuilder,
 103                                                   const std::string&        identifier) const
 104 {
 105     writeKvtCheckpointValue(stepsSinceLastCalculation_, stepsSinceLastCalculationName_, identifier,
 106                             kvtBuilder);
 107     writeKvtCheckpointValue(adaptiveForceConstantScale_, adaptiveForceConstantScaleName_,
 108                             identifier, kvtBuilder);
 109
 110     KeyValueTreeObjectBuilder exponentialMovingAverageKvtEntry =
 111             kvtBuilder.addObject(identifier + "-" + exponentialMovingAverageStateName_);
 112     exponentialMovingAverageStateAsKeyValueTree(exponentialMovingAverageKvtEntry,
 113                                                 exponentialMovingAverageState_);
 114 }
 115
 116 void DensityFittingForceProviderState::readState(const KeyValueTreeObject& kvtData,
 117                                                  const std::string&        identifier)
 118 {
 119     readKvtCheckpointValue(compat::make_not_null(&stepsSinceLastCalculation_),
 120                            stepsSinceLastCalculationName_, identifier, kvtData);
 121     readKvtCheckpointValue(compat::make_not_null(&adaptiveForceConstantScale_),
 122                            adaptiveForceConstantScaleName_, identifier, kvtData);
 123
 124     if (kvtData.keyExists(identifier + "-" + exponentialMovingAverageStateName_))
 125     {
 126         exponentialMovingAverageState_ = exponentialMovingAverageStateFromKeyValueTree(
 127                 kvtData[identifier + "-" + exponentialMovingAverageStateName_].asObject());
 128     }
 129 }
 130
 131 void DensityFittingForceProviderState::broadcastState(MPI_Comm communicator, bool isParallelRun)
 132 {
 133     if (isParallelRun)
 134     {
 135         block_bc(communicator, stepsSinceLastCalculation_);
 136         block_bc(communicator, adaptiveForceConstantScale_);
 137         block_bc(communicator, exponentialMovingAverageState_);
 138     }
 139 }
 140
 141
 142 /********************************************************************
 143  * DensityFittingForceProvider::Impl
 144  */
 145
 146 class DensityFittingForceProvider::Impl
 147 {
 148 public:
 149     //! \copydoc DensityFittingForceProvider(const DensityFittingParameters &parameters)
 150     Impl(const DensityFittingParameters&             parameters,
 151          basic_mdspan<const float, dynamicExtents3D> referenceDensity,
 152          const TranslateAndScale&                    transformationToDensityLattice,
 153          const LocalAtomSet&                         localAtomSet,
 154          PbcType                                     pbcType,
 155          double                                      simulationTimeStep,
 156          const DensityFittingForceProviderState&     state);
 157     ~Impl();
 158     void calculateForces(const ForceProviderInput& forceProviderInput,
 159                          ForceProviderOutput*      forceProviderOutput);
 160     const DensityFittingForceProviderState& stateToCheckpoint();
 161
 162 private:
 163     DensityFittingForceProviderState state();
 164     const DensityFittingParameters&  parameters_;
 165     DensityFittingForceProviderState state_;
 166     DensityFittingForceProviderState stateToCheckpoint_;
 167     LocalAtomSet                     localAtomSet_;
 168
 169     GaussianSpreadKernelParameters::Shape spreadKernel_;
 170     GaussTransform3D                      gaussTransform_;
 171     DensitySimilarityMeasure              measure_;
 172     DensityFittingForce                   densityFittingForce_;
 173     //! the local atom coordinates transformed into the grid coordinate system
 174     std::vector<RVec>             transformedCoordinates_;
 175     std::vector<RVec>             forces_;
 176     DensityFittingAmplitudeLookup amplitudeLookup_;
 177     TranslateAndScale             transformationToDensityLattice_;
 178     RVec                          referenceDensityCenter_;
 179     PbcType                       pbcType_;
 180
 181     //! Optionally scale the force according to a moving average of the similarity
 182     std::optional<ExponentialMovingAverage> expAverageSimilarity_;
 183
 184     //! Optionally translate the structure
 185     std::optional<AffineTransformation> affineTransformation_;
 186 };
 187
 188 DensityFittingForceProvider::Impl::~Impl() = default;
 189
 190 DensityFittingForceProvider::Impl::Impl(const DensityFittingParameters&             parameters,
 191                                         basic_mdspan<const float, dynamicExtents3D> referenceDensity,
 192                                         const TranslateAndScale& transformationToDensityLattice,
 193                                         const LocalAtomSet&      localAtomSet,
 194                                         PbcType                  pbcType,
 195                                         double                   simulationTimeStep,
 196                                         const DensityFittingForceProviderState& state) :
 197     parameters_(parameters),
 198     state_(state),
 199     localAtomSet_(localAtomSet),
 200     spreadKernel_(makeSpreadKernel(parameters_.gaussianTransformSpreadingWidth_,
 201                                    parameters_.gaussianTransformSpreadingRangeInMultiplesOfWidth_,
 202                                    transformationToDensityLattice.scaleOperationOnly())),
 203     gaussTransform_(referenceDensity.extents(), spreadKernel_),
 204     measure_(parameters.similarityMeasureMethod_, referenceDensity),
 205     densityFittingForce_(spreadKernel_),
 206     transformedCoordinates_(localAtomSet_.numAtomsLocal()),
 207     amplitudeLookup_(parameters_.amplitudeLookupMethod_),
 208     transformationToDensityLattice_(transformationToDensityLattice),
 209     pbcType_(pbcType),
 210     expAverageSimilarity_(std::nullopt)
 211 {
 212     if (parameters_.adaptiveForceScaling_)
 213     {
 214         GMX_ASSERT(simulationTimeStep > 0,
 215                    "Simulation time step must be larger than zero for adaptive for scaling.");
 216         expAverageSimilarity_.emplace(ExponentialMovingAverage(
 217                 parameters_.adaptiveForceScalingTimeConstant_
 218                         / (simulationTimeStep * parameters_.calculationIntervalInSteps_),
 219                 state.exponentialMovingAverageState_));
 220     }
 221
 222     // set up optional coordinate translation if the translation string contains a vector
 223     const std::optional<std::array<real, 3>> translationParametersAsArray =
 224             parsedArrayFromInputString<real, 3>(parameters_.translationString_);
 225     // set up optional coordinate transformation if the transformation string contains data
 226     const std::optional<std::array<real, 9>> transformationMatrixParametersAsArray =
 227             parsedArrayFromInputString<real, 9>(parameters_.transformationMatrixString_);
 228     if (translationParametersAsArray || transformationMatrixParametersAsArray)
 229     {
 230         Matrix3x3 translationMatrix = transformationMatrixParametersAsArray.has_value()
 231                                               ? *transformationMatrixParametersAsArray
 232                                               : identityMatrix<real, 3>();
 233         RVec translationVector = translationParametersAsArray.has_value()
 234                                          ? RVec((*translationParametersAsArray)[XX],
 235                                                 (*translationParametersAsArray)[YY],
 236                                                 (*translationParametersAsArray)[ZZ])
 237                                          : RVec(0, 0, 0);
 238         affineTransformation_.emplace(translationMatrix.asConstView(), translationVector);
 239     }
 240
 241     referenceDensityCenter_ = { real(referenceDensity.extent(XX)) / 2,
 242                                 real(referenceDensity.extent(YY)) / 2,
 243                                 real(referenceDensity.extent(ZZ)) / 2 };
 244     transformationToDensityLattice_.scaleOperationOnly().inverseIgnoringZeroScale(&referenceDensityCenter_);
 245     // correct the reference density center for a shift
 246     // if the reference density does not have its origin at (0,0,0)
 247     RVec referenceDensityOriginShift(0, 0, 0);
 248     transformationToDensityLattice_(&referenceDensityOriginShift);
 249     transformationToDensityLattice_.scaleOperationOnly().inverseIgnoringZeroScale(&referenceDensityOriginShift);
 250     referenceDensityCenter_ -= referenceDensityOriginShift;
 251 }
 252
 253 void DensityFittingForceProvider::Impl::calculateForces(const ForceProviderInput& forceProviderInput,
 254                                                         ForceProviderOutput* forceProviderOutput)
 255 {
 256     // TODO change if checkpointing moves to the start of the md loop
 257     stateToCheckpoint_ = state();
 258     // do nothing but count number of steps when not in density fitting step
 259     if (state_.stepsSinceLastCalculation_ % parameters_.calculationIntervalInSteps_ != 0)
 260     {
 261         ++(state_.stepsSinceLastCalculation_);
 262         return;
 263     }
 264
 265     state_.stepsSinceLastCalculation_ = 1;
 266
 267     transformedCoordinates_.resize(localAtomSet_.numAtomsLocal());
 268     // pick and copy atom coordinates
 269     std::transform(std::cbegin(localAtomSet_.localIndex()), std::cend(localAtomSet_.localIndex()),
 270                    std::begin(transformedCoordinates_),
 271                    [&forceProviderInput](int index) { return forceProviderInput.x_[index]; });
 272
 273     // apply additional structure transformations
 274     if (affineTransformation_)
 275     {
 276         (*affineTransformation_)(transformedCoordinates_);
 277     }
 278
 279     // pick periodic image that is closest to the center of the reference density
 280     {
 281         t_pbc pbc;
 282         set_pbc(&pbc, pbcType_, forceProviderInput.box_);
 283         for (RVec& x : transformedCoordinates_)
 284         {
 285             rvec dx;
 286             pbc_dx(&pbc, x, referenceDensityCenter_, dx);
 287             x = referenceDensityCenter_ + dx;
 288         }
 289     }
 290
 291     // transform local atom coordinates to density grid coordinates
 292     transformationToDensityLattice_(transformedCoordinates_);
 293
 294     // spread atoms on grid
 295     gaussTransform_.setZero();
 296
 297     std::vector<real> amplitudes =
 298             amplitudeLookup_(forceProviderInput.mdatoms_, localAtomSet_.localIndex());
 299
 300     if (parameters_.normalizeDensities_)
 301     {
 302         real sum = std::accumulate(std::begin(amplitudes), std::end(amplitudes), 0.);
 303         if (havePPDomainDecomposition(&forceProviderInput.cr_))
 304         {
 305             gmx_sum(1, &sum, &forceProviderInput.cr_);
 306         }
 307         for (real& amplitude : amplitudes)
 308         {
 309             amplitude /= sum;
 310         }
 311     }
 312
 313     auto amplitudeIterator = amplitudes.cbegin();
 314
 315     for (const auto& r : transformedCoordinates_)
 316     {
 317         gaussTransform_.add({ r, *amplitudeIterator });
 318         ++amplitudeIterator;
 319     }
 320
 321     // communicate grid
 322     if (havePPDomainDecomposition(&forceProviderInput.cr_))
 323     {
 324         // \todo update to real once GaussTransform class returns real
 325         gmx_sumf(gaussTransform_.view().mapping().required_span_size(),
 326                  gaussTransform_.view().data(), &forceProviderInput.cr_);
 327     }
 328
 329     // calculate grid derivative
 330     const DensitySimilarityMeasure::density& densityDerivative =
 331             measure_.gradient(gaussTransform_.constView());
 332     // calculate forces
 333     forces_.resize(localAtomSet_.numAtomsLocal());
 334     std::transform(
 335             std::begin(transformedCoordinates_), std::end(transformedCoordinates_), std::begin(amplitudes),
 336             std::begin(forces_), [&densityDerivative, this](const RVec r, real amplitude) {
 337                 return densityFittingForce_.evaluateForce({ r, amplitude }, densityDerivative);
 338             });
 339
 340     transformationToDensityLattice_.scaleOperationOnly().inverseIgnoringZeroScale(forces_);
 341
 342     auto       densityForceIterator = forces_.cbegin();
 343     const real effectiveForceConstant = state_.adaptiveForceConstantScale_ * parameters_.calculationIntervalInSteps_
 344                                         * parameters_.forceConstant_;
 345     for (const auto localAtomIndex : localAtomSet_.localIndex())
 346     {
 347         forceProviderOutput->forceWithVirial_.force_[localAtomIndex] +=
 348                 effectiveForceConstant * *densityForceIterator;
 349         ++densityForceIterator;
 350     }
 351
 352     const float similarity = measure_.similarity(gaussTransform_.constView());
 353     if (MASTER(&(forceProviderInput.cr_)))
 354     {
 355         // calculate corresponding potential energy
 356         const real energy = -similarity * parameters_.forceConstant_ * state_.adaptiveForceConstantScale_;
 357         forceProviderOutput->enerd_.term[F_DENSITYFITTING] += energy;
 358     }
 359
 360     if (expAverageSimilarity_.has_value())
 361     {
 362         expAverageSimilarity_->updateWithDataPoint(similarity);
 363
 364         if (expAverageSimilarity_->increasing())
 365         {
 366             state_.adaptiveForceConstantScale_ /= 1._real + expAverageSimilarity_->inverseTimeConstant();
 367         }
 368         else
 369         {
 370             state_.adaptiveForceConstantScale_ *=
 371                     1._real + 2 * expAverageSimilarity_->inverseTimeConstant();
 372         }
 373     }
 374 }
 375
 376 DensityFittingForceProviderState DensityFittingForceProvider::Impl::state()
 377 {
 378     if (expAverageSimilarity_.has_value())
 379     {
 380         state_.exponentialMovingAverageState_ = expAverageSimilarity_->state();
 381     }
 382     return state_;
 383 }
 384
 385 const DensityFittingForceProviderState& DensityFittingForceProvider::Impl::stateToCheckpoint()
 386 {
 387     return stateToCheckpoint_;
 388 }
 389 /********************************************************************
 390  * DensityFittingForceProvider
 391  */
 392
 393 DensityFittingForceProvider::~DensityFittingForceProvider() = default;
 394
 395 DensityFittingForceProvider::DensityFittingForceProvider(const DensityFittingParameters& parameters,
 396                                                          basic_mdspan<const float, dynamicExtents3D> referenceDensity,
 397                                                          const TranslateAndScale& transformationToDensityLattice,
 398                                                          const LocalAtomSet& localAtomSet,
 399                                                          PbcType             pbcType,
 400                                                          double              simulationTimeStep,
 401                                                          const DensityFittingForceProviderState& state) :
 402     impl_(new Impl(parameters, referenceDensity, transformationToDensityLattice, localAtomSet, pbcType, simulationTimeStep, state))
 403 {
 404 }
 405
 406 void DensityFittingForceProvider::calculateForces(const ForceProviderInput& forceProviderInput,
 407                                                   ForceProviderOutput*      forceProviderOutput)
 408 {
 409     impl_->calculateForces(forceProviderInput, forceProviderOutput);
 410 }
 411
 412 void DensityFittingForceProvider::writeCheckpointData(MdModulesWriteCheckpointData checkpointWriting,
 413                                                       const std::string&           moduleName)
 414 {
 415     impl_->stateToCheckpoint().writeState(checkpointWriting.builder_, moduleName);
 416 }
 417
 418 } // namespace gmx