Project 4 - Visualizing Drug Binding Sites

import pyrosetta
from pyrosetta import *
from pyrosetta.toolbox import *

pyrosetta.init(extra_options="-ignore_unrecognized_res false -load_PDB_components true -mute all")
print("PyRosetta ready!")

import urllib.request
from pyrosetta import pose_from_pdb

# Download COX-2 + Ibuprofen
url = "https://files.rcsb.org/download/4PH9.pdb"
urllib.request.urlretrieve(url, "4PH9.pdb")
pose = pose_from_pdb("4PH9.pdb")

print(f"Full structure: {pose.total_residue()} residues")

# Find all residues: identify IBP (Ibuprofen) and chain A protein
ligand_position = None
keep_positions  = []

for i in range(1, pose.total_residue() + 1):
    res   = pose.residue(i)
    chain = pose.pdb_info().chain(i)

    if res.name() == "pdb_IBP":
        ligand_position = i
        keep_positions.append(i)
        print(f"Found Ibuprofen at Rosetta position {i}")
    elif chain == "A" and res.is_protein():
        keep_positions.append(i)

# Build clean pose: chain A protein + Ibuprofen only
clean_pose = pose.clone()
delete_positions = [i for i in range(1, pose.total_residue() + 1)
                    if i not in keep_positions]
for i in reversed(delete_positions):
    clean_pose.delete_residue_slow(i)

print(f"Clean structure: {clean_pose.total_residue()} residues")

# Find Ibuprofen in the clean pose
for i in range(1, clean_pose.total_residue() + 1):
    if not clean_pose.residue(i).is_protein():
        ligand_pos_clean = i
        print(f"Ibuprofen in clean pose: position {ligand_pos_clean}")
        break

import numpy as np

def get_residue_centroid(pose, res_num):
    """Calculate the geometric centre of all atoms in a residue."""
    res    = pose.residue(res_num)
    coords = [np.array([res.atom(j).xyz().x,
                        res.atom(j).xyz().y,
                        res.atom(j).xyz().z])
              for j in range(1, res.natoms() + 1)]
    return np.mean(coords, axis=0)

def min_distance_to_ligand(pose, protein_res, ligand_res):
    """
    Calculate the minimum atom-atom distance between a protein residue
    and the drug molecule (minimum is more accurate than centroid-centroid).
    """
    lig  = pose.residue(ligand_res)
    prot = pose.residue(protein_res)
    min_dist = float("inf")

    for j in range(1, lig.natoms() + 1):
        lig_atom = np.array([lig.atom(j).xyz().x,
                             lig.atom(j).xyz().y,
                             lig.atom(j).xyz().z])
        for k in range(1, prot.natoms() + 1):
            prot_atom = np.array([prot.atom(k).xyz().x,
                                  prot.atom(k).xyz().y,
                                  prot.atom(k).xyz().z])
            dist = np.linalg.norm(lig_atom - prot_atom)
            if dist < min_dist:
                min_dist = dist

    return min_dist

# Distance cutoff for "contact"
CUTOFF = 5.0  # Ångströms

contact_residues = []

for i in range(1, clean_pose.total_residue() + 1):
    if i == ligand_pos_clean:
        continue  # skip the drug itself
    if not clean_pose.residue(i).is_protein():
        continue

    dist = min_distance_to_ligand(clean_pose, i, ligand_pos_clean)
    if dist <= CUTOFF:
        res       = clean_pose.residue(i)
        pdb_num   = clean_pose.pdb_info().number(i)
        aa_name   = res.name1()     # single-letter code
        aa_full   = res.name3()     # three-letter code
        contact_residues.append({
            "Position (Rosetta)": i,
            "Position (PDB)":     pdb_num,
            "AA (1-letter)":      aa_name,
            "AA (3-letter)":      aa_full,
            "Min Distance (Å)":   round(dist, 2)
        })

print(f"Found {len(contact_residues)} residues within {CUTOFF} Å of Ibuprofen:\n")
import pandas as pd
df_contacts = pd.DataFrame(contact_residues).sort_values("Min Distance (Å)")
print(df_contacts.to_string(index=False))

