gmx_MMPBSA_residual_comparative.py

#!/usr/bin/env python3
import re
import os
import logging
import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
from gmx_MMPBSA_residues_decomposition import parse_decomp_energy_terms

def plot_protein_delta_decomposition_heatmap_comparative(
    df_dict,
    wt_system_name='WT',
    output_file_path=None,
    output_file='delta_decomposition_heatmap_comparative.png',
    energy_threshold=2.0,
    figsize=None,  # Will auto-calculate
    dpi=600,
    square_cells=True  # New parameter
):
    """
    Creates a comparative heatmap showing energy values across systems and amino acids.
    Amino acids are selected based on the wild-type (WT) dataframe.
    
    Parameters:
    -----------
    df_dict : dict
        Dictionary mapping system names to DataFrames.
    
    wt_system_name : str
        Name of the wild-type system. Default: 'WT'
    
    output_file_path : str, optional
        Directory path to save the plot.
    
    output_file : str
        Filename for saving the plot.
    
    energy_threshold : float
        Absolute energy threshold (kcal/mol). Default: 2.0
    
    figsize : tuple, optional
        Figure size. If None, auto-calculated based on data dimensions.
    
    dpi : int
        Resolution. Default: 600
    
    square_cells : bool
        If True, cells are square. Default: True
    
    Returns:
    --------
    pd.DataFrame : Heatmap data.
    """
    
    plot_columns = [
        'Residue', 'Internal_Avg', 'Van_der_Waals_Avg', 'Electrostatic_Avg',
        'Polar_Solv_Avg', 'Non_Polar_Solv_Avg', 'Total_Avg'
    ]
    
    if wt_system_name not in df_dict:
        raise ValueError(f"'{wt_system_name}' system not found in df_dict. Available: {list(df_dict.keys())}")
    
    processed_dfs = {}
    
    for system_name, df in df_dict.items():
        get_data_protein = df[df['Residue'].str.startswith("R")]
        plot_df = pd.DataFrame(get_data_protein, columns=plot_columns)
        plot_df['Amino_Acid'] = plot_df['Residue'].str[2:]
        
        for col in plot_columns[1:]:
            plot_df[col] = pd.to_numeric(plot_df[col], errors='coerce', downcast='float')
        
        updated_plot_df = plot_df[
            ['Amino_Acid', 'Internal_Avg', 'Van_der_Waals_Avg',
             'Electrostatic_Avg', 'Polar_Solv_Avg',
             'Non_Polar_Solv_Avg', 'Total_Avg']
        ]
        updated_plot_df.set_index('Amino_Acid', inplace=True)
        
        updated_plot_df = updated_plot_df[
            (updated_plot_df['Total_Avg'] > energy_threshold) |
            (updated_plot_df['Total_Avg'] < -energy_threshold)
        ]
        
        processed_dfs[system_name] = updated_plot_df
    
    # Extract amino acids from WT
    wt_amino_acids = sorted(list(processed_dfs[wt_system_name].index))
    wt_amino_acids_sorted = residue_numeric_sort(wt_amino_acids)

    logging.info(f"\n--- Heatmap Configuration ---")
    logging.info(f"Reference system (WT): {wt_system_name}")
    logging.info(f"Total amino acids in WT (filtered |E| >= {energy_threshold}): {len(wt_amino_acids)}")
    logging.info(f"Systems to compare: {list(processed_dfs.keys())}\n")
    
    all_amino_acids = set()
    for df in processed_dfs.values():
        all_amino_acids.update(df.index)
    logging.info(f"Total unique amino acids across all systems: {len(all_amino_acids)}")
    # Create comparison matrix
    
    comparison_data = []
    for aa in all_amino_acids:
        row = []
        for system_name in df_dict.keys():
            if aa in processed_dfs[system_name].index:
                value = processed_dfs[system_name].loc[aa, 'Total_Avg']
            else:
                value = np.nan
            row.append(value)
        comparison_data.append(row)
    
    
    heatmap_df = pd.DataFrame(
    comparison_data,
    index=[aa.replace(":", " ") for aa in all_amino_acids],
    columns=list(df_dict.keys())
)

    # REINDEX: Use the sorted WT amino acid list to order the DataFrame
    # If your dataframe's index uses labels like ':LYS:717',
    # make sure they match exactly or adjust the format:
    heatmap_df = heatmap_df.reindex([aa.replace(":", " ") for aa in wt_amino_acids_sorted])

    # Auto-calculate figsize if not provided
    if figsize is None and square_cells:
        n_residues = len(processed_dfs)
        n_systems = len(heatmap_df.columns)
        # Aim for ~0.5 inches per cell
        cell_size = 0.5
        fig_width = max(8, n_systems * cell_size + 2)
        fig_height = max(8, n_residues * cell_size + 2)
        figsize = (fig_width, fig_height)
    elif figsize is None:
        figsize = (12, 8)
    
    sns.set_theme(style="white")
    
    fig, ax = plt.subplots(figsize=figsize, dpi=dpi)
    
    # Create heatmap
    # After creating and reindexing 'heatmap_df' as rows = amino acids, columns = systems
    # Transpose the DataFrame to switch rows and columns
    heatmap_df_T = heatmap_df.T

    # Then plot heatmap with 'heatmap_df_T'
    sns.heatmap(
        heatmap_df_T,
        annot=True,
        fmt='.1f',
        cmap='RdBu_r',
        center=0,
        ax=ax,
        cbar_kws={
            'label': 'Δ Energy (kcal/mol)',
            'shrink': 0.3
        },
        linewidths=1.0,
        linecolor='white',
        square=square_cells,
        annot_kws={'size': 10, 'weight': 'bold'},
        vmin=-max(heatmap_df_T.max().max(), -heatmap_df_T.min().min()),
        vmax=max(heatmap_df_T.max().max(), -heatmap_df_T.min().min())
    )

    # Update axis labels to reflect new layout
    ax.set_xlabel('Amino Acid', fontsize=12, fontweight='bold')
    ax.set_ylabel('System', fontsize=12, fontweight='bold')

    # Adjust tick parameters accordingly
    ax.tick_params(axis='y', labelrotation=0, labelsize=11)  # y-axis now systems
    ax.tick_params(axis='x', labelrotation=45, labelsize=11)  # x-axis now amino acids


    
    plt.tight_layout()
    
    if output_file_path is None:
        output_file_path = './'
    
    output_path = os.path.join(output_file_path, output_file)
    plt.savefig(
        output_path,
        dpi=600,
        bbox_inches='tight',
        facecolor='white'
    )
    logging.info(f"✓ Heatmap saved as: {output_path}")
    plt.close()
    
    return heatmap_df

