SASA.py

# %%
import seaborn as sns
import mdtraj as md
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Patch
import os
#%%
# -----------------------------
# First protein (apomon)
# -----------------------------
os.chdir('/home/pghw87/Documents/md-sim/5ue6/MDAniA/data/')

# Target residues with fixed colors
target_residues = [
    ("HIS", 86, 'gold'),
    ("HIS", 91, 'darkkhaki'),
    ("HIS", 126, 'beige'),
    ("CYS", 127, 'rebeccapurple'),
    ("HIS", 135, 'olive'),
    ("MET", 140, 'teal'),
    ("HIS", 281, 'goldenrod')
]
hatched_residues = {"HIS91", "HIS126", "HIS281"}

# Define residue groups (equivalent residues across chains)
residue_groups = [
    [("HIS", 86, 'gold'), ("HIS", 396, 'gold'), ("HIS", 706, 'gold')],
    [("HIS", 91, 'darkkhaki'), ("HIS", 401, 'darkkhaki'), ("HIS", 711, 'darkkhaki')],
    [("HIS", 126, 'beige'), ("HIS", 436, 'beige'), ("HIS", 746, 'beige')],
    [("CYS", 127, 'rebeccapurple'), ("CYS", 437, 'rebeccapurple'), ("CYS", 747, 'rebeccapurple')],
    [("HIS", 135, 'olive'), ("HIS", 445, 'olive'), ("HIS", 755, 'olive')],
    [("MET", 140, 'teal'), ("MET", 450, 'teal'), ("MET", 760, 'teal')],
    [("HIS", 281, 'goldenrod'), ("HIS", 591, 'goldenrod'), ("HIS", 901, 'goldenrod')]
]

hatched_residues = {"HIS91", "HIS126", "HIS281", "HIS401", "HIS711", "HIS436", "HIS746", "HIS591", "HIS901"}

def sasa_per_residue(traj):
    sasa = md.shrake_rupley(traj, mode='residue')
    sasa_mean = sasa.mean(axis=0)
    sasa_std = sasa.std(axis=0)
    
    selected_indices = []
    residues_sel = []
    
    for (name, pdb_idx, _) in target_residues:
        found = False
        for res in traj.topology.residues:
            if res.name == name and res.resSeq == pdb_idx:
                selected_indices.append(res.index)
                residues_sel.append(f"{res.name}{res.resSeq}")
                found = True
                break
        if not found:
            print(f"Residue {name}{pdb_idx} not found in holomon!")

    sasa_mean_sel = sasa_mean[selected_indices]
    sasa_std_sel = sasa_std[selected_indices]

    return sasa_mean_sel, sasa_std_sel, residues_sel

def mean_multichain_sasa(traj_files, top_files):
    sasa_dict = {f"{name}{idx}": [] for (name, idx, _) in target_residues}

    for traj_file, top_file in zip(traj_files, top_files):
        traj = md.load(traj_file, top=top_file)
        sasa = md.shrake_rupley(traj, mode='residue')
        for res in traj.topology.residues:
            for (name, pdb_idx, _) in target_residues:
                if res.name == name and res.resSeq == pdb_idx:
                    sasa_dict[f"{name}{pdb_idx}"].append(sasa[:, res.index])

    residue_means = []
    residue_stds = []
    residue_labels = []

    for label, arrays in sasa_dict.items():
        if len(arrays) == 0:
            print(f"Residue {label} not found!")
            continue
        stacked = np.vstack(arrays)
        all_vals = stacked.flatten()
        residue_means.append(all_vals.mean())
        residue_stds.append(all_vals.std())
        residue_labels.append(label)

    return residue_means, residue_stds, residue_labels

