"""
Plot interactions in 3D --- :mod:`prolif.plotting.complex3d`
============================================================
.. versionadded:: 2.0.0
.. autoclass:: Complex3D
:members:
"""
from __future__ import annotations
from contextlib import suppress
from copy import deepcopy
from typing import (
TYPE_CHECKING,
Any,
ClassVar,
Literal,
)
import numpy as np
import py3Dmol
from rdkit import Chem
from rdkit.Geometry import Point3D
from prolif.exceptions import RunRequiredError
from prolif.plotting.utils import metadata_iterator, separated_interaction_colors
from prolif.utils import get_centroid, get_residues_near_ligand, requires
with suppress(ModuleNotFoundError):
from IPython.display import Javascript, display
if TYPE_CHECKING:
from prolif.fingerprint import Fingerprint
from prolif.ifp import IFP
from prolif.molecule import Molecule
from prolif.residue import Residue, ResidueId
[docs]class Complex3D:
"""Creates a py3Dmol plot of interactions.
Parameters
----------
ifp : IFP
The interaction dictionary for a single frame.
lig_mol : Molecule
The ligand molecule to display.
prot_mol : Molecule
The protein molecule to display.
water_mol : Optional[Molecule]
Additional molecule (e.g. waters) to display.
Attributes
----------
COLORS : dict
Dictionnary of colors used in the plot for interactions.
LIGAND_STYLE : dict[str, dict] = {"stick": {"colorscheme": "cyanCarbon"}}
Style object passed to ``3Dmol.js`` for the ligand.
RESIDUES_STYLE : dict[str, dict] = {"stick": {}}
Style object passed to ``3Dmol.js`` for the protein residues involved in
interactions.
PROTEIN_STYLE : dict[str, dict] = {"cartoon": {"style": "edged"}}
Style object passed to ``3Dmol.js`` for the entire protein.
PEPTIDE_STYLE : dict[str, dict] = "cartoon": {"style": "edged", "colorscheme": "cyanCarbon"}
Style object passed to ``3Dmol.js`` for the ligand as a peptide if appropriate.
PEPTIDE_THRESHOLD : int = 2
Ligands with this number of residues or more will be displayed using
``PEPTIDE_STYLE`` in addition to the ``LIGAND_STYLE``.
LIGAND_DISPLAYED_ATOM : dict[str, int]
Which atom should be used to display an atom-to-atom interaction for the ligand.
Refers to the order defined in the SMARTS pattern used in interaction
definition. Interactions not specified here use ``0`` by default.
PROTEIN_DISPLAYED_ATOM : dict[str, int]
Same as :attr:`LIGAND_DISPLAYED_ATOM` for the protein.
LIGAND_RING_INTERACTIONS : set[str]
Which interactions should be displayed using the centroid instead of using
:attr:`LIGAND_DISPLAYED_ATOM` for the ligand.
PROTEIN_RING_INTERACTIONS : set[str]
Which interactions should be displayed using the centroid instead of using
:attr:`PROTEIN_DISPLAYED_ATOM` for the protein.
BRIDGED_INTERACTIONS : dict[str, str]
For bridged-interactions such as WaterBridge. The key is the interaction name,
and the value is the name of the molecule in the metadata indices dictionary.
RESIDUE_HOVER_CALLBACK : str
JavaScript callback executed when hovering a residue involved in an interaction.
INTERACTION_HOVER_CALLBACK : str
JavaScript callback executed when hovering an interaction line.
DISABLE_HOVER_CALLBACK : str
JavaScript callback executed when the hovering event is finished.
.. versionchanged:: 2.1.0
Added ``water_mol`` parameter to the constructor to display waters
involved in WaterBridge interactions. Added ``save_png`` method to save the
current state of the 3D viewer to a PNG. Added ``remove_hydrogens`` parameter
to the ``display`` and ``compare`` methods to remove non-polar hydrogens that
aren't involved in an interaction. Added ``only_interacting`` parameter to the
``display`` and ``compare`` methods to show all protein residues in the
vicinity of the ligand, or only the ones participating in an interaction.
""" # noqa: E501
COLORS: ClassVar[dict[str, str]] = {**separated_interaction_colors}
LIGAND_STYLE: ClassVar[dict] = {"stick": {"colorscheme": "cyanCarbon"}}
RESIDUES_STYLE: ClassVar[dict] = {"stick": {}}
PROTEIN_STYLE: ClassVar[dict] = {"cartoon": {"style": "edged"}}
PEPTIDE_STYLE: ClassVar[dict] = {
"cartoon": {"style": "edged", "colorscheme": "cyanCarbon"},
}
PEPTIDE_THRESHOLD: ClassVar[int] = 5
LIGAND_DISPLAYED_ATOM: ClassVar[dict] = {
"HBDonor": 1,
"XBDonor": 1,
}
PROTEIN_DISPLAYED_ATOM: ClassVar[dict] = {
"HBAcceptor": 1,
"XBAcceptor": 1,
}
RING_SYSTEMS: ClassVar[set[str]] = {
"PiStacking",
"EdgeToFace",
"FaceToFace",
}
LIGAND_RING_INTERACTIONS: ClassVar[set[str]] = {*RING_SYSTEMS, "PiCation"}
PROTEIN_RING_INTERACTIONS: ClassVar[set[str]] = {*RING_SYSTEMS, "CationPi"}
BRIDGED_INTERACTIONS: ClassVar[dict[str, str]] = {"WaterBridge": "water"}
RESIDUE_HOVER_CALLBACK: ClassVar[str] = """
function(atom,viewer) {
if(!atom.label) {
atom.label = viewer.addLabel('%s:'+atom.atom+atom.serial,
{position: atom, backgroundColor: 'mintcream', fontColor:'black'});
}
}"""
INTERACTION_HOVER_CALLBACK: ClassVar[str] = """
function(shape,viewer) {
if(!shape.label) {
shape.label = viewer.addLabel(shape.interaction,
{position: shape, backgroundColor: 'black', fontColor:'white'});
}
}"""
DISABLE_HOVER_CALLBACK: ClassVar[str] = """
function(obj,viewer) {
if(obj.label) {
viewer.removeLabel(obj.label);
delete obj.label;
}
}"""
def __init__(
self,
ifp: IFP,
lig_mol: Molecule,
prot_mol: Molecule,
water_mol: Molecule | None = None,
) -> None:
self.ifp = ifp
self.lig_mol = lig_mol
self.prot_mol = prot_mol
self.water_mol = water_mol
self._view: py3Dmol.view | None = None # type: ignore[no-any-unimported]
[docs] @classmethod
def from_fingerprint(
cls,
fp: Fingerprint,
lig_mol: Molecule,
prot_mol: Molecule,
water_mol: Molecule | None = None,
*,
frame: int,
) -> Complex3D:
"""Creates a py3Dmol plot of interactions.
Parameters
----------
fp : prolif.fingerprint.Fingerprint
The fingerprint object already executed using one of the ``run`` or
``run_from_iterable`` methods.
frame : int
The frame number chosen to select which interactions are going to be
displayed.
lig_mol : Molecule
The ligand molecule to display.
prot_mol : Molecule
The protein molecule to display.
water_mol : Optional[Molecule]
Additional molecule (e.g. waters) to display.
"""
if not hasattr(fp, "ifp"):
raise RunRequiredError(
"Please run the fingerprint analysis before attempting to display"
" results.",
)
ifp = fp.ifp[frame]
return cls(ifp, lig_mol, prot_mol, water_mol)
[docs] @staticmethod
def get_ring_centroid(mol: Molecule, indices: tuple[int, ...]) -> Point3D:
"""Get the centroid of a ring system."""
centroid = mol.xyz[list(indices)].mean(axis=0)
return Point3D(*centroid)
[docs] def display(
self,
size: tuple[int, int] = (650, 600),
display_all: bool = False,
only_interacting: bool = True,
remove_hydrogens: bool | Literal["ligand", "protein", "water"] = True,
) -> Complex3D:
"""Display as a py3Dmol widget view.
Parameters
----------
size: tuple[int, int] = (650, 600)
The size of the py3Dmol widget view.
display_all : bool = False
Display all occurences for a given pair of residues and interaction, or only
the shortest one. Not relevant if ``count=False`` in the ``Fingerprint``
object.
only_interacting : bool = True
Whether to show all protein residues in the vicinity of the ligand, or
only the ones participating in an interaction.
remove_hydrogens: bool | Literal["ligand", "protein", "water"] = True
Whether to remove non-polar hydrogens (unless they are involved in an
interaction).
.. versionchanged:: 2.1.0
Added ``only_interacting=True`` and ``remove_hydrogens=True`` parameters.
Non-polar hydrogen atoms that aren't involved in interactions are now
hidden. Added support for waters involved in WaterBridge interactions.
"""
v = py3Dmol.view(width=size[0], height=size[1], viewergrid=(1, 1), linked=False)
v.removeAllModels()
self._populate_view(
v,
position=(0, 0),
display_all=display_all,
only_interacting=only_interacting,
remove_hydrogens=remove_hydrogens,
)
self._view = v
return self
[docs] def compare(
self,
other: Complex3D,
*,
size: tuple[int, int] = (900, 600),
display_all: bool = False,
linked: bool = True,
color_unique: str | None = "magentaCarbon",
only_interacting: bool = True,
remove_hydrogens: bool | Literal["ligand", "protein", "water"] = True,
) -> Complex3D:
"""Displays the initial complex side-by-side with a second one for easier
comparison.
Parameters
----------
other: Complex3D
Other ``Complex3D`` object to compare to.
size: tuple[int, int] = (900, 600)
The size of the py3Dmol widget view.
display_all : bool = False
Display all occurences for a given pair of residues and interaction, or only
the shortest one. Not relevant if ``count=False`` in the ``Fingerprint``
object.
linked: bool = True
Link mouse interactions (pan, zoom, translate) on both views.
color_unique: str | None = "magentaCarbon",
Which color to use for residues that have interactions that are found in one
complex but not the other. Use ``None`` to disable the color override.
only_interacting : bool = True
Whether to show all protein residues in the vicinity of the ligand, or
only the ones participating in an interaction.
remove_hydrogens: bool | Literal["ligand", "protein", "water"] = True
Whether to remove non-polar hydrogens (unless they are involved in an
interaction).
.. versionadded:: 2.0.1
.. versionchanged:: 2.1.0
Added ``only_interacting=True`` and ``remove_hydrogens=True`` parameters.
Non-polar hydrogen atoms that aren't involved in interactions are now
hidden. Added support for waters involved in WaterBridge interactions.
"""
v = py3Dmol.view(
width=size[0],
height=size[1],
linked=linked,
viewergrid=(1, 2),
)
v.removeAllModels()
# get set of interactions for both poses
interactions1 = {
(resid[1], i)
for resid, interactions in self.ifp.items()
for i in interactions
}
interactions2 = {
(resid[1], i)
for resid, interactions in other.ifp.items()
for i in interactions
}
# get residues with interactions specific to pose 1
highlights = (
{r[0]: color_unique for r in interactions1 - interactions2}
if color_unique
else {}
)
self._populate_view(
v,
position=(0, 0),
display_all=display_all,
colormap=highlights,
only_interacting=only_interacting,
remove_hydrogens=remove_hydrogens,
)
# get residues with interactions specific to pose 2
highlights = (
{r[0]: color_unique for r in interactions2 - interactions1}
if color_unique
else {}
)
other._populate_view(
v,
position=(0, 1),
display_all=display_all,
colormap=highlights,
only_interacting=only_interacting,
remove_hydrogens=remove_hydrogens,
)
self._view = v
return self
def _populate_view( # type: ignore[no-any-unimported] # noqa: PLR0912
self,
view: py3Dmol.view | Complex3D,
position: tuple[int, int] = (0, 0),
display_all: bool = False,
colormap: dict[ResidueId, str] | None = None,
only_interacting: bool = True,
remove_hydrogens: bool | Literal["ligand", "protein", "water"] = True,
) -> None:
if isinstance(view, Complex3D):
# backwards compatibility for when display/compare used to return the view
if view._view is None:
raise ValueError(
"View not initialized, did you call `display`/`compare` first?",
)
v = view._view
else:
v = view
self._colormap = {} if colormap is None else colormap
self._models: dict[ResidueId, int] = {}
self._mid = -1
self._interacting_atoms: dict[str, set[int]] = {
"ligand": set(),
"protein": set(),
"water": set(),
}
# show all interacting residues
for (lresid, presid), interactions in self.ifp.items():
lres = self.lig_mol[lresid]
pres = self.prot_mol[presid]
# set model ids for reusing later
for resid, res, style in [
(lresid, lres, self.LIGAND_STYLE),
(presid, pres, self.RESIDUES_STYLE),
]:
if resid not in self._models:
self._add_residue_to_view(v, position, res, style)
for interaction, metadata_tuple in interactions.items():
# whether to display all interactions or only the one with the shortest
# distance
for metadata in metadata_iterator(metadata_tuple, display_all):
# record indices of atoms interacting
self._interacting_atoms["ligand"].update(
metadata["parent_indices"]["ligand"]
)
self._interacting_atoms["protein"].update(
metadata["parent_indices"]["protein"]
)
if interaction in self.BRIDGED_INTERACTIONS and self.water_mol:
self._interacting_atoms["water"].update(
metadata["parent_indices"][
self.BRIDGED_INTERACTIONS[interaction]
]
)
for wresid in metadata["water_residues"]:
wres = self.water_mol[wresid]
if wresid not in self._models:
self._add_residue_to_view(
v, position, wres, self.RESIDUES_STYLE
)
# show cylinders for WaterBridge
distances = [d for d in metadata if d.startswith("distance_")]
for distlabel in distances:
_, src, dest = distlabel.split("_")
if src == "ligand":
p1 = lres.GetConformer().GetAtomPosition(
metadata["indices"]["ligand"][
self.LIGAND_DISPLAYED_ATOM.get(interaction, 0)
],
)
else:
p1 = (
self.water_mol[src]
.GetConformer()
.GetAtomPosition(metadata["indices"][src][0])
)
if dest == "protein":
p2 = pres.GetConformer().GetAtomPosition(
metadata["indices"]["protein"][
self.PROTEIN_DISPLAYED_ATOM.get(interaction, 0)
],
)
else:
p2 = (
self.water_mol[dest]
.GetConformer()
.GetAtomPosition(metadata["indices"][dest][0])
)
self._add_interaction(
v,
position,
lresid,
interaction,
metadata[distlabel],
p1,
p2,
)
else:
# get coordinates for both points of the interaction
if interaction in self.LIGAND_RING_INTERACTIONS:
p1 = self.get_ring_centroid(
lres, metadata["indices"]["ligand"]
)
else:
p1 = lres.GetConformer().GetAtomPosition(
metadata["indices"]["ligand"][
self.LIGAND_DISPLAYED_ATOM.get(interaction, 0)
],
)
if interaction in self.PROTEIN_RING_INTERACTIONS:
p2 = self.get_ring_centroid(
pres,
metadata["indices"]["protein"],
)
else:
p2 = pres.GetConformer().GetAtomPosition(
metadata["indices"]["protein"][
self.PROTEIN_DISPLAYED_ATOM.get(interaction, 0)
],
)
# add interaction line
dist = metadata.get("distance", float("nan"))
self._add_interaction(
v, position, lresid, interaction, dist, p1, p2
)
# show "protein" residues that are close to the "ligand"
if not only_interacting:
pocket_residues = get_residues_near_ligand(self.lig_mol, self.prot_mol)
pocket_residues = set(pocket_residues).difference(self._models)
for resid in pocket_residues:
res = self.prot_mol[resid]
self._add_residue_to_view(v, position, res, self.RESIDUES_STYLE)
# hide non-polar hydrogens (except if they are involved in an interaction)
if remove_hydrogens:
to_remove = []
if remove_hydrogens in {"ligand", True}:
to_remove.append(("ligand", self.lig_mol))
if remove_hydrogens in {"protein", True}:
to_remove.append(("protein", self.prot_mol))
if remove_hydrogens in {"water", True} and self.water_mol:
to_remove.append(("water", self.water_mol))
for resid in self._models:
for moltype, mol in to_remove:
try:
modelID = self._models[resid]
res = mol[resid]
except KeyError:
continue
model = v.getModel(modelID, viewer=position)
int_atoms = self._interacting_atoms[moltype]
hide = [
a.GetIdx()
for a in res.GetAtoms()
if a.GetAtomicNum() == 1
and a.GetUnsignedProp("mapindex") not in int_atoms
and all(n.GetAtomicNum() in {1, 6} for n in a.GetNeighbors())
]
model.setStyle({"index": hide}, {"stick": {"hidden": True}})
# show protein
mol = Chem.RemoveAllHs(self.prot_mol, sanitize=False)
pdb = Chem.MolToPDBBlock(mol, flavor=0x20 | 0x10)
v.addModel(pdb, "pdb", viewer=position)
model = v.getModel(viewer=position)
model.setStyle({}, self.PROTEIN_STYLE)
# do the same for ligand if multiple residues
if self.lig_mol.n_residues >= self.PEPTIDE_THRESHOLD:
mol = Chem.RemoveAllHs(self.lig_mol, sanitize=False)
pdb = Chem.MolToPDBBlock(mol, flavor=0x20 | 0x10)
v.addModel(pdb, "pdb", viewer=position)
model = v.getModel(viewer=position)
model.setStyle({}, self.PEPTIDE_STYLE)
v.zoomTo({"model": list(self._models.values())}, viewer=position)
def _add_interaction( # type: ignore[no-any-unimported]
self,
v: py3Dmol.view,
position: tuple[int, int],
resid: ResidueId,
interaction: str,
distance: float,
p1: Point3D,
p2: Point3D,
) -> None:
"""Add an interaction line (and hover label) to the view."""
v.addCylinder(
{
"start": {"x": p1.x, "y": p1.y, "z": p1.z},
"end": {"x": p2.x, "y": p2.y, "z": p2.z},
"color": self.COLORS.get(interaction, "grey"),
"radius": 0.15,
"dashed": True,
"fromCap": 1,
"toCap": 1,
},
viewer=position,
)
# add label when hovering the middle of the dashed line by adding a dummy atom
arr = np.array([list(p1), list(p2)], dtype=float)
c = Point3D(*get_centroid(arr))
modelID = self._models[resid]
model = v.getModel(modelID, viewer=position)
interaction_label = f"{interaction}: {distance:.2f}Å"
model.addAtoms(
[
{
"elem": "Z",
"x": c.x,
"y": c.y,
"z": c.z,
"interaction": interaction_label,
},
],
)
model.setStyle(
{"interaction": interaction_label},
{"clicksphere": {"radius": 0.5}},
)
model.setHoverable(
{"interaction": interaction_label},
True,
self.INTERACTION_HOVER_CALLBACK,
self.DISABLE_HOVER_CALLBACK,
)
def _add_residue_to_view( # type: ignore[no-any-unimported]
self,
v: py3Dmol.view,
position: tuple[int, int],
res: Residue,
style: dict,
) -> None:
"""Add a residue to the view."""
self._mid += 1
resid = res.resid
v.addModel(Chem.MolToMolBlock(res), "sdf", viewer=position)
model = v.getModel(viewer=position)
if resid in self._colormap:
resid_style = deepcopy(style)
for key in resid_style:
resid_style[key]["colorscheme"] = self._colormap[resid]
else:
resid_style = style
model.setStyle({}, resid_style)
# add residue label
model.setHoverable(
{},
True,
self.RESIDUE_HOVER_CALLBACK % resid,
self.DISABLE_HOVER_CALLBACK,
)
self._models[resid] = self._mid
[docs] @requires("IPython.display")
def save_png(self) -> None:
"""Saves the current state of the 3D viewer to a PNG. Not available outside of a
notebook.
.. versionadded:: 2.1.0
"""
if self._view is None:
raise ValueError(
"View not initialized, did you call `display`/`compare` first?",
)
uid = self._view.uniqueid
display(
Javascript(f"""
var png = viewer_{uid}.pngURI()
var a = document.createElement('a')
a.href = png
a.download = "prolif-3d.png"
a.click()
a.remove()
"""),
)
def __getattr__(self, name: str) -> Any:
"""Get an attribute from the py3Dmol view."""
if self._view is None:
raise ValueError(
"View not initialized, did you call `display`/`compare` first?",
)
return getattr(self._view, name)
def _repr_html_(self) -> str | None:
if self._view:
return self._view._repr_html_() # type: ignore[no-any-return]
return None