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

@@ -63,6 +63,11 @@ Enhancements
    (PR #5038)
  * Moved distopia checking function to common import location in
    MDAnalysisTest.util (PR #5038)
+ * Added new method '_conclude_non_linear` to calculate time-averaged
+   mean square displacement in `package/MDAnalysis/analysis/msd.py`,
+   new bool keyword argument `non_linear=True` under `EinsteinMSD`
+   enables execution of this method.
+   (Issue #5028, PR #5066)
 
 Changes
  * Refactored the RDKit converter code to move the inferring code in a separate

package/MDAnalysis/analysis/msd.py

@@ -173,7 +173,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 +245,7 @@
 from .base import AnalysisBase
 from ..core import groups
 from tqdm import tqdm
+import collections
 
 logger = logging.getLogger("MDAnalysis.analysis.msd")
 
@@ -281,6 +282,10 @@ 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``.
 
     Attributes
     ----------
@@ -290,6 +295,8 @@ 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.
     ag : :class:`AtomGroup`
         The :class:`AtomGroup` resulting from your selection
     n_frames : int
@@ -298,10 +305,12 @@ class EinsteinMSD(AnalysisBase):
         Number of particles MSD was calculated over.
 
 
-    .. versionadded:: 2.0.0
+    .. versionadded:: 2.10.0
     """
 
-    def __init__(self, u, select="all", msd_type="xyz", fft=True, **kwargs):
+    def __init__(
+        self, u, select="all", msd_type="xyz", fft=True, non_linear=False, **kwargs
+    ):
         r"""
         Parameters
         ----------
@@ -316,11 +325,12 @@ def __init__(self, u, select="all", msd_type="xyz", fft=True, **kwargs):
             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`.
+        non_linear : bool
+            If ``True``, calculates MSD for trajectory where frames are
+            non-linearly dumped. To use this set `fft=False`.
         """
         if isinstance(u, groups.UpdatingAtomGroup):
-            raise TypeError(
-                "UpdatingAtomGroups are not valid for MSD " "computation"
-            )
+            raise TypeError("UpdatingAtomGroups are not valid for MSD " "computation")
 
         super(EinsteinMSD, self).__init__(u.universe.trajectory, **kwargs)
 
@@ -329,6 +339,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,16 +349,13 @@ 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
         # these need to be zeroed prior to each run() call
-        self.results.msds_by_particle = np.zeros(
-            (self.n_frames, self.n_particles)
-        )
-        self._position_array = np.zeros(
-            (self.n_frames, self.n_particles, self.dim_fac)
-        )
+        self.results.msds_by_particle = np.zeros((self.n_frames, self.n_particles))
+        self._position_array = np.zeros((self.n_frames, self.n_particles, self.dim_fac))
         # self.results.timeseries not set here
 
     def _parse_msd_type(self):
@@ -378,15 +386,16 @@ def _single_frame(self):
         r"""Constructs array of positions for MSD calculation."""
         # shape of position array set here, use span in last dimension
         # from this point on
-        self._position_array[self._frame_index] = self.ag.positions[
-            :, self._dim
-        ]
+        self._position_array[self._frame_index] = self.ag.positions[:, self._dim]
 
     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."""
@@ -417,7 +426,43 @@ def _conclude_fft(self):  # with FFT, np.float64 bit prescision required.
 
         positions = self._position_array.astype(np.float64)
         for n in tqdm(range(self.n_particles)):
-            self.results.msds_by_particle[:, n] = tidynamics.msd(
-                positions[:, n, :]
-            )
+            self.results.msds_by_particle[:, n] = tidynamics.msd(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, ...]}
+
+        # 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)
+
+        msd_dict[0] = [0]
+
+        # Prepare averaged results
+        delta_t_values = []
+        avg_msds = []
+        for delta_t in sorted(msd_dict.keys()):
+            msd_list = msd_dict[delta_t]
+            avg_msd = np.mean(msd_list)
+            delta_t_values.append(delta_t)
+            avg_msds.append(avg_msd)
+
+        self.results.timeseries = np.array(avg_msds, dtype=np.float64)
+        self.results.delta_t_values = np.array(delta_t_values, dtype=np.float64)

testsuite/MDAnalysisTests/analysis/test_msd.py

@@ -28,7 +28,13 @@
 from numpy.testing import assert_almost_equal, assert_equal
 import numpy as np
 
-from MDAnalysisTests.datafiles import PSF, DCD, RANDOM_WALK, RANDOM_WALK_TOPO
+from MDAnalysisTests.datafiles import (
+    PSF,
+    DCD,
+    RANDOM_WALK,
+    RANDOM_WALK_TOPO,
+    LAMMPSDUMP_non_linear,
+)
 from MDAnalysisTests.util import block_import, import_not_available
 
 import pytest
@@ -124,9 +130,7 @@ def test_msdtype_error(self, u, SELECTION, msdtype):
             ("z", 1),
         ],
     )
-    def test_simple_step_traj_all_dims(
-        self, step_traj, NSTEP, dim, dim_factor
-    ):
+    def test_simple_step_traj_all_dims(self, step_traj, NSTEP, dim, dim_factor):
         # testing the "simple" algorithm on constant velocity trajectory
         # should fit the polynomial y=dim_factor*x**2
         m_simple = MSD(step_traj, "all", msd_type=dim, fft=False)
@@ -146,18 +150,14 @@ def test_simple_step_traj_all_dims(
             ("z", 1),
         ],
     )
-    def test_simple_start_stop_step_all_dims(
-        self, step_traj, NSTEP, dim, dim_factor
-    ):
+    def test_simple_start_stop_step_all_dims(self, step_traj, NSTEP, dim, dim_factor):
         # testing the "simple" algorithm on constant velocity trajectory
         # test start stop step is working correctly
         m_simple = MSD(step_traj, "all", msd_type=dim, fft=False)
         m_simple.run(start=10, stop=1000, step=10)
         poly = characteristic_poly(NSTEP, dim_factor)
         # polynomial must take offset start into account
-        assert_almost_equal(
-            m_simple.results.timeseries, poly[0:990:10], decimal=4
-        )
+        assert_almost_equal(m_simple.results.timeseries, poly[0:990:10], decimal=4)
 
     def test_random_walk_u_simple(self, random_walk_u):
         # regress against random_walk test data
@@ -252,18 +252,14 @@ def test_fft_step_traj_all_dims(self, step_traj, NSTEP, dim, dim_factor):
             ("z", 1),
         ],
     )
-    def test_fft_start_stop_step_all_dims(
-        self, step_traj, NSTEP, dim, dim_factor
-    ):
+    def test_fft_start_stop_step_all_dims(self, step_traj, NSTEP, dim, dim_factor):
         # testing the fft algorithm on constant velocity trajectory
         # test start stop step is working correctly
         m_simple = MSD(step_traj, "all", msd_type=dim, fft=True)
         m_simple.run(start=10, stop=1000, step=10)
         poly = characteristic_poly(NSTEP, dim_factor)
         # polynomial must take offset start into account
-        assert_almost_equal(
-            m_simple.results.timeseries, poly[0:990:10], decimal=3
-        )
+        assert_almost_equal(m_simple.results.timeseries, poly[0:990:10], decimal=3)
 
     def test_random_walk_u_fft(self, random_walk_u):
         # regress against random_walk test data
@@ -272,3 +268,144 @@ def test_random_walk_u_fft(self, random_walk_u):
         norm = np.linalg.norm(msd_rw.results.timeseries)
         val = 3932.39927487146
         assert_almost_equal(norm, val, decimal=5)
+
+
+def test_msd_non_linear():
+    u = mda.Universe(LAMMPSDUMP_non_linear, format="LAMMPSDUMP")
+
+    msd = MSD(u, select="all", msd_type="xyz", non_linear=True)
+    msd.run()
+
+    result_msd = msd.results.timeseries
+    result_delta_t = msd.results.delta_t_values
+
+    assert result_msd.ndim == 1
+    assert result_msd.shape[0] > 0
+
+    assert result_delta_t.ndim == 1
+    assert result_delta_t.shape[0] > 0
+
+    expected_msd = np.array(
+        [
+            0.00000000e00,
+            7.70976963e-05,
+            2.90842662e-04,
+            6.55040347e-04,
+            1.20610926e-03,
+            2.52547250e-03,
+            3.31645965e-03,
+            5.38852795e-03,
+            1.01941562e-02,
+            1.24745603e-02,
+            1.35380300e-02,
+            1.57475527e-02,
+            2.85165801e-02,
+            3.50591021e-02,
+            3.81292797e-02,
+            3.96176470e-02,
+            3.83551274e-02,
+            5.51041371e-02,
+            5.95049433e-02,
+            6.07026502e-02,
+            6.14434181e-02,
+            6.19512436e-02,
+            6.61293773e-02,
+            9.46607497e-02,
+            1.01300585e-01,
+            9.96583811e-02,
+            9.81112279e-02,
+            9.72780657e-02,
+            9.69221886e-02,
+            1.29442431e-01,
+            1.80752226e-01,
+            1.86358673e-01,
+            1.98140564e-01,
+            2.00603000e-01,
+            1.99094789e-01,
+            1.97272787e-01,
+            1.96156023e-01,
+            2.67664446e-01,
+            4.50987076e-01,
+            4.02344442e-01,
+            3.91458056e-01,
+            4.10370922e-01,
+            4.22997445e-01,
+            4.26217251e-01,
+            4.26484034e-01,
+            4.26360794e-01,
+            6.91315347e-01,
+            9.94317423e-01,
+            1.19622365e00,
+            1.04919180e00,
+            1.06437594e00,
+            1.09426432e00,
+            1.10194082e00,
+            1.10275424e00,
+            1.10383947e00,
+            1.10493159e00,
+        ]
+    )
+
+    expected_delta_t = np.array(
+        [
+            0.000e00,
+            1.000e00,
+            2.000e00,
+            3.000e00,
+            4.000e00,
+            6.000e00,
+            7.000e00,
+            8.000e00,
+            1.200e01,
+            1.400e01,
+            1.500e01,
+            1.600e01,
+            2.400e01,
+            2.800e01,
+            3.000e01,
+            3.100e01,
+            3.200e01,
+            4.800e01,
+            5.600e01,
+            6.000e01,
+            6.200e01,
+            6.300e01,
+            6.400e01,
+            9.600e01,
+            1.120e02,
+            1.200e02,
+            1.240e02,
+            1.260e02,
+            1.270e02,
+            1.280e02,
+            1.920e02,
+            2.240e02,
+            2.400e02,
+            2.480e02,
+            2.520e02,
+            2.540e02,
+            2.550e02,
+            2.560e02,
+            3.840e02,
+            4.480e02,
+            4.800e02,
+            4.960e02,
+            5.040e02,
+            5.080e02,
+            5.100e02,
+            5.110e02,
+            5.120e02,
+            7.680e02,
+            8.960e02,
+            9.600e02,
+            9.920e02,
+            1.008e03,
+            1.016e03,
+            1.020e03,
+            1.022e03,
+            1.023e03,
+        ]
+    )
+
+    np.testing.assert_allclose(result_msd, expected_msd, rtol=1e-5)
+    np.testing.assert_allclose(result_delta_t, expected_delta_t, rtol=1e-5)