def sasa_per_trimer_residue(traj):
    sasa = md.shrake_rupley(traj, mode='residue')  # shape: (n_frames, n_residues)

    group_means = []
    group_stds = []
    group_labels = []

    for group in residue_groups:
        arrays = []
        for (name, pdb_idx, color) in group:
            for res in traj.topology.residues:
                if res.name == name and res.resSeq == pdb_idx:
                    arrays.append(sasa[:, res.index])  # (n_frames,)
                    break
        if not arrays:
            print(f"Residue group {group} not found!")
            continue

        stacked = np.vstack(arrays)  # shape (n_chains, n_frames)
        all_vals = stacked.flatten()  # combine frames × chains

        group_means.append(all_vals.mean())
        group_stds.append(all_vals.std())
        group_labels.append(f"{group[0][0]}{group[0][1]}")  # e.g. HIS82

    return group_means, group_stds, group_labels

# %%
def plot_sasa_bar(protein_data):
    color_map = {f"{name}{idx}": color for (name, idx, color) in target_residues}
    fig, ax = plt.subplots(figsize=(10, 6))

    x_positions = np.arange(len(protein_data))  
    bar_width = 0.5

    for x, (title, means, stds, labels) in zip(x_positions, protein_data):
        bottom = 0
        for mean, std, label in zip(means, stds, labels):
            hatch_style = 'xx' if label in hatched_residues else None
            ax.bar(
                x, mean, bottom=bottom,
                color=color_map[label], alpha=0.7,
                hatch=hatch_style, edgecolor='black', width=bar_width
            )
            y_center = bottom + mean / 2
            ax.errorbar(x, y_center, yerr=std/2, fmt='k_', capsize=3)
            bottom += mean

    # X-axis with protein names
    ax.set_xticks(x_positions)
    ax.set_xticklabels([title for title, _, _, _ in protein_data], ha='center', rotation=30)

    # Labels
    ax.set_ylabel("SASA (nm$^2$)")

    unique_labels = [f"{group[0][0]}{group[0][1]}" for group in residue_groups]  # HIS82, HIS87, ...
    legend_elements = []
    for lbl in unique_labels:
        hatch_style = 'xx' if lbl in hatched_residues else None
        legend_elements.append(
            Patch(facecolor=color_map[lbl], edgecolor='black',
                  hatch=hatch_style, label=lbl, alpha=0.7)
        )

    ax.legend(handles=legend_elements, bbox_to_anchor=(1.02, 1),
              loc='upper left', title="Residues")
    ymin, ymax = ax.get_ylim()
    # ax.set_ylim(ymin, ymax * 1.3)  
    plt.title("Trimer as a whole SASA per Residue")
    plt.tight_layout()
    plt.show()
