2 * This file is part of the GROMACS molecular simulation package.
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.
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],
83 sigmaInLatticeCoordinates[ZZ] },
89 /********************************************************************
90 * DensityFittingForceProviderState
93 const std::string DensityFittingForceProviderState::adaptiveForceConstantScaleName_ =
94 "adaptiveForceConstantScale";
96 const std::string DensityFittingForceProviderState::exponentialMovingAverageStateName_ =
97 "exponentialMovingAverageState";
99 const std::string DensityFittingForceProviderState::stepsSinceLastCalculationName_ =
100 "stepsSinceLastCalculation";
102 void DensityFittingForceProviderState::writeState(KeyValueTreeObjectBuilder kvtBuilder,
103 const std::string& identifier) const
105 writeKvtCheckpointValue(stepsSinceLastCalculation_, stepsSinceLastCalculationName_, identifier,
107 writeKvtCheckpointValue(adaptiveForceConstantScale_, adaptiveForceConstantScaleName_,
108 identifier, kvtBuilder);
110 KeyValueTreeObjectBuilder exponentialMovingAverageKvtEntry =
111 kvtBuilder.addObject(identifier + "-" + exponentialMovingAverageStateName_);
112 exponentialMovingAverageStateAsKeyValueTree(exponentialMovingAverageKvtEntry,
113 exponentialMovingAverageState_);
116 void DensityFittingForceProviderState::readState(const KeyValueTreeObject& kvtData,
117 const std::string& identifier)
119 readKvtCheckpointValue(compat::make_not_null(&stepsSinceLastCalculation_),
120 stepsSinceLastCalculationName_, identifier, kvtData);
121 readKvtCheckpointValue(compat::make_not_null(&adaptiveForceConstantScale_),
122 adaptiveForceConstantScaleName_, identifier, kvtData);
124 if (kvtData.keyExists(identifier + "-" + exponentialMovingAverageStateName_))
126 exponentialMovingAverageState_ = exponentialMovingAverageStateFromKeyValueTree(
127 kvtData[identifier + "-" + exponentialMovingAverageStateName_].asObject());
131 void DensityFittingForceProviderState::broadcastState(MPI_Comm communicator, bool isParallelRun)
135 block_bc(communicator, stepsSinceLastCalculation_);
136 block_bc(communicator, adaptiveForceConstantScale_);
137 block_bc(communicator, exponentialMovingAverageState_);
142 /********************************************************************
143 * DensityFittingForceProvider::Impl
146 class DensityFittingForceProvider::Impl
149 //! \copydoc DensityFittingForceProvider::DensityFittingForceProvider
150 Impl(const DensityFittingParameters& parameters,
151 basic_mdspan<const float, dynamicExtents3D> referenceDensity,
152 const TranslateAndScale& transformationToDensityLattice,
153 const LocalAtomSet& localAtomSet,
155 double simulationTimeStep,
156 const DensityFittingForceProviderState& state);
158 void calculateForces(const ForceProviderInput& forceProviderInput,
159 ForceProviderOutput* forceProviderOutput);
160 const DensityFittingForceProviderState& stateToCheckpoint();
163 DensityFittingForceProviderState state();
164 const DensityFittingParameters& parameters_;
165 DensityFittingForceProviderState state_;
166 DensityFittingForceProviderState stateToCheckpoint_;
167 LocalAtomSet localAtomSet_;
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_;
181 //! Optionally scale the force according to a moving average of the similarity
182 std::optional<ExponentialMovingAverage> expAverageSimilarity_;
184 //! Optionally translate the structure
185 std::optional<AffineTransformation> affineTransformation_;
188 DensityFittingForceProvider::Impl::~Impl() = default;
190 DensityFittingForceProvider::Impl::Impl(const DensityFittingParameters& parameters,
191 basic_mdspan<const float, dynamicExtents3D> referenceDensity,
192 const TranslateAndScale& transformationToDensityLattice,
193 const LocalAtomSet& localAtomSet,
195 double simulationTimeStep,
196 const DensityFittingForceProviderState& state) :
197 parameters_(parameters),
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),
210 expAverageSimilarity_(std::nullopt)
212 if (parameters_.adaptiveForceScaling_)
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_));
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)
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])
238 affineTransformation_.emplace(translationMatrix.asConstView(), translationVector);
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;
253 void DensityFittingForceProvider::Impl::calculateForces(const ForceProviderInput& forceProviderInput,
254 ForceProviderOutput* forceProviderOutput)
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)
261 ++(state_.stepsSinceLastCalculation_);
265 state_.stepsSinceLastCalculation_ = 1;
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]; });
273 // apply additional structure transformations
274 if (affineTransformation_)
276 (*affineTransformation_)(transformedCoordinates_);
279 // pick periodic image that is closest to the center of the reference density
282 set_pbc(&pbc, pbcType_, forceProviderInput.box_);
283 for (RVec& x : transformedCoordinates_)
286 pbc_dx(&pbc, x, referenceDensityCenter_, dx);
287 x = referenceDensityCenter_ + dx;
291 // transform local atom coordinates to density grid coordinates
292 transformationToDensityLattice_(transformedCoordinates_);
294 // spread atoms on grid
295 gaussTransform_.setZero();
297 std::vector<real> amplitudes =
298 amplitudeLookup_(forceProviderInput.mdatoms_, localAtomSet_.localIndex());
300 if (parameters_.normalizeDensities_)
302 real sum = std::accumulate(std::begin(amplitudes), std::end(amplitudes), 0.);
303 if (havePPDomainDecomposition(&forceProviderInput.cr_))
305 gmx_sum(1, &sum, &forceProviderInput.cr_);
307 for (real& amplitude : amplitudes)
313 auto amplitudeIterator = amplitudes.cbegin();
315 for (const auto& r : transformedCoordinates_)
317 gaussTransform_.add({ r, *amplitudeIterator });
322 if (havePPDomainDecomposition(&forceProviderInput.cr_))
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_);
329 // calculate grid derivative
330 const DensitySimilarityMeasure::density& densityDerivative =
331 measure_.gradient(gaussTransform_.constView());
333 forces_.resize(localAtomSet_.numAtomsLocal());
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);
340 transformationToDensityLattice_.scaleOperationOnly().inverseIgnoringZeroScale(forces_);
342 auto densityForceIterator = forces_.cbegin();
343 const real effectiveForceConstant = state_.adaptiveForceConstantScale_ * parameters_.calculationIntervalInSteps_
344 * parameters_.forceConstant_;
345 for (const auto localAtomIndex : localAtomSet_.localIndex())
347 forceProviderOutput->forceWithVirial_.force_[localAtomIndex] +=
348 effectiveForceConstant * *densityForceIterator;
349 ++densityForceIterator;
352 const float similarity = measure_.similarity(gaussTransform_.constView());
353 if (MASTER(&(forceProviderInput.cr_)))
355 // calculate corresponding potential energy
356 const real energy = -similarity * parameters_.forceConstant_ * state_.adaptiveForceConstantScale_;
357 forceProviderOutput->enerd_.term[F_DENSITYFITTING] += energy;
360 if (expAverageSimilarity_.has_value())
362 expAverageSimilarity_->updateWithDataPoint(similarity);
364 if (expAverageSimilarity_->increasing())
366 state_.adaptiveForceConstantScale_ /= 1._real + expAverageSimilarity_->inverseTimeConstant();
370 state_.adaptiveForceConstantScale_ *=
371 1._real + 2 * expAverageSimilarity_->inverseTimeConstant();
376 DensityFittingForceProviderState DensityFittingForceProvider::Impl::state()
378 if (expAverageSimilarity_.has_value())
380 state_.exponentialMovingAverageState_ = expAverageSimilarity_->state();
385 const DensityFittingForceProviderState& DensityFittingForceProvider::Impl::stateToCheckpoint()
387 return stateToCheckpoint_;
389 /********************************************************************
390 * DensityFittingForceProvider
393 DensityFittingForceProvider::~DensityFittingForceProvider() = default;
395 DensityFittingForceProvider::DensityFittingForceProvider(const DensityFittingParameters& parameters,
396 basic_mdspan<const float, dynamicExtents3D> referenceDensity,
397 const TranslateAndScale& transformationToDensityLattice,
398 const LocalAtomSet& localAtomSet,
400 double simulationTimeStep,
401 const DensityFittingForceProviderState& state) :
402 impl_(new Impl(parameters, referenceDensity, transformationToDensityLattice, localAtomSet, pbcType, simulationTimeStep, state))
406 void DensityFittingForceProvider::calculateForces(const ForceProviderInput& forceProviderInput,
407 ForceProviderOutput* forceProviderOutput)
409 impl_->calculateForces(forceProviderInput, forceProviderOutput);
412 void DensityFittingForceProvider::writeCheckpointData(MdModulesWriteCheckpointData checkpointWriting,
413 const std::string& moduleName)
415 impl_->stateToCheckpoint().writeState(checkpointWriting.builder_, moduleName);