| |
| """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 |
|
|
| |
| 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 |
| 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 |
|
|
| 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() |
|
|
