demo.py

#!/usr/bin/env python3
# Real-data demo: validate a protein structure with a Ramachandran plot.
# Computes backbone phi/psi dihedral angles for every residue of a REAL PDB
# structure (1CRN, crambin) and plots them against the sterically allowed
# regions. Downloads live from RCSB; falls back to the bundled real copy.
import os
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
from rama_utils import download_pdb, phi_psi

PDB_ID = "1CRN"
BUNDLED = "data/1crn.pdb"

# Approximate favoured (phi,psi) regions, drawn as guide boxes.
REGIONS = {
    "right-handed alpha": dict(xy=(-160, -70), w=140, h=90, color="#4C72B0"),
    "beta sheet":         dict(xy=(-180, 90),  w=130, h=90, color="#55A868"),
    "left-handed alpha":  dict(xy=(20, 10),    w=70,  h=70, color="#C44E52"),
}


def classify(phi, psi):
    for name, r in REGIONS.items():
        x, y = r["xy"]
        if x <= phi <= x + r["w"] and y <= psi <= y + r["h"]:
            return name
    return "other/outlier"


def get_structure():
    live = download_pdb(PDB_ID, "data/_live_%s.pdb" % PDB_ID)
    return live or BUNDLED


def main():
    os.makedirs("figures", exist_ok=True)
    os.makedirs("results", exist_ok=True)

    data = phi_psi(get_structure())
    resid = [d[0] for d in data]; resn = [d[1] for d in data]
    phi = np.array([d[2] for d in data]); psi = np.array([d[3] for d in data])
    regions = [classify(p, s) for p, s in zip(phi, psi)]

    fig, ax = plt.subplots(1, 2, figsize=(13, 6), gridspec_kw={"width_ratios": [1.4, 1]})

    # Panel 1: the Ramachandran plot.
    a = ax[0]
    for name, r in REGIONS.items():
        a.add_patch(Rectangle(r["xy"], r["w"], r["h"], color=r["color"], alpha=0.15, zorder=0))
    is_gly = np.array([n == "GLY" for n in resn])
    is_pro = np.array([n == "PRO" for n in resn])
    gen = ~(is_gly | is_pro)
    a.scatter(phi[gen], psi[gen], s=35, color="#333333", label="general", zorder=3)
    a.scatter(phi[is_gly], psi[is_gly], s=45, color="#DD8452", marker="^", label="glycine", zorder=3)
    a.scatter(phi[is_pro], psi[is_pro], s=45, color="#8172B3", marker="s", label="proline", zorder=3)
    a.axhline(0, color="grey", lw=0.5); a.axvline(0, color="grey", lw=0.5)
    a.set_xlim(-180, 180); a.set_ylim(-180, 180)
    a.set_xticks(range(-180, 181, 90)); a.set_yticks(range(-180, 181, 90))
    a.set_xlabel("phi (degrees)"); a.set_ylabel("psi (degrees)")
    a.set_title(f"Ramachandran plot — {PDB_ID} (crambin)"); a.legend(fontsize=8, loc="upper right")

    # Panel 2: residues per region.
    a = ax[1]
    order = ["right-handed alpha", "beta sheet", "left-handed alpha", "other/outlier"]
    counts = [regions.count(k) for k in order]
    colors = [REGIONS.get(k, {}).get("color", "#B0B0B0") for k in order]
    a.barh(order[::-1], counts[::-1], color=colors[::-1])
    a.set_xlabel("number of residues")
    a.set_title("Backbone conformations by region")
    for i, c in enumerate(counts[::-1]):
        a.text(c + 0.2, i, str(c), va="center", fontsize=9)

    fig.suptitle("Ramachandran structure validation (real PDB data)", fontsize=15)
    fig.tight_layout(rect=[0, 0, 1, 0.95])
    fig.savefig("figures/demo.png", dpi=120)

    pd.DataFrame({"residue_number": resid, "residue": resn,
                  "phi": phi, "psi": psi, "region": regions}).to_csv(
        "results/summary.csv", index=False)

    outliers = sum(1 for r in regions if r == "other/outlier")
    print(f"Residues with phi/psi: {len(data)}")
    print(f"Alpha: {regions.count('right-handed alpha')}, "
          f"Beta: {regions.count('beta sheet')}, "
          f"Left-handed: {regions.count('left-handed alpha')}, Other: {outliers}")
    print("Wrote figures/demo.png and results/summary.csv")


if __name__ == "__main__":
    main()