def plot_RNA_delta_decomposition_heatmap_comparative(
    df_dict,
    wt_system_name='WT',
    output_file_path=None,
    output_file='delta_decomposition_RNA_heatmap_comparative.png',
    energy_threshold=2.0,
    figsize=None,  # Will auto-calculate
    dpi=600,
    square_cells=True  # New parameter
):
    """
    Creates a comparative heatmap showing energy values across systems and amino acids.
    Amino acids are selected based on the wild-type (WT) dataframe.
    
    Parameters:
    -----------
    df_dict : dict
        Dictionary mapping system names to DataFrames.
    
    wt_system_name : str
        Name of the wild-type system. Default: 'WT'
    
    output_file_path : str, optional
        Directory path to save the plot.
    
    output_file : str
        Filename for saving the plot.
    
    energy_threshold : float
        Absolute energy threshold (kcal/mol). Default: 2.0
    
    figsize : tuple, optional
        Figure size. If None, auto-calculated based on data dimensions.
    
    dpi : int
        Resolution. Default: 600
    
    square_cells : bool
        If True, cells are square. Default: True
    
    Returns:
    --------
    pd.DataFrame : Heatmap data.
    """
    
    plot_columns = [
        'Residue', 'Internal_Avg', 'Van_der_Waals_Avg', 'Electrostatic_Avg',
        'Polar_Solv_Avg', 'Non_Polar_Solv_Avg', 'Total_Avg'
    ]
    
    if wt_system_name not in df_dict:
        raise ValueError(f"'{wt_system_name}' system not found in df_dict. Available: {list(df_dict.keys())}")
    
    processed_dfs = {}
    
    for system_name, df in df_dict.items():
        get_data_rna = df[df['Residue'].str.startswith("L")]
        plot_df = pd.DataFrame(get_data_rna, columns=plot_columns)
        
        # Extract base names
        plot_df['Base'] = plot_df['Residue'].str[3:10]
        
        for col in plot_columns[1:]:
            plot_df[col] = pd.to_numeric(plot_df[col], errors='coerce', downcast='float')
        
        updated_plot_df = plot_df[
            ['Base', 'Internal_Avg', 'Van_der_Waals_Avg',
             'Electrostatic_Avg', 'Polar_Solv_Avg',
             'Non_Polar_Solv_Avg', 'Total_Avg']
        ]
        updated_plot_df.set_index('Base', inplace=True)
        
        updated_plot_df = updated_plot_df[
            (updated_plot_df['Total_Avg'] > energy_threshold) |
            (updated_plot_df['Total_Avg'] < -energy_threshold)
        ]
        
        processed_dfs[system_name] = updated_plot_df
    
    # Extract amino acids from WT
    wt_nts = sorted(list(processed_dfs[wt_system_name].index))
    wt_nts_sorted = residue_numeric_sort(wt_nts)

    logging.info(f"\n--- Heatmap Configuration ---")
    logging.info(f"Reference system (WT): {wt_system_name}")
    logging.info(f"Total nucleotides in WT (filtered |E| >= {energy_threshold}): {len(wt_nts_sorted)}")
    logging.info(f"Systems to compare: {list(processed_dfs.keys())}\n")
    
    all_nts = set()
    for df in processed_dfs.values():
        all_nts.update(df.index)
    logging.info(f"Total unique nucleotides across all systems: {len(all_nts)}")
    # Create comparison matrix
    
    comparison_data = []
    for nt in all_nts:
        row = []
        for system_name in df_dict.keys():
            if nt in processed_dfs[system_name].index:
                value = processed_dfs[system_name].loc[nt, 'Total_Avg']
            else:
                value = np.nan
            row.append(value)
        comparison_data.append(row)
    
    heatmap_df = pd.DataFrame(
    comparison_data,
    index=[nt.replace(":", " ") for nt in all_nts],
    columns=list(df_dict.keys())
)

    # REINDEX: Use the sorted WT amino acid list to order the DataFrame
    # If your dataframe's index uses labels like ':LYS:717',
    # make sure they match exactly or adjust the format:
    heatmap_df = heatmap_df.reindex([nt.replace(":", " ") for nt in wt_nts_sorted])

    # Auto-calculate figsize if not provided
    if figsize is None and square_cells:
        n_residues = len(processed_dfs)
        n_systems = len(heatmap_df.columns)
        # Aim for ~0.5 inches per cell
        cell_size = 0.5
        fig_width = max(8, n_systems * cell_size + 2)
        fig_height = max(8, n_residues * cell_size + 2)
        figsize = (fig_width, fig_height)
    elif figsize is None:
        figsize = (12, 8)
    # print heatmap shape
    logging.info(f"Heatmap shape: {heatmap_df.shape} (residues × systems)\n")
    
    sns.set_theme(style="white")
    
    fig, ax = plt.subplots(figsize=figsize, dpi=dpi)
    
    # Create heatmap
    # After creating and reindexing 'heatmap_df' as rows = amino acids, columns = systems
    # Transpose the DataFrame to switch rows and columns
    heatmap_df_T = heatmap_df.T

    # Then plot heatmap with 'heatmap_df_T'
    sns.heatmap(
        heatmap_df_T,
        annot=True,
        fmt='.1f',
        cmap='RdBu_r',
        center=0,
        ax=ax,
        cbar_kws={
            'label': 'Δ Energy (kcal/mol)',
            'shrink': 0.3
        },
        linewidths=1.0,
        linecolor='white',
        square=square_cells,
        annot_kws={'size': 10, 'weight': 'bold'},
        vmin=-max(heatmap_df_T.max().max(), -heatmap_df_T.min().min()),
        vmax=max(heatmap_df_T.max().max(), -heatmap_df_T.min().min())
    )

    # Update axis labels to reflect new layout
    ax.set_xlabel('Nucleotides', fontsize=12, fontweight='bold')
    ax.set_ylabel('System', fontsize=12, fontweight='bold')

    # Adjust tick parameters accordingly
    ax.tick_params(axis='y', labelrotation=0, labelsize=11)  # y-axis now systems
    ax.tick_params(axis='x', labelrotation=45, labelsize=11)  # x-axis now amino acids


    
    plt.tight_layout()
    
    if output_file_path is None:
        output_file_path = './'
    
    output_path = os.path.join(output_file_path, output_file)
    plt.savefig(
        output_path,
        dpi=600,
        bbox_inches='tight',
        facecolor='white'
    )
    logging.info(f"✓ Heatmap saved as: {output_path}")
    plt.close()
    
    return heatmap_df

