add to_barcode_plot method

CHANGELOG.md

@@ -7,7 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 
 ### Added
-- Added a `Barcode` class for plotting interactions.
+- Added a `Barcode` class for plotting interactions. Added the
+  `Fingerprint.to_barcode_plot` method to generate the plot directly from an FP object.
 - Added a `count` argument in `Fingerprint`. If `count=True`, enumerates all groups of
   atoms that satisfy interaction constraints (instead of stopping at the first one),
   allowing users to generate a count-fingerprint. The `Fingerprint.to_dataframe` method

docs/notebooks/protein-protein_interactions.ipynb

@@ -131,6 +131,22 @@
     "g.loc[g > 0.3]"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Note: Add `%matplotlib ipympl` at the top of the cell for an interactive visualization:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fp.to_barcode_plot()"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -194,7 +210,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.5"
+   "version": "3.9.13"
   },
   "toc": {
    "base_numbering": 1,

docs/notebooks/quickstart.ipynb

@@ -225,7 +225,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Here's a simple example to plot the interactions over time"
+    "Here's a simple example to plot the interactions over time. Add `%matplotlib ipympl` at the top of the cell for an interactive visualization:"
    ]
   },
   {
@@ -234,9 +234,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from prolif.plotting.barcode import Barcode\n",
-    "\n",
-    "Barcode.from_fingerprint(fp).display();"
+    "fp.to_barcode_plot()"
    ]
   },
   {

docs/source/modules/plotting.rst

@@ -2,3 +2,5 @@ Plotting
 ========
 
 .. automodule:: prolif.plotting.network
+
+.. automodule:: prolif.plotting.barcode

prolif/fingerprint.py

@@ -26,6 +26,7 @@
 """
 from collections.abc import Sized
 from functools import wraps
+from typing import Literal, Optional, Tuple
 
 import dill
 import multiprocess as mp
@@ -38,6 +39,7 @@
 from prolif.interactions.base import _BASE_INTERACTIONS, _INTERACTIONS
 from prolif.molecule import Molecule
 from prolif.parallel import MolIterablePool, TrajectoryPool
+from prolif.plotting.utils import IS_NOTEBOOK
 from prolif.utils import (
     get_residues_near_ligand,
     to_bitvectors,
@@ -893,3 +895,63 @@ def to_ligplot(
         raise RunRequiredError(
             "Please run the fingerprint analysis before attempting to display results."
         )
+
+    def to_barcode_plot(
+        self,
+        figsize: Tuple[int, int] = (8, 10),
+        dpi: int = 100,
+        interactive: bool = IS_NOTEBOOK,
+        n_frame_ticks: int = 10,
+        residues_tick_location: Literal["top", "bottom"] = "top",
+        xlabel: str = "Frame",
+        subplots_kwargs: Optional[dict] = None,
+        tight_layout_kwargs: Optional[dict] = None,
+    ):
+        """Generate and display a :class:`~prolif.plotting.barcode.Barcode` plot from
+        a fingerprint object that has been used to run an analysis.
+
+        Parameters
+        ----------
+        figsize: Tuple[int, int] = (8, 10)
+            Size of the matplotlib figure.
+        dpi: int = 100
+            DPI used for the matplotlib figure.
+        interactive: bool
+            Add hover interactivity to the plot (only relevant for notebooks). You may
+            need to add ``%matplotlib notebook`` or ``%matplotlib ipympl`` for it to
+            work as expected.
+        n_frame_ticks: int = 10
+            Number of ticks on the X axis. May use ±1 tick to have them evenly spaced.
+        residues_tick_location: Literal["top", "bottom"] = "top"
+            Whether the Y ticks appear at the top or at the bottom of the series of
+            interactions of each residue.
+        xlabel: str = "Frame"
+            Label displayed for the X axis.
+        subplots_kwargs: Optional[dict] = None
+            Other parameters passed to :func:`matplotlib.pyplot.subplots`.
+        tight_layout_kwargs: Optional[dict] = None
+            Other parameters passed to :meth:`matplotlib.figure.Figure.tight_layout`.
+
+        See Also
+        --------
+        :class:`prolif.plotting.barcode.Barcode`
+
+        .. versionadded:: 2.0.0
+        """
+        if hasattr(self, "ifp"):
+            from prolif.plotting.barcode import Barcode
+
+            barcode = Barcode.from_fingerprint(self)
+            return barcode.display(
+                figsize=figsize,
+                dpi=dpi,
+                interactive=interactive,
+                n_frame_ticks=n_frame_ticks,
+                residues_tick_location=residues_tick_location,
+                xlabel=xlabel,
+                subplots_kwargs=subplots_kwargs,
+                tight_layout_kwargs=tight_layout_kwargs,
+            )
+        raise RunRequiredError(
+            "Please run the fingerprint analysis before attempting to display results."
+        )

prolif/plotting/barcode.py

@@ -8,8 +8,7 @@
    :members:
 
 """
-import builtins
-from typing import ClassVar, Literal, Optional
+from typing import ClassVar, List, Literal, Optional, Tuple
 
 import numpy as np
 import pandas as pd
@@ -18,9 +17,7 @@
 from matplotlib.patches import Patch
 
 from prolif.fingerprint import Fingerprint
-from prolif.plotting.colors import separated_interaction_colors
-
-_IS_NOTEBOOK = hasattr(builtins, "__IPYTHON__")
+from prolif.plotting.utils import IS_NOTEBOOK, separated_interaction_colors
 
 
 class Barcode:
@@ -44,37 +41,50 @@ class Barcode:
     }
 
     def __init__(self, df: pd.DataFrame) -> None:
