src/gromacs/applied_forces/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
  48 #include "gromacs/compat/optional.h"
  49 #include "gromacs/gmxlib/network.h"
  50 #include "gromacs/math/densityfit.h"
  51 #include "gromacs/math/densityfittingforce.h"
  52 #include "gromacs/math/exponentialmovingaverage.h"
  53 #include "gromacs/math/gausstransform.h"
  54 #include "gromacs/mdtypes/commrec.h"
  55 #include "gromacs/mdtypes/enerdata.h"
  56 #include "gromacs/mdtypes/forceoutput.h"
  57 #include "gromacs/mdtypes/iforceprovider.h"
  58 #include "gromacs/pbcutil/pbc.h"
  59
  60 #include "densityfittingamplitudelookup.h"
  61 #include "densityfittingparameters.h"
  62
  63 namespace gmx
  64 {
  65
  66 namespace
  67 {
  68
  69 /*! \internal \brief Generate the spread kernel from Gaussian parameters.
  70  *
  71  * \param[in] sigma the width of the Gaussian to be spread
  72  * \param[in] nSigma the range of the Gaussian in multiples of sigma
  73  * \param[in] scaleToLattice the coordinate transformation into the spreading lattice
  74  * \returns A Gauss-transform kernel shape
  75  */
  76 GaussianSpreadKernelParameters::Shape makeSpreadKernel(real sigma, real nSigma, const ScaleCoordinates& scaleToLattice)
  77 {
  78     RVec sigmaInLatticeCoordinates{ sigma, sigma, sigma };
  79     scaleToLattice({ &sigmaInLatticeCoordinates, &sigmaInLatticeCoordinates + 1 });
  80     return { DVec{ sigmaInLatticeCoordinates[XX], sigmaInLatticeCoordinates[YY],
  81                    sigmaInLatticeCoordinates[ZZ] },
  82              nSigma };
  83 }
  84
  85 } // namespace
  86
  87 /********************************************************************
  88  * DensityFittingForceProvider::Impl
  89  */
  90
  91 class DensityFittingForceProvider::Impl
  92 {
  93 public:
  94     //! \copydoc DensityFittingForceProvider(const DensityFittingParameters &parameters)
  95     Impl(const DensityFittingParameters&             parameters,
  96          basic_mdspan<const float, dynamicExtents3D> referenceDensity,
  97          const TranslateAndScale&                    transformationToDensityLattice,
  98          const LocalAtomSet&                         localAtomSet,
  99          int                                         pbcType,
 100          double                                      simulationTimeStep,
 101          const DensityFittingForceProviderState&     state);
 102     ~Impl();
 103     void calculateForces(const ForceProviderInput& forceProviderInput,
 104                          ForceProviderOutput*      forceProviderOutput);
 105     const DensityFittingForceProviderState& stateToCheckpoint();
 106
 107 private:
 108     DensityFittingForceProviderState state();
 109     const DensityFittingParameters&  parameters_;
 110     DensityFittingForceProviderState state_;
 111     DensityFittingForceProviderState stateToCheckpoint_;
 112     LocalAtomSet                     localAtomSet_;
 113
 114     GaussianSpreadKernelParameters::Shape spreadKernel_;
 115     GaussTransform3D                      gaussTransform_;
 116     DensitySimilarityMeasure              measure_;
 117     DensityFittingForce                   densityFittingForce_;
 118     //! the local atom coordinates transformed into the grid coordinate system
 119     std::vector<RVec>             transformedCoordinates_;
 120     std::vector<RVec>             forces_;
 121     DensityFittingAmplitudeLookup amplitudeLookup_;
 122     TranslateAndScale             transformationToDensityLattice_;
 123     RVec                          referenceDensityCenter_;
 124     int                           pbcType_;
 125
 126     //! Optionally scale the force according to a moving average of the similarity
 127     compat::optional<ExponentialMovingAverage> expAverageSimilarity_;
 128 };
 129
 130 DensityFittingForceProvider::Impl::~Impl() = default;
 131
 132 DensityFittingForceProvider::Impl::Impl(const DensityFittingParameters&             parameters,
 133                                         basic_mdspan<const float, dynamicExtents3D> referenceDensity,
 134                                         const TranslateAndScale& transformationToDensityLattice,
 135                                         const LocalAtomSet&      localAtomSet,
 136                                         int                      pbcType,
 137                                         double                   simulationTimeStep,
 138                                         const DensityFittingForceProviderState& state) :
 139     parameters_(parameters),
 140     state_(state),
 141     localAtomSet_(localAtomSet),
 142     spreadKernel_(makeSpreadKernel(parameters_.gaussianTransformSpreadingWidth_,
 143                                    parameters_.gaussianTransformSpreadingRangeInMultiplesOfWidth_,
 144                                    transformationToDensityLattice.scaleOperationOnly())),
 145     gaussTransform_(referenceDensity.extents(), spreadKernel_),
 146     measure_(parameters.similarityMeasureMethod_, referenceDensity),
 147     densityFittingForce_(spreadKernel_),
 148     transformedCoordinates_(localAtomSet_.numAtomsLocal()),
 149     amplitudeLookup_(parameters_.amplitudeLookupMethod_),
 150     transformationToDensityLattice_(transformationToDensityLattice),
 151     pbcType_(pbcType),
 152     expAverageSimilarity_(compat::nullopt)
 153 {
 154     if (parameters_.adaptiveForceScaling_)
 155     {
 156         GMX_ASSERT(simulationTimeStep > 0,
 157                    "Simulation time step must be larger than zero for adaptive for scaling.");
 158         expAverageSimilarity_.emplace(ExponentialMovingAverage(
 159                 parameters_.adaptiveForceScalingTimeConstant_
 160                         / (simulationTimeStep * parameters_.calculationIntervalInSteps_),
 161                 state.exponentialMovingAverageState_));
 162     }
 163     referenceDensityCenter_ = { real(referenceDensity.extent(XX)) / 2,
 164                                 real(referenceDensity.extent(YY)) / 2,
 165                                 real(referenceDensity.extent(ZZ)) / 2 };
 166     transformationToDensityLattice_.scaleOperationOnly().inverseIgnoringZeroScale(
 167             { &referenceDensityCenter_, &referenceDensityCenter_ + 1 });
 168     // correct the reference density center for a shift
 169     // if the reference density does not have its origin at (0,0,0)
 170     RVec referenceDensityOriginShift(0, 0, 0);
 171     transformationToDensityLattice_({ &referenceDensityOriginShift, &referenceDensityOriginShift + 1 });
 172     transformationToDensityLattice_.scaleOperationOnly().inverseIgnoringZeroScale(
 173             { &referenceDensityOriginShift, &referenceDensityOriginShift + 1 });
 174     referenceDensityCenter_ -= referenceDensityOriginShift;
 175 }
 176
 177 void DensityFittingForceProvider::Impl::calculateForces(const ForceProviderInput& forceProviderInput,
 178                                                         ForceProviderOutput* forceProviderOutput)
 179 {
 180     // TODO change if checkpointing moves to the start of the md loop
 181     stateToCheckpoint_ = state();
 182     // do nothing but count number of steps when not in density fitting step
 183     if (state_.stepsSinceLastCalculation_ % parameters_.calculationIntervalInSteps_ != 0)
 184     {
 185         ++(state_.stepsSinceLastCalculation_);
 186         return;
 187     }
 188
 189     state_.stepsSinceLastCalculation_ = 1;
 190
 191     transformedCoordinates_.resize(localAtomSet_.numAtomsLocal());
 192     // pick and copy atom coordinates
 193     std::transform(std::cbegin(localAtomSet_.localIndex()), std::cend(localAtomSet_.localIndex()),
 194                    std::begin(transformedCoordinates_),
 195                    [&forceProviderInput](int index) { return forceProviderInput.x_[index]; });
 196
 197     // pick periodic image that is closest to the center of the reference density
 198     {
 199         t_pbc pbc;
 200         set_pbc(&pbc, pbcType_, forceProviderInput.box_);
 201         for (RVec& x : transformedCoordinates_)
 202         {
 203             rvec dx;
 204             pbc_dx(&pbc, x, referenceDensityCenter_, dx);
 205             x = referenceDensityCenter_ + dx;
 206         }
 207     }
 208
 209     // transform local atom coordinates to density grid coordinates
 210     transformationToDensityLattice_(transformedCoordinates_);
 211
 212     // spread atoms on grid
 213     gaussTransform_.setZero();
 214
 215     std::vector<real> amplitudes =
 216             amplitudeLookup_(forceProviderInput.mdatoms_, localAtomSet_.localIndex());
 217
 218     if (parameters_.normalizeDensities_)
 219     {
 220         real sum = std::accumulate(std::begin(amplitudes), std::end(amplitudes), 0.);
 221         if (havePPDomainDecomposition(&forceProviderInput.cr_))
 222         {
 223             gmx_sum(1, &sum, &forceProviderInput.cr_);
 224         }
 225         for (real& amplitude : amplitudes)
 226         {
 227             amplitude /= sum;
 228         }
 229     }
 230
 231     auto amplitudeIterator = amplitudes.cbegin();
 232
 233     for (const auto& r : transformedCoordinates_)
 234     {
 235         gaussTransform_.add({ r, *amplitudeIterator });
 236         ++amplitudeIterator;
 237     }
 238
 239     // communicate grid
 240     if (havePPDomainDecomposition(&forceProviderInput.cr_))
 241     {
 242         // \todo update to real once GaussTransform class returns real
 243         gmx_sumf(gaussTransform_.view().mapping().required_span_size(),
 244                  gaussTransform_.view().data(), &forceProviderInput.cr_);
 245     }
 246
 247     // calculate grid derivative
 248     const DensitySimilarityMeasure::density& densityDerivative =
 249             measure_.gradient(gaussTransform_.constView());
 250     // calculate forces
 251     forces_.resize(localAtomSet_.numAtomsLocal());
 252     std::transform(
 253             std::begin(transformedCoordinates_), std::end(transformedCoordinates_), std::begin(amplitudes),
 254             std::begin(forces_), [&densityDerivative, this](const RVec r, real amplitude) {
 255                 return densityFittingForce_.evaluateForce({ r, amplitude }, densityDerivative);
 256             });
 257
 258     transformationToDensityLattice_.scaleOperationOnly().inverseIgnoringZeroScale(forces_);
 259
 260     auto       densityForceIterator = forces_.cbegin();
 261     const real effectiveForceConstant = state_.adaptiveForceConstantScale_ * parameters_.calculationIntervalInSteps_
 262                                         * parameters_.forceConstant_;
 263     for (const auto localAtomIndex : localAtomSet_.localIndex())
 264     {
 265         forceProviderOutput->forceWithVirial_.force_[localAtomIndex] +=
 266                 effectiveForceConstant * *densityForceIterator;
 267         ++densityForceIterator;
 268     }
 269
 270     const float similarity = measure_.similarity(gaussTransform_.constView());
 271     if (MASTER(&(forceProviderInput.cr_)))
 272     {
 273         // calculate corresponding potential energy
 274         const real energy = -similarity * parameters_.forceConstant_ * state_.adaptiveForceConstantScale_;
 275         forceProviderOutput->enerd_.term[F_DENSITYFITTING] += energy;
 276     }
 277
 278     if (expAverageSimilarity_.has_value())
 279     {
 280         expAverageSimilarity_->updateWithDataPoint(similarity);
 281
 282         if (expAverageSimilarity_->increasing())
 283         {
 284             state_.adaptiveForceConstantScale_ /= 1._real + expAverageSimilarity_->inverseTimeConstant();
 285         }
 286         else
 287         {
 288             state_.adaptiveForceConstantScale_ *=
 289                     1._real + 2 * expAverageSimilarity_->inverseTimeConstant();
 290         }
 291     }
 292 }
 293
 294 DensityFittingForceProviderState DensityFittingForceProvider::Impl::state()
 295 {
 296     if (expAverageSimilarity_.has_value())
 297     {
 298         state_.exponentialMovingAverageState_ = expAverageSimilarity_->state();
 299     }
 300     return state_;
 301 }
 302
 303 const DensityFittingForceProviderState& DensityFittingForceProvider::Impl::stateToCheckpoint()
 304 {
 305     return stateToCheckpoint_;
 306 }
 307 /********************************************************************
 308  * DensityFittingForceProvider
 309  */
 310
 311 DensityFittingForceProvider::~DensityFittingForceProvider() = default;
 312
 313 DensityFittingForceProvider::DensityFittingForceProvider(const DensityFittingParameters& parameters,
 314                                                          basic_mdspan<const float, dynamicExtents3D> referenceDensity,
 315                                                          const TranslateAndScale& transformationToDensityLattice,
 316                                                          const LocalAtomSet& localAtomSet,
 317                                                          int                 pbcType,
 318                                                          double              simulationTimeStep,
 319                                                          const DensityFittingForceProviderState& state) :
 320     impl_(new Impl(parameters, referenceDensity, transformationToDensityLattice, localAtomSet, pbcType, simulationTimeStep, state))
 321 {
 322 }
 323
 324 void DensityFittingForceProvider::calculateForces(const ForceProviderInput& forceProviderInput,
 325                                                   ForceProviderOutput*      forceProviderOutput)
 326 {
 327     impl_->calculateForces(forceProviderInput, forceProviderOutput);
 328 }
 329
 330 const DensityFittingForceProviderState& DensityFittingForceProvider::stateToCheckpoint()
 331 {
 332     return impl_->stateToCheckpoint();
 333 }
 334
 335 } // namespace gmx