diff --git a/CHANGELOG.md b/CHANGELOG.md index e3ad721..9fcd954 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -29,6 +29,8 @@ The rules for this file: +- Added widget for MSD Analysis (PR #35) + ### Fixed diff --git a/mdadash/backend/analyses/__init__.py b/mdadash/backend/analyses/__init__.py index 97daa12..6921d8f 100644 --- a/mdadash/backend/analyses/__init__.py +++ b/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", ] diff --git a/mdadash/backend/analyses/msd.py b/mdadash/backend/analyses/msd.py new file mode 100644 index 0000000..a9988cc --- /dev/null +++ b/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, + ) diff --git a/mdadash/backend/tests/test_server.py b/mdadash/backend/tests/test_server.py index cb0d937..3f30970 100644 --- a/mdadash/backend/tests/test_server.py +++ b/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("") diff --git a/mdadash/backend/tests/utils.py b/mdadash/backend/tests/utils.py index 72e73ac..e252adf 100644 --- a/mdadash/backend/tests/utils.py +++ b/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"] diff --git a/mdadash/backend/widgets/base.py b/mdadash/backend/widgets/base.py index 461f9a0..28bea4f 100644 --- a/mdadash/backend/widgets/base.py +++ b/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