beae64d65e997c07c119c284dab0d0da6cf5f7d5
[alexxy/gromacs.git] / src / gromacs / trajectoryanalysis / modules / msd.cpp
1 /*
2  * This file is part of the GROMACS molecular simulation package.
3  *
4  * Copyright (c) 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  * Defines the trajectory analysis module for mean squared displacement calculations.
38  *
39  * \author Kevin Boyd <kevin44boyd@gmail.com>
40  * \ingroup module_trajectoryanalysis
41  */
42
43 #include "gmxpre.h"
44
45 #include "msd.h"
46
47 #include <numeric>
48 #include <optional>
49
50 #include "gromacs/analysisdata/analysisdata.h"
51 #include "gromacs/analysisdata/modules/average.h"
52 #include "gromacs/analysisdata/modules/plot.h"
53 #include "gromacs/analysisdata/paralleloptions.h"
54 #include "gromacs/fileio/oenv.h"
55 #include "gromacs/fileio/trxio.h"
56 #include "gromacs/math/functions.h"
57 #include "gromacs/math/vectypes.h"
58 #include "gromacs/options/basicoptions.h"
59 #include "gromacs/options/filenameoption.h"
60 #include "gromacs/options/ioptionscontainer.h"
61 #include "gromacs/pbcutil/pbc.h"
62 #include "gromacs/selection/selectionoption.h"
63 #include "gromacs/statistics/statistics.h"
64 #include "gromacs/trajectoryanalysis/analysissettings.h"
65 #include "gromacs/trajectoryanalysis/topologyinformation.h"
66 #include "gromacs/trajectory/trajectoryframe.h"
67 #include "gromacs/utility/stringutil.h"
68 #include "gromacs/utility.h"
69
70 namespace gmx
71 {
72 namespace analysismodules
73 {
74
75 namespace
76 {
77
78 //! Convert nm^2/ps to 10e-5 cm^2/s
79 constexpr double c_diffusionConversionFactor = 1000.0;
80 //! Three dimensional diffusion coefficient multiplication constant.
81 constexpr double c_3DdiffusionDimensionFactor = 6.0;
82 //! Two dimensional diffusion coefficient multiplication constant.
83 constexpr double c_2DdiffusionDimensionFactor = 4.0;
84 //! One dimensional diffusion coefficient multiplication constant.
85 constexpr double c_1DdiffusionDimensionFactor = 2.0;
86
87
88 /*! \brief Mean Squared Displacement data accumulator
89  *
90  * This class is used to accumulate individual MSD data points
91  * and emit tau-averaged results once data is finished collecting. Displacements at each observed
92  * time difference (tau) are recorded from the trajectory. Because it is not known in advance which
93  * time differences will be observed from the trajectory, this data structure is built adaptively.
94  * New columns corresponding to observed time differences are added as needed, and additional
95  * observations at formerly observed time differences are added to those columns. Separate time lags
96  * will likely have differing total data points.
97  *
98  * Data columns per tau are accessed via operator[], which always guarantees
99  * a column is initialized and returns an MsdColumProxy to the column that can push data.
100  */
101 class MsdData
102 {
103 public:
104     //! Proxy to a MsdData tau column vector. Supports only push_back.
105     class MsdColumnProxy
106     {
107     public:
108         MsdColumnProxy(std::vector<double>* column) : column_(column) {}
109
110         void push_back(double value) { column_->push_back(value); }
111
112     private:
113         std::vector<double>* column_;
114     };
115     //! Returns a proxy to the column for the given tau index. Guarantees that the column is initialized.
116     MsdColumnProxy operator[](size_t index)
117     {
118         if (msds_.size() <= index)
119         {
120             msds_.resize(index + 1);
121         }
122         return MsdColumnProxy(&msds_[index]);
123     }
124     /*! \brief Compute per-tau MSDs averaged over all added points.
125      *
126      * The resulting vector is size(max tau index). Any indices
127      * that have no data points have MSD set to 0.
128      *
129      * \return Average MSD per tau
130      */
131     [[nodiscard]] std::vector<real> averageMsds() const;
132
133 private:
134     //! Results - first indexed by tau, then data points
135     std::vector<std::vector<double>> msds_;
136 };
137
138
139 std::vector<real> MsdData::averageMsds() const
140 {
141     std::vector<real> msdSums;
142     msdSums.reserve(msds_.size());
143     for (gmx::ArrayRef<const double> msdValues : msds_)
144     {
145         if (msdValues.empty())
146         {
147             msdSums.push_back(0.0);
148             continue;
149         }
150         msdSums.push_back(std::accumulate(msdValues.begin(), msdValues.end(), 0.0, std::plus<>())
151                           / msdValues.size());
152     }
153     return msdSums;
154 }
155
156 /*! \brief Calculates 1,2, or 3D distance for two vectors.
157  *
158  * \todo Remove NOLINTs once clang-tidy is updated to v11, it should be able to handle constexpr.
159  *
160  * \tparam x If true, calculate x dimension of displacement
161  * \tparam y If true, calculate y dimension of displacement
162  * \tparam z If true, calculate z dimension of displacement
163  * \param[in] c1 First point
164  * \param[in] c2 Second point
165  * \return Euclidian distance for the given dimension.
166  */
167 template<bool x, bool y, bool z>
168 inline double calcSingleSquaredDistance(const RVec c1, const RVec c2)
169 {
170     static_assert(x || y || z, "zero-dimensional MSD selected");
171     const DVec firstCoords  = c1.toDVec();
172     const DVec secondCoords = c2.toDVec();
173     double     result       = 0;
174     if constexpr (x)
175     {
176         result += (firstCoords[XX] - secondCoords[XX]) * (firstCoords[XX] - secondCoords[XX]);
177     }
178     // NOLINTNEXTLINE(readability-misleading-indentation) remove when clang-tidy-13 is required
179     if constexpr (y)
180     {
181         result += (firstCoords[YY] - secondCoords[YY]) * (firstCoords[YY] - secondCoords[YY]);
182     }
183     // NOLINTNEXTLINE(readability-misleading-indentation) remove when clang-tidy-13 is required
184     if constexpr (z)
185     {
186         result += (firstCoords[ZZ] - secondCoords[ZZ]) * (firstCoords[ZZ] - secondCoords[ZZ]);
187     }
188     // NOLINTNEXTLINE(readability-misleading-indentation) remove when clang-tidy-13 is required
189     return result;
190 }
191
192 /*! \brief Calculate average displacement between sets of points
193  *
194  * Each displacement c1[i] - c2[i] is calculated and the distances
195  * are averaged.
196  *
197  * \tparam x If true, calculate x dimension of displacement
198  * \tparam y If true, calculate y dimension of displacement
199  * \tparam z If true, calculate z dimension of displacement
200  * \param[in] c1 First vector
201  * \param[in] c2 Second vector
202  * \return Per-particle averaged distance
203  */
204 template<bool x, bool y, bool z>
205 double calcAverageDisplacement(ArrayRef<const RVec> c1, ArrayRef<const RVec> c2)
206 {
207     double result = 0;
208     for (size_t i = 0; i < c1.size(); i++)
209     {
210         result += calcSingleSquaredDistance<x, y, z>(c1[i], c2[i]);
211     }
212     return result / c1.size();
213 }
214
215
216 //! Describes 1D MSDs, in the given dimension.
217 enum class SingleDimDiffType : int
218 {
219     X = 0,
220     Y,
221     Z,
222     Unused,
223     Count,
224 };
225
226 //! Describes 2D MSDs, in the plane normal to the given dimension.
227 enum class TwoDimDiffType : int
228 {
229     NormalToX = 0,
230     NormalToY,
231     NormalToZ,
232     Unused,
233     Count,
234 };
235
236 /*! \brief Removes jumps across periodic boundaries for currentFrame, based on the positions in
237  * previousFrame. Updates currentCoords in place.
238  */
239 void removePbcJumps(ArrayRef<RVec> currentCoords, ArrayRef<const RVec> previousCoords, t_pbc* pbc)
240 {
241     // There are two types of "pbc removal" in gmx msd. The first happens in the trajectoryanalysis
242     // framework, which makes molecules whole across periodic boundaries and is done
243     // automatically where the inputs support it. This lambda performs the second PBC correction, where
244     // any "jump" across periodic boundaries BETWEEN FRAMES is put back. The order of these
245     // operations is important - since the first transformation may only apply to part of a
246     // molecule (e.g., one half in/out of the box is put on one side of the box), the
247     // subsequent step needs to be applied to the molecule COM rather than individual atoms, or
248     // we'd have a clash where the per-mol PBC removal moves an atom that gets put back into
249     // it's original position by the second transformation. Therefore, this second transformation
250     // is applied *after* per molecule coordinates have been consolidated into COMs.
251     auto pbcRemover = [pbc](RVec in, RVec prev) {
252         rvec dx;
253         pbc_dx(pbc, in, prev, dx);
254         return prev + dx;
255     };
256     std::transform(
257             currentCoords.begin(), currentCoords.end(), previousCoords.begin(), currentCoords.begin(), pbcRemover);
258 }
259
260 //! Holds data needed for MSD calculations for a single molecule, if requested.
261 struct MoleculeData
262 {
263     //! Number of atoms in the molecule.
264     int atomCount = 0;
265     //! Total mass.
266     double mass = 0;
267     //! MSD accumulator and calculator for the molecule
268     MsdData msdData;
269     //! Calculated diffusion coefficient
270     real diffusionCoefficient = 0;
271 };
272
273 /*! \brief Handles coordinate operations for MSD calculations.
274  *
275  * Can be used to hold coordinates for individual atoms as well as molecules COMs. Handles PBC
276  * jump removal between consecutive frames.
277  *
278  * Only previous_ contains valid data between calls to buildCoordinates(), although both vectors
279  * are maintained at the correct size for the number of coordinates needed.
280  */
281 class MsdCoordinateManager
282 {
283 public:
284     MsdCoordinateManager(const int                    numAtoms,
285                          ArrayRef<const MoleculeData> molecules,
286                          ArrayRef<const int>          moleculeIndexMapping) :
287         current_(numAtoms), previous_(numAtoms), molecules_(molecules), moleculeIndexMapping_(moleculeIndexMapping)
288     {
289     }
290     /*! \brief Prepares coordinates for the current frame.
291      *
292      * Reads in selection data, and returns an ArrayRef of the particle positions or molecule
293      * centers of mass (if the molecules input is not empty). Removes jumps across periodic
294      * boundaries based on the previous frame coordinates, except for the first frame built with
295      * builtCoordinates(), which has no previous frame as a reference.
296      *
297      * \param[in] sel                   The selection object which holds coordinates
298      * \param[in] pbc                   Information about periodicity.
299      * \returns                         The current frames coordinates in proper format.
300      */
301     ArrayRef<const RVec> buildCoordinates(const Selection& sel, t_pbc* pbc);
302
303 private:
304     //! The current coordinates.
305     std::vector<RVec> current_;
306     //! The previous frame's coordinates.
307     std::vector<RVec> previous_;
308     //! Molecule data.
309     ArrayRef<const MoleculeData> molecules_;
310     //! Mapping of atom indices to molecle indices;
311     ArrayRef<const int> moleculeIndexMapping_;
312     //! Tracks if a previous frame exists to compare with for PBC handling.
313     bool containsPreviousFrame_ = false;
314 };
315
316 ArrayRef<const RVec> MsdCoordinateManager::buildCoordinates(const Selection& sel, t_pbc* pbc)
317 {
318
319     if (molecules_.empty())
320     {
321         std::copy(sel.coordinates().begin(), sel.coordinates().end(), current_.begin());
322     }
323     else
324     {
325         // Prepare for molecule COM calculation, then sum up all positions per molecule.
326
327         std::fill(current_.begin(), current_.end(), RVec(0, 0, 0));
328         gmx::ArrayRef<const real> masses = sel.masses();
329         for (int i = 0; i < sel.posCount(); i++)
330         {
331             const int moleculeIndex = moleculeIndexMapping_[i];
332             // accumulate ri * mi, and do division at the end to minimize number of divisions.
333             current_[moleculeIndex] += RVec(sel.position(i).x()) * masses[i];
334         }
335         // Divide accumulated mass * positions to get COM.
336         std::transform(current_.begin(),
337                        current_.end(),
338                        molecules_.begin(),
339                        current_.begin(),
340                        [](const RVec& position, const MoleculeData& molecule) -> RVec {
341                            return position / molecule.mass;
342                        });
343     }
344
345     if (containsPreviousFrame_)
346     {
347         removePbcJumps(current_, previous_, pbc);
348     }
349     else
350     {
351         containsPreviousFrame_ = true;
352     }
353
354     // Previous is no longer needed, swap with current and return "current" coordinates which
355     // now reside in previous.
356     current_.swap(previous_);
357     return previous_;
358 }
359
360
361 //! Holds per-group coordinates, analysis, and results.
362 struct MsdGroupData
363 {
364     explicit MsdGroupData(const Selection&             inputSel,
365                           ArrayRef<const MoleculeData> molecules,
366                           ArrayRef<const int>          moleculeAtomMapping) :
367         sel(inputSel),
368         coordinateManager_(molecules.empty() ? sel.posCount() : molecules.size(), molecules, moleculeAtomMapping)
369     {
370     }
371
372     //! Selection associated with this group.
373     const Selection& sel;
374
375     //! Stored coordinates, indexed by frame then atom number.
376     std::vector<std::vector<RVec>> frames;
377
378     //! MSD result accumulator
379     MsdData msds;
380     //! Coordinate handler for the group.
381     MsdCoordinateManager coordinateManager_;
382     //! Collector for processed MSD averages per tau
383     std::vector<real> msdSums;
384     //! Fitted diffusion coefficient
385     real diffusionCoefficient = 0.0;
386     //! Uncertainty of diffusion coefficient
387     double sigma = 0.0;
388 };
389
390 } // namespace
391
392 /*! \brief Implements the gmx msd module
393  *
394  * \todo Implement -(no)mw. Right now, all calculations are mass-weighted with -mol, and not otherwise
395  * \todo Implement -tensor for full MSD tensor calculation
396  * \todo Implement -rmcomm for total-frame COM removal
397  * \todo Implement -pdb for molecule B factors
398  * \todo Implement -maxtau option proposed at https://gitlab.com/gromacs/gromacs/-/issues/3870
399  * \todo Update help text as options are added and clarifications decided on at https://gitlab.com/gromacs/gromacs/-/issues/3869
400  */
401 class Msd : public TrajectoryAnalysisModule
402 {
403 public:
404     Msd();
405
406     void initOptions(IOptionsContainer* options, TrajectoryAnalysisSettings* settings) override;
407     void optionsFinished(TrajectoryAnalysisSettings* settings) override;
408     void initAfterFirstFrame(const TrajectoryAnalysisSettings& settings, const t_trxframe& fr) override;
409     void initAnalysis(const TrajectoryAnalysisSettings& settings, const TopologyInformation& top) override;
410     void analyzeFrame(int                           frameNumber,
411                       const t_trxframe&             frame,
412                       t_pbc*                        pbc,
413                       TrajectoryAnalysisModuleData* pdata) override;
414     void finishAnalysis(int nframes) override;
415     void writeOutput() override;
416
417 private:
418     //! Selections for MSD output
419     SelectionList selections_;
420
421     //! MSD type information, for -type {x,y,z}
422     SingleDimDiffType singleDimType_ = SingleDimDiffType::Unused;
423     //! MSD type information, for -lateral {x,y,z}
424     TwoDimDiffType twoDimType_ = TwoDimDiffType::Unused;
425     //! Diffusion coefficient conversion factor
426     double diffusionCoefficientDimensionFactor_ = c_3DdiffusionDimensionFactor;
427     //! Method used to calculate MSD - changes based on dimensonality.
428     std::function<double(ArrayRef<const RVec>, ArrayRef<const RVec>)> calcMsd_ =
429             calcAverageDisplacement<true, true, true>;
430
431     //! Picoseconds between restarts
432     double trestart_ = 10.0;
433     //! Initial time
434     double t0_ = 0;
435     //! Inter-frame delta-t
436     std::optional<double> dt_ = std::nullopt;
437
438     //! First tau value to fit from for diffusion coefficient, defaults to 0.1 * max tau
439     real beginFit_ = -1.0;
440     //! Final tau value to fit to for diffusion coefficient, defaults to 0.9 * max tau
441     real endFit_ = -1.0;
442
443     //! All selection group-specific data stored here.
444     std::vector<MsdGroupData> groupData_;
445
446     //! Time of all frames.
447     std::vector<double> times_;
448     //! Taus for output - won't know the size until the end.
449     std::vector<double> taus_;
450     //! Tau indices for fitting.
451     size_t beginFitIndex_ = 0;
452     size_t endFitIndex_   = 0;
453
454     // MSD per-molecule stuff
455     //! Are we doing molecule COM-based MSDs?
456     bool molSelected_ = false;
457     //! Per molecule topology information and MSD accumulators.
458     std::vector<MoleculeData> molecules_;
459     //! Atom index -> mol index map
460     std::vector<int> moleculeIndexMappings_;
461
462     // Output stuff
463     AnalysisData msdPlotData_;
464     AnalysisData msdMoleculePlotData_;
465
466     AnalysisDataPlotSettings plotSettings_;
467     //! Per-tau MSDs for each selected group
468     std::string output_;
469     //! Per molecule diffusion coefficients if -mol is selected.
470     std::string moleculeOutput_;
471 };
472
473
474 Msd::Msd() = default;
475
476 void Msd::initOptions(IOptionsContainer* options, TrajectoryAnalysisSettings* settings)
477 {
478     static const char* const desc[] = {
479         "[THISMODULE] computes the mean square displacement (MSD) of atoms from",
480         "a set of initial positions. This provides an easy way to compute",
481         "the diffusion constant using the Einstein relation.",
482         "The time between the reference points for the MSD calculation",
483         "is set with [TT]-trestart[tt].",
484         "The diffusion constant is calculated by least squares fitting a",
485         "straight line (D*t + c) through the MSD(t) from [TT]-beginfit[tt] to",
486         "[TT]-endfit[tt] (note that t is time from the reference positions,",
487         "not simulation time). An error estimate given, which is the difference",
488         "of the diffusion coefficients obtained from fits over the two halves",
489         "of the fit interval.[PAR]",
490         "There are three, mutually exclusive, options to determine different",
491         "types of mean square displacement: [TT]-type[tt], [TT]-lateral[tt]",
492         "and [TT]-ten[tt]. Option [TT]-ten[tt] writes the full MSD tensor for",
493         "each group, the order in the output is: trace xx yy zz yx zx zy.[PAR]",
494         "If [TT]-mol[tt] is set, [THISMODULE] plots the MSD for individual molecules",
495         "(including making molecules whole across periodic boundaries): ",
496         "for each individual molecule a diffusion constant is computed for ",
497         "its center of mass. The chosen index group will be split into ",
498         "molecules. With -mol, only one index group can be selected.[PAR]",
499         "The diffusion coefficient is determined by linear regression of the MSD.",
500         "When [TT]-beginfit[tt] is -1, fitting starts at 10%",
501         "and when [TT]-endfit[tt] is -1, fitting goes to 90%.",
502         "Using this option one also gets an accurate error estimate",
503         "based on the statistics between individual molecules.",
504         "Note that this diffusion coefficient and error estimate are only",
505         "accurate when the MSD is completely linear between",
506         "[TT]-beginfit[tt] and [TT]-endfit[tt].[PAR]",
507     };
508     settings->setHelpText(desc);
509
510     // Selections
511     options->addOption(SelectionOption("sel")
512                                .storeVector(&selections_)
513                                .required()
514                                .onlyStatic()
515                                .multiValue()
516                                .description("Selections to compute MSDs for from the reference"));
517
518     // Select MSD type - defaults to 3D if neither option is selected.
519     EnumerationArray<SingleDimDiffType, const char*> enumTypeNames    = { "x", "y", "z", "unused" };
520     EnumerationArray<TwoDimDiffType, const char*>    enumLateralNames = { "x", "y", "z", "unused" };
521     options->addOption(EnumOption<SingleDimDiffType>("type")
522                                .enumValue(enumTypeNames)
523                                .store(&singleDimType_)
524                                .defaultValue(SingleDimDiffType::Unused));
525     options->addOption(EnumOption<TwoDimDiffType>("lateral")
526                                .enumValue(enumLateralNames)
527                                .store(&twoDimType_)
528                                .defaultValue(TwoDimDiffType::Unused));
529
530     options->addOption(DoubleOption("trestart")
531                                .description("Time between restarting points in trajectory (ps)")
532                                .defaultValue(10.0)
533                                .store(&trestart_));
534     options->addOption(
535             RealOption("beginfit").description("Time point at which to start fitting.").store(&beginFit_));
536     options->addOption(RealOption("endfit").description("End time for fitting.").store(&endFit_));
537
538     // Output options
539     options->addOption(FileNameOption("o")
540                                .filetype(OptionFileType::Plot)
541                                .outputFile()
542                                .store(&output_)
543                                .defaultBasename("msdout")
544                                .description("MSD output"));
545     options->addOption(
546             FileNameOption("mol")
547                     .filetype(OptionFileType::Plot)
548                     .outputFile()
549                     .store(&moleculeOutput_)
550                     .storeIsSet(&molSelected_)
551                     .defaultBasename("diff_mol")
552                     .description("Report diffusion coefficients for each molecule in selection"));
553 }
554
555 void Msd::optionsFinished(TrajectoryAnalysisSettings gmx_unused* settings)
556 {
557     if (singleDimType_ != SingleDimDiffType::Unused && twoDimType_ != TwoDimDiffType::Unused)
558     {
559         std::string errorMessage =
560                 "Options -type and -lateral are mutually exclusive. Choose one or neither (for 3D "
561                 "MSDs).";
562         GMX_THROW(InconsistentInputError(errorMessage.c_str()));
563     }
564     if (selections_.size() > 1 && molSelected_)
565     {
566         std::string errorMessage =
567                 "Cannot have multiple groups selected with -sel when using -mol.";
568         GMX_THROW(InconsistentInputError(errorMessage.c_str()));
569     }
570 }
571
572
573 void Msd::initAnalysis(const TrajectoryAnalysisSettings& settings, const TopologyInformation& top)
574 {
575     plotSettings_ = settings.plotSettings();
576
577     // Enumeration helpers for dispatching the right MSD calculation type.
578     const EnumerationArray<SingleDimDiffType, decltype(&calcAverageDisplacement<true, true, true>)>
579             oneDimensionalMsdFunctions = { &calcAverageDisplacement<true, false, false>,
580                                            &calcAverageDisplacement<false, true, false>,
581                                            &calcAverageDisplacement<false, false, true> };
582     const EnumerationArray<TwoDimDiffType, decltype(&calcAverageDisplacement<true, true, true>)>
583             twoDimensionalMsdFunctions = { calcAverageDisplacement<false, true, true>,
584                                            calcAverageDisplacement<true, false, true>,
585                                            calcAverageDisplacement<true, true, false> };
586
587     // Parse dimensionality and assign the MSD calculating function.
588     // Note if we don't hit either of these cases, we're computing 3D MSDs.
589     if (singleDimType_ != SingleDimDiffType::Unused)
590     {
591         calcMsd_                             = oneDimensionalMsdFunctions[singleDimType_];
592         diffusionCoefficientDimensionFactor_ = c_1DdiffusionDimensionFactor;
593     }
594     else if (twoDimType_ != TwoDimDiffType::Unused)
595     {
596         calcMsd_                             = twoDimensionalMsdFunctions[twoDimType_];
597         diffusionCoefficientDimensionFactor_ = c_2DdiffusionDimensionFactor;
598     }
599
600     // TODO validate that we have mol info and not atom only - and masses, and topology.
601     if (molSelected_)
602     {
603         Selection& sel  = selections_[0];
604         const int  nMol = sel.initOriginalIdsToGroup(top.mtop(), INDEX_MOL);
605
606         gmx::ArrayRef<const int> mappedIds = selections_[0].mappedIds();
607         moleculeIndexMappings_.resize(selections_[0].posCount());
608         std::copy(mappedIds.begin(), mappedIds.end(), moleculeIndexMappings_.begin());
609
610         // Precalculate each molecules mass for speeding up COM calculations.
611         ArrayRef<const real> masses = sel.masses();
612
613         molecules_.resize(nMol);
614         for (int i = 0; i < sel.posCount(); i++)
615         {
616             molecules_[mappedIds[i]].atomCount++;
617             molecules_[mappedIds[i]].mass += masses[i];
618         }
619     }
620
621     // Accumulated frames and results
622     for (const Selection& sel : selections_)
623     {
624         groupData_.emplace_back(sel, molecules_, moleculeIndexMappings_);
625     }
626 }
627
628 void Msd::initAfterFirstFrame(const TrajectoryAnalysisSettings gmx_unused& settings, const t_trxframe& fr)
629 {
630     t0_ = std::round(fr.time);
631 }
632
633
634 void Msd::analyzeFrame(int gmx_unused                frameNumber,
635                        const t_trxframe&             frame,
636                        t_pbc*                        pbc,
637                        TrajectoryAnalysisModuleData* pdata)
638 {
639     const real time = std::round(frame.time);
640     // Need to populate dt on frame 2;
641     if (!dt_.has_value() && !times_.empty())
642     {
643         dt_ = time - times_[0];
644     }
645
646     // Each frame gets an entry in times, but frameTimes only updates if we're at a restart.
647     times_.push_back(time);
648
649     // Each frame will get a tau between it and frame 0, and all other frame combos should be
650     // covered by this.
651     // \todo this will no longer hold exactly when maxtau is added
652     taus_.push_back(time - times_[0]);
653
654     for (MsdGroupData& msdData : groupData_)
655     {
656         const Selection& sel = pdata->parallelSelection(msdData.sel);
657
658         ArrayRef<const RVec> coords = msdData.coordinateManager_.buildCoordinates(sel, pbc);
659
660         // For each preceding frame, calculate tau and do comparison.
661         for (size_t i = 0; i < msdData.frames.size(); i++)
662         {
663             // If dt > trestart, the tau increment will be dt rather than trestart.
664             double  tau      = time - (t0_ + std::max<double>(trestart_, *dt_) * i);
665             int64_t tauIndex = gmx::roundToInt64(tau / *dt_);
666             msdData.msds[tauIndex].push_back(calcMsd_(coords, msdData.frames[i]));
667
668             for (size_t molInd = 0; molInd < molecules_.size(); molInd++)
669             {
670                 molecules_[molInd].msdData[tauIndex].push_back(
671                         calcMsd_(arrayRefFromArray(&coords[molInd], 1),
672                                  arrayRefFromArray(&msdData.frames[i][molInd], 1)));
673             }
674         }
675
676
677         // We only store the frame for the future if it's a restart per -trestart.
678         if (bRmod(time, t0_, trestart_))
679         {
680             msdData.frames.emplace_back(coords.begin(), coords.end());
681         }
682     }
683 }
684
685 //! Calculate the tau index for fitting. If userFitTau < 0, uses the default fraction of max tau.
686 static size_t calculateFitIndex(const int    userFitTau,
687                                 const double defaultTauFraction,
688                                 const int    numTaus,
689                                 const double dt)
690 {
691     if (userFitTau < 0)
692     {
693         return gmx::roundToInt((numTaus - 1) * defaultTauFraction);
694     }
695     return std::min<size_t>(numTaus - 1, gmx::roundToInt(static_cast<double>(userFitTau) / dt));
696 }
697
698
699 void Msd::finishAnalysis(int gmx_unused nframes)
700 {
701     static constexpr double c_defaultStartFitIndexFraction = 0.1;
702     static constexpr double c_defaultEndFitIndexFraction   = 0.9;
703     beginFitIndex_ = calculateFitIndex(beginFit_, c_defaultStartFitIndexFraction, taus_.size(), *dt_);
704     endFitIndex_   = calculateFitIndex(endFit_, c_defaultEndFitIndexFraction, taus_.size(), *dt_);
705     const int numTausForFit = 1 + endFitIndex_ - beginFitIndex_;
706
707     // These aren't used, except for correlationCoefficient, which is used to estimate error if
708     // enough points are available.
709     real b = 0.0, correlationCoefficient = 0.0, chiSquared = 0.0;
710
711     for (MsdGroupData& msdData : groupData_)
712     {
713         msdData.msdSums = msdData.msds.averageMsds();
714
715         if (numTausForFit >= 4)
716         {
717             const int halfNumTaus         = numTausForFit / 2;
718             const int secondaryStartIndex = beginFitIndex_ + halfNumTaus;
719             // Split the fit in 2, and compare the results of each fit;
720             real a = 0.0, a2 = 0.0;
721             lsq_y_ax_b_xdouble(halfNumTaus,
722                                &taus_[beginFitIndex_],
723                                &msdData.msdSums[beginFitIndex_],
724                                &a,
725                                &b,
726                                &correlationCoefficient,
727                                &chiSquared);
728             lsq_y_ax_b_xdouble(halfNumTaus,
729                                &taus_[secondaryStartIndex],
730                                &msdData.msdSums[secondaryStartIndex],
731                                &a2,
732                                &b,
733                                &correlationCoefficient,
734                                &chiSquared);
735             msdData.sigma = std::abs(a - a2);
736         }
737         lsq_y_ax_b_xdouble(numTausForFit,
738                            &taus_[beginFitIndex_],
739                            &msdData.msdSums[beginFitIndex_],
740                            &msdData.diffusionCoefficient,
741                            &b,
742                            &correlationCoefficient,
743                            &chiSquared);
744         msdData.diffusionCoefficient *= c_diffusionConversionFactor / diffusionCoefficientDimensionFactor_;
745         msdData.sigma *= c_diffusionConversionFactor / diffusionCoefficientDimensionFactor_;
746     }
747
748     for (MoleculeData& molecule : molecules_)
749     {
750         std::vector<real> msds = molecule.msdData.averageMsds();
751         lsq_y_ax_b_xdouble(numTausForFit,
752                            &taus_[beginFitIndex_],
753                            &msds[beginFitIndex_],
754                            &molecule.diffusionCoefficient,
755                            &b,
756                            &correlationCoefficient,
757                            &chiSquared);
758         molecule.diffusionCoefficient *= c_diffusionConversionFactor / diffusionCoefficientDimensionFactor_;
759     }
760 }
761
762 void Msd::writeOutput()
763 {
764     // AnalysisData currently doesn't support changing column counts after analysis has started.
765     // We can't determine the number of tau values until the trajectory is fully read, so analysis
766     // data construction and plotting are done here.
767     AnalysisDataPlotModulePointer msdPlotModule(new AnalysisDataPlotModule(plotSettings_));
768     msdPlotModule->setFileName(output_);
769     msdPlotModule->setTitle("Mean Squared Displacement");
770     msdPlotModule->setXLabel("tau (ps)");
771     msdPlotModule->setYLabel(R"(MSD (nm\\S2\\N))");
772     msdPlotModule->setYFormat(10, 6, 'g');
773     for (const auto& group : groupData_)
774     {
775         const real D = group.diffusionCoefficient;
776         if (D > 0.01 && D < 1e4)
777         {
778             msdPlotModule->appendLegend(formatString(
779                     "D[%10s] = %.4f (+/- %.4f) (1e-5 cm^2/s)", group.sel.name(), D, group.sigma));
780         }
781         else
782         {
783             msdPlotModule->appendLegend(formatString(
784                     "D[%10s] = %.4g (+/- %.4f) (1e-5 cm^2/s)", group.sel.name(), D, group.sigma));
785         }
786     }
787     msdPlotData_.addModule(msdPlotModule);
788     msdPlotData_.setDataSetCount(groupData_.size());
789     for (size_t i = 0; i < groupData_.size(); i++)
790     {
791         msdPlotData_.setColumnCount(i, 1);
792     }
793     AnalysisDataHandle dh = msdPlotData_.startData({});
794     for (size_t tauIndex = 0; tauIndex < taus_.size(); tauIndex++)
795     {
796         dh.startFrame(tauIndex, taus_[tauIndex]);
797         for (size_t dataSetIndex = 0; dataSetIndex < groupData_.size(); dataSetIndex++)
798         {
799             dh.selectDataSet(dataSetIndex);
800             dh.setPoint(0, groupData_[dataSetIndex].msdSums[tauIndex]);
801         }
802         dh.finishFrame();
803     }
804     dh.finishData();
805
806     if (molSelected_)
807     {
808         AnalysisDataPlotModulePointer molPlotModule(new AnalysisDataPlotModule(plotSettings_));
809         molPlotModule->setFileName(moleculeOutput_);
810         molPlotModule->setTitle("Mean Squared Displacement / Molecule");
811         molPlotModule->setXLabel("Molecule");
812         molPlotModule->setYLabel("D(1e-5 cm^2/s)");
813         molPlotModule->setYFormat(10, 0, 'g');
814         msdMoleculePlotData_.addModule(molPlotModule);
815         msdMoleculePlotData_.setDataSetCount(1);
816         msdMoleculePlotData_.setColumnCount(0, 1);
817         AnalysisDataHandle molDh = msdMoleculePlotData_.startData({});
818         for (size_t moleculeIndex = 0; moleculeIndex < molecules_.size(); moleculeIndex++)
819         {
820             molDh.startFrame(moleculeIndex, moleculeIndex);
821             molDh.setPoint(0, molecules_[moleculeIndex].diffusionCoefficient);
822             molDh.finishFrame();
823         }
824         molDh.finishData();
825     }
826 }
827
828
829 const char                      MsdInfo::name[]             = "msd";
830 const char                      MsdInfo::shortDescription[] = "Compute mean squared displacements";
831 TrajectoryAnalysisModulePointer MsdInfo::create()
832 {
833     return TrajectoryAnalysisModulePointer(new Msd);
834 }
835
836
837 } // namespace analysismodules
838 } // namespace gmx