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