File size: 5,049 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 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 | #!/usr/bin/env python
"""One-page summary figure for a single replica of a protac_md.py run: the
radius of gyration, potential energy and density traces of the production run,
with the PROTAC 2D structure beside them when --smiles is given. Writes
<rep>/<name>_summary.png.
Usage:
python plot_results.py --rep egfrvhl_rank4/chloroform/rep1
python plot_results.py --rep egfrvhl_rank4/chloroform/rep1 --smiles "<SMILES>"
"""
import argparse, csv, glob, io, os
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
import mdtraj as md
from rdkit import Chem
from rdkit.Chem import AllChem
from rdkit.Chem.Draw import rdMolDraw2D
from PIL import Image
# Okabe-Ito colorblind-safe hues
BLUE, VERM, GREEN = "#0072B2", "#D55E00", "#009E73"
INK, MUTED, GRID = "#222222", "#666666", "#DDDDDD"
def load_csv(path):
with open(path) as fh:
rows = list(csv.reader(fh))
header = [h.strip().strip('#"').strip() for h in rows[0]]
data = np.array([[float(x) for x in r] for r in rows[1:] if r])
col = {h: data[:, i] for i, h in enumerate(header)}
def find(key):
return next(c for h, c in col.items() if key.lower() in h.lower())
t = find("time") / 1000.0 # ps -> ns
return t - t[0], find("potential"), find("temperature"), find("density")
def load_rg(rep_dir, name):
"""Rg of the PROTAC (nm), same selection as analyze_conformations.py."""
tops = glob.glob(os.path.join(rep_dir, os.pardir, "*_solvated.pdb"))
traj = md.load(os.path.join(rep_dir, f"{name}_prod.dcd"), top=tops[0])
lig = [a.index for r in traj.topology.residues
if not r.is_water and r.n_atoms > 5 for a in r.atoms]
sub = traj.atom_slice(lig)
sub.center_coordinates()
return md.compute_rg(sub)
def mol_image(smiles, px=700):
mol = Chem.MolFromSmiles(smiles)
AllChem.Compute2DCoords(mol)
d = rdMolDraw2D.MolDraw2DCairo(px, px)
d.drawOptions().bondLineWidth = 2
d.drawOptions().padding = 0.06
rdMolDraw2D.PrepareAndDrawMolecule(d, mol)
d.FinishDrawing()
return Image.open(io.BytesIO(d.GetDrawingText()))
def style(ax, title, ylabel):
ax.set_title(title, fontsize=11, color=INK, loc="left", pad=6, weight="bold")
ax.set_ylabel(ylabel, fontsize=9, color=MUTED)
ax.set_xlabel("production time (ns)", fontsize=9, color=MUTED)
ax.grid(True, color=GRID, lw=0.6, alpha=0.8)
ax.set_axisbelow(True)
for s in ("top", "right"):
ax.spines[s].set_visible(False)
for s in ("left", "bottom"):
ax.spines[s].set_color(GRID)
ax.tick_params(colors=MUTED, labelsize=8, length=3)
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("--smiles", default=None, help="draw this structure in the left panel")
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)))
solvent = os.path.basename(os.path.dirname(rep))
t, pe, temp, dens = load_csv(os.path.join(rep, f"{name}_prod.csv"))
rg = load_rg(rep, name)
rg_t = np.arange(len(rg)) * a.report_ps / 1000.0 # frame index -> ns
fig = plt.figure(figsize=(13, 7.5), dpi=150)
fig.patch.set_facecolor("white")
ncol = 3 if a.smiles else 2
gs = GridSpec(2, ncol, figure=fig, hspace=0.42, wspace=0.30,
left=0.06, right=0.98, top=0.90, bottom=0.09)
off = 0
if a.smiles:
axm = fig.add_subplot(gs[:, 0])
axm.axis("off")
axm.imshow(mol_image(a.smiles))
axm.set_title(name, fontsize=11, color=INK, loc="left", weight="bold", pad=6)
off = 1
ax1 = fig.add_subplot(gs[0, off:])
ax1.plot(rg_t, rg, color=BLUE, lw=1.2)
m = rg.mean()
ax1.axhline(m, color=MUTED, lw=1.0, ls="--")
ax1.annotate(f"mean {m:.2f} nm", (rg_t[-1], m), color=MUTED, fontsize=8,
ha="right", va="bottom", xytext=(0, 3), textcoords="offset points")
style(ax1, "Radius of gyration", "Rg (nm)")
ax2 = fig.add_subplot(gs[1, off])
ax2.plot(t, pe / 1000.0, color=VERM, lw=1.0)
style(ax2, "Potential energy", "PE (10^3 kJ/mol)")
ax3 = fig.add_subplot(gs[1, off + 1])
ax3.plot(t, dens, color=GREEN, lw=1.0)
style(ax3, "System density", "density (g/mL)")
fig.suptitle(f"Free-PROTAC MD summary - {name} in {solvent} "
f"({t[-1]:.0f} ns, 300 K / 1 bar)",
fontsize=13, color=INK, x=0.06, ha="left", weight="bold")
out = os.path.join(rep, f"{name}_summary.png")
fig.savefig(out, facecolor="white")
print("wrote", out)
print(f"Rg {rg.mean():.3f}±{rg.std():.3f} nm | T {temp.mean():.1f} K | "
f"rho {dens.mean():.3f} g/mL")
if __name__ == "__main__":
main()
|