Ramachandran plot (plotting the torsion angles of protein backbone)

Question

I would like to have some insights on the code I created. It would be really insightful for me to get some feedback about the code, what I should do differently, if there is an easier, more elegant way, etc. Thanks!

Code:

#!/usr/bin/python3

import math
import sys
import os

import matplotlib.pyplot as plt
import numpy as np
from Bio import PDB
from matplotlib import colors

"""
The preferences were calculated from the following artice:
Lovell et al. Structure validation by Calpha geometry: phi,psi and Cbeta deviation. 2003
DOI: 10.1002/prot.10286
"""

# General variable for the background preferences
rama_preferences = {
    "General": {
        "file": "pref_general.data",
        "cmap": colors.ListedColormap(['#FFFFFF', '#B3E8FF', '#7FD9FF']),
        "bounds": [0, 0.0005, 0.02, 1],
    },
    "GLY": {
        "file": "pref_glycine.data",
        "cmap": colors.ListedColormap(['#FFFFFF', '#FFE8C5', '#FFCC7F']),
        "bounds": [0, 0.002, 0.02, 1],
    },
    "PRO": {
        "file": "pref_proline.data",
        "cmap": colors.ListedColormap(['#FFFFFF', '#D0FFC5', '#7FFF8C']),
        "bounds": [0, 0.002, 0.02, 1],
    },
    "PRE-PRO": {
        "file": "pref_preproline.data",
        "cmap": colors.ListedColormap(['#FFFFFF', '#B3E8FF', '#7FD9FF']),
        "bounds": [0, 0.002, 0.02, 1],
    }
}

# Read in the expected torsion angles
__location__ = os.path.realpath(os.path.join(os.getcwd(), os.path.dirname(__file__)))
rama_pref_values = {}
for key, val in rama_preferences.items():
    rama_pref_values[key] = np.full((360, 360), 0, dtype=np.float64)
    with open(os.path.join(__location__, val["file"])) as fn:
        for line in fn:
            if not line.startswith("#"):
                # Preference file has values for every second position only
                rama_pref_values[key][int(float(line.split()[1])) + 180][int(float(line.split()[0])) + 180] = float(
                    line.split()[2])
                rama_pref_values[key][int(float(line.split()[1])) + 179][int(float(line.split()[0])) + 179] = float(
                    line.split()[2])
                rama_pref_values[key][int(float(line.split()[1])) + 179][int(float(line.split()[0])) + 180] = float(
                    line.split()[2])
                rama_pref_values[key][int(float(line.split()[1])) + 180][int(float(line.split()[0])) + 179] = float(
                    line.split()[2])

normals = {}
outliers = {}
for key, val in rama_preferences.items():
    normals[key] = {"x": [], "y": []}
    outliers[key] = {"x": [], "y": []}

# Calculate the torsion angle of the inputs
for inp in sys.argv[1:]:
    if not os.path.isfile(inp):
        print("{} not found!".format(inp))
        continue
    structure = PDB.PDBParser().get_structure('input_structure', inp)
    for model in structure:
        for chain in model:
            polypeptides = PDB.PPBuilder().build_peptides(chain)
            for poly_index, poly in enumerate(polypeptides):
                phi_psi = poly.get_phi_psi_list()
                for res_index, residue in enumerate(poly):
                    res_name = "{}".format(residue.resname)
                    res_num = residue.id[1]
                    phi, psi = phi_psi[res_index]
                    if phi and psi:
                        aa_type = ""
                        if str(poly[res_index + 1].resname) == "PRO":
                            aa_type = "PRE-PRO"
                        elif res_name == "PRO":
                            aa_type = "PRO"
                        elif res_name == "GLY":
                            aa_type = "GLY"
                        else:
                            aa_type = "General"
                        if rama_pref_values[aa_type][int(math.degrees(psi)) + 180][int(math.degrees(phi)) + 180] < \
                                rama_preferences[aa_type]["bounds"][1]:
                            print("{} {} {} {}{} is an outlier".format(inp, model, chain, res_name, res_num))
                            outliers[aa_type]["x"].append(math.degrees(phi))
                            outliers[aa_type]["y"].append(math.degrees(psi))
                        else:
                            normals[aa_type]["x"].append(math.degrees(phi))
                            normals[aa_type]["y"].append(math.degrees(psi))

# Generate the plots
for idx, (key, val) in enumerate(sorted(rama_preferences.items(), key=lambda x: x[0].lower())):
    plt.subplot(2, 2, idx + 1)
    plt.title(key)
    plt.imshow(rama_pref_values[key], cmap=rama_preferences[key]["cmap"],
               norm=colors.BoundaryNorm(rama_preferences[key]["bounds"], rama_preferences[key]["cmap"].N),
               extent=(-180, 180, 180, -180))
    plt.scatter(normals[key]["x"], normals[key]["y"])
    plt.scatter(outliers[key]["x"], outliers[key]["y"], color="red")
    plt.xlim([-180, 180])
    plt.ylim([-180, 180])
    plt.plot([-180, 180], [0, 0], color="black")
    plt.plot([0, 0], [-180, 180], color="black")
    plt.locator_params(axis='x', nbins=7)
    plt.xlabel(r'$\phi$')
    plt.ylabel(r'$\psi$')
    plt.grid()

plt.tight_layout()
# plt.savefig("asd.png", dpi=300)
plt.show()

Sample output:

Answer 1

Strongly-typed preference structures

Avoid using a dictionary for rama_preferences. We aren't in JavaScript: Python has actual classes, and there are light-weight options such as @dataclass. That last option in particular supports good type hinting, which will lend your program better structure and testability.

The code directly after Read in the expected torsion angles can likely turn into a class method.

Reserved names

"Dunder" (double-underscore) names are reserved for special purposes and should not be used for general variables, so __location__ needs to be renamed, probably just to location.

Functions

All of your code is in global scope. Consider moving it into functions for better testability and more meaningful stack traces.

Shared axes

Consider sharing your axis ticks, labels and title for both horizontal and vertical. Vertical (psi) will only have these on the extreme left, and horizontal (phi) will only have these on the extreme bottom.

Refer to sharex, sharey : Axes, optional or get_shared_y_axes in the documentation.