rama_utils.py

#!/usr/bin/env python3
# Backbone dihedral (phi/psi) computation from a real PDB structure.
# Standard library + numpy only. Downloads from RCSB or reads a bundled file.
import os
import urllib.request
import numpy as np

RCSB_URL = "https://files.rcsb.org/download/{pdb_id}.pdb"


def download_pdb(pdb_id, dest, timeout=30):
    try:
        os.makedirs(os.path.dirname(dest), exist_ok=True)
        with urllib.request.urlopen(RCSB_URL.format(pdb_id=pdb_id.upper()), timeout=timeout) as r:
            open(dest, "w").write(r.read().decode("utf-8"))
        return dest
    except Exception as e:
        print(f"  live download failed ({e}); using bundled copy")
        return None


def parse_backbone(path):
    # Return per-residue backbone atoms {resseq: {"N":xyz,"CA":xyz,"C":xyz, "res":name}}.
    res = {}
    order = []
    with open(path) as fh:
        for line in fh:
            if not line.startswith("ATOM"):
                continue
            atom = line[12:16].strip()
            if atom not in ("N", "CA", "C"):
                continue
            rid = int(line[22:26])
            if rid not in res:
                res[rid] = {"res": line[17:20].strip()}
                order.append(rid)
            res[rid][atom] = np.array([float(line[30:38]), float(line[38:46]), float(line[46:54])])
    return res, order


def dihedral(p0, p1, p2, p3):
    b0, b1, b2 = p0 - p1, p2 - p1, p3 - p2
    b1n = b1 / np.linalg.norm(b1)
    v = b0 - np.dot(b0, b1n) * b1n
    w = b2 - np.dot(b2, b1n) * b1n
    x = np.dot(v, w)
    y = np.dot(np.cross(b1n, v), w)
    return np.degrees(np.arctan2(y, x))


def phi_psi(path):
    res, order = parse_backbone(path)
    out = []
    for i, rid in enumerate(order):
        r = res[rid]
        if not all(k in r for k in ("N", "CA", "C")):
            continue
        phi = psi = None
        if i > 0 and "C" in res[order[i - 1]]:
            phi = dihedral(res[order[i - 1]]["C"], r["N"], r["CA"], r["C"])
        if i < len(order) - 1 and "N" in res[order[i + 1]]:
            psi = dihedral(r["N"], r["CA"], r["C"], res[order[i + 1]]["N"])
        if phi is not None and psi is not None:
            out.append((rid, r["res"], phi, psi))
    return out


if __name__ == "__main__":
    print("residues with phi/psi:", len(phi_psi("data/1crn.pdb")))