ipqb bootcamp programming assignment

rama.py

from __future__ import division, print_function
import sys
if sys.version_info[0] < 3:
    print("Python 3 required")
    sys.exit(1)
import math, urllib.request, re
try:
    import matplotlib.pyplot as plt
    import seaborn as sns
except ImportError:
    print("Requires matplotlib and seaborn for plotting")
    sys.exit(1)


def dot(a, b):
    return sum([p*q for p, q in zip(a, b)])


def add(a, b):
    return [p+q for p, q in zip(a, b)]


def sub(a, b):
    return [p-q for p, q in zip(a, b)]


def norm(a):
    length = math.sqrt(sum([p**2 for p in a]))
    return [p/length for p in a]


def cross(a, b):
    return [a[1]*b[2] - a[2]*b[1], a[2]*b[0] - a[0]*b[2],
            a[0]*b[1] - a[1]*b[0]]


if __name__ == "__main__":
    try:
        pdb_id = sys.argv[1]
    except IndexError as e:
        print("Requires PDB id")
    regex = r"ATOM\s*\d*\s*([CA|N]*)\s*([A-Z]{3})\s[A-Z]\s*([-0-9]*)\s*([-0-9,\.]*)\s*([-0-9,\.]*)\s*([0-9,\.]*)"
    regex = re.compile(regex, re.MULTILINE)

    url = "http://www.rcsb.org/pdb/files/%s.pdb"%pdb_id
    response = "".join(urllib.request.urlopen(url).read().decode("utf-8").split("\n"))

    atom_types, res_names, pos, xs, ys, zs = zip(*re.findall(regex, response))
    num_residues = len(atom_types)//3
    coords = list(zip(map(float, xs), map(float, ys), map(float, zs)))

    assert len(atom_types)%3 is 0, "Some residues may be missing backbone atoms"

    angles = [[], []]
    for angle_idx, angle_offset in enumerate([2, 0]):
        for res_pos in range(len(atom_types)//3 - 1):
            atom_set = [coords[res_pos*3+angle_offset-3+j] for j in range(4)]
            b1, b2, b3 = (sub(atom_set[i], atom_set[i+1]) for i in range(3))
            n1, n2 = norm(cross(b1, b2)), norm(cross(b2, b3))

            x = dot(n1, n2)
            y = dot(cross(n1, norm(b2)), n2)

            angles[angle_idx].append(math.degrees(math.atan2(y, x)))

    sns.set_style("whitegrid")
    plt.title(pdb_id)
    plt.axes().set_aspect("equal")
    plt.xticks([-180, -90, 0, 90, 180])
    plt.yticks([-180, -90, 0, 90, 180])
    plt.scatter(*angles, color=sns.color_palette()[0])
    plt.xlim((-180, 180))
    plt.ylim((-180, 180))
    sns.axlabel(r"$\phi$", r"$\psi$")
    plt.savefig("%s.png"%pdb_id)

    print("Saved Ramachandran plot to %s.png in the current working directory"%pdb_id)