PvsCOM.py

import os
import sparse
import yaml
import numpy as np
import matplotlib.pyplot as plt
import MDAnalysis as mda
import MDAnalysis.lib.distances

# --- Configuration ---
# Updated BASE_DIR to match your latest version
BASE_DIR = "/ceph/users/akv16273/34_Elcock_13abundant_inEColi/21_Analysis_Datastructure/"
MIN_SAMPLE_SIZE = 5

FORCE_FIELD_CONFIG = {
    "AMBER": {
        "path": os.path.join(BASE_DIR, "20_AMBER/"),
        "pdb_name": "md_0_1002ns.pdb",
        "color": "olive",
        "linestyle": "-"
    },
    "CHARMM": {
        "path": os.path.join(BASE_DIR, "12_AA/1_monomer/"),
        "pdb_name": "md_0_1002ns.pdb",
        "color": "cyan",
        "linestyle": "-"
    },
    "SIRAH": {
        "path": os.path.join(BASE_DIR, "13_SIRAH/"),
        "pdb_name": "md_0_1Mu.pdb",
        "color": "brown",
        "linestyle": "-"
    },
    "M3": {
        "path": os.path.join(BASE_DIR, "9_Martini3/3_303K/"),
        "pdb_name": "md_0_1Mu.pdb",
        "color": "magenta",
        "linestyle": "--"
    },
    "M2": {
        "path": os.path.join(BASE_DIR, "10_Martini2_nonPW/3_303K/"),
        "pdb_name": "md_0_1Mu.pdb",
        "color": "orange",
        "linestyle": "--"
    },
    "M2 (PW)": {
        "path": os.path.join(BASE_DIR, "11_Martini2_PW/3_303K/"),
        "pdb_name": "md_0_1Mu.pdb",
        "color": "blue",
        "linestyle": "--"
    }
}


def get_protein_map_from_yaml(yaml_path):
    """
    This function was designed to load a YAML file and return the mapping
    of chain IDs to protein names.
    """
    try:
        with open(yaml_path, 'r') as f:
            data = yaml.safe_load(f)
        return data.get('chain_labels', {})
    except (FileNotFoundError, yaml.YAMLError) as e:
        print(f"  -> Warning: Could not read or parse YAML file at {yaml_path}. {e}")
        return {}