#%%
apotri_individual_sasa, apotri_individual_std, apotri_individual_labels = mean_multichain_sasa(
    ['trimer/chainA_3us.xtc', 'trimer/chainB_3us.xtc', 'trimer/chainC_3us.xtc'],
    top_files=['trimer/chainA.gro', 'trimer/chainB.gro', 'trimer/chainC.gro']
)
#%%
holotri6_empty_hyd_individual_sasa, holotri6_empty_hyd_individual_std, holotri6noCU_empty_hyd_individual_labels = mean_multichain_sasa(
    ['holotrimer6-hyd/bonded/chainA_noCU.xtc', 'holotrimer6-hyd/bonded/chainB_noCU.xtc', 'holotrimer6-hyd/bonded/chainC_noCU.xtc'],
    top_files=['holotrimer6-hyd/bonded/chainA_noCU_renamed.gro', 'holotrimer6-hyd/bonded/chainB_noCU_renamed.gro', 'holotrimer6-hyd/bonded/chainC_noCU_renamed.gro']
)
#%%
holotri6_loaded_hyd_individual_sasa, holotri6_loaded_hyd_individual_std, holotri6_loaded_hyd_individual_labels = mean_multichain_sasa(
    ['holotrimer6-hyd/bonded/chainA.xtc', 'holotrimer6-hyd/bonded/chainB.xtc', 'holotrimer6-hyd/bonded/chainC.xtc'],
    top_files=['holotrimer6-hyd/bonded/chainA_renamed.gro', 'holotrimer6-hyd/bonded/chainB_renamed.gro', 'holotrimer6-hyd/bonded/chainC_renamed.gro']
)
#%%
holotri6_empty_wat_individual_sasa, holotri6_empty_wat_individual_std, holotri6noCU_empty_wat_individual_labels = mean_multichain_sasa(
    ['holotrimer6/bonded/chainA_noCU_3us.xtc', 'holotrimer6/bonded/chainB_noCU_3us.xtc', 'holotrimer6/bonded/chainC_noCU_3us.xtc'],
    top_files=['holotrimer6/bonded/chainA_noCU_renamed.gro', 'holotrimer6/bonded/chainB_noCU_renamed.gro', 'holotrimer6/bonded/chainC_noCU_renamed.gro']
)
#%%
holotri6_loaded_wat_individual_sasa, holotri6_loaded_wat_individual_std, holotri6_loaded_wat_individual_labels = mean_multichain_sasa(
    ['holotrimer6/bonded/chainA_3us.xtc', 'holotrimer6/bonded/chainB_3us.xtc', 'holotrimer6/bonded/chainC_3us.xtc'],
    top_files=['holotrimer6/bonded/chainA_renamed.gro', 'holotrimer6/bonded/chainB_renamed.gro', 'holotrimer6/bonded/chainC_renamed.gro']
)
#%%
holotri3_empty_individual_sasa, holotri3_empty_individual_std, holotri3_empty_individual_labels = mean_multichain_sasa(
    ['holotrimer3/bonded/chainA_noCU_3us.xtc', 'holotrimer3/bonded/chainB_noCU_3us.xtc', 'holotrimer3/bonded/chainC_noCU_3us.xtc'],
    top_files=['holotrimer3/bonded/chainA_noCU_renamed.gro', 'holotrimer3/bonded/chainB_noCU_renamed.gro', 'holotrimer3/bonded/chainC_noCU_renamed.gro']
)
#%%
holotri3_loaded_individual_sasa, holotri3_loaded_individual_std, holotri3_loaded_individual_labels = mean_multichain_sasa(
    ['holotrimer3/bonded/chainA_3us.xtc', 'holotrimer3/bonded/chainB_3us.xtc', 'holotrimer3/bonded/chainC_3us.xtc'],
    top_files=['holotrimer3/bonded/chainA_renamed.gro', 'holotrimer3/bonded/chainB_renamed.gro', 'holotrimer3/bonded/chainC_renamed.gro']
)
#%%
apomon_sasa, apomon_std, apomon_labels = sasa_per_residue(
    md.load('apomon/apomon_3us.xtc', top='apomon/apomon.gro')
)
#%%
holomon_loaded_sasa, holomon_loaded_std, holomon_loaded_labels = sasa_per_residue(
    md.load('holomon/holomon_3us.xtc', top='holomon/holomon_renamed.gro')
)
#%%
holomon_empty_sasa, holomon_empty_std, holomon_empty_labels = sasa_per_residue(
    md.load('holomon/holomon_noCU_3us.xtc', top='holomon/holomon_noCU_renamed.gro')
)
#%%
apotri_whole_sasa, apotri_whole_std, apotri_whole_labels = sasa_per_trimer_residue(
    md.load("trimer/trimer_3us.xtc", top="trimer/trimer.gro")
)
#%%
holotri3_empty_whole_sasa, holotri3_empty_whole_std, holotri3_empty_whole_labels = sasa_per_trimer_residue(
    md.load("holotrimer3/bonded/trimer_noCU_3us.xtc", top="holotrimer3/bonded/trimer_noCU_renamed.gro")
)
#%%
holotri3_loaded_whole_sasa, holotri3_loaded_whole_std, holotri3_loaded_whole_labels = sasa_per_trimer_residue(
    md.load("holotrimer3/bonded/trimer_3us.xtc", top="holotrimer3/bonded/trimer_renamed.gro")
)
#%%
holotri6_empty_hyd_whole_sasa, holotri6_empty_hyd_whole_std, holotri6_empty_hyd_whole_labels = sasa_per_trimer_residue(
    md.load("holotrimer6-hyd/bonded/trimer_noCU.xtc", top="holotrimer6-hyd/bonded/trimer_noCU_renamed.gro")
)
#%%
holotri6_loaded_hyd_whole_sasa, holotri6_loaded_hyd_whole_std, holotri6_loaded_hyd_whole_labels = sasa_per_trimer_residue(
    md.load("holotrimer6-hyd/bonded/trimer.xtc", top="holotrimer6-hyd/bonded/trimer_renamed.gro")
)
#%%
holotri6_empty_wat_whole_sasa, holotri6_empty_wat_whole_std, holotri6_empty_wat_whole_labels = sasa_per_trimer_residue(
    md.load("holotrimer6/bonded/trimer_noCU_3us.xtc", top="holotrimer6/bonded/trimer_noCU_renamed.gro")
)
#%%
holotri6_loaded_wat_whole_sasa, holotri6_loaded_wat_whole_std, holotri6_loaded_wat_whole_labels = sasa_per_trimer_residue(
    md.load("holotrimer6/bonded/trimer_3us.xtc", top="holotrimer6/bonded/trimer_renamed.gro")
)
#%%
protein_data = [
    ### Monomer ###
    # ("apomon", apomon_sasa, apomon_std, apomon_labels),
    # ("Cu Type1", holomon_loaded_sasa, holomon_loaded_std, holomon_loaded_labels),
    # ("Cu Type1 removed", holomon_empty_sasa, holomon_empty_std, holomon_empty_labels),

    # ### Trimer Mean of Chains ###
    # ("apo trimer", apotri_individual_sasa, apotri_individual_std, apotri_individual_labels),
    # ("Cu Type1", holotri3_loaded_individual_sasa, holotri3_loaded_individual_std, holotri3_loaded_individual_labels),
    # ("Cu Type1 removed", holotri3_empty_individual_sasa, holotri3_empty_individual_std, holotri3_empty_individual_labels),
    # ("Cu Type1&2 removed, Water", holotri6_empty_wat_individual_sasa, holotri6_empty_wat_individual_std, holotri6noCU_empty_wat_individual_labels),
    # ("Cu Type1&2, Water", holotri6_loaded_wat_individual_sasa, holotri6_loaded_wat_individual_std, holotri6_loaded_wat_individual_labels),
    # ("Cu Type1&2 removed, Hyd", holotri6_empty_hyd_individual_sasa, holotri6_empty_hyd_individual_std, holotri6noCU_empty_hyd_individual_labels),
    # ("Cu Type1&2, Hyd", holotri6_loaded_hyd_individual_sasa, holotri6_loaded_hyd_individual_std, holotri6_loaded_hyd_individual_labels),

    ### Trimer as a whole ####
    ("apo trimer", apotri_whole_sasa, apotri_whole_std, apotri_whole_labels),
    ("Cu Type1", holotri3_loaded_whole_sasa, holotri3_loaded_whole_std, holotri3_loaded_whole_labels),
    # ("Cu Type1 removed", holotri3_empty_whole_sasa, holotri3_empty_whole_std, holotri3_empty_whole_labels),
    ("Cu Type1&2", holotri6_loaded_wat_whole_sasa, holotri6_loaded_wat_whole_std, holotri6_loaded_wat_whole_labels),
    # ("Cu Type1&2 removed", holotri6_empty_wat_whole_sasa, holotri6_empty_wat_whole_std, holotri6_empty_wat_whole_labels),
    # ("Cu Type1&2, Hyd", holotri6_loaded_hyd_whole_sasa, holotri6_loaded_hyd_whole_std, holotri6_loaded_hyd_whole_labels),
    # ("Cu Type1&2 removed, Hyd", holotri6_empty_hyd_whole_sasa, holotri6_empty_hyd_whole_std, holotri6_empty_hyd_whole_labels),
]
plot_sasa_bar(protein_data)
# %%