Add Contacts within cutoff widget (#60)

Author: PardhavMaradani

CHANGELOG.md

@@ -31,6 +31,7 @@ The rules for this file:
 
 - Added widget for RMSD (PR #54)
 - Added widget for Native Contacts (PR #59)
+- Added widget for Contacts within cutoff (PR #60)
 
 ### Fixed
 

mdadash/backend/analyses/__init__.py

@@ -5,6 +5,7 @@
 from . import (
     acf,
     com_distance,
+    contacts,
     custom_code,
     dssp,
     energies,
@@ -19,6 +20,7 @@
 __all__ = [
     "acf",
     "com_distance",
+    "contacts",
     "custom_code",
     "dssp",
     "energies",

mdadash/backend/analyses/contacts.py

@@ -0,0 +1,232 @@
+"""
+Contacts within a cutoff
+"""
+
+import logging
+from collections import deque
+from typing import ClassVar
+
+import matplotlib.pyplot as plt
+from IPython.display import display
+from joblib import delayed
+from MDAnalysis.lib.distances import capped_distance
+
+from mdadash.backend.widgets.base import WidgetBase
+
+logger = logging.getLogger(__name__)
+
+
+class Contacts(WidgetBase):
+    """
+
+    **Contacts within a cutoff**
+
+    This widget uses `MDAnalysis.lib.distances.capped_distance`_ to calculate number
+    of contacts with a given cutoff between two groups.
+
+    .. _MDAnalysis.lib.distances.capped_distance: https://docs.mdanalysis.org/stable/
+        documentation_pages/lib/distances.html#MDAnalysis.lib.distances.capped_distance
+
+    """
+
+    name = "Contacts"
+    description = "Contacts within a cutoff"
+
+    _inputs: ClassVar = [
+        {
+            "attribute": "_run_frequency",
+            "name": "Run frequency",
+            "description": "The frequency with which the widget is run",
+            "type": "select",
+            "items": [
+                "every-frame",
+                "batch",
+            ],
+        },
+        {
+            "attribute": "_run_mode",
+            "name": "Run mode",
+            "description": "The mode in which the widget is run",
+            "type": "select",
+            "items": [
+                "serial",
+                "parallel",
+            ],
+        },
+        {
+            "attribute": "selection1",
+            "name": "Contacting Group 1",
+            "description": "MDAnalysis selection phrase of first group",
+            "type": "str",
+            "validations": ["required"],
+        },
+        {
+            "attribute": "selection2",
+            "name": "Contacting Group 2",
+            "description": "MDAnalysis selection phrase of second group",
+            "type": "str",
+            "validations": ["required"],
+        },
+        {
+            "attribute": "radius",
+            "name": "Radius",
+            "description": "Radius within which contacts exist",
+            "type": "float",
+        },
+        {
+            "attribute": "custom_title",
+            "name": "Custom title",
+            "description": "Custom title for the plot",
+            "type": "str",
+        },
+        {
+            "attribute": "maxlen",
+            "name": "Max values",
+            "description": "Max values to show in plot",
+            "type": "int",
+        },
+        {
+            "attribute": "x_type",
+            "name": "X-axis",
+            "type": "toggle",
+            "options": [
+                {"name": "Time", "value": "time"},
+                {"name": "Step", "value": "step"},
+            ],
+        },
+    ]
+
+    def __init__(self):
+        super().__init__()
+        self.selection1 = "(resname ASP GLU) and (name OE* OD*)"
+        self.selection2 = "(resname ARG LYS) and (name NH* NZ)"
+        self.radius = 4.5
+        self.ag1 = None
+        self.ag2 = None
+        self.title = "Contacts within cutoff"
+        self.custom_title = None
+        self.default_maxlen = 100
+        self.maxlen = self.default_maxlen
+        self.x_type = "time"
+        self.x_values = None
+        self._setup_plot()
+        self._reset_plot_values()
+
+    def _setup_plot(self):
+        """Setup matplotlib plot"""
+        self.fig, self.ax = plt.subplots()
+        (self.plot,) = self.ax.plot([], [])
+        self.ax.set_ylabel("Number of contacts")
+        self.ax.grid(True)
+        self._set_title()
+
+    def _reset_plot_values(self):
+        """Reset plot values"""
+        self.steps = deque(maxlen=self.maxlen)
+        self.times = deque(maxlen=self.maxlen)
+        self.y_values = deque(maxlen=self.maxlen)
+        self._set_x_values()
+
+    def _set_title(self):
+        """Set plot title"""
+        self.ax.set_title(
+            self.custom_title.replace("\\n", "\n") if self.custom_title else self.title
+        )
+
+    def _set_x_values(self):
+        """Set the values for the x-axis"""
+        if self.x_type == "step":
+            x_label = "Step"
+            self.x_values = self.steps
+        else:
+            x_label = "Time (ps)"
+            self.x_values = self.times
+        self.ax.set_xlabel(x_label)
+
+    def _update_selections(self):
+        """Update atom groups when selection phrases change"""
+        self.ag1 = self.u.select_atoms(self.selection1)
+        self.ag2 = self.u.select_atoms(self.selection2)
+        self.title = f"Contacts between\n'{self.selection1}' and '{self.selection2}'"
+        self._set_title()
+
+    def on_post_create(self):
+        """on_post_create handler"""
+        self._set_title()
+        self._reset_plot_values()
+
+    def on_post_connect(self):
+        """on_post_connect handler"""
+        self._update_selections()
+
+    def on_input_change(self, attribute, _old_value, new_value):
+        """on_input_change handler"""
+        if attribute == "maxlen":
+            if new_value < 0:
+                self.maxlen = self.default_maxlen
+            self._reset_plot_values()
+        elif attribute == "x_type":
+            self._set_x_values()
+        elif attribute == "custom_title":
+            self._set_title()
+        elif attribute in ("selection1", "selection2", "radius"):
+            self._reset_plot_values()
+            self._update_selections()
+
+    def _compute_current_frame(self):
+        """Compute values for current frame"""
+        pairs = capped_distance(
+            self.ag1.positions,
+            self.ag2.positions,
+            max_cutoff=self.radius,
+            box=self.u.dimensions,
+            return_distances=False,
+        )
+        return (
+            self.u.trajectory.ts.data["step"],
+            self.u.trajectory.ts.data["time"],
+            len(pairs),
+        )
+
+    def _compute_batch(self):
+        """Compute values for current batch"""
+        values = []
+        for i in range(self.u.trajectory.buffer_size):
+            _ = self.u.trajectory[i]
+            values.append(self._compute_current_frame())
+        return values
+
+    def _update_plot(self, values):
+        """Append values and update plot"""
+        if isinstance(values, tuple):
+            values = [values]
+        # update plot points
+        for value in values:
+            (steps, times, v) = value
+            self.steps.append(steps)
+            self.times.append(times)
+            self.y_values.append(v)
+        # update plot
+        self.plot.set_data(self.x_values, self.y_values)
+        self.ax.relim()
+        self.ax.autoscale_view()
+        self.fig.canvas.draw()
+        display(self.fig)
+
+    def run_every_frame(self):
+        """every-frame run handler"""
+        self._update_plot(self._compute_current_frame())
+
+    def run_batch(self):
+        """batch run handler"""
+        self._update_plot(self._compute_batch())
+
+    def get_parallel_job(self):
+        """get parallel job handler"""
+        if self._run_frequency == "batch":
+            return delayed(self._compute_batch)()
+        return delayed(self._compute_current_frame)()
+
+    def apply_parallel_results(self, values):
+        """apply parallel results handler"""
+        self._update_plot(values)

mdadash/backend/tests/test_server.py

@@ -661,6 +661,65 @@ async def test_widget_run_native_contacts_parallel_batch(_client, imd_server):
     await disconnect_from_simulation()
 
 
+async def test_widget_run_contacts_serial_every_frame(_client, imd_server):
+    uuid = await add_widget("Contacts")
+    await connect_to_simulation(imd_server)
+    inputs = [
+        ("selection1", "protein"),
+        ("selection2", "resid 1:10"),
+        ("maxlen", -1),
+        ("x_type", "step"),
+        ("custom_title", "Title"),
+    ]
+    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_contacts_serial_batch(_client, imd_server):
+    uuid = await add_widget("Contacts")
+    await connect_to_simulation(imd_server)
+    inputs = [
+        ("_run_frequency", "batch"),
+    ]
+    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_contacts_parallel_every_frame(_client, imd_server):
+    uuid = await add_widget("Contacts")
+    inputs = [
+        ("_run_mode", "parallel"),
+    ]
+    await check_input_changes(uuid, inputs)
+    await connect_to_simulation(imd_server)
+    await resume_simulation(imd_server)
+    timeout = 30 if sys.platform == "win32" else 20
+    assert await sio_event_emitted(sio, "widgets:output", n=1, timeout=timeout)
+    await remove_widget(uuid)
+    await disconnect_from_simulation()
+
+
+async def test_widget_run_contacts_parallel_batch(_client, imd_server):
+    uuid = await add_widget("Contacts")
+    inputs = [
+        ("_run_frequency", "batch"),
+        ("_run_mode", "parallel"),
+    ]
+    await check_input_changes(uuid, inputs)
+    await connect_to_simulation(imd_server)
+    await resume_simulation(imd_server)
+    timeout = 30 if sys.platform == "win32" else 20
+    assert await sio_event_emitted(sio, "widgets:output", n=1, timeout=timeout)
+    await remove_widget(uuid)
+    await disconnect_from_simulation()
+
+
 async def test_widget_run_dssp_serial_every_frame(_client, imd_server):
     await connect_to_simulation(imd_server)
     uuid = await add_widget("DSSP Analysis")