Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Implementation of Water Bridge Interaction Calculation #225

Closed
wants to merge 12 commits into from
Prev Previous commit
Next Next commit
chore: allow WaterBridge in interactions list
  • Loading branch information
cbouy committed Nov 18, 2024
commit 8844af766ae605c2909ec5ac732d49ae93ccf8ba
82 changes: 57 additions & 25 deletions prolif/fingerprint.py
Original file line number Diff line number Diff line change
@@ -207,7 +207,6 @@ def __init__(
"PiCation",
"VdWContact",
]
self.interactions = interactions
self.count = count
self._set_interactions(interactions, parameters)
self.vicinity_cutoff = vicinity_cutoff
@@ -218,6 +217,19 @@ 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:
pass
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")
@@ -475,22 +487,49 @@ def run(
if converter_kwargs is not None and len(converter_kwargs) != 2:
raise ValueError("converter_kwargs must be a list of 2 dicts")

# setup defaults
converter_kwargs = converter_kwargs or ({}, {})
if n_jobs is None:
n_jobs = int(os.environ.get("PROLIF_N_JOBS", 0)) or None
if residues == "all":
residues = list(Molecule.from_mda(prot, **converter_kwargs[1]).residues)
if n_jobs != 1:
return self._run_parallel(

if self.interactions:
if n_jobs == 1:
ifp = self._run_serial(
traj,
lig,
prot,
residues=residues,
converter_kwargs=converter_kwargs,
progress=progress,
)
else:
ifp = self._run_parallel(
traj,
lig,
prot,
residues=residues,
converter_kwargs=converter_kwargs,
progress=progress,
n_jobs=n_jobs,
)
self.ifp = ifp

if water_bridge_params := getattr(self, "_water_bridge_parameters", None):
self.run_bridged_analysis(
traj,
lig,
prot,
**water_bridge_params,
residues=residues,
converter_kwargs=converter_kwargs,
progress=progress,
n_jobs=n_jobs,
)
return self

def _run_serial(self, traj, lig, prot, *, residues, converter_kwargs, progress):
"""Serial implementation for trajectories."""
iterator = tqdm(traj) if progress else traj
ifp = {}
for ts in iterator:
@@ -502,8 +541,7 @@ def run(
residues=residues,
metadata=True,
)
self.ifp = ifp
return self
return ifp

def _run_parallel(
self,
@@ -539,8 +577,7 @@ def _run_parallel(
for ifp_data_chunk in pool.process(args_iterable):
ifp.update(ifp_data_chunk)

self.ifp = ifp
return self
return ifp

def run_from_iterable(
self,
@@ -1070,16 +1107,17 @@ def run_bridged_analysis(self, traj, lig, prot, water, **kwargs):
TODO
"""
kwargs.pop("n_jobs", None)
residues = kwargs.pop("residues", None)
set_converter_cache_size(3)
fp = Fingerprint(
interactions=["HBDonor", "HBAcceptor"], parameters=self.parameters
)

# 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)
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", {})
@@ -1132,15 +1170,9 @@ def run_bridged_analysis(self, traj, lig, prot, water, **kwargs):
)

# 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]}
ifp = self.ifp.setdefault(frame, IFP())
ifp.setdefault((data1.ligand, data2.protein), {}).setdefault(
"WaterBridge", []
).append(metadata)

return self
6 changes: 6 additions & 0 deletions tests/test_fingerprint.py
Original file line number Diff line number Diff line change
@@ -333,3 +333,9 @@ def test_interaction_params(self):
assert fp.hydrophobic.distance == 1.0
fp = Fingerprint()
assert fp.hydrophobic.distance == 4.5

def test_water_bridge_instance_without_params_raises_error(self):
with pytest.raises(
ValueError, match="Must specify settings for the `WaterBridge` interaction"
):
Fingerprint(["WaterBridge"])