diff --git a/CHANGELOG.md b/CHANGELOG.md index 9fcd954..403a25e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -30,6 +30,7 @@ The rules for this file: - Added widget for MSD Analysis (PR #35) +- Added widget for Velocity Autocorrelation Function (VACF) (PR #36) ### Fixed diff --git a/mdadash/backend/analyses/__init__.py b/mdadash/backend/analyses/__init__.py index 6921d8f..365105b 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, msd, ramachandran, rog +from . import acf, com_distance, dssp, energies, janin, msd, ramachandran, rog __all__ = [ "energies", @@ -12,4 +12,5 @@ "ramachandran", "janin", "msd", + "acf", ] diff --git a/mdadash/backend/analyses/acf.py b/mdadash/backend/analyses/acf.py new file mode 100644 index 0000000..66beb40 --- /dev/null +++ b/mdadash/backend/analyses/acf.py @@ -0,0 +1,340 @@ +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 scipy import integrate + +from mdadash.backend.widgets.base import WidgetBase + +logger = logging.getLogger(__name__) + + +class ACFAnalysis(WidgetBase): + name = "ACF" + description = "Autocorrelation Function" + + _inputs = [ + { + "attribute": "_run_mode", + "name": "Run mode", + "description": "The mode in which the widget is run", + "type": "select", + "items": [ + "serial", + "parallel", + ], + }, + { + "attribute": "physical_property", + "name": "Physical property", + "description": "Physical property to analyze", + "type": "select", + "items": [ + "velocity", + "position", + "force", + ], + }, + { + "attribute": "selection", + "name": "Selection", + "description": "MDAnalysis selection phrase", + "type": "str", + }, + { + "attribute": "dim_type", + "name": "Dimension type", + "description": "Desired dimensions to be included in the ACF", + "type": "select", + "items": [ + "xyz", + "xy", + "yz", + "xz", + "x", + "y", + "z", + ], + }, + { + "attribute": "centered", + "name": "Centered", + "description": ( + "Use mean subtracted values to calculate ACF. " + "A running updated mean based on data processed so far is used. " + "The number of data samples must be much greater than the lag-time " + "window for this to be accurate." + ), + "type": "bool", + }, + { + "attribute": "show_running_integral", + "name": "Show running integral", + "description": "Show running integral of the ACF", + "type": "bool", + }, + { + "attribute": "show_particle_acfs", + "name": "Show particle ACFs", + "description": "Show ACFs for individual particles of the selection", + "type": "bool", + }, + { + "attribute": "normalized", + "name": "Normalize", + "description": "Normalize ACF values", + "type": "bool", + }, + { + "attribute": "custom_title", + "name": "Custom title", + "description": "Custom title for the plot", + "type": "str", + }, + ] + + def __init__(self): + super().__init__() + self.acf = None + self.physical_property = "velocity" + self.selection = "all" + self.dim_type = "xyz" + self.centered = False + self.show_running_integral = False + self.show_particle_acfs = False + self.normalized = False + self.custom_title = None + self._setup_plot() + + def _setup_plot(self): + """Setup matplotlib plot""" + self.fig, self.ax = plt.subplots() + (self.plot,) = self.ax.plot([], [], 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"Time (ps)") + self.ax.grid(True, linestyle="--", alpha=0.6) + self._set_title() + self._set_y_label() + + def _set_title(self): + """Set plot title""" + if self.show_running_integral: + title = ( + f"Running integral of {self.physical_property.title()} " + f"ACF of '{self.selection}'" + ) + else: + title = f"{self.physical_property.title()} ACF of '{self.selection}'" + self.ax.set_title(self.custom_title if self.custom_title else title) + + def _set_y_label(self): + """Set plot y label""" + if self.show_running_integral: + self.ax.set_ylabel( + f"Running integral of {self.physical_property.title()} ACF" + ) + else: + self.ax.set_ylabel( + f"{self.physical_property.title()} Autocorrelation Function" + ) + + def _create_acf(self): + """Create acf instance""" + self.acf = SlidingWindowACF( + self.u, + select=self.selection, + physical_property=self.physical_property, + dim_type=self.dim_type, + centered=self.centered, + show_running_integral=self.show_running_integral, + show_particle_acfs=self.show_particle_acfs, + ) + self._set_title() + self._set_y_label() + + def on_post_create(self): + """on_post_create handler""" + self._set_title() + self._set_y_label() + + def on_post_connect(self): + """on_post_connect handler""" + self._create_acf() + + def on_input_change(self, attribute, _old_value, new_value): + """on_input_change handler""" + if attribute == "custom_title": + self._set_title() + elif attribute in ("normalized", "_run_mode"): + pass + else: + self._create_acf() + + def _compute(self, normalized: bool = False, parallel: bool = False): + """Run ACF for the current timesteps window""" + return self.acf.run(normalized=normalized, 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 y2 is not None 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(normalized=self.normalized) + self._update_plot(x, y1, y2) + + def get_parallel_job(self): + """get parallel job handler""" + return delayed(self._compute)(normalized=self.normalized, parallel=True) + + def apply_parallel_results(self, values): + """apply parallel results handler""" + x, y1, y2, (v1, v2, v3, v4, v5, v6) = values + self._update_plot(x, y1, y2) + # update acf state + self.acf.acf_sums = v1 + self.acf.acf_counts = v2 + self.acf.running_sum = v3 + self.acf.running_count = v4 + if v5 is not None: + self.acf.particle_acf_sums = v5 + self.acf.particle_acf_counts = v6 + + +class SlidingWindowACF: + """Sliding Window ACF + + Calculate ACF for a sliding window of frames + + """ + + # pylint: disable=too-many-arguments,too-many-positional-arguments + def __init__( + self, + u: mda.Universe, + physical_property: str = "velocity", + select: str = "all", + dim_type: str = "xyz", + centered: bool = False, + show_running_integral: bool = False, + show_particle_acfs: bool = False, + ): + self.u = u + property_map = { + "velocity": "velocities", + "position": "positions", + "force": "forces", + } + self.physical_property = property_map[physical_property] + self.select = select + self.dim_type = dim_type + self.centered = centered + self.show_running_integral = show_running_integral + self.show_particle_acfs = (not show_running_integral) and show_particle_acfs + self._parse_dim_type() + self.ag = u.select_atoms(self.select) + self.n_atoms = self.ag.atoms.n_atoms + self.n_lags = u.trajectory.buffer_size + self.frame_dt = None + self.running_sum = np.zeros_like( + getattr(self.ag, self.physical_property)[:, self._dim], dtype=np.float64 + ) + self.running_count = 0 + self.acf_sums = np.zeros(self.n_lags) + self.acf_counts = np.zeros(self.n_lags, dtype=int) + if self.show_particle_acfs: + self.particle_acf_sums = np.zeros((self.n_lags, self.n_atoms)) + self.particle_acf_counts = np.zeros((self.n_lags, self.n_atoms), dtype=int) + + def _parse_dim_type(self): + """Sets up the desired dimensionality.""" + keys = { + "x": [0], + "y": [1], + "z": [2], + "xy": [0, 1], + "xz": [0, 2], + "yz": [1, 2], + "xyz": [0, 1, 2], + } + self._dim = keys[self.dim_type.lower()] + + # pylint: disable=too-many-locals + def run(self, normalized: bool = False, parallel: bool = False) -> tuple: + """Run ACF for the current window""" + + n = len(self.u.trajectory) # buffer / window might not be full yet + current = getattr(self.ag, self.physical_property)[:, self._dim] + self.running_sum += current + self.running_count += 1 + mu = self.running_sum / self.running_count + for i in range(n): + lag = n - 1 - i + _ = self.u.trajectory[i] # set trajectory to past frame + previous = getattr(self.ag, self.physical_property) + acf = current * previous[:, self._dim] + if self.centered: + acf = acf - (mu**2) + acf_sum = np.sum(acf, axis=-1) + self.acf_sums[lag] += np.mean(acf_sum) + self.acf_counts[lag] += 1 + if self.show_particle_acfs: + self.particle_acf_sums[lag, :] += acf_sum + self.particle_acf_counts[lag, :] += 1 + + if self.frame_dt is None: + # 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 + self.frame_dt = self.u.trajectory[1].time - self.u.trajectory[0].time + delta_t_values = np.arange(n) * self.frame_dt + avg_acfs = self.acf_sums[:n] / self.acf_counts[:n] + + if normalized: + avg_acfs = avg_acfs / avg_acfs[0] + + acfs_by_particle_lines = None + if self.show_particle_acfs: + acfs_by_particle_array = ( + self.particle_acf_sums[:n, :] / self.particle_acf_counts[:n, :] + ) + if normalized: + acfs_by_particle_array = ( + acfs_by_particle_array / acfs_by_particle_array[0] + ) + acfs_by_particle_lines = np.empty((self.n_atoms, n, 2)) + acfs_by_particle_lines[:, :, 0] = delta_t_values + acfs_by_particle_lines[:, :, 1] = acfs_by_particle_array.T + + if self.show_running_integral: + running_integral = integrate.cumulative_trapezoid( + avg_acfs, + delta_t_values, + initial=0, + ) / len(self._dim) + + return ( + delta_t_values, + running_integral if self.show_running_integral else avg_acfs, + acfs_by_particle_lines, + ( + self.acf_sums, + self.acf_counts, + self.running_sum, + self.running_count, + self.particle_acf_sums if self.show_particle_acfs else None, + self.particle_acf_counts if self.show_particle_acfs else None, + ) + if parallel + else (None,) * 4, + ) diff --git a/mdadash/backend/analyses/msd.py b/mdadash/backend/analyses/msd.py index a9988cc..98eaac1 100644 --- a/mdadash/backend/analyses/msd.py +++ b/mdadash/backend/analyses/msd.py @@ -34,8 +34,8 @@ class MSDAnalysis(WidgetBase): "type": "str", }, { - "attribute": "msd_type", - "name": "MSD type", + "attribute": "dim_type", + "name": "Dimension type", "description": "Desired dimensions to be included in the MSD", "type": "select", "items": [ @@ -48,6 +48,12 @@ class MSDAnalysis(WidgetBase): "z", ], }, + { + "attribute": "show_diffusion_coefficient", + "name": "Show diffusion coefficient", + "description": "Show self-diffusion coefficient calculated from MSD", + "type": "bool", + }, { "attribute": "show_particle_msds", "name": "Show particle MSDs", @@ -72,12 +78,13 @@ def __init__(self): super().__init__() self.msd = None self.selection = "all" - self.msd_type = "xyz" + self.dim_type = "xyz" self.log_scale = False + self.show_diffusion_coefficient = False self.show_particle_msds = False - self.title = "MSD" self.custom_title = None self._setup_plot() + self._set_y_label() def _setup_plot(self): """Setup matplotlib plot""" @@ -87,15 +94,26 @@ def _setup_plot(self): (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.set_xlabel(r"Time (ps)") self.ax.grid(True, linestyle="--", alpha=0.6) self._set_title() + self._set_y_label() 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) + if self.show_diffusion_coefficient: + title = f"Diffusion coefficient of '{self.selection}'" + else: + title = f"MSD of '{self.selection}'" + self.ax.set_title(self.custom_title if self.custom_title else title) + + def _set_y_label(self): + """Set plot y label""" + if self.show_diffusion_coefficient: + self.ax.set_ylabel(r"Diffusion Coefficient (${\AA}^2$/ps)") + else: + self.ax.set_ylabel(r"MSD ($\AA^2$)") def _set_axes_scale(self): """Set axes scale""" @@ -107,15 +125,17 @@ def _create_msd(self): self.msd = SlidingWindowMSD( self.u, select=self.selection, - msd_type=self.msd_type, + dim_type=self.dim_type, + show_diffusion_coefficient=self.show_diffusion_coefficient, show_particle_msds=self.show_particle_msds, ) - self.title = f"MSD of '{self.selection}'" self._set_title() + self._set_y_label() def on_post_create(self): """on_post_create handler""" self._set_title() + self._set_y_label() self._set_axes_scale() def on_post_connect(self): @@ -140,7 +160,7 @@ def _compute(self, parallel: bool = False): 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.lc.set_segments(y2 if y2 is not None else []) self.ax.relim() self.ax.autoscale_view() self.fig.canvas.draw() @@ -162,7 +182,7 @@ def apply_parallel_results(self, values): # update msd state self.msd.msd_sums = v1 self.msd.msd_counts = v2 - if self.show_particle_msds: + if v3 is not None: self.msd.particle_msd_sums = v3 self.msd.particle_msd_counts = v4 @@ -178,14 +198,18 @@ def __init__( self, u: mda.Universe, select: str = "all", - msd_type: str = "xyz", + dim_type: str = "xyz", + show_diffusion_coefficient: bool = False, 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.dim_type = dim_type + self.show_diffusion_coefficient = show_diffusion_coefficient + self.show_particle_msds = ( + not show_diffusion_coefficient + ) and show_particle_msds + self._parse_dim_type() self.ag = u.select_atoms(self.select) self.n_atoms = self.ag.atoms.n_atoms self.n_lags = u.trajectory.buffer_size @@ -197,8 +221,8 @@ def __init__( 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.""" + def _parse_dim_type(self): + """Sets up the desired dimensionality.""" keys = { "x": [0], "y": [1], @@ -208,7 +232,7 @@ def _parse_msd_type(self): "yz": [1, 2], "xyz": [0, 1, 2], } - self._dim = keys[self.msd_type.lower()] + self._dim = keys[self.dim_type.lower()] def run(self, parallel: bool = False) -> tuple: """Run MSD for the current window""" @@ -242,9 +266,14 @@ def run(self, parallel: bool = False) -> tuple: msds_by_particle_lines[:, :, 0] = delta_t_values msds_by_particle_lines[:, :, 1] = msds_by_particle_array.T + if self.show_diffusion_coefficient: + diffusion_coefficient = np.gradient(avg_msds, delta_t_values) / ( + 2 * len(self._dim) + ) + return ( delta_t_values, - avg_msds, + diffusion_coefficient if self.show_diffusion_coefficient else avg_msds, msds_by_particle_lines, ( self.msd_sums, diff --git a/mdadash/backend/tests/data/adk_oplsaa.trr b/mdadash/backend/tests/data/adk_oplsaa.trr new file mode 100644 index 0000000..bde5942 Binary files /dev/null and b/mdadash/backend/tests/data/adk_oplsaa.trr differ diff --git a/mdadash/backend/tests/data/files.py b/mdadash/backend/tests/data/files.py index 8a7d482..587122a 100644 --- a/mdadash/backend/tests/data/files.py +++ b/mdadash/backend/tests/data/files.py @@ -10,6 +10,7 @@ __all__ = [ "TPR", + "TRR", "XTC", ] @@ -18,4 +19,5 @@ data_directory = importlib.resources.files("mdadash") / "backend" / "tests" / "data" TPR = data_directory / "adk_oplsaa.tpr" +TRR = data_directory / "adk_oplsaa.trr" XTC = data_directory / "adk_oplsaa.xtc" diff --git a/mdadash/backend/tests/test_server.py b/mdadash/backend/tests/test_server.py index 61d702c..60be1f0 100644 --- a/mdadash/backend/tests/test_server.py +++ b/mdadash/backend/tests/test_server.py @@ -12,7 +12,7 @@ from mdadash.backend.kernel.core import BufferedTrajectory from mdadash.backend.main import MDADash, app, sio, start_server from mdadash.backend.state.manager import StateManager -from mdadash.backend.tests.data.files import TPR, XTC +from mdadash.backend.tests.data.files import TPR, TRR, XTC from mdadash.backend.widgets.base import WidgetBase, WidgetManager from .utils import ( @@ -54,6 +54,20 @@ def imd_server_fixture(): server.cleanup() +@pytest.fixture(name="imd_server_trr") +def imd_server_fixture_trr(): + u = mda.Universe(TPR, TRR) + server = InThreadIMDServer(u.trajectory) + info = create_default_imdsinfo_v3() + info.velocities = True + info.forces = False + info.box = True + server.set_imdsessioninfo(info) + server.handshake_sequence("localhost", first_frame=True) + yield server + server.cleanup() + + def test_start_server(mocker): # mock the command line params mocker.patch.object( @@ -636,6 +650,69 @@ async def test_widget_run_msd_parallel(_client, imd_server): await disconnect_from_simulation() +async def test_widget_run_msd_diffusion_coefficient(_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_diffusion_coefficient", True), + ] + 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_vacf_serial(_client, imd_server_trr): + uuid = await add_widget("ACF") + await connect_to_simulation(imd_server_trr, step=1, batch_size=2) + inputs = [ + ("physical_property", "velocity"), + ("selection", "resid 1"), + ("custom_title", ""), + ("show_particle_acfs", True), + ("centered", True), + ("normalized", True), + ] + await check_input_changes(uuid, inputs) + await resume_simulation(imd_server_trr, n_frames=5) + assert await sio_event_emitted(sio, "widgets:output", n=1) + await remove_widget(uuid) + await disconnect_from_simulation() + + +async def test_widget_run_vacf_parallel(_client, imd_server_trr): + uuid = await add_widget("ACF") + await connect_to_simulation(imd_server_trr, step=1, batch_size=2) + inputs = [ + ("physical_property", "velocity"), + ("selection", "resid 1"), + ("show_particle_acfs", True), + ("_run_mode", "parallel"), + ] + await check_input_changes(uuid, inputs) + await resume_simulation(imd_server_trr, n_frames=5) + assert await sio_event_emitted(sio, "widgets:output", n=1) + await remove_widget(uuid) + await disconnect_from_simulation() + + +async def test_widget_run_vacf_running_integral(_client, imd_server_trr): + uuid = await add_widget("ACF") + await connect_to_simulation(imd_server_trr, step=1, batch_size=2) + inputs = [ + ("physical_property", "velocity"), + ("selection", "resid 1"), + ("show_running_integral", True), + ] + await check_input_changes(uuid, inputs) + await resume_simulation(imd_server_trr, n_frames=5) + 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 9dd1fc9..8f661ee 100644 --- a/mdadash/backend/tests/utils.py +++ b/mdadash/backend/tests/utils.py @@ -32,7 +32,7 @@ async def sio_event_emitted(_sio, event, n=1, timeout=5.0): args, _ = call if args[0] == event: count += 1 - if count == n: + if count >= n: emitted = True break await asyncio.sleep(0) @@ -73,13 +73,13 @@ async def disconnect_from_simulation(): assert response["status"] == "ok" -async def resume_simulation(imd_server): +async def resume_simulation(imd_server, n_frames=10): sio.emit.reset_mock() # clear emit.await_args_list handler = sio.handlers["/"]["resume_simulations"] response = await run_task_until_done(handler("_sid")) assert response["status"] == "ok" # send the frames needed by imdclient here - imd_server.send_frames(1, 10) + imd_server.send_frames(1, n_frames) async def pause_simulation():