Water bridge 2nd Order implementation (#230)

* 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 <[email protected]>

.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: |

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": <AtomGroup or Molecule>}}`'
-                )
-            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,24 +270,29 @@ 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
         ----------
         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(

prolif/interactions/__init__.py

@@ -23,3 +23,4 @@
     XBAcceptor,
     XBDonor,
 )
+from prolif.interactions.water_bridge import WaterBridge

prolif/interactions/base.py

@@ -6,18 +6,21 @@
 """
 
 import warnings
+from abc import abstractmethod
 from itertools import product
 from math import degrees, radians
 
 import numpy as np
 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