interaction_Grid.py

import os
import yaml
import numpy as np
import sparse
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from collections import Counter

# --- Configuration ---
y_label = ["AMBER", "CHARMM", "SIRAH hybrid", "M3", "M2", "M2PW"]
folder_paths = [
    "/ceph/users/akv16273/34_Elcock_13abundant_inEColi/21_Analysis_Datastructure/20_AMBER/",
    "/ceph/users/akv16273/34_Elcock_13abundant_inEColi/21_Analysis_Datastructure/12_AA/1_monomer/",
    "/ceph/users/akv16273/34_Elcock_13abundant_inEColi/21_Analysis_Datastructure/13_SIRAH/",
    "/ceph/users/akv16273/34_Elcock_13abundant_inEColi/21_Analysis_Datastructure/9_Martini3/3_303K/",
    "/ceph/users/akv16273/34_Elcock_13abundant_inEColi/21_Analysis_Datastructure/10_Martini2_nonPW/3_303K/",
    "/ceph/users/akv16273/34_Elcock_13abundant_inEColi/21_Analysis_Datastructure/11_Martini2_PW/3_303K/"
]


def plot_top_interacting_pairs_grid(folder_paths, n, y_label, output_filename="top_interactions_grid.png"):
    """
    Analyzed and plotted top protein interactions from multiple simulations.

    Interactions seen across multiple force fields were colored uniquely for comparison.
    Interactions that appeared in only one force field's top 'n' list were colored
    grey. The final plot was saved to a file.

    Args:
        folder_paths (list): A list of paths to folders containing simulation data.
        n (int): The number of top interactions to display per subplot.
        y_label (list): A list of titles for each subplot.
        output_filename (str): The path and name for the output image file.
    """
    if len(folder_paths) > 6:
        print("Warning: This function is designed for up to 6 plots. Only the first 6 were processed.")
        folder_paths = folder_paths[:6]
        y_label = y_label[:6]

    # --- Step 1: Pre-scan data and COUNT occurrences of top pairs ---
    pair_counter = Counter()
    precomputed_data = []

    print("Step 1: Pre-scanning all folders to identify and count top interaction pairs...")
    for dir_path in folder_paths:
        matrices = []
        for file in os.listdir(dir_path):
            if file.endswith(".npz"):
                sparse_matrix = sparse.load_npz(os.path.join(dir_path, file))
                matrices.append(np.sum(np.sign(sparse_matrix.todense()), axis=2))

        mean_matrix = np.mean(matrices, axis=0)
        std_matrix = np.std(matrices, axis=0)

        upper_triangle_mean = np.triu(mean_matrix, k=1)
        upper_triangle_std = np.triu(std_matrix, k=1)
        flat_indices = np.argsort(upper_triangle_mean, axis=None)[-n:]
        i_indices, j_indices = np.unravel_index(flat_indices, upper_triangle_mean.shape)

        with open(os.path.join(dir_path, "protein_names.yaml"), "r") as g:
            protein_names_mapping = yaml.safe_load(g)["chain_labels"]

        top_values_ns = upper_triangle_mean[i_indices, j_indices]
        sorted_indices = np.argsort(top_values_ns)[::-1]

        sorted_pairs = []
        for i in sorted_indices:
            protein_i = protein_names_mapping[chr(ord('A') + i_indices[i])]
            protein_j = protein_names_mapping[chr(ord('A') + j_indices[i])]
            canonical_pair = tuple(sorted((protein_i, protein_j)))
            sorted_pairs.append(canonical_pair)

        pair_counter.update(set(sorted_pairs))

        precomputed_data.append({
            "values_ns": top_values_ns[sorted_indices],
            "stds_ns": upper_triangle_std[i_indices, j_indices][sorted_indices],
            "pairs": sorted_pairs
        })

    # --- Step 2: Create a grey-free color map for recurring pairs ---
    print("Step 2: Creating a grey-free color map...")
    color_map = {}

    recurring_pairs = sorted([p for p, c in pair_counter.items() if c > 1])
    num_recurring = len(recurring_pairs)

    colors_set1 = plt.get_cmap('Set1').colors
    colors_paired = plt.get_cmap('Paired').colors
    custom_colors = list(colors_set1) + list(colors_paired)

    if num_recurring > len(custom_colors):
        print(
            f"Warning: {num_recurring} recurring pairs found, which is more than the {len(custom_colors)} available distinct colors. Some colors will be reused.")

    for i, pair in enumerate(recurring_pairs):
        color_map[pair] = custom_colors[i % len(custom_colors)]

    for pair, count in pair_counter.items():
        if count == 1:
            color_map[pair] = '#a9a9a9'

    # --- Step 3: Create the 3x2 grid of subplots ---
    print("Step 3: Generating plots...")
    fig, axes = plt.subplots(3, 2, figsize=(14, 9), sharey=True)
    axes = axes.flatten()

    for i, data in enumerate(precomputed_data):
        ax = axes[i]
        values_us = data["values_ns"] / 1000.0
        stds_us = data["stds_ns"] / 1000.0
        bar_colors = [color_map[pair] for pair in data["pairs"]]

        ax.bar(range(n), values_us, yerr=stds_us, align='center', alpha=0.8, capsize=5, color=bar_colors)

        ax.set_title(y_label[i], fontweight='bold', fontsize=16)

        # A y-axis label was added to every subplot.
        ax.set_ylabel('Mean Interaction Time (µs)', fontsize=12)

        ax.set_ylim([0, 1.3])
        ax.set_yticks([0, 0.4, 0.8, 1.2])
        ax.grid(True, axis='y', linestyle='--', alpha=0.6)

        xtick_labels = [f'{p[0]}-{p[1]}' for p in data["pairs"]]
        ax.set_xticks(range(n))
        ax.set_xticklabels(xtick_labels, rotation=60, ha='right')
        ax.tick_params(axis='x', labelsize=10)

    for i in range(len(precomputed_data), len(axes)):
        axes[i].axis('off')

    # --- Step 4: Add a global legend and main title ---
    legend_patches = [mpatches.Patch(color=color_map[pair], label=f'{pair[0]}-{pair[1]}') for pair in recurring_pairs]
    if any(c == 1 for c in pair_counter.values()):
        legend_patches.append(mpatches.Patch(color='#a9a9a9', label='Single-Occurrence Pair'))

    fig.legend(handles=legend_patches, bbox_to_anchor=(1.01, 0.95), loc='upper left')

    # The layout was adjusted to make room for the new main title.
    plt.tight_layout()

    # --- Step 5: Save the figure ---
    if output_filename:
        print(f"Step 5: Saving the figure to {output_filename}...")
        plt.savefig(output_filename, dpi=600, bbox_inches='tight')
        print("Figure saved successfully.")
    else:
        plt.show()

    plt.close(fig)

    return None


if __name__ == "__main__":
    plot_top_interacting_pairs_grid(
        folder_paths=folder_paths,
        n=10,
        y_label=y_label,
        output_filename="Top_10_Interacting_Pairs_Grid_Final.png"
    )