+        # mapping interaction type (HBond...etc.) to an arbitrary value which
+        # corresponds to a color
         self.color_mapper = {
             interaction: value for value, interaction in enumerate(self.COLORS)
         }
+        # reverse: map value to interaction type
         self.inv_color_mapper = {
             value: interaction for interaction, value in self.color_mapper.items()
         }
+        # matplotlib colormap
         self.cmap = ListedColormap(list(self.COLORS.values()))
 
+        # drop ligand level if single residue
+        # else concatenate ligand with protein and drop ligand if peptide
+        n_ligand_residues = len(np.unique(df.columns.get_level_values("ligand")))
+        if n_ligand_residues == 1:
+            df = df.droplevel("ligand", axis=1)
+        else:
+            df.columns = pd.MultiIndex.from_tuples(
+                [(f"{items[0]}-{items[1]}", items[2]) for items in df.columns],
+                names=["protein", "interaction"],
+            )
+
         def _bit_to_color_value(s: pd.Series) -> pd.Series:
+            """Replaces a bit value with it's corresponding color value"""
             interaction = s.name[-1]
             return s.apply(
                 lambda v: self.color_mapper[interaction]
                 if v
                 else self.color_mapper[None]
             )
 
-        self.df = (
-            df.droplevel("ligand", axis=1)
-            .astype(np.uint8)
-            .T.apply(_bit_to_color_value, axis=1)
-        )
+        self.df = df.astype(np.uint8).T.apply(_bit_to_color_value, axis=1)
 
     @classmethod
     def from_fingerprint(cls, fp: Fingerprint) -> "Barcode":
