Feature non linear time averaged msd

package/AUTHORS

@@ -259,7 +259,7 @@ Chronological list of authors
   - Tulga-Erdene Sodjargal
   - Gareth Elliott
   - Marc Schuh
-
+  - Sirsha Ganguly
 
 External code
 -------------

package/CHANGELOG

@@ -40,6 +40,9 @@ Fixes
    directly passing them. (Issue #3520, PR #5006)
 
 Enhancements
+ * Added capability to calculate MSD from frames with irregular (non-linear)
+   time spacing in analysis.msd.EinsteinMSD with keyword argument
+   `non_linear=True` (Issue #5028, PR #5066)
  * Added `TemplateInferrer` and `RDKitInferrer` dataclasses to the
    `RDKitInferring` module to be used by the RDKit converter. (PR #4305)
  * Improve speed of GROMOS11 (TRC) reader (Issue #5079, PR #5080)

package/MDAnalysis/analysis/msd.py

@@ -103,8 +103,7 @@
 
     msd =  MSD.results.timeseries
 
-Visual inspection of the MSD is important, so let's take a look at it with a
- simple plot.
+Visual inspection of the MSD is important, so let's take a look at it with a simple plot.
 
 .. code-block:: python
 
@@ -173,7 +172,7 @@
     start_index = int(start_time/timestep)
     end_time = 60
     linear_model = linregress(lagtimes[start_index:end_index],
-    				  		  msd[start_index:end_index])
+                                                  msd[start_index:end_index])
     slope = linear_model.slope
     error = linear_model.stderr
     # dim_fac is 3 as we computed a 3D msd with 'xyz'
@@ -245,6 +244,7 @@
 from .base import AnalysisBase
 from ..core import groups
 from tqdm import tqdm
+import collections
 
 logger = logging.getLogger("MDAnalysis.analysis.msd")
 
@@ -281,6 +281,13 @@ class EinsteinMSD(AnalysisBase):
         the MSD. Otherwise, use the simple "windowed" algorithm.
         The tidynamics package is required for `fft=True`.
         Defaults to ``True``.
+    non_linear : bool
+        If ``True``, calculates MSD for trajectory where frames are
+        non-linearly dumped. To use this set `fft=False`.
+        Defaults to ``False``.
+
+        .. versionadded:: 2.10.0
+
 
     Attributes
     ----------
@@ -290,6 +297,12 @@ class EinsteinMSD(AnalysisBase):
         The averaged MSD over all the particles with respect to lag-time.
     results.msds_by_particle : :class:`numpy.ndarray`
         The MSD of each individual particle with respect to lag-time.
+    results.delta_t_values : :class:`numpy.ndarray`
+        Array of unique Δt (time differences) at which time-averaged MSD values are
+        computed.
+
+        .. versionadded:: 2.10.0
+
     ag : :class:`AtomGroup`
         The :class:`AtomGroup` resulting from your selection
     n_frames : int
@@ -299,24 +312,20 @@ class EinsteinMSD(AnalysisBase):
 
 
     .. versionadded:: 2.0.0
+    .. versionchanged:: 2.10.0
+       Added ability to calculate MSD from samples that are not linearly spaced with the
+       new `non_linear` keyword argument.
     """
 
-    def __init__(self, u, select="all", msd_type="xyz", fft=True, **kwargs):
-        r"""
-        Parameters
-        ----------
-        u : Universe or AtomGroup
-            An MDAnalysis :class:`Universe` or :class:`AtomGroup`.
-        select : str
-            A selection string. Defaults to "all" in which case
-            all atoms are selected.
-        msd_type : {'xyz', 'xy', 'yz', 'xz', 'x', 'y', 'z'}
-            Desired dimensions to be included in the MSD.
-        fft : bool
-            If ``True``, uses a fast FFT based algorithm for computation of
-            the MSD. Otherwise, use the simple "windowed" algorithm.
-            The tidynamics package is required for `fft=True`.
-        """
+    def __init__(
+        self,
+        u,
+        select="all",
+        msd_type="xyz",
+        fft=True,
+        non_linear=False,
+        **kwargs,
+    ):
         if isinstance(u, groups.UpdatingAtomGroup):
             raise TypeError(
                 "UpdatingAtomGroups are not valid for MSD " "computation"
@@ -329,6 +338,7 @@ def __init__(self, u, select="all", msd_type="xyz", fft=True, **kwargs):
         self.msd_type = msd_type
         self._parse_msd_type()
         self.fft = fft
+        self.non_linear = non_linear
 
         # local
         self.ag = u.select_atoms(self.select)
@@ -338,6 +348,7 @@ def __init__(self, u, select="all", msd_type="xyz", fft=True, **kwargs):
         # result
         self.results.msds_by_particle = None
         self.results.timeseries = None
+        self.results.delta_t_values = None
 
     def _prepare(self):
         # self.n_frames only available here
@@ -383,10 +394,13 @@ def _single_frame(self):
         ]
 
     def _conclude(self):
-        if self.fft:
-            self._conclude_fft()
+        if self.non_linear:
+            self._conclude_non_linear()
         else:
-            self._conclude_simple()
+            if self.fft:
+                self._conclude_fft()
+            else:
+                self._conclude_simple()
 
     def _conclude_simple(self):
         r"""Calculates the MSD via the simple "windowed" algorithm."""
@@ -421,3 +435,48 @@ def _conclude_fft(self):  # with FFT, np.float64 bit prescision required.
                 positions[:, n, :]
             )
         self.results.timeseries = self.results.msds_by_particle.mean(axis=1)
+
+    def _conclude_non_linear(self):
+
+        dump_times = self.times
+        n_frames = self.n_frames
+        n_atoms = self.n_particles
+        positions = self._position_array.astype(np.float64)
+
+        msd_dict = collections.defaultdict(
+            list
+        )  # Dictionary to collect MSDs: {Δt: [msd1, msd2, ...]}
+        msds_by_particle_dict = collections.defaultdict(list)
+
+        # TODO: optimize the code
+        # Looping over all the frames as if the referenced gets shifted frame to frame
+        for i in range(n_frames):
+            for j in range(i + 1, n_frames):
+                delta_t = dump_times[j] - dump_times[i]
+                # Compute displacement and squared displacement
+                disp = positions[j] - positions[i]
+                squared_disp = np.sum(disp**2, axis=1)
+                msd = np.mean(squared_disp)
+                # Store MSD under corresponding Δt
+                msd_dict[delta_t].append(msd)
+                msds_by_particle_dict[delta_t].append(squared_disp)
+
+        msd_dict[0] = [0]
+        msds_by_particle_dict[0.0] = [np.zeros(n_atoms)]
+
+        # For each delta_t, stacked all squared_disp arrays and averaging over axis=0 (time origins)
+        delta_t_values = sorted(msd_dict.keys())
+        avg_msds = [np.mean(msd_dict[dt]) for dt in delta_t_values]
+        msds_by_particle_array = np.zeros((len(delta_t_values), n_atoms))
+        for idx, dt in enumerate(delta_t_values):
+            # Stack list of arrays like -- (n_time_origins, n_atoms)
+            arr = np.vstack(msds_by_particle_dict[dt])
+            msds_by_particle_array[idx, :] = np.mean(arr, axis=0)
+
+        self.results.timeseries = np.array(avg_msds, dtype=np.float64)
+        self.results.delta_t_values = np.array(
+            delta_t_values, dtype=np.float64
+        )
+        self.results.msds_by_particle = np.array(
+            msds_by_particle_array, dtype=np.float64
+        )