Implementation of Water Bridge Interaction Calculation

prolif/fingerprint.py

@@ -34,6 +34,7 @@
 import dill
 import multiprocess as mp
 import numpy as np
+from MDAnalysis.converters.RDKit import set_converter_cache_size
 from rdkit import Chem
 from tqdm.auto import tqdm
 
@@ -210,6 +211,7 @@ def __init__(
         self.count = count
         self._set_interactions(interactions, parameters)
         self.vicinity_cutoff = vicinity_cutoff
+        self.parameters = parameters
 
     def _set_interactions(self, interactions, parameters):
         # read interactions to compute
@@ -1062,3 +1064,83 @@ def plot_3d(
             only_interacting=only_interacting,
             remove_hydrogens=remove_hydrogens,
         )
+
+    def run_bridged_analysis(self, traj, lig, prot, water, **kwargs):
+        """
+        TODO
+        """
+        kwargs.pop("n_jobs", None)
+        set_converter_cache_size(3)
+
+        # run analysis twice, once on ligand-water, then on water-prot
+        ifp_stores: list[dict[int, IFP]] = []
+        for pair in [(lig, water), (water, prot)]:
+            fp = Fingerprint(
+                interactions=["HBDonor", "HBAcceptor"], parameters=self.parameters
+            )
+            fp.run(traj, *pair, n_jobs=1, **kwargs)
+            ifp_stores.append(fp.ifp)
+
+        # 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
+                    if frame not in self.ifp:
+                        ifp = self.ifp[frame] = IFP()
+                    ifp = self.ifp[frame]
+                    if int_data := ifp.get((data1.ligand, data2.protein)):
+                        if "WaterBridge" in int_data:
+                            int_data["WaterBridge"].append(metadata)
+                        else:
+                            int_data["WaterBridge"] = [metadata]
+                    else:
+                        ifp[data1.ligand, data2.protein] = {"WaterBridge": [metadata]}
+
+        return self