def plot_initial_distance_vs_interaction_prob(config):
    """
    This analysis calculated and plotted the probability of subsequent protein
    interaction based on the initial, PBC-corrected COM distance at the
    second simulation frame. A final summary line, aggregating all force
    field data, was also generated. Points are not plotted if the sample size
    for a bin is below MIN_SAMPLE_SIZE.
    """
    plt.figure(figsize=(7, 5), dpi=600)

    all_ff_initial_distances = []
    all_ff_ever_interacted_flags = []

    for ff_name, ff_config in config.items():
        dir_path = ff_config["path"]
        print(f"Processing model: {ff_name}...")

        yaml_path = os.path.join(dir_path, 'protein_names.yaml')
        protein_map = get_protein_map_from_yaml(yaml_path)
        if not protein_map:
            print(f"Skipping {ff_name} due to missing protein map.")
            continue

        chain_ids_ordered = sorted(protein_map.keys())
        num_proteins = len(chain_ids_ordered)

        try:
            take_folders = sorted([os.path.join(dir_path, d) for d in os.listdir(dir_path) if
                                   d.startswith('Take_') and os.path.isdir(os.path.join(dir_path, d))])
            npz_files = sorted([os.path.join(dir_path, f) for f in os.listdir(dir_path) if f.endswith('.npz')])
            if len(take_folders) != len(npz_files) or not take_folders:
                print(f"  -> Warning: Mismatch or no data found for replicas in {dir_path}. Skipping.")
                continue
        except FileNotFoundError:
            print(f"  -> Warning: Directory not found: {dir_path}. Skipping.")
            continue

        initial_distances = []
        ever_interacted_flags = []

        for take_folder, npz_path in zip(take_folders, npz_files):
            pdb_path = os.path.join(take_folder, ff_config["pdb_name"])
            xtc_path = os.path.join(take_folder, "md_noWater_PBC_unwrapped_1000FRAMES.xtc")

            if not all(os.path.exists(p) for p in [pdb_path, xtc_path, npz_path]):
                continue

            print(f"  - Analyzing replica: {os.path.basename(take_folder)}")

            binary_matrix = np.sign(sparse.load_npz(npz_path).todense())
            universe = mda.Universe(pdb_path, xtc_path)

            if len(universe.trajectory) < 2:
                print("    -> FAILED: Trajectory has fewer than 2 frames. Skipping replica.")
                continue

            universe.trajectory[1]
            box = universe.dimensions
            proteins = [universe.select_atoms(f'segid {chain_id}') for chain_id in chain_ids_ordered]

            for p1_idx in range(num_proteins):
                for p2_idx in range(p1_idx + 1, num_proteins):
                    com1 = proteins[p1_idx].center_of_mass()
                    com2 = proteins[p2_idx].center_of_mass()

                    initial_dist = MDAnalysis.lib.distances.calc_bonds(com1, com2, box=box) / 10.0
                    initial_distances.append(initial_dist)

                    interaction_timeseries = binary_matrix[p1_idx, p2_idx, 1:]
                    did_interact = np.any(interaction_timeseries)
                    ever_interacted_flags.append(1 if did_interact else 0)

        if not initial_distances:
            print(f"Warning: No data successfully processed for {ff_name}.")
            continue

        dist_arr = np.array(initial_distances)
        flag_arr = np.array(ever_interacted_flags)

        # Using the bin_width from your provided script
        bin_width = 0.2
        bins = np.arange(0, np.max(dist_arr) + bin_width, bin_width)

        total_counts, _ = np.histogram(dist_arr, bins=bins)
        interaction_counts, _ = np.histogram(dist_arr, bins=bins, weights=flag_arr)

        probability = np.zeros_like(interaction_counts, dtype=float)
        non_zero_mask = total_counts > 0
        probability[non_zero_mask] = interaction_counts[non_zero_mask] / total_counts[non_zero_mask]

        probability[total_counts < MIN_SAMPLE_SIZE] = np.nan

        bin_centers = (bins[:-1] + bins[1:]) / 2

        plt.plot(bin_centers, probability, label=ff_name, color=ff_config["color"], linestyle=ff_config["linestyle"])

        all_ff_initial_distances.extend(initial_distances)
        all_ff_ever_interacted_flags.extend(ever_interacted_flags)

    if all_ff_initial_distances:
        print("\nProcessing combined data for all force fields...")

        all_dist_arr = np.array(all_ff_initial_distances)
        all_flag_arr = np.array(all_ff_ever_interacted_flags)

        bins = np.arange(0, np.max(all_dist_arr) + bin_width, bin_width)

        total_counts, _ = np.histogram(all_dist_arr, bins=bins)
        interaction_counts, _ = np.histogram(all_dist_arr, bins=bins, weights=all_flag_arr)

        comb_probability = np.zeros_like(interaction_counts, dtype=float)
        non_zero_mask = total_counts > 0
        comb_probability[non_zero_mask] = interaction_counts[non_zero_mask] / total_counts[non_zero_mask]

        comb_probability[total_counts < MIN_SAMPLE_SIZE] = np.nan

        bin_centers = (bins[:-1] + bins[1:]) / 2

        plt.plot(bin_centers, comb_probability, label='All Combined', color='black', linewidth=2.5, zorder=10)

    # --- Final Plot Configuration ---
    plt.xlabel('Initial COM Distance (nm)')
    plt.ylabel('Probability of Subsequent Interaction')
    # plt.title('Interaction Probability vs. Initial Distance (PBC Corrected)')
    # Using the x-limits from your provided script
    plt.xlim(4.2, 14.6)
    plt.ylim(0, 1.05)
    plt.grid(True, linestyle='--', alpha=0.6)
    plt.tight_layout()
    plt.savefig("Initial_Distance_vs_Interaction_Prob_PBC_Combined.png", dpi=600)

    # --- Histogram Plot ---
    if all_ff_initial_distances:
        plt.figure(figsize=(7, 4), dpi=600)
        plt.hist(all_ff_initial_distances, bins=bins, color='gray', edgecolor='black')
        # plt.title('Distribution of Initial Distances')
        plt.xlabel('Initial COM Distance (nm)')
        plt.ylabel('Number of Protein Pairs (Frequency)')
        plt.yscale('log')
        plt.grid(True, linestyle='--', alpha=0.6)
        plt.xlim(2.5, 17.5)

        # --- MODIFICATION: Add requested annotation lines ---
        # Get the current x-axis limits to draw segmented lines
        x_min, x_max = plt.xlim()

        # Add a horizontal line at y=10, but only outside the range [4.2, 14.6]
        plt.hlines(y=10, xmin=x_min, xmax=4.2, color='red', linestyle='--', linewidth=2)
        plt.hlines(y=10, xmin=14.6, xmax=x_max, color='red', linestyle='--', linewidth=2)

        # Add vertical lines at x=4.2 and x=14.6
        plt.axvline(x=4.2, color='red', linestyle='--', linewidth=2)
        plt.axvline(x=14.6, color='red', linestyle='--', linewidth=2)

        plt.tight_layout()
        plt.savefig("Initial_Distance_Distribution.png", dpi=600)

    plt.show()
    return None


if __name__ == "__main__":
    plot_initial_distance_vs_interaction_prob(FORCE_FIELD_CONFIG)