diff --git a/environment.yml b/environment.yml index 08dd07d..fecaecb 100644 --- a/environment.yml +++ b/environment.yml @@ -10,6 +10,7 @@ dependencies: - openff-units - pip - tqdm + - prolif - pyyaml # for testing - coverage diff --git a/pyproject.toml b/pyproject.toml index f430640..1c38e3e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,6 +21,7 @@ dependencies = [ 'spyrmsd', 'rdkit', 'matplotlib', + 'prolif', ] description="Trajectory analysis of free energy calculations." readme="README.md" diff --git a/src/openfe_analysis/prolif.py b/src/openfe_analysis/prolif.py new file mode 100644 index 0000000..09a4cfa --- /dev/null +++ b/src/openfe_analysis/prolif.py @@ -0,0 +1,253 @@ +from __future__ import annotations + +import numpy as np +from typing import Any, Dict, Optional, Sequence, Tuple, Literal +import warnings + +import MDAnalysis as mda + +import prolif as plf + +from .utils.universe_utils import guess_ligand_bonds + + +class ProLIFAnalysis: + """ + ProLIF interaction fingerprint analysis for an OpenFEReader Universe. + + Parameters + ---------- + universe + MDAnalysis Universe containing topology and trajectory. + ligand_ag + mda.AtomGroup representing the ligand. + water_order + Maximum WaterBridge interaction order (water-water interaction). + Only used if "WaterBridge" is tracked. + protein_cutoff + Distance cutoff in angstrom used to define the protein pocket + around the ligand. + water_cutoff + Distance cutoff in angstrom used to define waters considered + around the ligand/protein pocket. + interactions + Which interactions to track: + - None: ProLIF defaults (Hydrophobic, HBDonor, HBAcceptor, + PiStacking, Anionic, Cationic, CationPi, PiCation, VdWContact); + see ProLIF's ``DEFAULT_INTERACTIONS``: + https://github.com/chemosim-lab/ProLIF/blob/6d993eb1b54cd20cc160461dba1ee5e775cb4037/prolif/fingerprint.py#L97 + - "all": every available interaction, including bridged ones + (e.g. WaterBridge) + - Sequence[str]: explicit list like ["VdWContact", "HBDonor"] + guess_bonds + If True, guess bonds for (protein, ligand, water) so ProLIF can + recognize donors/acceptors and bonded hydrogens. + """ + + def __init__( + self, + universe: mda.Universe, + ligand_ag: mda.AtomGroup, + water_order: int = 3, + protein_cutoff: float = 12.0, + water_cutoff: float = 8.0, + interactions: Optional[Sequence[str] | str] = None, + guess_bonds: bool = True, + ) -> None: + self.universe = universe + self.ligand_ag = ligand_ag + self.water_order = water_order + + self.frames: Optional[np.ndarray] = None + self.times: Optional[np.ndarray] = None + self.n_frames: Optional[int] = None + + if guess_bonds: + self._guess_bonds() + + self._setup_selections(protein_cutoff, water_cutoff) + + self.fp = self._build_fingerprint(interactions) + + def _guess_bonds(self) -> None: + """ + Guess bonds for the protein, ligand and water selections in-place so + RDKit/ProLIF can detect donors/acceptors and bonded hydrogens. + """ + # Protein: guess on the full protein so any pocket residue later has bonds + guess_ligand_bonds(self.universe.select_atoms("protein")) + + # Ligand: stable group + guess_ligand_bonds(self.ligand_ag) + + # Water: only if water-mediated interactions are of interest + water = self.universe.select_atoms("water") + if water.n_atoms: + guess_ligand_bonds(water) + + def _setup_selections(self, protein_cutoff: float, water_cutoff: float) -> None: + """ + Build the updating pocket and water selections around the ligand. + """ + self.protein_ag = self.universe.select_atoms( + f"protein and byres around {protein_cutoff} group ligand", + ligand=self.ligand_ag, + updating=True, + ) + self.water_ag = self.universe.select_atoms( + f"water and byres around {water_cutoff} (group ligand or group pocket)", + ligand=self.ligand_ag, + pocket=self.protein_ag, + updating=True, + ) + + def _build_fingerprint( + self, interactions: Optional[Sequence[str] | str] + ) -> plf.Fingerprint: + """ + Resolve the requested interactions and construct the ProLIF Fingerprint. + + Configures WaterBridge parameters when it is requested and waters are + present, and drops WaterBridge (with a warning) when they are not. + """ + available = plf.Fingerprint.list_available(show_bridged=True) + + fp_interactions: Optional[list[str]] + if interactions is None: + fp_interactions = None + + elif interactions == "all": + # ProLIF's "all" excludes bridged interactions (e.g. WaterBridge) + fp_interactions = list(available) + + else: + fp_interactions = list(interactions) + + self._parameters: Optional[dict] = None + if fp_interactions is not None and "WaterBridge" in fp_interactions: + if self.water_ag.n_atoms == 0: + warnings.warn( + "WaterBridge selected but water selection is empty at the initial " + "frame; removing WaterBridge from the requested interactions.", + UserWarning, + stacklevel=3, + ) + fp_interactions = [ + interaction + for interaction in fp_interactions + if interaction != "WaterBridge" + ] + else: + self._parameters = { + "WaterBridge": {"water": self.water_ag, "order": self.water_order} + } + + if not fp_interactions: + return plf.Fingerprint(parameters=self._parameters) + return plf.Fingerprint( + interactions=fp_interactions, + parameters=self._parameters, + ) + + def run( + self, + *, + start: Optional[int] = None, + stop: Optional[int] = None, + step: Optional[int] = None, + residues: Optional[Literal["all"] | Sequence[str | int]] = None, + progress: bool = True, + n_jobs: Optional[int] = None, + parallel_strategy: Optional[Literal["chunk", "queue"]] = None, + converter_kwargs: Optional[Tuple[Dict[str, Any], Dict[str, Any]]] = None, + ) -> "ProLIFAnalysis": + """ + Run the fingerprint calculation over a slice of the trajectory. + + Parameters + ---------- + start, stop, step + Trajectory slicing parameters. + residues + Passed to ProLIF: ``"all"`` to track every residue, or an explicit + sequence of residue identifiers. If None, ProLIF's default is used + and interactions with atoms are identified. + progress + Show progress bar. + n_jobs + Number of workers for parallel execution.). + parallel_strategy + ProLIF parallel strategy. If None, this wrapper sets: + - "chunk" for n_jobs None/1 + - "queue" for n_jobs > 1 + converter_kwargs + Two dicts: (ligand_kwargs, protein_kwargs) forwarded to the MDAnalysis→RDKit + converter. If None, we default to: + - ligand: {"inferrer": None, "implicit_hydrogens": False} (avoid valence issues) + - protein: {"implicit_hydrogens": False} (use topology bonds) + + Returns + ------- + self + Returned for fluent chaining. + """ + # Due to FEReader trajectory only certain strategies work with the format + if parallel_strategy is None: + # avoid ProLIF trying to pickle FEReader/netCDF trajectory to auto-pick strategy + parallel_strategy = "chunk" if (n_jobs is None or n_jobs == 1) else "queue" + + _slice = slice(start, stop, step) + traj = self.universe.trajectory[_slice] + + try: + n_total = len(self.universe.trajectory) + s0, s1, s2 = _slice.indices(n_total) + self.frames = np.arange(s0, s1, s2, dtype=int) + self.n_frames = len(traj) + + if ( + hasattr(self.universe.trajectory, "times") + and self.universe.trajectory.times is not None + ): + self.times = np.asarray(self.universe.trajectory.times)[self.frames] + elif getattr(self.universe.trajectory, "dt", None) is not None: + self.times = self.frames * self.universe.trajectory.dt + else: + self.times = None + except Exception: + self.frames = None + self.times = None + self.n_frames = None + + if converter_kwargs is None: + # Avoid Valence errors + converter_kwargs = ( + {"inferrer": None, "implicit_hydrogens": False}, # ligand + {"implicit_hydrogens": False}, # protein + ) + + self.fp.run( + traj, + self.ligand_ag, + self.protein_ag, + residues=residues, + converter_kwargs=converter_kwargs, + progress=progress, + n_jobs=n_jobs, + parallel_strategy=parallel_strategy, + ) + + return self + + @property + def ifp(self): + """ + Convenience accessor for underlying ProLIF fingerprint results. + """ + return getattr(self.fp, "ifp", None) + + def to_dataframe(self, **kwargs): + """ + Transform fingerprint results to pd.DataFrame. + """ + return self.fp.to_dataframe(**kwargs) \ No newline at end of file diff --git a/src/openfe_analysis/tests/test_prolif.py b/src/openfe_analysis/tests/test_prolif.py new file mode 100644 index 0000000..0bee946 --- /dev/null +++ b/src/openfe_analysis/tests/test_prolif.py @@ -0,0 +1,366 @@ +import MDAnalysis as mda +import numpy as np +import pytest +from rdkit.Chem import Lipinski + +from openfe_analysis.reader import FEReader +from openfe_analysis.prolif import ProLIFAnalysis + + +def test_prolifanalysis_runs_vdwcontact( + simulation_skipped_nc, hybrid_system_skipped_pdb +): + """ + Test for identification of interactions + """ + u = mda.Universe( + hybrid_system_skipped_pdb, simulation_skipped_nc, format=FEReader, index=0 + ) + ligand_ag = u.select_atoms("resname UNK") + + analysis = ProLIFAnalysis( + u, ligand_ag, interactions=["VdWContact"], guess_bonds=True + ) + analysis.run(stop=5, step=1, n_jobs=1, progress=False) + + df = analysis.to_dataframe(dtype=np.uint8) + assert df.shape[0] == 5 + assert hasattr(analysis.fp, "ifp") + assert len(analysis.fp.ifp) == 5 + + assert analysis.ifp is analysis.fp.ifp + assert len(analysis.ifp) == 5 + + # Ensure there is at least one detected interaction across all processed frames + assert sum(len(v) for v in analysis.fp.ifp.values()) > 0 + + +def test_guess_bonds_enables_protein_chemistry( + simulation_skipped_nc, hybrid_system_skipped_pdb +): + """ + Test for protein connectivity + """ + u = mda.Universe( + hybrid_system_skipped_pdb, simulation_skipped_nc, format=FEReader, index=0 + ) + ligand_ag = u.select_atoms("resname UNK") + + analysis = ProLIFAnalysis( + u, ligand_ag, interactions=["VdWContact"], guess_bonds=True + ) + + # pick a residue from the pocket and check it has connectivity in RDKit + u.trajectory[0] + res_atoms = analysis.protein_ag.residues[0].atoms + res_mol = res_atoms.convert_to("RDKIT", implicit_hydrogens=False) + assert res_mol.GetNumBonds() > 0 + + # ensure the protein donors/acceptors exist + prot_mol = analysis.protein_ag.convert_to("RDKIT", implicit_hydrogens=False) + assert Lipinski.NumHDonors(prot_mol) + Lipinski.NumHAcceptors(prot_mol) > 0 + + +def test_prolifanalysis_accepts_all_keyword( + simulation_skipped_nc, hybrid_system_skipped_pdb +): + """ + The string "all" should be accepted as the special keyword for + all available ProLIF interactions. + """ + u = mda.Universe( + hybrid_system_skipped_pdb, simulation_skipped_nc, format=FEReader, index=0 + ) + ligand_ag = u.select_atoms("resname UNK") + + analysis = ProLIFAnalysis(u, ligand_ag, interactions="all", guess_bonds=True) + + assert analysis.fp is not None + + +def test_default_interactions_are_prolif_defaults( + simulation_skipped_nc, hybrid_system_skipped_pdb +): + """interactions=None should track ProLIF's DEFAULT_INTERACTIONS.""" + u = mda.Universe( + hybrid_system_skipped_pdb, simulation_skipped_nc, format=FEReader, index=0 + ) + ligand_ag = u.select_atoms("resname UNK") + + analysis = ProLIFAnalysis(u, ligand_ag, interactions=None) + + expected = { + "Hydrophobic", + "HBDonor", + "HBAcceptor", + "PiStacking", + "Anionic", + "Cationic", + "CationPi", + "PiCation", + "VdWContact", + } + assert set(analysis.fp.interactions) == expected + + +def test_waterbridge_empty_selection_warns_and_skips_parameters( + simulation_skipped_nc, hybrid_system_skipped_pdb, monkeypatch +): + """ + Requesting WaterBridge with an empty water selection should warn + instead of raising, and should not configure WaterBridge parameters. + """ + u = mda.Universe( + hybrid_system_skipped_pdb, simulation_skipped_nc, format=FEReader, index=0 + ) + ligand_ag = u.select_atoms("resname UNK") + + original_select_atoms = u.select_atoms + + def patched_select_atoms(selection, *args, **kwargs): + if selection == "water and byres around 8 (group ligand or group pocket)": + return u.atoms[[]] + return original_select_atoms(selection, *args, **kwargs) + + monkeypatch.setattr(u, "select_atoms", patched_select_atoms) + + with pytest.warns(UserWarning, match="WaterBridge selected"): + analysis = ProLIFAnalysis( + u, + ligand_ag, + interactions=["WaterBridge"], + guess_bonds=True, + ) + + assert analysis._parameters is None + + +def test_waterbridge_with_water_sets_parameters( + simulation_skipped_nc, hybrid_system_skipped_pdb +): + """ + Requesting WaterBridge with waters present should configure the + WaterBridge parameters. + """ + u = mda.Universe( + hybrid_system_skipped_pdb, simulation_skipped_nc, format=FEReader, index=0 + ) + ligand_ag = u.select_atoms("resname UNK") + + analysis = ProLIFAnalysis(u, ligand_ag, interactions=["WaterBridge"]) + + assert analysis._parameters is not None + assert "WaterBridge" in analysis._parameters + + +def test_guess_bonds_false_skips_guessing( + simulation_skipped_nc, hybrid_system_skipped_pdb +): + """guess_bonds=False should skip bond guessing and still build a fingerprint.""" + u = mda.Universe( + hybrid_system_skipped_pdb, simulation_skipped_nc, format=FEReader, index=0 + ) + ligand_ag = u.select_atoms("resname UNK") + + analysis = ProLIFAnalysis( + u, ligand_ag, interactions=["VdWContact"], guess_bonds=False + ) + + assert analysis.fp is not None + + +def test_plot_prolif_lignetwork_builds_ligand_mol_and_delegates(monkeypatch): + """ + plot_prolif_lignetwork builds a ligand molecule when one is not provided + and delegates to the fingerprint's plot_lignetwork. + """ + from openfe_analysis.utils.plotting import plot_prolif_lignetwork + + calls = {} + + class DummyTrajectory: + def __init__(self): + self.last_frame = None + + def __getitem__(self, frame): + self.last_frame = frame + return None + + traj = DummyTrajectory() + ligand_ag = type("AG", (), {"universe": type("U", (), {"trajectory": traj})()})() + + class DummyFP: + ifp = {0: {"dummy": []}} + use_segid = False + + def plot_lignetwork(self, ligand_mol, **kwargs): + calls["plot_lignetwork"] = (ligand_mol, kwargs) + return "fake-view" + + fp = DummyFP() + fake_ligand_mol = object() + + def fake_from_mda(atomgroup, **kwargs): + calls["from_mda"] = (atomgroup, kwargs) + return fake_ligand_mol + + monkeypatch.setattr( + "openfe_analysis.utils.plotting.plf.Molecule.from_mda", + fake_from_mda, + ) + + view = plot_prolif_lignetwork(fp, ligand_ag, frame=0, kind="frame") + + assert view == "fake-view" + assert calls["from_mda"][0] is ligand_ag + assert calls["from_mda"][1]["inferrer"] is None + assert calls["from_mda"][1]["implicit_hydrogens"] is False + assert calls["from_mda"][1]["use_segid"] == fp.use_segid + assert calls["plot_lignetwork"][0] is fake_ligand_mol + assert calls["plot_lignetwork"][1]["frame"] == 0 + assert calls["plot_lignetwork"][1]["kind"] == "frame" + assert traj.last_frame == 0 + + +def test_plot_prolif_3d_builds_mols_and_delegates(monkeypatch): + """plot_prolif_3d builds ligand/protein/water mols and delegates to fp.plot_3d.""" + from openfe_analysis.utils.plotting import plot_prolif_3d + + calls = {} + traj = {0: None} + + def ag(n): + return type( + "AG", + (), + {"n_atoms": n, "universe": type("U", (), {"trajectory": traj})()}, + )() + + ligand_ag, protein_ag, water_ag = ag(10), ag(100), ag(3) + + class DummyFP: + ifp = {0: {"x": []}} + use_segid = False + + def plot_3d(self, lig, prot, **kw): + calls.update(args=(lig, prot), kw=kw) + return "fake-3d" + + fp = DummyFP() + + made = [] + monkeypatch.setattr( + "openfe_analysis.utils.plotting.plf.Molecule.from_mda", + lambda ag, **kw: made.append(ag) or object(), + ) + + assert plot_prolif_3d(fp, ligand_ag, protein_ag, water_ag, frame=0) == "fake-3d" + assert made == [ligand_ag, protein_ag, water_ag] + assert calls["kw"]["frame"] == 0 + + +def test_plot_prolif_functions_raise_without_ifp(): + """Plotting before the fingerprint is run raises a clear RuntimeError.""" + from openfe_analysis.utils.plotting import ( + plot_prolif_3d, + plot_prolif_barcode, + plot_prolif_lignetwork, + ) + + fp = type("FP", (), {"ifp": {}})() + ag = object() + with pytest.raises(RuntimeError, match="No ProLIF fingerprint data"): + plot_prolif_lignetwork(fp, ag) + with pytest.raises(RuntimeError, match="No ProLIF fingerprint data"): + plot_prolif_3d(fp, ag, ag) + with pytest.raises(RuntimeError, match="No ProLIF fingerprint data"): + plot_prolif_barcode(fp) + + +def test_plot_prolif_lignetwork_invalid_frame_raises(): + """kind='frame' with a frame not in the results raises ValueError.""" + from openfe_analysis.utils.plotting import plot_prolif_lignetwork + + fp = type("FP", (), {"ifp": {0: {"x": []}}, "use_segid": False})() + ag = type("AG", (), {"universe": type("U", (), {"trajectory": {0: None}})()})() + with pytest.raises(ValueError, match="not present"): + plot_prolif_lignetwork(fp, ag, frame=99, kind="frame") + + +def test_plot_prolif_lignetwork_auto_picks_first_frame(monkeypatch): + """With no frame given, the first available frame is used.""" + from openfe_analysis.utils.plotting import plot_prolif_lignetwork + + calls = {} + ag = type("AG", (), {"universe": type("U", (), {"trajectory": {5: None}})()})() + + class DummyFP: + ifp = {5: {"x": []}} + use_segid = False + + def plot_lignetwork(self, ligand_mol, **kwargs): + calls["kw"] = kwargs + return "fake-view" + + monkeypatch.setattr( + "openfe_analysis.utils.plotting.plf.Molecule.from_mda", + lambda ag, **kw: object(), + ) + + assert plot_prolif_lignetwork(DummyFP(), ag) == "fake-view" + assert calls["kw"]["frame"] == 5 + + +def test_plot_prolif_barcode_delegates(): + """plot_prolif_barcode delegates to the fingerprint's plot_barcode.""" + from openfe_analysis.utils.plotting import plot_prolif_barcode + + calls = {} + + class DummyFP: + ifp = {0: {"x": []}} + + def plot_barcode(self, **kwargs): + calls["kw"] = kwargs + return "fake-barcode" + + assert plot_prolif_barcode(DummyFP()) == "fake-barcode" + assert calls["kw"]["xlabel"] == "Frame" + + +def test_plot_prolif_3d_invalid_frame_raises(): + """plot_prolif_3d with a frame not in the results raises ValueError.""" + from openfe_analysis.utils.plotting import plot_prolif_3d + + fp = type("FP", (), {"ifp": {0: {"x": []}}, "use_segid": False})() + ag = type("AG", (), {"universe": type("U", (), {"trajectory": {0: None}})()})() + with pytest.raises(ValueError, match="not present"): + plot_prolif_3d(fp, ag, ag, frame=99) + + +def test_plot_prolif_3d_skips_water_when_absent(monkeypatch): + """With water_ag=None, no water molecule is built and water_mol stays None.""" + from openfe_analysis.utils.plotting import plot_prolif_3d + + calls = {} + + def ag(): + return type("AG", (), {"universe": type("U", (), {"trajectory": {0: None}})()})() + + class DummyFP: + ifp = {0: {"x": []}} + use_segid = False + + def plot_3d(self, lig, prot, **kw): + calls["kw"] = kw + return "fake-3d" + + made = [] + monkeypatch.setattr( + "openfe_analysis.utils.plotting.plf.Molecule.from_mda", + lambda a, **kw: made.append(a) or object(), + ) + + assert plot_prolif_3d(DummyFP(), ag(), ag(), water_ag=None, frame=0) == "fake-3d" + assert len(made) == 2 # ligand + protein only + assert calls["kw"]["water_mol"] is None \ No newline at end of file diff --git a/src/openfe_analysis/utils/plotting.py b/src/openfe_analysis/utils/plotting.py index 18f20df..9058377 100644 --- a/src/openfe_analysis/utils/plotting.py +++ b/src/openfe_analysis/utils/plotting.py @@ -1,7 +1,10 @@ # This code is part of OpenFE and is licensed under the MIT license. # For details, see https://github.com/OpenFreeEnergy/openfe +from typing import Literal, Optional + import matplotlib.pyplot as plt import numpy as np +import prolif as plf def plot_2D_rmsd(data: list[np.ndarray] | list[list[float]], vmax: float = 5.0) -> plt.Figure: @@ -137,3 +140,193 @@ def plot_ligand_RMSD(time: list[float], data: list[np.ndarray]) -> plt.Figure: ylabel=r"RMSD ($\AA$)", title="Ligand RMSD", ) + + +def plot_prolif_lignetwork( + fingerprint, + ligand_ag, + *, + ligand_mol=None, + frame: Optional[int] = None, + kind: Literal["aggregate", "frame"] = "frame", + display_all: bool = False, + threshold: float = 0.3, + use_coordinates: bool = True, + flatten_coordinates: bool = True, + kekulize: bool = False, + molsize: int = 35, + rotation: float = 0, + carbon: float = 0.16, + width: str = "100%", + height: str = "500px", + fontsize: int = 20, + show_interaction_data: bool = False, +): + """ + 2D ProLIF ligand-network visualization. + + Parameters + ---------- + fingerprint : prolif.Fingerprint + A fingerprint that has already been run. + ligand_ag : mda.AtomGroup + Ligand atoms used to build the 2D depiction; its universe is advanced + to ``frame`` before rendering. + """ + if not getattr(fingerprint, "ifp", None): + raise RuntimeError( + "No ProLIF fingerprint data found; run the fingerprint first." + ) + + available_frames = list(fingerprint.ifp.keys()) + + if frame is None: + frame = available_frames[0] + + if kind == "frame" and frame not in fingerprint.ifp: + preview = available_frames[:10] + suffix = " ..." if len(available_frames) > 10 else "" + raise ValueError( + f"frame={frame} not present in fingerprint results. " + f"Available frames: {preview}{suffix}" + ) + + ligand_ag.universe.trajectory[frame] + + if ligand_mol is None: + ligand_mol = plf.Molecule.from_mda( + ligand_ag, + inferrer=None, + implicit_hydrogens=False, + use_segid=fingerprint.use_segid, + ) + + return fingerprint.plot_lignetwork( + ligand_mol, + kind=kind, + frame=frame, + display_all=display_all, + threshold=threshold, + use_coordinates=use_coordinates, + flatten_coordinates=flatten_coordinates, + kekulize=kekulize, + molsize=molsize, + rotation=rotation, + carbon=carbon, + width=width, + height=height, + fontsize=fontsize, + show_interaction_data=show_interaction_data, + ) + + +def plot_prolif_barcode( + fingerprint, + *, + figsize: tuple[int, int] = (8, 10), + dpi: int = 100, + interactive: bool = True, + 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, +): + """ + Barcode plot of interactions across frames. + + Parameters + ---------- + fingerprint : prolif.Fingerprint + A fingerprint that has already been run. + + Returns + ------- + matplotlib.figure.Figure + """ + if not getattr(fingerprint, "ifp", None): + raise RuntimeError( + "No ProLIF fingerprint data found; run the fingerprint first." + ) + + return fingerprint.plot_barcode( + 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, + ) + + +def plot_prolif_3d( + fingerprint, + ligand_ag, + protein_ag, + water_ag=None, + *, + ligand_mol=None, + protein_mol=None, + water_mol=None, + frame: int = 0, + size: tuple[int, int] = (650, 600), + display_all: bool = False, + only_interacting: bool = True, + remove_hydrogens: bool | Literal["ligand", "protein", "water"] = True, +): + """ + 3D ProLIF interaction visualization using py3Dmol. + + Parameters + ---------- + fingerprint : prolif.Fingerprint + A fingerprint that has already been run. + ligand_ag, protein_ag : mda.AtomGroup + Ligand and pocket atoms used to build the 3D depiction. + water_ag : mda.AtomGroup, optional + Water atoms for water-mediated interactions; ignored if None/empty. + """ + if not getattr(fingerprint, "ifp", None): + raise RuntimeError( + "No ProLIF fingerprint data found; run the fingerprint first." + ) + + if frame not in fingerprint.ifp: + raise ValueError(f"frame={frame} not present in fingerprint results.") + + ligand_ag.universe.trajectory[frame] + + if ligand_mol is None: + ligand_mol = plf.Molecule.from_mda( + ligand_ag, + inferrer=None, + implicit_hydrogens=False, + use_segid=fingerprint.use_segid, + ) + + if protein_mol is None: + protein_mol = plf.Molecule.from_mda( + protein_ag, + implicit_hydrogens=False, + use_segid=fingerprint.use_segid, + ) + + if water_mol is None and water_ag is not None and water_ag.n_atoms: + water_mol = plf.Molecule.from_mda( + water_ag, + implicit_hydrogens=False, + use_segid=fingerprint.use_segid, + ) + + return fingerprint.plot_3d( + ligand_mol, + protein_mol, + water_mol=water_mol, + frame=frame, + size=size, + display_all=display_all, + only_interacting=only_interacting, + remove_hydrogens=remove_hydrogens, + )