From ca4c9da2935d5686ef8e796944d2128e3d591ae9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?C=C3=A9dric=20Bouysset?= Date: Mon, 28 Aug 2023 12:02:54 +0100 Subject: [PATCH] add to_barcode_plot method --- CHANGELOG.md | 3 +- .../protein-protein_interactions.ipynb | 18 ++- docs/notebooks/quickstart.ipynb | 6 +- docs/source/modules/plotting.rst | 2 + prolif/fingerprint.py | 62 +++++++++ prolif/plotting/barcode.py | 118 ++++++++++-------- tests/plotting/test_barcode.py | 23 +++- 7 files changed, 174 insertions(+), 58 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2d47df8..e3fb033 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/docs/notebooks/protein-protein_interactions.ipynb b/docs/notebooks/protein-protein_interactions.ipynb index 7806991..d79cee6 100644 --- a/docs/notebooks/protein-protein_interactions.ipynb +++ b/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, diff --git a/docs/notebooks/quickstart.ipynb b/docs/notebooks/quickstart.ipynb index 07a7e10..6337aae 100644 --- a/docs/notebooks/quickstart.ipynb +++ b/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()" ] }, { diff --git a/docs/source/modules/plotting.rst b/docs/source/modules/plotting.rst index fe8a9af..990f365 100644 --- a/docs/source/modules/plotting.rst +++ b/docs/source/modules/plotting.rst @@ -2,3 +2,5 @@ Plotting ======== .. automodule:: prolif.plotting.network + +.. automodule:: prolif.plotting.barcode \ No newline at end of file diff --git a/prolif/fingerprint.py b/prolif/fingerprint.py index 86f166d..ff3797d 100644 --- a/prolif/fingerprint.py +++ b/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." + ) diff --git a/prolif/plotting/barcode.py b/prolif/plotting/barcode.py index 63553e4..1154497 100644 --- a/prolif/plotting/barcode.py +++ b/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,15 +41,31 @@ 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] @@ -60,21 +73,18 @@ def _bit_to_color_value(s: pd.Series) -> pd.Series: 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 diff --git a/tests/plotting/test_barcode.py b/tests/plotting/test_barcode.py index b1ae7f4..d9cc531 100644 --- a/tests/plotting/test_barcode.py +++ b/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)