| |
| """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 = {"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 |
| |
| 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) |
| n = len(rg); early = slice(0, n // 3); late = slice(2 * n // 3, n) |
| fe = np.arange(n)[early][np.argmax(rg[early])] |
| fc = np.arange(n)[late][np.argmin(rg[late])] |
| xyz_A = sub.xyz * 10.0 |
| 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() |
|
|
