mechanisms.py

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from augment_omics.style import Style, set_pub_style, wrap_labels
import os


set_pub_style()

style_colors = [Style.blue, Style.teal, Style.purple, Style.orange, Style.green,
                Style.red, Style.yellow, Style.brown, Style.olive, Style.silver]
model_color_map = {
    'CLIP':        Style.blue,
    'FedCoder':    Style.teal,
    'Transformer': Style.purple,
    'MLP':         Style.orange,
    'Ridge':       Style.green,
    'Proteomics':  Style.red,
}

name_dict = {
    'clip': 'CLIP',
    'clip_selected': 'CLIP_Selected',
    'fedcoder': 'FedCoder',
    'transformer': 'Transformer',
    'mlp': 'MLP',
    'cov': 'Covariates',
    'clip_concat_raw': 'CLIP+Raw',
    'ridge': 'Ridge',
    'proteomics': 'Proteomics',
}

PROT_VS_MET_PATH = 'results/data/survival_metrics/paired_proteomics_vs_metabolomics_with_cov.csv'

no_cov_model_dict = {
    'CLIP':        'results/data/survival_metrics/paired_clip_vs_baseline_no_cov.csv',
    'FedCoder':    'results/data/survival_metrics/paired_fedcoder_vs_baseline_no_cov.csv',
    'Transformer': 'results/data/survival_metrics/paired_transformer_translation_vs_baseline_no_cov.csv',
    'MLP':         'results/data/survival_metrics/paired_mlp_vs_baseline_no_cov.csv',
    'Ridge':       'results/data/survival_metrics/paired_ridge_vs_baseline_no_cov.csv',
    'Proteomics':  'results/data/survival_metrics/paired_proteomics_vs_baseline_no_cov.csv',
}

correlated_proteins = {
    'other acute ischaemic\nheart diseases':
        ['PLTP', 'KITLG', 'GIP', 'FGF21', 'HSD11B1',
         'MFGE8', 'LDLR', 'IGFBP1', 'PCSK9', 'TGFA'],
    'chronic nephritic\nsyndrome':
        ['PLA2G7', 'FGFBP1', 'ANGPTL4', 'LDLR', 'VMO1',
         'ANGPTL3', 'APOC1', 'MFGE8', 'TNFSF11', 'HRC'],
    'vascular dementia':
        ['FGF21', 'NRP1', 'ANGPTL3', 'FST', 'PLA2G10',
         'ANGPTL4', 'FABP4', 'APOA1', 'LCAT', 'IGFBP2'],
    'alcoholic liver\ndisease':
        ['PLA2G7', 'KITLG', 'FGF21', 'LDLR', 'ANGPTL4',
         'IGFBP1', 'FGL1', 'IL1RL2', 'MVK', 'CHI3L1'],
}

# ---- Disease subsets ----
prot_vs_met_df = pd.read_csv(PROT_VS_MET_PATH)
mask = prot_vs_met_df['cindex_diff_mean'] > 0.01
potential_codes = set(prot_vs_met_df[mask]['code'].values)
rest_codes = set(prot_vs_met_df[~mask]['code'].values)

# ---- Panel a: Correlation box plot data ----
scores = {}
for file in os.listdir('results/data/correlation'):
    if not file.endswith('.csv'):
        continue
    model_name = file.split('_')[0].lower()
    model_name = name_dict.get(model_name, model_name)
    df = pd.read_csv(f'results/data/correlation/{file}')
    scores[model_name] = df['pearson_r'].values

corr_model_names = sorted(scores.keys(), key=lambda m: np.median(scores[m]))
corr_plot_data = [scores[m] for m in corr_model_names]
corr_colors = [model_color_map.get(m, style_colors[i % len(style_colors)])
               for i, m in enumerate(corr_model_names)]

# ---- Panels b & c: Delta C-index data ----
def build_delta_data(disease_codes):
    delta_scores = {}
    for model_name, csv_path in no_cov_model_dict.items():
        df = pd.read_csv(csv_path)
        df = df[df['code'].isin(disease_codes)]
        delta_scores[model_name] = (df['cindex_1_mean'] - df['cindex_2_mean']).values
    model_names = sorted(delta_scores.keys(), key=lambda m: np.median(delta_scores[m]))
    plot_data = [delta_scores[m] for m in model_names]
    colors = [model_color_map.get(m, style_colors[i % len(style_colors)])
              for i, m in enumerate(model_names)]
    return model_names, plot_data, colors

prot_better_names, prot_better_data, prot_better_colors = build_delta_data(potential_codes)
rest_names, rest_data, rest_colors = build_delta_data(rest_codes)

# ---- Panel d-g: CLIP correlation bar chart data ----
df_clip_corr = pd.read_csv('results/data/correlation/clip_protein_correlations.csv')

# ---- Figure layout: 2 rows, shared 12-column grid ----
# Row 1: 3 box plots, each spanning 4 cols
fig = plt.figure(figsize=(Style.width * 1.5, Style.height))
gs = gridspec.GridSpec(2, 12, figure=fig, hspace=0.55, wspace=1.2)

ax_corr    = fig.add_subplot(gs[0, 0:4])
ax_prot    = fig.add_subplot(gs[0, 4:8])
ax_rest    = fig.add_subplot(gs[0, 8:12])

def draw_boxplot(ax, plot_data, model_names, colors, ylabel):
    bp = ax.boxplot(plot_data, patch_artist=True, widths=0.5,
                    medianprops=dict(color='black', linewidth=1),
                    meanprops=dict(marker='D', markerfacecolor='black', markersize=3),
                    whiskerprops=dict(linewidth=0.8), capprops=dict(linewidth=0.8),
                    flierprops=dict(marker='o', markersize=1.5, alpha=0.5),
                    showmeans=True)
    for patch, color in zip(bp['boxes'], colors):
        patch.set_facecolor(color)
        patch.set_alpha(0.7)
    ax.set_xticks(range(1, len(model_names) + 1))
    ax.set_xticklabels(wrap_labels(model_names, width=12), rotation=45, ha='right')
    ax.set_ylabel(ylabel)
    
    if ylabel != 'Pearson correlation':
        ax.set_ylim([-0.35, 0.35])
    ax.tick_params(axis='y', pad=1)  # default is ~4, lower = closer


draw_boxplot(ax_corr, corr_plot_data, corr_model_names, corr_colors, 'Pearson correlation')

draw_boxplot(ax_prot, prot_better_data, prot_better_names, prot_better_colors,
             'Δ C-Index (Model − Baseline)')
ax_prot.axhline(y=0, color='gray', linestyle='--', linewidth=1)

draw_boxplot(ax_rest, rest_data, rest_names, rest_colors,
             'Δ C-Index (Model − Baseline)')
ax_rest.axhline(y=0, color='gray', linestyle='--', linewidth=1)


os.makedirs('results/main/mechanism', exist_ok=True)
plt.savefig('results/main/mechanism/mechanism_panels_abc.png', dpi=300, bbox_inches='tight')