# Standard amino acid chemical classifications
HYDROPHOBIC = set("AVILMFYW")   # Ala, Val, Ile, Leu, Met, Phe, Tyr, Trp
POLAR       = set("STNQ")        # Ser, Thr, Asn, Gln
CHARGED_POS = set("RKH")         # Arg, Lys, His (positive at physiological pH)
CHARGED_NEG = set("DE")          # Asp, Glu (negative at physiological pH)
SPECIAL     = set("CGP")         # Cys, Gly, Pro (unique properties)

def classify_aa(aa1):
    if aa1 in HYDROPHOBIC:  return "Hydrophobic"
    if aa1 in POLAR:        return "Polar"
    if aa1 in CHARGED_POS:  return "Charged (+)"
    if aa1 in CHARGED_NEG:  return "Charged (−)"
    if aa1 in SPECIAL:      return "Special"
    return "Unknown"

df_contacts["Chemical Class"] = df_contacts["AA (1-letter)"].apply(classify_aa)

# Summary by class
print("BINDING SITE CHEMICAL PROFILE")
print("=" * 45)
class_counts = df_contacts["Chemical Class"].value_counts()
total = len(df_contacts)
for cls, count in class_counts.items():
    pct = 100 * count / total
    print(f"  {cls:<18} {count} residues ({pct:.0f}%)")
print()
print(df_contacts[["AA (3-letter)", "Position (PDB)", "Min Distance (Å)", "Chemical Class"]]
      .sort_values("Min Distance (Å)").to_string(index=False))

import matplotlib.pyplot as plt
import matplotlib.patches as mpatches

# Colour by chemical class
CLASS_COLORS = {
    "Hydrophobic":  "#e07b39",
    "Polar":        "#4a90d9",
    "Charged (+)":  "#2ecc71",
    "Charged (−)":  "#e74c3c",
    "Special":      "#9b59b6",
    "Unknown":      "#95a5a6"
}

df_plot = df_contacts.sort_values("Min Distance (Å)")
colors  = df_plot["Chemical Class"].map(CLASS_COLORS)

fig, ax = plt.subplots(figsize=(12, 6))
bars = ax.barh(
    df_plot["AA (3-letter)"] + " " + df_plot["Position (PDB)"].astype(str),
    df_plot["Min Distance (Å)"],
    color=colors, height=0.6
)

# Reference line at 5 Å cutoff
ax.axvline(x=CUTOFF, color="black", linestyle="--", linewidth=1,
           label=f"{CUTOFF} Å cutoff")

ax.set_xlabel("Minimum Distance to Ibuprofen (Å)", fontsize=12)
ax.set_title("COX-2 Binding Site Contact Map\n(Ibuprofen — 4PH9)",
             fontsize=13, fontweight="bold")

# Legend for chemical classes
patches = [mpatches.Patch(color=v, label=k) for k, v in CLASS_COLORS.items()
           if k in df_plot["Chemical Class"].values]
ax.legend(handles=patches, loc="lower right", fontsize=9)

plt.tight_layout()
plt.savefig("binding_site_contacts.png", dpi=150, bbox_inches="tight")
plt.show()
print("Chart saved as binding_site_contacts.png")

# Full summary table
print("BINDING SITE REPORT — COX-2 + Ibuprofen (4PH9)")
print("=" * 60)
print(f"Drug: Ibuprofen (IBP)")
print(f"Protein: COX-2 (PDB 4PH9, Chain A)")
print(f"Distance cutoff: {CUTOFF} Å\n")

print(f"Total contact residues: {len(df_contacts)}")
print(f"Closest residue: {df_contacts.sort_values('Min Distance (Å)').iloc[0]['AA (3-letter)']} "
      f"at {df_contacts['Min Distance (Å)'].min():.2f} Å")
print(f"Chemical composition:")
for cls, count in df_contacts["Chemical Class"].value_counts().items():
    print(f"   {cls}: {count}")

# Save to CSV for use in Projects 5 and 6
df_contacts.to_csv("binding_site_report.csv", index=False)
print("\nFull report saved to binding_site_report.csv")
print("You will use this file in Project 5 (Alanine Scanning).")