Add widget for MSD analysis (#35)

- Add a new `SlidingWindowMSD` class that is O(n) for each window
- Support to compute and display particle MSDs
- Support for configuring selection, dimension type and log scale

Author: PardhavMaradani

CHANGELOG.md

@@ -29,6 +29,8 @@ The rules for this file:
 
 <!-- New added features -->
 
+- Added widget for MSD Analysis (PR #35)
+
 ### Fixed
 
 <!-- Bug fixes -->

mdadash/backend/analyses/__init__.py

@@ -2,7 +2,7 @@
 Module that has all the analyses widgets
 """
 
-from . import com_distance, dssp, energies, janin, ramachandran, rog
+from . import com_distance, dssp, energies, janin, msd, ramachandran, rog
 
 __all__ = [
     "energies",
@@ -11,4 +11,5 @@
     "dssp",
     "ramachandran",
     "janin",
+    "msd",
 ]

mdadash/backend/analyses/msd.py

@@ -0,0 +1,257 @@
+import logging
+
+import matplotlib.pyplot as plt
+import MDAnalysis as mda
+import numpy as np
+from IPython.display import display
+from joblib import delayed
+from matplotlib.collections import LineCollection
+
+from mdadash.backend.widgets.base import WidgetBase
+
+logger = logging.getLogger(__name__)
+
+
+class MSDAnalysis(WidgetBase):
+    name = "MSD Analysis"
+    description = "Mean squared displacement analysis"
+
+    _inputs = [
+        {
+            "attribute": "_run_mode",
+            "name": "Run mode",
+            "description": "The mode in which the widget is run",
+            "type": "select",
+            "items": [
+                "serial",
+                "parallel",
+            ],
+        },
+        {
+            "attribute": "selection",
+            "name": "Selection",
+            "description": "MDAnalysis selection phrase",
+            "type": "str",
+        },
+        {
+            "attribute": "msd_type",
+            "name": "MSD type",
+            "description": "Desired dimensions to be included in the MSD",
+            "type": "select",
+            "items": [
+                "xyz",
+                "xy",
+                "yz",
+                "xz",
+                "x",
+                "y",
+                "z",
+            ],
+        },
+        {
+            "attribute": "show_particle_msds",
+            "name": "Show particle MSDs",
+            "description": "Show MSDs for individual particles of the selection",
+            "type": "bool",
+        },
+        {
+            "attribute": "log_scale",
+            "name": "Log scale",
+            "description": "Use a log scale for the axes",
+            "type": "bool",
+        },
+        {
+            "attribute": "custom_title",
+            "name": "Custom title",
+            "description": "Custom title for the plot",
+            "type": "str",
+        },
+    ]
+
+    def __init__(self):
+        super().__init__()
+        self.msd = None
+        self.selection = "all"
+        self.msd_type = "xyz"
+        self.log_scale = False
+        self.show_particle_msds = False
+        self.title = "MSD"
+        self.custom_title = None
+        self._setup_plot()
+
+    def _setup_plot(self):
+        """Setup matplotlib plot"""
+        self.fig, self.ax = plt.subplots()
+        # use non-empty values to prevent initial exception
+        # if widget is configured to use log scale
+        (self.plot,) = self.ax.plot([1], [1], color="red", zorder=2)
+        self.lc = LineCollection([], colors="gray", alpha=0.2, lw=0.5, zorder=1)
+        self.ax.add_collection(self.lc)
+        self.ax.set_xlabel(r"Lag time $\Delta$t (ps)")
+        self.ax.set_ylabel(r"MSD ($\AA^2$)")
+        self.ax.grid(True, linestyle="--", alpha=0.6)
+        self._set_title()
+        self._set_axes_scale()
+
+    def _set_title(self):
+        """Set plot title"""
+        self.ax.set_title(self.custom_title if self.custom_title else self.title)
+
+    def _set_axes_scale(self):
+        """Set axes scale"""
+        self.ax.set_xscale("log" if self.log_scale else "linear")
+        self.ax.set_yscale("log" if self.log_scale else "linear")
+
+    def _create_msd(self):
+        """Create msd instance"""
+        self.msd = SlidingWindowMSD(
+            self.u,
+            select=self.selection,
+            msd_type=self.msd_type,
+            show_particle_msds=self.show_particle_msds,
+        )
+        self.title = f"MSD of '{self.selection}'"
+        self._set_title()
+
+    def on_post_create(self):
+        """on_post_create handler"""
+        self._set_title()
+        self._set_axes_scale()
+
+    def on_post_connect(self):
+        """on_post_connect handler"""
+        self._create_msd()
+
+    def on_input_change(self, attribute, _old_value, new_value):
+        """on_input_change handler"""
+        if attribute == "custom_title":
+            self._set_title()
+        elif attribute == "log_scale":
+            self._set_axes_scale()
+        elif attribute == "_run_mode":
+            pass
+        else:
+            self._create_msd()
+
+    def _compute(self, parallel: bool = False):
+        """Run MSD for the current timesteps window"""
+        return self.msd.run(parallel=parallel)
+
+    def _update_plot(self, x, y1, y2):
+        """Update plot with computed values"""
+        self.plot.set_data(x, y1)
+        self.lc.set_segments(y2 if self.show_particle_msds else [])
+        self.ax.relim()
+        self.ax.autoscale_view()
+        self.fig.canvas.draw()
+        display(self.fig)
+
+    def run_every_frame(self):
+        """every-frame run handler"""
+        x, y1, y2, _ = self._compute()
+        self._update_plot(x, y1, y2)
+
+    def get_parallel_job(self):
+        """get parallel job handler"""
+        return delayed(self._compute)(parallel=True)
+
+    def apply_parallel_results(self, values):
+        """apply parallel results handler"""
+        x, y1, y2, (v1, v2, v3, v4) = values
+        self._update_plot(x, y1, y2)
+        # update msd state
+        self.msd.msd_sums = v1
+        self.msd.msd_counts = v2
+        if self.show_particle_msds:
+            self.msd.particle_msd_sums = v3
+            self.msd.particle_msd_counts = v4
+
+
+class SlidingWindowMSD:
+    """Sliding Window MSD
+
+    Calculate MSD for a sliding window of frames
+
+    """
+
+    def __init__(
+        self,
+        u: mda.Universe,
+        select: str = "all",
+        msd_type: str = "xyz",
+        show_particle_msds: bool = False,
+    ):
+        self.u = u
+        self.select = select
+        self.msd_type = msd_type
+        self.show_particle_msds = show_particle_msds
+        self._parse_msd_type()
+        self.ag = u.select_atoms(self.select)
+        self.n_atoms = self.ag.atoms.n_atoms
+        self.n_lags = u.trajectory.buffer_size
+        self.msd_sums = np.zeros(self.n_lags)
+        self.msd_counts = np.zeros(self.n_lags, dtype=int)
+        self.msd_counts[0] = 1
+        if self.show_particle_msds:
+            self.particle_msd_sums = np.zeros((self.n_lags, self.n_atoms))
+            self.particle_msd_counts = np.zeros((self.n_lags, self.n_atoms), dtype=int)
+            self.particle_msd_counts[0, :] = 1
+
+    def _parse_msd_type(self):
+        """Sets up the desired dimensionality of the MSD."""
+        keys = {
+            "x": [0],
+            "y": [1],
+            "z": [2],
+            "xy": [0, 1],
+            "xz": [0, 2],
+            "yz": [1, 2],
+            "xyz": [0, 1, 2],
+        }
+        self._dim = keys[self.msd_type.lower()]
+
+    def run(self, parallel: bool = False) -> tuple:
+        """Run MSD for the current window"""
+
+        n = len(self.u.trajectory)  # buffer / window might not be full yet
+        positions_current = self.ag.positions[:, self._dim]
+        for i in range(n - 1):
+            lag = n - 1 - i
+            _ = self.u.trajectory[i]  # set the buffered trajectory frame
+            disp = positions_current - self.ag.positions[:, self._dim]
+            squared_disp = np.sum(disp**2, axis=1)
+            msd = np.mean(squared_disp)
+            self.msd_sums[lag] += msd
+            self.msd_counts[lag] += 1
+            if self.show_particle_msds:
+                self.particle_msd_sums[lag, :] += squared_disp
+                self.particle_msd_counts[lag, :] += 1
+
+        # We will have at least 2 frames by the time we are here.
+        # frame_dt will ensure the delta_t is correct even if we have step
+        # value (other than 1) configured in the universe configuration
+        frame_dt = round(self.u.trajectory[1].time - self.u.trajectory[0].time, 2)
+        delta_t_values = np.arange(n) * frame_dt
+        avg_msds = self.msd_sums[:n] / self.msd_counts[:n]
+        msds_by_particle_lines = None
+        if self.show_particle_msds:
+            msds_by_particle_array = (
+                self.particle_msd_sums[:n, :] / self.particle_msd_counts[:n, :]
+            )
+            msds_by_particle_lines = np.empty((self.n_atoms, n, 2))
+            msds_by_particle_lines[:, :, 0] = delta_t_values
+            msds_by_particle_lines[:, :, 1] = msds_by_particle_array.T
+
+        return (
+            delta_t_values,
+            avg_msds,
+            msds_by_particle_lines,
+            (
+                self.msd_sums,
+                self.msd_counts,
+                self.particle_msd_sums if self.show_particle_msds else None,
+                self.particle_msd_counts if self.show_particle_msds else None,
+            )
+            if parallel
+            else (None,) * 4,
+        )

mdadash/backend/tests/test_server.py

@@ -604,6 +604,37 @@ async def test_widget_run_janin(_client, imd_server):
     await disconnect_from_simulation()
 
 
+async def test_widget_run_msd_serial(_client, imd_server):
+    uuid = await add_widget("MSD Analysis")
+    await connect_to_simulation(imd_server, step=1, batch_size=2)
+    inputs = [
+        ("selection", "resid 1"),
+        ("custom_title", ""),
+        ("show_particle_msds", True),
+        ("log_scale", False),
+    ]
+    await check_input_changes(uuid, inputs)
+    await resume_simulation(imd_server)
+    assert await sio_event_emitted(sio, "widgets:output", n=1)
+    await remove_widget(uuid)
+    await disconnect_from_simulation()
+
+
+async def test_widget_run_msd_parallel(_client, imd_server):
+    uuid = await add_widget("MSD Analysis")
+    await connect_to_simulation(imd_server, step=1, batch_size=2)
+    inputs = [
+        ("selection", "resid 1"),
+        ("show_particle_msds", True),
+        ("_run_mode", "parallel"),
+    ]
+    await check_input_changes(uuid, inputs)
+    await resume_simulation(imd_server)
+    assert await sio_event_emitted(sio, "widgets:output", n=1)
+    await remove_widget(uuid)
+    await disconnect_from_simulation()
+
+
 def test_state_load(tmp_path):
     # test with no state file
     sm = StateManager("")

mdadash/backend/tests/utils.py

@@ -52,13 +52,13 @@ async def check_input_changes(uuid, inputs, status="ok"):
         assert response["status"] == status
 
 
-async def connect_to_simulation(imd_server):
+async def connect_to_simulation(imd_server, step=2, batch_size=1):
     main.mdadash.sm.universe_configs[0].update(
         {
             "topology": str(TPR),
             "trajectory": f"imd://localhost:{imd_server.port}",
-            "step": 2,
-            "batch_size": 1,
+            "step": step,
+            "batch_size": batch_size,
         }
     )
     handler = sio.handlers["/"]["connect_to_simulations"]

mdadash/backend/widgets/base.py

@@ -556,8 +556,12 @@ def _run_parallel_jobs(self, parallel_widgets, parallel_results):
             func, args, kwargs = widget.get_parallel_job()
             parallel_jobs.append((self.with_reset_frame, (func,) + args, kwargs))
         try:
+            # without max_nbytes=None, np arrays passed / returned
+            # are marked read-only in subsequent calls (eg: msd case)
             results = Parallel(
-                n_jobs=self.n_jobs, initializer=WidgetManager._patch_IMDReader
+                n_jobs=self.n_jobs,
+                max_nbytes=None,
+                initializer=WidgetManager._patch_IMDReader,
             )(parallel_jobs)
             parallel_results.extend(results)
         # pylint: disable=broad-exception-caught