File size: 3,804 Bytes
f12d3b3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | #!/usr/bin/env python
"""Ball-and-stick renders of two PROTAC conformations from a replica trajectory:
the most extended early frame beside the most collapsed late frame. Both frames
are also written out as PDBs for PyMOL / VMD / ChimeraX.
Usage:
python render_3d.py --rep egfrvhl_rank4/chloroform/rep1
"""
import argparse, glob, os
import numpy as np
import mdtraj as md
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Line3DCollection
# CPK element colors and relative sizes
CPK = {"C": ("#4d4d4d", 55), "N": ("#3050f8", 70), "O": ("#ff2d2d", 70),
"S": ("#e6c200", 95), "Cl": ("#1fd01f", 95), "F": ("#6fd0c0", 60),
"H": ("#dddddd", 20)}
def draw(ax, xyz, syms, bonds, title):
xyz = xyz - xyz.mean(0)
segs = [[xyz[i], xyz[j]] for i, j in bonds]
ax.add_collection3d(Line3DCollection(segs, colors="#888888", linewidths=1.6))
for el in set(syms):
m = [k for k, s in enumerate(syms) if s == el]
c, sz = CPK.get(el, ("#cc44cc", 60))
ax.scatter(xyz[m, 0], xyz[m, 1], xyz[m, 2], c=c, s=sz,
edgecolors="white", linewidths=0.4, depthshade=True)
ax.set_title(title, fontsize=12, weight="bold", color="#222", pad=0)
r = np.abs(xyz).max() * 1.05
ax.set_xlim(-r, r); ax.set_ylim(-r, r); ax.set_zlim(-r, r)
ax.set_box_aspect((1, 1, 1)); ax.axis("off")
def main():
p = argparse.ArgumentParser()
p.add_argument("--rep", required=True, help="replica directory, e.g. <name>/<solvent>/rep1")
p.add_argument("--name", default=None, help="run name (default: the run directory's name)")
p.add_argument("--report-ps", type=float, default=10.0, help="DCD frame interval")
a = p.parse_args()
rep = a.rep.rstrip("/")
name = a.name or os.path.basename(os.path.dirname(os.path.dirname(rep)))
solvated = glob.glob(os.path.join(rep, os.pardir, "*_solvated.pdb"))[0]
traj = md.load(os.path.join(rep, f"{name}_prod.dcd"), top=solvated)
top = traj.topology
# the PROTAC is the one non-water residue with many atoms; water and ions are tiny
lig = [at.index for res in top.residues
if not res.is_water and res.n_atoms > 5 for at in res.atoms]
heavy = [i for i in lig if top.atom(i).element.symbol != "H"]
lset = set(heavy)
remap = {g: k for k, g in enumerate(heavy)}
bonds = [(remap[b0.index], remap[b1.index]) for b0, b1 in top.bonds
if b0.index in lset and b1.index in lset]
syms = [top.atom(i).element.symbol for i in heavy]
sub = traj.atom_slice(heavy)
rg = md.compute_rg(sub) # nm
n = len(rg); early = slice(0, n // 3); late = slice(2 * n // 3, n)
fe = np.arange(n)[early][np.argmax(rg[early])] # extended
fc = np.arange(n)[late][np.argmin(rg[late])] # collapsed
xyz_A = sub.xyz * 10.0 # nm -> Angstrom
ns = a.report_ps / 1000.0
fig = plt.figure(figsize=(13, 6.5), dpi=150); fig.patch.set_facecolor("white")
ax1 = fig.add_subplot(121, projection="3d")
draw(ax1, xyz_A[fe], syms, bonds, f"Extended (t={fe*ns:.1f} ns, Rg={rg[fe]:.2f} nm)")
ax2 = fig.add_subplot(122, projection="3d")
draw(ax2, xyz_A[fc], syms, bonds, f"Collapsed (t={fc*ns:.1f} ns, Rg={rg[fc]:.2f} nm)")
fig.suptitle(f"Free-PROTAC conformations from MD - {name} (heavy atoms, H hidden)",
fontsize=13, weight="bold", color="#222", y=0.97)
out = os.path.join(rep, f"{name}_3d.png")
fig.savefig(out, facecolor="white", bbox_inches="tight"); print("wrote", out)
sub[fe].save_pdb(os.path.join(rep, f"{name}_extended.pdb"))
sub[fc].save_pdb(os.path.join(rep, f"{name}_collapsed.pdb"))
print(f"extended frame {fe} (Rg {rg[fe]:.2f}), collapsed frame {fc} (Rg {rg[fc]:.2f})")
if __name__ == "__main__":
main()
|