src/gromacs/applied_forces/densityfittingforceprovider.cpp

   1 /*
   2  * This file is part of the GROMACS molecular simulation package.
   3  *
   4  * Copyright (c) 2019, 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
 106     DensityFittingForceProviderState state();
 107
 108 private:
 109     const DensityFittingParameters&  parameters_;
 110     DensityFittingForceProviderState state_;
 111     LocalAtomSet                     localAtomSet_;
 112
 113     GaussianSpreadKernelParameters::Shape spreadKernel_;
 114     GaussTransform3D                      gaussTransform_;
 115     DensitySimilarityMeasure              measure_;
 116     DensityFittingForce                   densityFittingForce_;
 117     //! the local atom coordinates transformed into the grid coordinate system
 118     std::vector<RVec>             transformedCoordinates_;
 119     std::vector<RVec>             forces_;
 120     DensityFittingAmplitudeLookup amplitudeLookup_;
 121     TranslateAndScale             transformationToDensityLattice_;
 122     RVec                          referenceDensityCenter_;
 123     int                           pbcType_;
 124
 125     //! Optionally scale the force according to a moving average of the similarity
 126     compat::optional<ExponentialMovingAverage> expAverageSimilarity_;
 127 };
 128
 129 DensityFittingForceProvider::Impl::~Impl() = default;
 130
 131 DensityFittingForceProvider::Impl::Impl(const DensityFittingParameters&             parameters,
 132                                         basic_mdspan<const float, dynamicExtents3D> referenceDensity,
 133                                         const TranslateAndScale& transformationToDensityLattice,
 134                                         const LocalAtomSet&      localAtomSet,
 135                                         int                      pbcType,
 136                                         double                   simulationTimeStep,
 137                                         const DensityFittingForceProviderState& state) :
 138     parameters_(parameters),
 139     state_(state),
 140     localAtomSet_(localAtomSet),
 141     spreadKernel_(makeSpreadKernel(parameters_.gaussianTransformSpreadingWidth_,
 142                                    parameters_.gaussianTransformSpreadingRangeInMultiplesOfWidth_,
 143                                    transformationToDensityLattice.scaleOperationOnly())),
 144     gaussTransform_(referenceDensity.extents(), spreadKernel_),
 145     measure_(parameters.similarityMeasureMethod_, referenceDensity),
 146     densityFittingForce_(spreadKernel_),
 147     transformedCoordinates_(localAtomSet_.numAtomsLocal()),
 148     amplitudeLookup_(parameters_.amplitudeLookupMethod_),
 149     transformationToDensityLattice_(transformationToDensityLattice),
 150     pbcType_(pbcType),
 151     expAverageSimilarity_(compat::nullopt)
 152 {
 153     if (parameters_.adaptiveForceScaling_)
 154     {
 155         GMX_ASSERT(simulationTimeStep > 0,
 156                    "Simulation time step must be larger than zero for adaptive for scaling.");
 157         expAverageSimilarity_.emplace(ExponentialMovingAverage(
 158                 parameters_.adaptiveForceScalingTimeConstant_
 159                         / (simulationTimeStep * parameters_.calculationIntervalInSteps_),
 160                 state.exponentialMovingAverageState_));
 161     }
 162     referenceDensityCenter_ = { real(referenceDensity.extent(XX)) / 2,
 163                                 real(referenceDensity.extent(YY)) / 2,
 164                                 real(referenceDensity.extent(ZZ)) / 2 };
 165     transformationToDensityLattice_.scaleOperationOnly().inverseIgnoringZeroScale(
 166             { &referenceDensityCenter_, &referenceDensityCenter_ + 1 });
 167     // correct the reference density center for a shift
 168     // if the reference density does not have its origin at (0,0,0)
 169     RVec referenceDensityOriginShift(0, 0, 0);
 170     transformationToDensityLattice_({ &referenceDensityOriginShift, &referenceDensityOriginShift + 1 });
 171     transformationToDensityLattice_.scaleOperationOnly().inverseIgnoringZeroScale(
 172             { &referenceDensityOriginShift, &referenceDensityOriginShift + 1 });
 173     referenceDensityCenter_ -= referenceDensityOriginShift;
 174 }
 175
 176 void DensityFittingForceProvider::Impl::calculateForces(const ForceProviderInput& forceProviderInput,
 177                                                         ForceProviderOutput* forceProviderOutput)
 178 {
 179     // do nothing but count number of steps when not in density fitting step
 180     if (state_.stepsSinceLastCalculation_ % parameters_.calculationIntervalInSteps_ != 0)
 181     {
 182         ++(state_.stepsSinceLastCalculation_);
 183         return;
 184     }
 185
 186     state_.stepsSinceLastCalculation_ = 1;
 187
 188     // do nothing if there are no density fitting atoms on this node
 189     if (localAtomSet_.numAtomsLocal() == 0)
 190     {
 191         return;
 192     }
 193     transformedCoordinates_.resize(localAtomSet_.numAtomsLocal());
 194     // pick and copy atom coordinates
 195     std::transform(std::cbegin(localAtomSet_.localIndex()), std::cend(localAtomSet_.localIndex()),
 196                    std::begin(transformedCoordinates_),
 197                    [&forceProviderInput](int index) { return forceProviderInput.x_[index]; });
 198
 199     // pick periodic image that is closest to the center of the reference density
 200     {
 201         t_pbc pbc;
 202         set_pbc(&pbc, pbcType_, forceProviderInput.box_);
 203         for (RVec& x : transformedCoordinates_)
 204         {
 205             rvec dx;
 206             pbc_dx(&pbc, x, referenceDensityCenter_, dx);
 207             x = referenceDensityCenter_ + dx;
 208         }
 209     }
 210
 211     // transform local atom coordinates to density grid coordinates
 212     transformationToDensityLattice_(transformedCoordinates_);
 213
 214     // spread atoms on grid
 215     gaussTransform_.setZero();
 216
 217     std::vector<real> amplitudes =
 218             amplitudeLookup_(forceProviderInput.mdatoms_, localAtomSet_.localIndex());
 219
 220     if (parameters_.normalizeDensities_)
 221     {
 222         real sum = std::accumulate(std::begin(amplitudes), std::end(amplitudes), 0.);
 223         if (PAR(&forceProviderInput.cr_))
 224         {
 225             gmx_sum(1, &sum, &forceProviderInput.cr_);
 226         }
 227         for (real& amplitude : amplitudes)
 228         {
 229             amplitude /= sum;
 230         }
 231     }
 232
 233     auto amplitudeIterator = amplitudes.cbegin();
 234
 235     for (const auto& r : transformedCoordinates_)
 236     {
 237         gaussTransform_.add({ r, *amplitudeIterator });
 238         ++amplitudeIterator;
 239     }
 240
 241     // communicate grid
 242     if (PAR(&forceProviderInput.cr_))
 243     {
 244         // \todo update to real once GaussTransform class returns real
 245         gmx_sumf(gaussTransform_.view().mapping().required_span_size(),
 246                  gaussTransform_.view().data(), &forceProviderInput.cr_);
 247     }
 248
 249     // calculate grid derivative
 250     const DensitySimilarityMeasure::density& densityDerivative =
 251             measure_.gradient(gaussTransform_.constView());
 252     // calculate forces
 253     forces_.resize(localAtomSet_.numAtomsLocal());
 254     std::transform(
 255             std::begin(transformedCoordinates_), std::end(transformedCoordinates_), std::begin(amplitudes),
 256             std::begin(forces_), [&densityDerivative, this](const RVec r, real amplitude) {
 257                 return densityFittingForce_.evaluateForce({ r, amplitude }, densityDerivative);
 258             });
 259
 260     transformationToDensityLattice_.scaleOperationOnly().inverseIgnoringZeroScale(forces_);
 261
 262     auto       densityForceIterator = forces_.cbegin();
 263     const real effectiveForceConstant = state_.adaptiveForceConstantScale_ * parameters_.calculationIntervalInSteps_
 264                                         * parameters_.forceConstant_;
 265     for (const auto localAtomIndex : localAtomSet_.localIndex())
 266     {
 267         forceProviderOutput->forceWithVirial_.force_[localAtomIndex] +=
 268                 effectiveForceConstant * *densityForceIterator;
 269         ++densityForceIterator;
 270     }
 271
 272     // calculate corresponding potential energy
 273     const float similarity = measure_.similarity(gaussTransform_.constView());
 274     const real energy = -similarity * parameters_.forceConstant_ * state_.adaptiveForceConstantScale_;
 275     forceProviderOutput->enerd_.term[F_DENSITYFITTING] += energy;
 276
 277     if (expAverageSimilarity_.has_value())
 278     {
 279         expAverageSimilarity_->updateWithDataPoint(similarity);
 280
 281         if (expAverageSimilarity_->increasing())
 282         {
 283             state_.adaptiveForceConstantScale_ /= 1._real + expAverageSimilarity_->inverseTimeConstant();
 284         }
 285         else
 286         {
 287             state_.adaptiveForceConstantScale_ *= 1._real + expAverageSimilarity_->inverseTimeConstant();
 288         }
 289     }
 290 }
 291
 292 DensityFittingForceProviderState DensityFittingForceProvider::Impl::state()
 293 {
 294     if (expAverageSimilarity_.has_value())
 295     {
 296         state_.exponentialMovingAverageState_ = expAverageSimilarity_->state();
 297     }
 298     return state_;
 299 }
 300
 301 /********************************************************************
 302  * DensityFittingForceProvider
 303  */
 304
 305 DensityFittingForceProvider::~DensityFittingForceProvider() = default;
 306
 307 DensityFittingForceProvider::DensityFittingForceProvider(const DensityFittingParameters& parameters,
 308                                                          basic_mdspan<const float, dynamicExtents3D> referenceDensity,
 309                                                          const TranslateAndScale& transformationToDensityLattice,
 310                                                          const LocalAtomSet& localAtomSet,
 311                                                          int                 pbcType,
 312                                                          double              simulationTimeStep,
 313                                                          const DensityFittingForceProviderState& state) :
 314     impl_(new Impl(parameters, referenceDensity, transformationToDensityLattice, localAtomSet, pbcType, simulationTimeStep, state))
 315 {
 316 }
 317
 318 void DensityFittingForceProvider::calculateForces(const ForceProviderInput& forceProviderInput,
 319                                                   ForceProviderOutput*      forceProviderOutput)
 320 {
 321     impl_->calculateForces(forceProviderInput, forceProviderOutput);
 322 }
 323
 324 DensityFittingForceProviderState DensityFittingForceProvider::state()
 325 {
 326     return impl_->state();
 327 }
 328
 329 } // namespace gmx