From 9cb60d2595c39a014a9c878af91c6e42dd274909 Mon Sep 17 00:00:00 2001 From: Valerij Talagayev <82884038+talagayev@users.noreply.github.com> Date: Sun, 29 Dec 2024 16:22:47 +0100 Subject: [PATCH] Water bridge 2nd Order implementation (#230) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Update fingerprint.py added code for 2nd order water bridges * refactor: move to separate class * fix: typing * feat: add implementation for any order * chore: fix CI warning * fix: stable order of water residues in metadata * tests: rearrange waterbridge tests * refactor: introduce base bridged class and interface * tests: test water bridge params and validation --------- Co-authored-by: Cédric Bouysset --- .github/workflows/ci.yml | 2 +- prolif/fingerprint.py | 140 +++++--------- prolif/interactions/__init__.py | 1 + prolif/interactions/base.py | 39 +++- prolif/interactions/water_bridge.py | 278 ++++++++++++++++++++++++++++ pyproject.toml | 3 +- tests/conftest.py | 21 +++ tests/test_fingerprint.py | 29 +-- tests/test_interactions.py | 48 +++++ 9 files changed, 434 insertions(+), 127 deletions(-) create mode 100644 prolif/interactions/water_bridge.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index ff80a16..937e86a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -64,7 +64,7 @@ jobs: python-version: ${{ matrix.python-version }} auto-update-conda: true use-mamba: true - miniforge-variant: Mambaforge + miniforge-variant: Miniforge3 - name: Check conda and pip run: | diff --git a/prolif/fingerprint.py b/prolif/fingerprint.py index e6fd339..edfe1ac 100644 --- a/prolif/fingerprint.py +++ b/prolif/fingerprint.py @@ -29,6 +29,7 @@ import warnings from collections.abc import Sized from functools import wraps +from inspect import signature from typing import Literal, Optional, Tuple, Union import dill @@ -39,7 +40,11 @@ from tqdm.auto import tqdm from prolif.ifp import IFP -from prolif.interactions.base import _BASE_INTERACTIONS, _INTERACTIONS +from prolif.interactions.base import ( + _BASE_INTERACTIONS, + _BRIDGED_INTERACTIONS, + _INTERACTIONS, +) from prolif.molecule import Molecule from prolif.parallel import MolIterablePool, TrajectoryPool from prolif.plotting.utils import IS_NOTEBOOK @@ -217,22 +222,11 @@ def _set_interactions(self, interactions, parameters): parameters = parameters or {} if interactions == "all": interactions = self.list_available() - # prepare water bridge interaction - try: - i = interactions.index("WaterBridge") - except ValueError: - self._water_bridge_parameters = None - else: - interactions.pop(i) - if "WaterBridge" not in parameters: - raise ValueError( - "Must specify settings for the `WaterBridge` interaction: try " - '`parameters={"WaterBridge": {"water": }}`' - ) - self._water_bridge_parameters = parameters.pop("WaterBridge") + # sanity check self._check_valid_interactions(interactions, "interactions") self._check_valid_interactions(parameters, "parameters") + # add interaction methods self.interactions = {} wrapper = all_occurences if self.count else first_occurence @@ -243,10 +237,28 @@ def _set_interactions(self, interactions, parameters): if name in interactions: self.interactions[name] = wrapper(interaction) + # add bridged interactions + self.bridged_interactions = {} + for name, interaction_cls in _BRIDGED_INTERACTIONS.items(): + if name in interactions: + params = parameters.get(name, {}) + if not params: + raise ValueError( + f"Must specify settings for bridged interaction {name!r}: try " + f'`parameters={{"{name}": {{...}}}}`' + ) + sig = signature(interaction_cls.__init__) + if "count" in sig.parameters: + params.setdefault("count", self.count) + interaction = interaction_cls(**params) + setattr(self, name.lower(), interaction) + self.bridged_interactions[name] = interaction + def _check_valid_interactions(self, interactions_iterable, varname): """Raises a NameError if an unknown interaction is given.""" unsafe = set(interactions_iterable) - unknown = unsafe.symmetric_difference(_INTERACTIONS.keys()) & unsafe + known = {*_INTERACTIONS, *_BRIDGED_INTERACTIONS} + unknown = unsafe.difference(known) if unknown: raise NameError( f"Unknown interaction(s) in {varname!r}: {', '.join(unknown)}", @@ -258,7 +270,7 @@ def __repr__(self): # pragma: no cover return f"<{name}: {params} at {id(self):#x}>" @staticmethod - def list_available(show_hidden=False): + def list_available(show_hidden=False, show_bridged=False): """List interactions available to the Fingerprint class. Parameters @@ -266,16 +278,21 @@ def list_available(show_hidden=False): show_hidden : bool Show hidden classes (base classes meant to be inherited from to create custom interactions). + show_bridged : bool + Show bridged interaction classes such as ``WaterBridge``. """ + interactions = sorted(_INTERACTIONS) + if show_bridged: + interactions.extend(sorted(_BRIDGED_INTERACTIONS)) if show_hidden: - return sorted(_BASE_INTERACTIONS) + sorted(_INTERACTIONS) - return sorted(_INTERACTIONS) + return sorted(_BASE_INTERACTIONS) + interactions + return interactions @property def _interactions_list(self): interactions = list(self.interactions) - if self._water_bridge_parameters: - interactions.append("WaterBridge") + if self.bridged_interactions: + interactions.extend(self.bridged_interactions) return interactions @property @@ -497,11 +514,11 @@ def run( # setup defaults converter_kwargs = converter_kwargs or ({}, {}) if ( - self._water_bridge_parameters + self.bridged_interactions and (maxsize := atomgroup_to_mol.cache_parameters()["maxsize"]) and maxsize <= 2 ): - set_converter_cache_size(3) + set_converter_cache_size(2 + len(self.bridged_interactions)) if n_jobs is None: n_jobs = int(os.environ.get("PROLIF_N_JOBS", "0")) or None if residues == "all": @@ -529,12 +546,11 @@ def run( ) self.ifp = ifp - if self._water_bridge_parameters: + if self.bridged_interactions: self._run_bridged_analysis( traj, lig, prot, - **self._water_bridge_parameters, residues=residues, converter_kwargs=converter_kwargs, progress=progress, @@ -712,7 +728,7 @@ def _run_iter_parallel( self.ifp = ifp return self - def _run_bridged_analysis(self, traj, lig, prot, water, **kwargs): + def _run_bridged_analysis(self, traj, lig, prot, **kwargs): """Implementation of the WaterBridge analysis for trajectories. Parameters @@ -724,79 +740,11 @@ def _run_bridged_analysis(self, traj, lig, prot, water, **kwargs): An MDAnalysis AtomGroup for the ligand prot : MDAnalysis.core.groups.AtomGroup An MDAnalysis AtomGroup for the protein (with multiple residues) - water: MDAnalysis.core.groups.AtomGroup - An MDAnalysis AtomGroup for the water molecules """ # noqa: E501 - kwargs.pop("n_jobs", None) - residues = kwargs.pop("residues", None) - fp = Fingerprint( - interactions=["HBDonor", "HBAcceptor"], - parameters=self.parameters, - count=self.count, - ) - - # run analysis twice, once on ligand-water, then on water-prot - ifp_stores: list[dict[int, IFP]] = [ - fp._run_serial(traj, lig, water, residues=None, **kwargs), - fp._run_serial(traj, water, prot, residues=residues, **kwargs), - ] - - # merge results from the 2 runs on matching water residues self.ifp = getattr(self, "ifp", {}) - for (frame, ifp1), ifp2 in zip(ifp_stores[0].items(), ifp_stores[1].values()): - # for each ligand-water interaction in ifp1 - for data1 in ifp1.interactions(): - # for each water-protein interaction in ifp2 where water1 == water2 - for data2 in [ - d2 for d2 in ifp2.interactions() if d2.ligand == data1.protein - ]: - # construct merged metadata - metadata = ( - { - "indices": { - "ligand": data1.metadata["indices"]["ligand"], - "protein": data2.metadata["indices"]["protein"], - "water": tuple( - set().union( - data1.metadata["indices"]["protein"], - data2.metadata["indices"]["ligand"], - ) - ), - }, - "parent_indices": { - "ligand": data1.metadata["parent_indices"]["ligand"], - "protein": data2.metadata["parent_indices"]["protein"], - "water": tuple( - set().union( - data1.metadata["parent_indices"]["protein"], - data2.metadata["parent_indices"]["ligand"], - ) - ), - }, - "water_residue": data1.protein, - "ligand_role": data1.interaction, - "protein_role": ( # invert role - "HBDonor" - if data2.interaction == "HBAcceptor" - else "HBAcceptor" - ), - **{ - f"{key}{suffix}": data.metadata[key] - for suffix, data in [ - ("_ligand_water", data1), - ("_water_protein", data2), - ] - for key in ["distance", "DHA_angle"] - }, - }, - ) - - # store metadata - ifp = self.ifp.setdefault(frame, IFP()) - ifp.setdefault((data1.ligand, data2.protein), {}).setdefault( - "WaterBridge", [] - ).append(metadata) - + for interaction in self.bridged_interactions.values(): + interaction.setup(ifp_store=self.ifp, **kwargs) + interaction.run(traj, lig, prot) return self def to_dataframe( diff --git a/prolif/interactions/__init__.py b/prolif/interactions/__init__.py index e3a3c81..f498687 100644 --- a/prolif/interactions/__init__.py +++ b/prolif/interactions/__init__.py @@ -23,3 +23,4 @@ XBAcceptor, XBDonor, ) +from prolif.interactions.water_bridge import WaterBridge diff --git a/prolif/interactions/base.py b/prolif/interactions/base.py index ecb8486..d083b99 100644 --- a/prolif/interactions/base.py +++ b/prolif/interactions/base.py @@ -6,6 +6,7 @@ """ import warnings +from abc import abstractmethod from itertools import product from math import degrees, radians @@ -13,11 +14,13 @@ from rdkit import Geometry from rdkit.Chem import MolFromSmarts +from prolif.ifp import IFP from prolif.interactions.utils import DISTANCE_FUNCTIONS, get_mapindex from prolif.utils import angle_between_limits, get_centroid, get_ring_normal_vector -_INTERACTIONS = {} -_BASE_INTERACTIONS = {} +_INTERACTIONS: dict[str, type["Interaction"]] = {} +_BRIDGED_INTERACTIONS: dict[str, type["BridgedInteraction"]] = {} +_BASE_INTERACTIONS: dict[str, type["Interaction"]] = {} class Interaction: @@ -111,6 +114,38 @@ def detect(self, ligand, residue): return inverted +class BridgedInteraction: + """Base class for bridged interactions.""" + + def __init_subclass__(cls): + super().__init_subclass__() + name = cls.__name__ + register = _BRIDGED_INTERACTIONS + if name in register: + warnings.warn( + f"The {name!r} interaction has been superseded by a " + f"new class with id {id(cls):#x}", + stacklevel=2, + ) + register[name] = cls + + def __init__(self): + self.ifp = {} + # force empty setup to initialize args with defaults + self.setup() + + def setup(self, ifp_store=None, **kwargs) -> None: + """Setup additional arguments passed at runtime to the fingerprint generator's + ``run`` method. + """ + self.ifp = ifp_store if ifp_store is not None else {} + self.kwargs = kwargs + + @abstractmethod + def run(self, traj, lig, prot) -> dict[int, IFP]: + raise NotImplementedError() + + class Distance(Interaction, is_abstract=True): """Generic class for distance-based interactions diff --git a/prolif/interactions/water_bridge.py b/prolif/interactions/water_bridge.py new file mode 100644 index 0000000..49dbd1c --- /dev/null +++ b/prolif/interactions/water_bridge.py @@ -0,0 +1,278 @@ +"""Module for the WaterBridge interaction implementation.""" + +import itertools as it +from collections import defaultdict +from typing import Iterator, Optional + +import networkx as nx + +from prolif.ifp import IFP, InteractionData +from prolif.interactions.base import BridgedInteraction + + +class WaterBridge(BridgedInteraction): + """Implementation of the WaterBridge analysis. + + Parameters + ---------- + water : MDAnalysis.core.groups.AtomGroup + An MDAnalysis AtomGroup containing the water molecules + order : int + Maximum number of water molecules that can be involved in a water-bridged + interaction. + min_order : int + Minimum number of water molecules that can be involved in a water-bridged + interaction. + hbdonor : Optional[dict] + Parameters for the HBDonor interaction passed to the underlying fingerprint + generator. See :class:`~prolif.interactions.interactions.HBDonor` for more + details. + hbacceptor : Optional[dict] + Same as above for :class:`~prolif.interactions.interactions.HBAcceptor`. + count : bool + Whether to generate a count fingerprint or just a binary one. + + Notes + ----- + This analysis currently only runs in serial. + """ + + def __init__( + self, + water, + order: int = 1, + min_order: int = 1, + hbdonor: Optional[dict] = None, + hbacceptor: Optional[dict] = None, + count: bool = False, + ): + # circular import + from prolif.fingerprint import Fingerprint + + super().__init__() + if order < 1: + raise ValueError("order must be greater than 0") + if min_order > order: + raise ValueError("min_order cannot be greater than order") + self.water = water + self.order = order + self.min_order = min_order + self.water_fp = Fingerprint( + interactions=["HBDonor", "HBAcceptor"], + parameters={"HBDonor": hbdonor or {}, "HBAcceptor": hbacceptor or {}}, + count=count, + ) + + def setup(self, ifp_store=None, **kwargs) -> None: + super().setup(ifp_store=ifp_store, **kwargs) + kwargs.pop("n_jobs", None) + self.residues = kwargs.pop("residues", None) + self.kwargs = kwargs + + def run(self, traj, lig, prot) -> dict[int, IFP]: + """Run the water bridge analysis. + + Parameters + ---------- + traj : MDAnalysis.coordinates.base.ProtoReader or MDAnalysis.coordinates.base.FrameIteratorSliced + Iterate over this Universe trajectory or sliced trajectory object + to extract the frames used for the fingerprint extraction + lig : MDAnalysis.core.groups.AtomGroup + An MDAnalysis AtomGroup for the ligand + prot : MDAnalysis.core.groups.AtomGroup + An MDAnalysis AtomGroup for the protein (with multiple residues) + """ # noqa: E501 + # Run analysis for ligand-water and water-protein interactions + lig_water_ifp: dict[int, IFP] = self.water_fp._run_serial( + traj, lig, self.water, residues=None, **self.kwargs + ) + water_prot_ifp: dict[int, IFP] = self.water_fp._run_serial( + traj, self.water, prot, residues=self.residues, **self.kwargs + ) + if self.order >= 2: + # Run water-water interaction analysis + water_ifp: dict[int, IFP] = self.water_fp._run_serial( + traj, self.water, self.water, residues=None, **self.kwargs + ) + + for frame in lig_water_ifp: + ifp_lw = lig_water_ifp[frame] # Ligand → Water + ifp_wp = water_prot_ifp[frame] # Water → Protein + self.ifp.setdefault(frame, IFP()) + + if self.order >= 2: + ifp_ww = water_ifp[frame] # WaterX -> WaterY + self._any_order(frame, ifp_lw, ifp_ww, ifp_wp) + + else: + self._first_order_only(frame, ifp_lw, ifp_wp) + + return self.ifp + + def _first_order_only( + self, frame: int, ifp_lw: IFP, ifp_wp: IFP + ) -> Iterator[tuple[InteractionData, InteractionData]]: + """Iterates over all relevant combinations of ligand-water-protein""" + # for each ligand-water interaction + for data_lw in ifp_lw.interactions(): + # for each water-protein interaction + for data_wp in ifp_wp.interactions(): + if data_lw.protein == data_wp.ligand: + self._merge_metadata(frame, data_lw, data_wp) + + def _any_order(self, frame: int, ifp_lw: IFP, ifp_ww: IFP, ifp_wp: IFP) -> None: + """Generate results for any order of water-bridge interactions. + + Constructs a graph to represent the water network and iterates over all paths + up to a given length (corresponding to ``order + 1``). + """ + # MultiGraph to allow the same pair of nodes to interact as both HBA and HBD + # and potentially multiple groups of atoms satisfying the constraints. + # Each of these interaction will have its own edge in the graph. + graph = nx.MultiGraph() + nodes = defaultdict(set) + + # construct graph of water interactions + for ifp, role in [(ifp_lw, "ligand"), (ifp_wp, "protein")]: + for data in ifp.interactions(): + graph.add_edge(data.ligand, data.protein, int_data=data) + # assign ligand and protein residue nodes to corresponding role + nodes[role].add(getattr(data, role)) + + # remove mirror interactions before adding water-water to the graph + deduplicated = { + frozenset((ligand_resid, protein_resid)) + for ligand_resid, protein_resid in ifp_ww + if ligand_resid != protein_resid + } + for ligand_resid, protein_resid in deduplicated: + ww_dict = ifp_ww.data[ligand_resid, protein_resid] + for int_name, metadata_tuple in ww_dict.items(): + for metadata in metadata_tuple: + data = InteractionData( + ligand=ligand_resid, + protein=protein_resid, + interaction=int_name, + metadata=metadata, + ) + graph.add_edge( + data.ligand, data.protein, int_data=data, water_only=True + ) + + # find all edge paths of length up to `order + 1` + for source in nodes["ligand"]: + targets = (t for t in nodes["protein"] if nx.has_path(graph, source, t)) + paths = nx.all_simple_edge_paths( + graph, source, targets, cutoff=self.order + 1 + ) + for path in paths: + if len(path) <= self.min_order: + continue + # path is a list[tuple[node_id1, node_id2, deduplication_key]] + # first element in path is lig-water1, last is waterN-prot + data_lw = graph.edges[path[0]]["int_data"] + data_wp = graph.edges[path[-1]]["int_data"] + ww_edges = [graph.edges[p] for p in path[1:-1]] + # only include if strictly passing through water (going back through + # ligand or protein is not a valid higher-order interaction) + if all(e.get("water_only") for e in ww_edges): + # reorder ligand and protein in InteractionData to be contiguous + # i.e. lig-w1, w1-w2, w2-prot instead of lig-w1, w2-w1, w2-prot + data_ww_list = [] + left = data_lw.protein + for e in ww_edges: + d = e["int_data"] + is_sorted = d.ligand == left + data_ww = InteractionData( + ligand=d.ligand if is_sorted else d.protein, + protein=d.protein if is_sorted else d.ligand, + # interaction name is not kept in final metadata + # so no need to invert it + interaction=d.interaction, + # the indices of "ligand" water and "protein" water are + # merged in the final metadata for water mols so no need to + # invert the roles in indices + metadata=d.metadata, + ) + data_ww_list.append(data_ww) + left = data_ww.protein + self._merge_metadata(frame, data_lw, data_wp, *data_ww_list) + + def _merge_metadata( + self, + frame: int, + data_lw: InteractionData, + data_wp: InteractionData, + *data_ww_args: InteractionData, + ) -> None: + """Merge results from all fingerprints on matching water residues""" + # get indices for individual water molecules + water_indices = defaultdict(set) + for data, role in [ + (data_lw, "protein"), + (data_wp, "ligand"), + *it.chain.from_iterable([ + [(data_ww, "ligand"), (data_ww, "protein")] for data_ww in data_ww_args + ]), + ]: + resid = getattr(data, role) + water_indices[str(resid)].update(data.metadata["indices"][role]) + # construct merged metadata + metadata = { + "indices": { + "ligand": data_lw.metadata["indices"]["ligand"], + "protein": data_wp.metadata["indices"]["protein"], + **{key: tuple(indices) for key, indices in water_indices.items()}, + }, + "parent_indices": { + "ligand": data_lw.metadata["parent_indices"]["ligand"], + "protein": data_wp.metadata["parent_indices"]["protein"], + "water": tuple( + set().union( + data_lw.metadata["parent_indices"]["protein"], + data_wp.metadata["parent_indices"]["ligand"], + *it.chain.from_iterable([ + [ + data_ww.metadata["parent_indices"]["ligand"], + data_ww.metadata["parent_indices"]["protein"], + ] + for data_ww in data_ww_args + ]), + ) + ), + }, + "water_residues": tuple( + dict.fromkeys( # uniquify but keep order + [ + data_lw.protein, + *it.chain.from_iterable([ + [data_ww.ligand, data_ww.protein] + for data_ww in data_ww_args + ]), + data_wp.ligand, + ] + ) + ), + "order": len(data_ww_args) + 1, + "ligand_role": data_lw.interaction, + "protein_role": ( # invert role + "HBDonor" if data_wp.interaction == "HBAcceptor" else "HBAcceptor" + ), + **{ + f"{key}{suffix}": data.metadata[key] + for suffix, data in [ + (f"_ligand_{data_lw.protein}", data_lw), + *it.chain.from_iterable([ + [(f"_{data_ww.ligand}_{data_ww.protein}", data_ww)] + for data_ww in data_ww_args + ]), + (f"_{data_wp.ligand}_protein", data_wp), + ] + for key in ["distance", "DHA_angle"] + }, + } + + # store metadata + self.ifp[frame].setdefault((data_lw.ligand, data_wp.protein), {}).setdefault( + "WaterBridge", [] + ).append(metadata) diff --git a/pyproject.toml b/pyproject.toml index 1d47d62..dd62795 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,6 +34,7 @@ dependencies = [ "numpy>=1.13.3", "scipy>=1.3.0", "mdanalysis>=2.2.0", + "networkx>=2.5.0", "tqdm", "multiprocess", "dill", @@ -62,7 +63,7 @@ tests = [ "pytest>=6.1.2", "pytest-cov", "ipython", - "prolif[plots]" + "prolif[plots,water]" ] dev = [ "prolif[build,tests]", diff --git a/tests/conftest.py b/tests/conftest.py index 73ed41d..3f1ca85 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -170,6 +170,27 @@ def metal_false(): return from_mol2("metal_false.mol2") +@pytest.fixture(scope="session") +def water_u(): + top_path = (datapath / "water_m2.pdb").as_posix() + traj_path = (datapath / "water_m2.xtc").as_posix() + return Universe(top_path, traj_path) + + +@pytest.fixture(scope="session") +def water_params(water_u): + ligand = water_u.select_atoms("resname QNB") + protein = water_u.select_atoms( + "protein and byres around 4 group ligand", ligand=ligand + ) + water = water_u.select_atoms( + "resname TIP3 and byres around 6 (group ligand or group pocket)", + ligand=ligand, + pocket=protein, + ) + return ligand, protein, water + + class BaseTestMixinRDKitMol: def test_init(self, mol): assert isinstance(mol, Chem.Mol) diff --git a/tests/test_fingerprint.py b/tests/test_fingerprint.py index f7efc7e..9d0ce6e 100644 --- a/tests/test_fingerprint.py +++ b/tests/test_fingerprint.py @@ -338,38 +338,13 @@ def test_interaction_params(self): fp = Fingerprint() assert fp.hydrophobic.distance == 4.5 - @pytest.fixture(scope="class") - def water_u(self): - top_path = (datapath / "water_m2.pdb").as_posix() - traj_path = (datapath / "water_m2.xtc").as_posix() - return mda.Universe(top_path, traj_path) - - @pytest.fixture(scope="class") - def water_params(self, water_u): - ligand = water_u.select_atoms("resname QNB") - protein = water_u.select_atoms( - "protein and byres around 4 group ligand", ligand=ligand - ) - water = water_u.select_atoms( - "resname TIP3 and byres around 4 group ligand", ligand=ligand - ) - return ligand, protein, water - def test_water_bridge_instance_without_params_raises_error(self): with pytest.raises( - ValueError, match="Must specify settings for the `WaterBridge` interaction" + ValueError, + match="Must specify settings for bridged interaction 'WaterBridge'", ): Fingerprint(["WaterBridge"]) - def test_direct_water_bridge(self, water_u, water_params): - ligand, protein, water = water_params - fp = Fingerprint(["WaterBridge"], parameters={"WaterBridge": {"water": water}}) - fp.run(water_u.trajectory[:1], ligand, protein) - int_data = next(fp.ifp[0].interactions()) - - assert int_data.interaction == "WaterBridge" - assert str(int_data.protein) == "TRP400.X" - def test_mix_water_bridge_and_other_interactions(self, water_u, water_params): ligand, protein, water = water_params fp = Fingerprint( diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 4216c8c..6b71dde 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -347,3 +347,51 @@ def test_edgetoface_phe331(self, ligand_mol, protein_mol): lig, phe331 = ligand_mol[0], protein_mol["PHE331.B"] assert next(fp.edgetoface(lig, phe331)) is True assert next(fp.pistacking(lig, phe331)) is True + + +class TestBridgedInteractions: + @pytest.mark.parametrize( + ("kwargs", "match"), + [ + ({"order": 0}, "order must be greater than 0"), + ({"order": 1, "min_order": 2}, "min_order cannot be greater than order"), + ], + ) + def test_water_bridge_validation(self, water_params, kwargs, match): + *_, water = water_params + with pytest.raises(ValueError, match=match): + Fingerprint( + ["WaterBridge"], + parameters={"WaterBridge": {"water": water, **kwargs}}, + ) + + def test_direct_water_bridge(self, water_u, water_params): + ligand, protein, water = water_params + fp = Fingerprint(["WaterBridge"], parameters={"WaterBridge": {"water": water}}) + fp.run(water_u.trajectory[:1], ligand, protein) + int_data = next(fp.ifp[0].interactions()) + + assert int_data.interaction == "WaterBridge" + assert str(int_data.protein) == "TRP400.X" + + @pytest.mark.parametrize( + ("kwargs", "num_expected"), + [ + ({}, 3), + ({"min_order": 2}, 2), + ], + ) + def test_higher_order_water_bridge(self, water_u, kwargs, num_expected): + ligand = water_u.select_atoms("resname QNB") + pocket = water_u.select_atoms("protein and resid 399:403") + water = water_u.select_atoms("segid WAT and (resid 17 or resid 83)") + fp = Fingerprint( + ["WaterBridge"], + parameters={"WaterBridge": {"water": water, "order": 2, **kwargs}}, + ) + fp.run(water_u.trajectory[:1], ligand, pocket) + all_int_data = list(fp.ifp[0].interactions()) + + assert len(all_int_data) == num_expected + int_data = all_int_data[-1] + assert "distance_TIP383.X_TIP317.X" in int_data.metadata