2 * This file is part of the GROMACS molecular simulation package.
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.
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.
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.
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.
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.
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.
37 * Implements force provider for density fitting
39 * \author Christian Blau <blau@kth.se>
40 * \ingroup module_applied_forces
44 #include "densityfittingforceprovider.h"
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"
62 #include "densityfittingamplitudelookup.h"
63 #include "densityfittingparameters.h"
71 /*! \internal \brief Generate the spread kernel from Gaussian parameters.
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
78 GaussianSpreadKernelParameters::Shape makeSpreadKernel(real sigma, real nSigma, const ScaleCoordinates& scaleToLattice)
80 RVec sigmaInLatticeCoordinates{ sigma, sigma, sigma };
81 scaleToLattice(&sigmaInLatticeCoordinates);
82 return { DVec{ sigmaInLatticeCoordinates[XX], sigmaInLatticeCoordinates[YY], sigmaInLatticeCoordinates[ZZ] },
88 /********************************************************************
89 * DensityFittingForceProviderState
92 const std::string DensityFittingForceProviderState::adaptiveForceConstantScaleName_ =
93 "adaptiveForceConstantScale";
95 const std::string DensityFittingForceProviderState::exponentialMovingAverageStateName_ =
96 "exponentialMovingAverageState";
98 const std::string DensityFittingForceProviderState::stepsSinceLastCalculationName_ =
99 "stepsSinceLastCalculation";
101 void DensityFittingForceProviderState::writeState(KeyValueTreeObjectBuilder kvtBuilder,
102 const std::string& identifier) const
104 writeKvtCheckpointValue(
105 stepsSinceLastCalculation_, stepsSinceLastCalculationName_, identifier, kvtBuilder);
106 writeKvtCheckpointValue(
107 adaptiveForceConstantScale_, adaptiveForceConstantScaleName_, identifier, kvtBuilder);
109 KeyValueTreeObjectBuilder exponentialMovingAverageKvtEntry =
110 kvtBuilder.addObject(identifier + "-" + exponentialMovingAverageStateName_);
111 exponentialMovingAverageStateAsKeyValueTree(exponentialMovingAverageKvtEntry,
112 exponentialMovingAverageState_);
115 void DensityFittingForceProviderState::readState(const KeyValueTreeObject& kvtData,
116 const std::string& identifier)
118 readKvtCheckpointValue(compat::make_not_null(&stepsSinceLastCalculation_),
119 stepsSinceLastCalculationName_,
122 readKvtCheckpointValue(compat::make_not_null(&adaptiveForceConstantScale_),
123 adaptiveForceConstantScaleName_,
127 if (kvtData.keyExists(identifier + "-" + exponentialMovingAverageStateName_))
129 exponentialMovingAverageState_ = exponentialMovingAverageStateFromKeyValueTree(
130 kvtData[identifier + "-" + exponentialMovingAverageStateName_].asObject());
134 void DensityFittingForceProviderState::broadcastState(MPI_Comm communicator, bool isParallelRun)
138 block_bc(communicator, stepsSinceLastCalculation_);
139 block_bc(communicator, adaptiveForceConstantScale_);
140 block_bc(communicator, exponentialMovingAverageState_);
145 /********************************************************************
146 * DensityFittingForceProvider::Impl
149 class DensityFittingForceProvider::Impl
152 //! \copydoc DensityFittingForceProvider::DensityFittingForceProvider
153 Impl(const DensityFittingParameters& parameters,
154 basic_mdspan<const float, dynamicExtents3D> referenceDensity,
155 const TranslateAndScale& transformationToDensityLattice,
156 const LocalAtomSet& localAtomSet,
158 double simulationTimeStep,
159 const DensityFittingForceProviderState& state);
161 void calculateForces(const ForceProviderInput& forceProviderInput,
162 ForceProviderOutput* forceProviderOutput);
163 const DensityFittingForceProviderState& stateToCheckpoint();
166 DensityFittingForceProviderState state();
167 const DensityFittingParameters& parameters_;
168 DensityFittingForceProviderState state_;
169 DensityFittingForceProviderState stateToCheckpoint_;
170 LocalAtomSet localAtomSet_;
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_;
184 //! Optionally scale the force according to a moving average of the similarity
185 std::optional<ExponentialMovingAverage> expAverageSimilarity_;
187 //! Optionally translate the structure
188 std::optional<AffineTransformation> affineTransformation_;
191 DensityFittingForceProvider::Impl::~Impl() = default;
193 DensityFittingForceProvider::Impl::Impl(const DensityFittingParameters& parameters,
194 basic_mdspan<const float, dynamicExtents3D> referenceDensity,
195 const TranslateAndScale& transformationToDensityLattice,
196 const LocalAtomSet& localAtomSet,
198 double simulationTimeStep,
199 const DensityFittingForceProviderState& state) :
200 parameters_(parameters),
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),
213 expAverageSimilarity_(std::nullopt)
215 if (parameters_.adaptiveForceScaling_)
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_));
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)
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])
241 affineTransformation_.emplace(translationMatrix.asConstView(), translationVector);
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;
256 void DensityFittingForceProvider::Impl::calculateForces(const ForceProviderInput& forceProviderInput,
257 ForceProviderOutput* forceProviderOutput)
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)
264 ++(state_.stepsSinceLastCalculation_);
268 state_.stepsSinceLastCalculation_ = 1;
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]; });
277 // apply additional structure transformations
278 if (affineTransformation_)
280 (*affineTransformation_)(transformedCoordinates_);
283 // pick periodic image that is closest to the center of the reference density
286 set_pbc(&pbc, pbcType_, forceProviderInput.box_);
287 for (RVec& x : transformedCoordinates_)
290 pbc_dx(&pbc, x, referenceDensityCenter_, dx);
291 x = referenceDensityCenter_ + dx;
295 // transform local atom coordinates to density grid coordinates
296 transformationToDensityLattice_(transformedCoordinates_);
298 // spread atoms on grid
299 gaussTransform_.setZero();
301 std::vector<real> amplitudes = amplitudeLookup_(
302 forceProviderInput.chargeA_, forceProviderInput.massT_, localAtomSet_.localIndex());
304 if (parameters_.normalizeDensities_)
306 real sum = std::accumulate(std::begin(amplitudes), std::end(amplitudes), 0.);
307 if (havePPDomainDecomposition(&forceProviderInput.cr_))
309 gmx_sum(1, &sum, &forceProviderInput.cr_);
311 for (real& amplitude : amplitudes)
317 auto amplitudeIterator = amplitudes.cbegin();
319 for (const auto& r : transformedCoordinates_)
321 gaussTransform_.add({ r, *amplitudeIterator });
326 if (havePPDomainDecomposition(&forceProviderInput.cr_))
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_);
334 // calculate grid derivative
335 const DensitySimilarityMeasure::density& densityDerivative =
336 measure_.gradient(gaussTransform_.constView());
338 forces_.resize(localAtomSet_.numAtomsLocal());
340 std::begin(transformedCoordinates_),
341 std::end(transformedCoordinates_),
342 std::begin(amplitudes),
344 [&densityDerivative, this](const RVec r, real amplitude) {
345 return densityFittingForce_.evaluateForce({ r, amplitude }, densityDerivative);
348 transformationToDensityLattice_.scaleOperationOnly().inverseIgnoringZeroScale(forces_);
350 auto densityForceIterator = forces_.cbegin();
351 const real effectiveForceConstant = state_.adaptiveForceConstantScale_ * parameters_.calculationIntervalInSteps_
352 * parameters_.forceConstant_;
353 for (const auto localAtomIndex : localAtomSet_.localIndex())
355 forceProviderOutput->forceWithVirial_.force_[localAtomIndex] +=
356 effectiveForceConstant * *densityForceIterator;
357 ++densityForceIterator;
360 const float similarity = measure_.similarity(gaussTransform_.constView());
361 if (MASTER(&(forceProviderInput.cr_)))
363 // calculate corresponding potential energy
364 const real energy = -similarity * parameters_.forceConstant_ * state_.adaptiveForceConstantScale_;
365 forceProviderOutput->enerd_.term[F_DENSITYFITTING] += energy;
368 if (expAverageSimilarity_.has_value())
370 expAverageSimilarity_->updateWithDataPoint(similarity);
372 if (expAverageSimilarity_->increasing())
374 state_.adaptiveForceConstantScale_ /= 1._real + expAverageSimilarity_->inverseTimeConstant();
378 state_.adaptiveForceConstantScale_ *=
379 1._real + 2 * expAverageSimilarity_->inverseTimeConstant();
384 DensityFittingForceProviderState DensityFittingForceProvider::Impl::state()
386 if (expAverageSimilarity_.has_value())
388 state_.exponentialMovingAverageState_ = expAverageSimilarity_->state();
393 const DensityFittingForceProviderState& DensityFittingForceProvider::Impl::stateToCheckpoint()
395 return stateToCheckpoint_;
397 /********************************************************************
398 * DensityFittingForceProvider
401 DensityFittingForceProvider::~DensityFittingForceProvider() = default;
403 DensityFittingForceProvider::DensityFittingForceProvider(const DensityFittingParameters& parameters,
404 basic_mdspan<const float, dynamicExtents3D> referenceDensity,
405 const TranslateAndScale& transformationToDensityLattice,
406 const LocalAtomSet& localAtomSet,
408 double simulationTimeStep,
409 const DensityFittingForceProviderState& state) :
410 impl_(new Impl(parameters, referenceDensity, transformationToDensityLattice, localAtomSet, pbcType, simulationTimeStep, state))
414 void DensityFittingForceProvider::calculateForces(const ForceProviderInput& forceProviderInput,
415 ForceProviderOutput* forceProviderOutput)
417 impl_->calculateForces(forceProviderInput, forceProviderOutput);
420 void DensityFittingForceProvider::writeCheckpointData(MDModulesWriteCheckpointData checkpointWriting,
421 const std::string& moduleName)
423 impl_->stateToCheckpoint().writeState(checkpointWriting.builder_, moduleName);