6363the normal MDAnalysis citations.
6464
6565.. warning::
66- To correctly compute the MSD using this analysis module, you must supply
67- coordinates in the **unwrapped** convention. That is, when atoms pass
68- the periodic boundary, they must not be **wrapped** back into the primary
69- simulation cell. MDAnalysis does not currently offer this functionality in
70- the ``MDAnalysis.transformations`` API despite having functions with
71- similar names. We plan to implement the appropriate transformations in the
72- future. In the meantime, various simulation packages provide utilities to
73- convert coordinates to the unwrapped convention. In GROMACS for example,
74- this can be done using ``gmx trjconv`` with the ``-pbc nojump`` flag.
66+ To correctly compute the MSD using this analysis module, you must supply
67+ coordinates in the **unwrapped** convention, also known as **no-jump**.
68+ That is, when atoms pass the periodic boundary, they must not be wrapped
69+ back into the primary simulation cell.
70+
71+ In MDAnalysis you can use the
72+ :class:`~MDAnalysis.transformations.nojump.NoJump`
73+ transformation.
74+
75+ In GROMACS, for example, this can be done using `gmx trjconv`_ with the
76+ ``-pbc nojump`` flag.
77+
78+ .. _`gmx trjconv`: https://manual.gromacs.org/current/onlinehelp/gmx-trjconv.html
79+
80+ .. SeeAlso::
81+ :mod:`MDAnalysis.transformations.nojump`
82+
7583
7684Computing an MSD
7785----------------
102110.. code-block:: python
103111
104112 msd = MSD.results.timeseries
113+ lagtimes = MSD.results.delta_t_values
105114
106- Visual inspection of the MSD is important, so let's take a look at it with a
107- simple plot.
115+ Visual inspection of the MSD is important, so let's take a look at it with a simple plot.
108116
109117.. code-block:: python
110118
111119 import matplotlib.pyplot as plt
112120 nframes = MSD.n_frames
113- timestep = 1 # this needs to be the actual time between frames
114- lagtimes = np.arange(nframes)*timestep # make the lag-time axis
115121 fig = plt.figure()
116122 ax = plt.axes()
117123 # plot the actual MSD
172178 start_time = 20
173179 start_index = int(start_time/timestep)
174180 end_time = 60
175- linear_model = linregress(lagtimes[start_index:end_index],
176- msd[start_index:end_index])
181+ linear_model = linregress(lagtimes[start_index:end_index], msd[start_index:end_index])
177182 slope = linear_model.slope
178183 error = linear_model.stderr
179184 # dim_fac is 3 as we computed a 3D msd with 'xyz'
245250from .base import AnalysisBase
246251from ..core import groups
247252from tqdm import tqdm
253+ import collections
248254
249255logger = logging .getLogger ("MDAnalysis.analysis.msd" )
250256
@@ -281,15 +287,32 @@ class EinsteinMSD(AnalysisBase):
281287 the MSD. Otherwise, use the simple "windowed" algorithm.
282288 The tidynamics package is required for `fft=True`.
283289 Defaults to ``True``.
290+ non_linear : bool
291+ If ``True``, calculates MSD for trajectory where frames are
292+ non-linearly dumped. To use this set `fft=False`.
293+ Defaults to ``False``.
294+
295+ .. versionadded:: 2.10.0
296+
284297
285298 Attributes
286299 ----------
287300 dim_fac : int
288301 Dimensionality :math:`d` of the MSD.
289302 results.timeseries : :class:`numpy.ndarray`
290- The averaged MSD over all the particles with respect to lag-time.
303+ The averaged MSD over all the particles with respect to constant lag-time or
304+ unique Δt intervals.
291305 results.msds_by_particle : :class:`numpy.ndarray`
292- The MSD of each individual particle with respect to lag-time.
306+ The MSD of each individual particle with respect to constant lag-time or
307+ unique Δt intervals.
308+ - for `non_linear=False`: a 2D array of shape (n_lagtimes, n_atoms)
309+ - for `non_linear=True`: a 2D array of shape (n_delta_t_values, n_atoms)
310+ results.delta_t_values : :class:`numpy.ndarray`
311+ Array of unique Δt (time differences) at which time-averaged MSD values are
312+ computed.
313+
314+ .. versionadded:: 2.10.0
315+
293316 ag : :class:`AtomGroup`
294317 The :class:`AtomGroup` resulting from your selection
295318 n_frames : int
@@ -299,27 +322,23 @@ class EinsteinMSD(AnalysisBase):
299322
300323
301324 .. versionadded:: 2.0.0
325+ .. versionchanged:: 2.10.0
326+ Added ability to calculate MSD from samples that are not linearly spaced with the
327+ new `non_linear` keyword argument.
302328 """
303329
304- def __init__ (self , u , select = "all" , msd_type = "xyz" , fft = True , ** kwargs ):
305- r"""
306- Parameters
307- ----------
308- u : Universe or AtomGroup
309- An MDAnalysis :class:`Universe` or :class:`AtomGroup`.
310- select : str
311- A selection string. Defaults to "all" in which case
312- all atoms are selected.
313- msd_type : {'xyz', 'xy', 'yz', 'xz', 'x', 'y', 'z'}
314- Desired dimensions to be included in the MSD.
315- fft : bool
316- If ``True``, uses a fast FFT based algorithm for computation of
317- the MSD. Otherwise, use the simple "windowed" algorithm.
318- The tidynamics package is required for `fft=True`.
319- """
330+ def __init__ (
331+ self ,
332+ u ,
333+ select = "all" ,
334+ msd_type = "xyz" ,
335+ fft = True ,
336+ non_linear = False ,
337+ ** kwargs ,
338+ ):
320339 if isinstance (u , groups .UpdatingAtomGroup ):
321340 raise TypeError (
322- "UpdatingAtomGroups are not valid for MSD " " computation"
341+ "UpdatingAtomGroups are not valid for MSD computation"
323342 )
324343
325344 super (EinsteinMSD , self ).__init__ (u .universe .trajectory , ** kwargs )
@@ -329,6 +348,7 @@ def __init__(self, u, select="all", msd_type="xyz", fft=True, **kwargs):
329348 self .msd_type = msd_type
330349 self ._parse_msd_type ()
331350 self .fft = fft
351+ self .non_linear = non_linear
332352
333353 # local
334354 self .ag = u .select_atoms (self .select )
@@ -338,6 +358,7 @@ def __init__(self, u, select="all", msd_type="xyz", fft=True, **kwargs):
338358 # result
339359 self .results .msds_by_particle = None
340360 self .results .timeseries = None
361+ self .results .delta_t_values = None
341362
342363 def _prepare (self ):
343364 # self.n_frames only available here
@@ -383,10 +404,13 @@ def _single_frame(self):
383404 ]
384405
385406 def _conclude (self ):
386- if self .fft :
387- self ._conclude_fft ()
407+ if self .non_linear :
408+ self ._conclude_non_linear ()
388409 else :
389- self ._conclude_simple ()
410+ if self .fft :
411+ self ._conclude_fft ()
412+ else :
413+ self ._conclude_simple ()
390414
391415 def _conclude_simple (self ):
392416 r"""Calculates the MSD via the simple "windowed" algorithm."""
@@ -397,6 +421,9 @@ def _conclude_simple(self):
397421 sqdist = np .square (disp ).sum (axis = - 1 )
398422 self .results .msds_by_particle [lag , :] = np .mean (sqdist , axis = 0 )
399423 self .results .timeseries = self .results .msds_by_particle .mean (axis = 1 )
424+ self .results .delta_t_values = np .arange (self .n_frames ) * (
425+ self .times [1 ] - self .times [0 ]
426+ )
400427
401428 def _conclude_fft (self ): # with FFT, np.float64 bit prescision required.
402429 r"""Calculates the MSD via the FCA fast correlation algorithm."""
@@ -421,3 +448,44 @@ def _conclude_fft(self): # with FFT, np.float64 bit prescision required.
421448 positions [:, n , :]
422449 )
423450 self .results .timeseries = self .results .msds_by_particle .mean (axis = 1 )
451+ self .results .delta_t_values = np .arange (self .n_frames ) * (
452+ self .times [1 ] - self .times [0 ]
453+ )
454+
455+ def _conclude_non_linear (self ):
456+
457+ n_frames = self .n_frames
458+ n_atoms = self .n_particles
459+ positions = self ._position_array .astype (np .float64 )
460+ # Dictionary to collect MSDs: {Δt: [msd1, msd2, ...]}
461+ msd_dict = collections .defaultdict (list )
462+ msds_by_particle_dict = collections .defaultdict (list )
463+
464+ # TODO: optimize the code
465+ # Looping over all the frames as if the referenced gets shifted frame to frame
466+ for i in range (n_frames ):
467+ for j in range (i + 1 , n_frames ):
468+ delta_t = self .times [j ] - self .times [i ]
469+ # Compute displacement and squared displacement
470+ disp = positions [j ] - positions [i ]
471+ squared_disp = np .sum (disp ** 2 , axis = 1 )
472+ msd = np .mean (squared_disp )
473+ # Store MSD under corresponding Δt
474+ msd_dict [delta_t ].append (msd )
475+ msds_by_particle_dict [delta_t ].append (squared_disp )
476+
477+ msd_dict [0 ] = [0 ]
478+ msds_by_particle_dict [0.0 ] = [np .zeros (n_atoms )]
479+
480+ # For each delta_t, stacked all squared_disp arrays and averaging over axis=0 (time origins)
481+ delta_t_values = sorted (msd_dict .keys ())
482+ avg_msds = [np .mean (msd_dict [dt ]) for dt in delta_t_values ]
483+ msds_by_particle_array = np .zeros ((len (delta_t_values ), n_atoms ))
484+ for idx , dt in enumerate (delta_t_values ):
485+ # Stack list of arrays like -- (n_time_origins, n_atoms)
486+ arr = np .vstack (msds_by_particle_dict [dt ])
487+ msds_by_particle_array [idx , :] = np .mean (arr , axis = 0 )
488+
489+ self .results .timeseries = np .array (avg_msds )
490+ self .results .delta_t_values = np .array (delta_t_values )
491+ self .results .msds_by_particle = msds_by_particle_array
0 commit comments