Skip to content

Commit

Permalink
Merge branch 'update_visualisation' into fix_MSRC_issues
Browse files Browse the repository at this point in the history
  • Loading branch information
@DomFijan
DomFijan committed Oct 8, 2023
2 parents d71a24f + 49cd8e1 commit 33c4e42
Show file tree
Hide file tree
Showing 2 changed files with 142 additions and 106 deletions.
244 changes: 140 additions & 104 deletions ConservedWaterSearch/utils.py
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from __future__ import annotations
import os
import platform
from re import L
import numpy as np
from typing import TYPE_CHECKING

Expand Down Expand Up @@ -239,6 +240,143 @@ def _add_crystal_waters(
os.remove(crystal_waters + ".cif")


def _add_hydrogen_and_bond(wname, Hpos, Hname, resn, resi):
from pymol import cmd

cmd.pseudoatom(
wname,
pos=[Hpos[0], Hpos[1], Hpos[2]],
name=Hname,
resn=resn,
resi=resi,
elem="H",
chain="W",
)
cmd.bond(f"{wname} and name O", f"{wname} and name {Hname}")


def _make_water_objects(water_type, waterO, waterH1, waterH2, output_file):
from pymol import cmd

cntr = {"FCW": 0, "WCW": 0, "HCW": 0}
ind = 0 # initialize index
while ind < len(water_type):
tip, Opos, H1pos, H2pos = (
water_type[ind],
waterO[ind],
waterH1[ind],
waterH2[ind],
)
cntr[tip] += 1
wname = tip + str(cntr[tip])
resis = cmd.identify("all", mode=0)
if len(resis) == 0:
highest_resi = 0
else:
highest_resi = np.max(resis)
if output_file is None or output_file.endswith(".pse"):
resn = "SOL"
elif output_file.endswith(".pdb"):
resn = f"{tip}"
cmd.create(wname, "none", source_state=0, target_state=0)
cmd.pseudoatom(
wname,
pos=[Opos[0], Opos[1], Opos[2]],
name="O",
resn=resn,
resi=highest_resi + 1,
elem="O",
chain="W",
)
if tip == "onlyO":
cmd.show("spheres", wname)
cmd.set("sphere_scale", 0.1, wname)
else:
_add_hydrogen_and_bond(wname, H1pos, "H1", resn, highest_resi + 1)
_add_hydrogen_and_bond(wname, H2pos, "H2", resn, highest_resi + 1)
# check future water type
add_ind = 0
ghind = 0
while ind + add_ind + 1 < len(water_type) and tip != "FCW":
if (
tip == "HCW"
and water_type[ind + add_ind + 1] == "HCW"
and np.array_equal(Opos, waterO[ind + add_ind + 1])
):
_add_hydrogen_and_bond(
wname,
waterH2[ind + add_ind + 1],
f"H{add_ind+3}",
resn,
highest_resi + 1,
)
add_ind += 1
elif (
tip == "WCW"
and water_type[ind + add_ind + 1] == "WCW"
and np.array_equal(Opos, waterO[ind + add_ind + 1])
):
hind = 0
# check for all previos Hs
H1eq = [
True
if np.array_equal(waterH1[ind], waterH1[ind + dind + 1])
else False
for dind in range(add_ind)
] + [
True
if np.array_equal(waterH2[ind], waterH1[ind + dind + 1])
else False
for dind in range(add_ind)
]
if H1eq.count(True) == 0:
_add_hydrogen_and_bond(
wname,
waterH1[ind + add_ind + 1],
f"H{add_ind+ghind+3}",
resn,
highest_resi + 1,
)
hind += 1
H2eq = [
True
if np.array_equal(waterH1[ind], waterH2[ind + dind + 1])
else False
for dind in range(add_ind)
] + [
True
if np.array_equal(waterH2[ind], waterH2[ind + dind + 1])
else False
for dind in range(add_ind)
]
if H2eq.count(True) == 0:
_add_hydrogen_and_bond(
wname,
waterH2[ind + add_ind + 1],
f"H{add_ind+hind+ghind+3}",
resn,
highest_resi + 1,
)
hind += 1
if hind > 0:
add_ind += 1
ghind += hind
else:
break
ind += 1 + add_ind
if output_file is None or output_file.endswith(".pse"):
cmd.show("lines", wname)
if tip == "FCW":
cmd.color("firebrick", wname)
elif tip == "HCW":
cmd.color("skyblue", wname)
if add_ind > 0:
cmd.color("red", "name H1 and " + wname)
elif tip == "WCW":
cmd.color("limegreen", wname)
cmd.show("sticks", wname)


def visualise_pymol(
water_type: list[str],
waterO: list[list[float]],
Expand Down Expand Up @@ -345,107 +483,7 @@ def visualise_pymol(
)
# protein surface
_make_protein_surface_with_ligand()
cntr = {"FCW": 0, "WCW": 0, "HCW": 0}
for tip, Opos, H1pos, H2pos in zip(water_type, waterO, waterH1, waterH2):
cntr[tip] += 1
wname = tip + str(cntr[tip])
resis = cmd.identify("all", mode=0)
if len(resis) == 0:
highest_resi = 0
else:
highest_resi = np.max(resis)
cmd.create(wname, "none", source_state=0, target_state=0)
cmd.pseudoatom(
wname,
pos=[0, 0, 0],
name="O",
resn="HOH",
resi=highest_resi + 1,
elem="O",
chain="W",
)
cmd.pseudoatom(
wname,
pos=[0.957, 0, 0],
name="H1",
resn="HOH",
resi=highest_resi + 1,
elem="H",
chain="W",
)
cmd.pseudoatom(
wname,
pos=[-0.239, 0.927, 0],
name="H2",
resn="HOH",
resi=highest_resi + 1,
elem="H",
chain="W",
)
cmd.bond(f"{wname} and name O", f"{wname} and name H1")
cmd.bond(f"{wname} and name O", f"{wname} and name H2")
cmd.alter_state(
0,
wname,
"(x,y,z)=(" + str(Opos[0]) + "," + str(Opos[1]) + "," + str(Opos[2]) + ")",
)
if tip == "onlyO":
cmd.delete(wname + "and elem H")
else:
indiciesH: list[int] = []
cmd.iterate_state(
-1,
wname + " and elem H",
"indiciesH.append(index)",
space={"indiciesH": indiciesH},
)
cmd.alter_state(
0,
wname + " and index " + str(indiciesH[0]),
"(x,y,z)=("
+ str(H1pos[0])
+ ","
+ str(H1pos[1])
+ ","
+ str(H1pos[2])
+ ")",
)
cmd.alter_state(
0,
wname + " and index " + str(indiciesH[1]),
"(x,y,z)=("
+ str(H2pos[0])
+ ","
+ str(H2pos[1])
+ ","
+ str(H2pos[2])
+ ")",
)
if output_file is None or output_file.endswith(".pse"):
cmd.alter_state(
0,
wname,
"resn='SOL'",
)
cmd.show("lines", wname)
if tip == "FCW":
cmd.color("firebrick", wname)
elif tip == "HCW":
cmd.color("skyblue", wname)
elif tip == "WCW":
cmd.color("limegreen", wname)
elif output_file.endswith(".pdb"):
cmd.alter_state(
0,
wname,
"resn='" + str(tip) + "'",
)

if tip == "onlyO":
cmd.show("spheres", wname)
cmd.set("sphere_scale", 0.1, wname)
else:
cmd.show("sticks", wname)
_make_water_objects(water_type, waterO, waterH1, waterH2, output_file)
waters: str = cmd.get_unused_name("waters")
cmd.select(waters, "SOL in (FCW* or WCW* or HCW*)")
if polar_contacts:
Expand Down Expand Up @@ -544,24 +582,22 @@ def visualise_pymol_from_pdb(
# conserved waters
conserved = cmd.get_unused_name("FCW_")
cmd.select(conserved, "resname FCW")
cmd.show("lines", conserved)
cmd.color("firebrick", conserved)
# half conserved waters
half_conserved = cmd.get_unused_name("HCW_")
cmd.select(half_conserved, "resname HCW")
cmd.show("lines", half_conserved)
cmd.color("skyblue", half_conserved)
# semi conserved waters
semi_conserved = cmd.get_unused_name("WCW_")
cmd.select(semi_conserved, "resname WCW")
cmd.show("lines", semi_conserved)
cmd.color("limegreen", semi_conserved)
# all waters
waters = cmd.get_unused_name("waters_")
cmd.select(
waters,
semi_conserved + " or " + half_conserved + " or " + conserved,
)
cmd.show("sticks", waters)
# Add crystal waters
if crystal_waters is not None:
_add_crystal_waters(
Expand Down
4 changes: 2 additions & 2 deletions ConservedWaterSearch/water_clustering.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -403,9 +403,9 @@ def __scan_clustering_params(
allow_single,
restart: bool = True,
):
found: bool = False if len(Odata) < self.nsnaps else True
for wt in whichH:
wta = [wt]
found: bool = False if len(Odata) < self.nsnaps else True
while found:
found = False
# loop over minsamps- from N(snapshots) to 0.75*N(snapshots)
Expand Down Expand Up @@ -482,7 +482,7 @@ def __scan_clustering_params(
):
found = True
allow_single = True
if len(Odata) < self.nsnaps or not restart:
if len(Odata) < self.nsnaps:
found = False
if (self.debugH == 1 or self.debugO == 1) and self.plotend:
plt = __check_mpl_installation()
Expand Down

0 comments on commit 33c4e42

Please sign in to comment.