plot_tm_funnels.py

import sys
from tqdm import tqdm
from util import compute_structure_similarity

#if len(sys.argv) < 2:
#    print("Usage: python script.py <pdb_id>")
#    sys.exit(1)

#pdb_id = sys.argv[1].upper()


pdb_ids = ['2jw1', '2lgk', '2lsk', '2kwv', '2law', '2mnu', '2mps', '5tp6', '5urn', '7ovc', '7jyn', '7jq8', '6h8c']
for pdb_id in pdb_ids:
    # Get CSP rank score file path
    csp_rank_score_file = f'./CSP_Rank_Scores/CSP_{pdb_id.lower()}_CSpred.csv'

    import os
    from os import listdir
    from os.path import isfile, join
    import pandas as pd

    # Read the CSP rank scores file
    try:
        df = pd.read_csv(csp_rank_score_file)
    except Exception as e:
        print(f"Error reading CSP rank scores file: {e}")
        sys.exit(1)

    # Define a function to determine the special case color
    def determine_special_case(holo_model_path):
        if 'exp_' + pdb_id.lower() in holo_model_path:
            return 'green'
        elif 'comp_' + pdb_id.lower() in holo_model_path:
            return 'cyan'
        elif 'v3_' in holo_model_path and 'dropout' in holo_model_path:
            return 'blue'
        elif 'v2_' in holo_model_path and 'dropout' in holo_model_path:
            return 'pink'
        elif 'dropout' in holo_model_path:
            return 'red'
        elif 'v2_' in holo_model_path:
            return 'purple'
        elif 'v3_' in holo_model_path:
            return 'purple'
        elif 'notemplate' in holo_model_path:
            return 'orange'
        elif 'multimer' in holo_model_path:
            return 'yellow'
        else:
            return 'NA'

    # Apply the function to each row in the dataframe
    df['special_cases'] = df['holo_model_path'].apply(determine_special_case)

    # Create scatter plot of consensus vs TM scores, colored by special cases
    import matplotlib.pyplot as plt

    plt.figure(figsize=(10, 6))

    # Define colors for each special case
    colors = {
        'green': 'NMR',
        'cyan': 'Baseline AF2',
        'blue': 'AFS v3',
        'pink': 'AFS v2',
        'red': 'AFS v1',
        'purple': 'AFS2 v2',
        'orange': 'AF ALT',
        'yellow': 'AFS2 v1/3',
        'gray': 'NA'
    }

    # Create a figure with two subplots side by side
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(20, 6))

    # Plot consensus scores on first subplot
    for case, color in colors.items():
        subset = df[df['special_cases'] == case]
        ax1.scatter(subset['tm_score'], subset['consensus'], alpha=0.5, label=None, color=case)

    ax1.set_ylabel('Consensus Score')
    ax1.set_xlabel('TM Score')
    ax1.set_title('Consensus Score vs TM Score')
    ax1.grid(True, linestyle='--', alpha=0.7)
    ax1.set_xlim(0, 1)
    ax1.set_ylim(0, 1)

    # Plot consensus*confidence scores on second subplot
    for case, color in colors.items():
        subset = df[df['special_cases'] == case]
        ax2.scatter(subset['tm_score'], subset['consensus'] * subset['Confidence'], alpha=0.5, label=color, color=case)

    ax2.set_ylabel('Bayes Score')
    ax2.set_xlabel('TM Score')
    ax2.set_title('Bayes Score vs TM Score')
    ax2.grid(True, linestyle='--', alpha=0.7)
    ax2.set_xlim(0, 1)
    ax2.set_ylim(0, 1)

    # Adjust layout and add legend
    plt.subplots_adjust(right=0.85)  # Make room for legend
    fig.legend(title='Special Cases', loc='center right', bbox_to_anchor=(0.98, 0.5))

    # Save the plot
    plt.savefig(f'./Figures/bayes_tm_scatter_{pdb_id.lower()}.png', bbox_inches='tight', dpi=300)
    plt.close()