def residue_numeric_sort(labels):
    # Extract numeric ending from each label
    return sorted(labels, key=lambda x: int(re.search(r'(\d+)$', x).group(1)))

### MAIN SCRIPT EXECUTION
def main():
	file_path = "/home/shrikant/my_work/OBJ4/SF3A1_UBL/ANALYSIS_SF3A1/Binding-Energy-Decomposition/"
	output_file_path = "/home/shrikant/my_work/OBJ4/SF3A1_UBL/ANALYSIS_SF3A1/Binding-Energy-Decomposition/plot/"
	# DECOMPOSED ENERGY FILE PATHS
	decomp_wt_fh = os.path.join(file_path, "wt/8ID2_WT_DECOMP_MMPBSA.dat")
	decomp_r788_df = os.path.join(file_path, "788A/8ID2_R788A_DECOMP_MMPBSA.dat")
	decomp_r791_df = os.path.join(file_path, "791A/8ID2_R791A_DECOMP_MMPBSA.dat")
	decomp_r788_791_df = os.path.join(file_path, "788A-791A/8ID2_R788A_R791A_DECOMP_MMPBSA.dat")
	decomp_e787a_df = os.path.join(file_path, "E787A/8ID2_E787A_DECOMP_MMPBSA.dat")

		# Parse the decomposed energy terms
	wt_delta_total, wt_delta_backbone, wt_delta_sidechain = parse_decomp_energy_terms(decomp_wt_fh, delta_energy=True)
	r788a_delta_total, r788a_delta_backbone, r788a_delta_sidechain = parse_decomp_energy_terms(decomp_r788_df, delta_energy=True)
	r791a_delta_total, r791a_delta_backbone, r791a_delta_sidechain = parse_decomp_energy_terms(decomp_r791_df, delta_energy=True)
	r788_791a_delta_total, r788_791a_delta_backbone, r788_791a_delta_sidechain = parse_decomp_energy_terms(decomp_r788_791_df, delta_energy=True)
	e787a_delta_total, e787a_delta_backbone, e787a_delta_sidechain = parse_decomp_energy_terms(decomp_e787a_df, delta_energy=True)

	# Create dictionary of systems
	df_dict = {
		'WT': wt_delta_total,
		'R788A': r788a_delta_total,
		'R791A': r791a_delta_total,
		'R788A_R791A': r788_791a_delta_total,
		'E787A': e787a_delta_total
	}


	## Plot heatmap using WT as reference
	heatmap_df = plot_protein_delta_decomposition_heatmap_comparative(
		df_dict=df_dict,
		wt_system_name='WT',  # Use WT amino acids as rows
		output_file_path=output_file_path,
		output_file='comparative_energy_heatmap_PROTEIN.png',
		energy_threshold=2.0,
		square_cells=True,  # Enable square cells
		dpi=600
	)

	heatmap_df = plot_RNA_delta_decomposition_heatmap_comparative(
		df_dict=df_dict,
		wt_system_name='WT',  # Use WT amino acids as rows
		output_file_path=output_file_path,
		output_file='comparative_energy_heatmap_RNA.png',
		energy_threshold=1.0,
		square_cells=True,  # Enable square cells
		dpi=600
	)

if __name__ == "__main__":
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        datefmt='%Y-%m-%d %H:%M:%S'
    )
    main()