+        """Creates a barcode object from a fingerprint."""
         return cls(fp.to_dataframe())
 
     def display(
         self,
-        figsize: tuple[int, int] = (8, 10),
+        figsize: Tuple[int, int] = (8, 10),
         dpi: int = 100,
-        interactive: bool = _IS_NOTEBOOK,
+        interactive: bool = IS_NOTEBOOK,
         n_frame_ticks: int = 10,
         residues_tick_location: Literal["top", "bottom"] = "top",
         xlabel: str = "Frame",
@@ -85,7 +95,7 @@ def display(
 
         Parameters
         ----------
-        figsize: tuple[int, int] = (8, 10)
+        figsize: Tuple[int, int] = (8, 10)
             Size of the matplotlib figure.
         dpi: int = 100
             DPI used for the matplotlib figure.
@@ -159,7 +169,7 @@ def display(
         ax.yaxis.set_ticks(indices, residues[indices])
 
         # legend
-        values: list[int] = np.unique(self.df.values).tolist()
+        values: List[int] = np.unique(self.df.values).tolist()
         values.pop(values.index(0))  # remove None color
         legend_colors = {
             self.inv_color_mapper[value]: im.cmap(value) for value in values
@@ -172,42 +182,52 @@ def display(
 
         # interactive
         if interactive:
-            annot = ax.annotate(
-                "",
-                xy=(0, 0),
-                xytext=(5, 5),
-                textcoords="offset points",
-                alpha=0.8,
-                bbox={"boxstyle": "round", "facecolor": "w"},
-                wrap=True,
+            self._add_interaction_callback(
+                fig,
+                ax,
+                im=im,
+                frames=frames,
+                residues=residues,
+                interactions=interactions,
             )
-            annot.set_visible(False)
-
-            def hover(event):
-                if (
-                    event.inaxes is ax
-                    and event.xdata is not None
-                    and event.ydata is not None
-                ):
-                    x, y = round(event.xdata), round(event.ydata)
-                    if self.df.values[y, x]:
-                        annot.xy = (x, y)
-                        frame = frames[x]
-                        interaction = interactions[y]
-                        residue = residues[y]
-                        annot.set_text(f"Frame {frame}: {residue}")
-                        color = im.cmap(self.color_mapper[interaction])
-                        annot.get_bbox_patch().set_facecolor(color)
-                        annot.set_visible(True)
-                        fig.canvas.draw_idle()
-                        return
-                if annot.get_visible():
-                    annot.set_visible(False)
-                    fig.canvas.draw_idle()
-
-            fig.canvas.mpl_connect("motion_notify_event", hover)
-            fig.canvas.header_visible = False
-            fig.canvas.footer_visible = False
 
         fig.tight_layout(**tight_layout_kwargs)
         return ax
+
+    def _add_interaction_callback(self, fig, ax, *, im, frames, residues, interactions):
+        annot = ax.annotate(
+            "",
+            xy=(0, 0),
+            xytext=(5, 5),
+            textcoords="offset points",
+            alpha=0.8,
+            bbox={"boxstyle": "round", "facecolor": "w"},
+            wrap=True,
+        )
+        annot.set_visible(False)
+
+        def hover_callback(event):
+            if (
+                event.inaxes is ax
+                and event.xdata is not None
+                and event.ydata is not None
+            ):
+                x, y = round(event.xdata), round(event.ydata)
+                if self.df.values[y, x]:
+                    annot.xy = (x, y)
+                    frame = frames[x]
+                    interaction = interactions[y]
+                    residue = residues[y]
+                    annot.set_text(f"Frame {frame}: {residue}")
+                    color = im.cmap(self.color_mapper[interaction])
+                    annot.get_bbox_patch().set_facecolor(color)
+                    annot.set_visible(True)
+                    fig.canvas.draw_idle()
+                    return
+            if annot.get_visible():
+                annot.set_visible(False)
+                fig.canvas.draw_idle()
+
+        fig.canvas.mpl_connect("motion_notify_event", hover_callback)
+        fig.canvas.header_visible = False
+        fig.canvas.footer_visible = False

tests/plotting/test_barcode.py

@@ -20,11 +20,15 @@ def fp(self, request: pytest.FixtureRequest) -> plf.Fingerprint:
         return request.getfixturevalue(request.param)
 
     @pytest.fixture(scope="class")
-    def barcode(
+    def fp_run(
         self, u: mda.Universe, ligand_ag, protein_ag, fp: plf.Fingerprint
-    ) -> Barcode:
+    ) -> plf.Fingerprint:
         fp.run(u.trajectory[0:2], ligand_ag, protein_ag)
-        return Barcode.from_fingerprint(fp)
+        return fp
+
+    @pytest.fixture(scope="class")
+    def barcode(self, fp_run: plf.Fingerprint) -> Barcode:
+        return Barcode.from_fingerprint(fp_run)
 
     def test_display(self, barcode: Barcode) -> None:
         ax = barcode.display()
@@ -42,3 +46,16 @@ def test_display_kwargs(self, barcode: Barcode) -> None:
             tight_layout_kwargs={"pad": 2},
         )
         assert isinstance(ax, plt.Axes)
+
+    def test_fp_to_barcode_plot(self, fp_run: plf.Fingerprint) -> None:
+        ax = fp_run.to_barcode_plot(
+            figsize=(1, 2),
+            dpi=200,
+            interactive=True,
+            n_frame_ticks=2,
+            residues_tick_location="bottom",
+            xlabel="foobar",
+            subplots_kwargs={},
+            tight_layout_kwargs={"pad": 2},
+        )
+        assert isinstance(ax, plt.Axes)