hpp_benchmark.py

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.lines as mlines
import pandas as pd
import numpy as np
from scipy import stats
from statsmodels.stats.multitest import multipletests
import os
import sys
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from src.augment_omics.style import Style, set_pub_style, wrap_labels


EXCLUDED = [
    'no medical diagnoses',
    'wear glasses',
    'surgery improve vision',
    'past fertility treatments',
]

METRICS = [
    ('roc_auc', 'ΔROC-AUC', '#dc0ab4', 'o'),   # magenta
    ('pr_auc',  'ΔPR-AUC',  '#ffa300', 's'),   # orange
    ('f2_score','ΔF2',       '#00bfa0', 'D'),   # teal
]

FDR_ALPHA = 0.05


# ── statistical helpers ────────────────────────────────────────────────────────

def compute_stats(df):
    """
    Per disease × metric: mean delta, 95 % CI, FDR-corrected p-value.
    Returns a dict keyed by (disease, metric_key).
    """
    records = []
    for disease, grp in df.groupby('disease'):
        for key, _, _, _ in METRICS:
            delta = grp[f'model_{key}'] - grp[f'base_{key}']
            n = len(delta)
            mean = delta.mean()
            sem  = delta.sem()
            _, pval = stats.wilcoxon(delta) if n > 1 else (None, 1.0)
            ci_half = stats.t.ppf(0.975, df=n - 1) * sem if n > 1 else 0.0
            records.append(dict(disease=disease, metric=key,
                                mean=mean, ci=ci_half, pval=pval,
                                delta_vals=delta.values))
    result_df = pd.DataFrame(records)

    # FDR correction (Benjamini–Hochberg) across all tests
    reject, qvals, _, _ = multipletests(result_df['pval'], alpha=FDR_ALPHA, method='fdr_bh')
    result_df['qval']   = qvals
    result_df['reject'] = reject
    result_df['sig_pos'] = result_df['reject'] & (result_df['mean'] > 0)
    result_df['sig_neg'] = result_df['reject'] & (result_df['mean'] < 0)
    
    # drop delta_vals
    result_df = result_df.drop('delta_vals', axis=1)
    
    return result_df


def fdr_stars(q):
    if q < 0.001:
        return '***'
    elif q < 0.01:
        return '**'
    elif q < 0.05:
        return '*'
    return ''


# ── sorting ────────────────────────────────────────────────────────────────────

def sort_diseases(stats_df):
    """Sort by (consensus_count DESC, mean ΔROC-AUC DESC)."""
    consensus = (
        stats_df[stats_df['sig_pos']]
        .groupby('disease')
        .size()
        .reindex(stats_df['disease'].unique(), fill_value=0)
    )
    roc_mean = (
        stats_df[stats_df['metric'] == 'roc_auc']
        .set_index('disease')['mean']
    )
    order = (
        pd.DataFrame({'consensus': consensus, 'roc': roc_mean})
        .sort_values(['consensus', 'roc'], ascending=[False, False])
        .index.tolist()
    )
    return order


# ── figure ────────────────────────────────────────────────────────────────────

def make_figure(df, stats_df, disease_order, case_counts, figures_dir, model_name):
    set_pub_style()

    n = len(disease_order)
    x_pos = np.arange(n)

    # Disease labels with case count
    def make_label(d):
        cnt = case_counts.get(d, '')
        name = d.strip().capitalize()
        return f'{name} ({cnt})' if cnt != '' else name

    xlabels = [make_label(d) for d in disease_order]
    
    # xlabels = [l for l in xlabels if l.startswith]
    # if the label is ADHD full name, change it to ADHD
    xlabels = [l.replace('Attention deficit hyperactivity disorder', 'ADHD') for l in xlabels]

    metric_keys = [k   for k, _, _, _ in METRICS]
    metric_lbls = [lbl for _, lbl, _, _ in METRICS]

    # ── layout: panel a (top dot plot) + panel c (bottom heatmap) ─────────────
    fig_w = max(Style.width, n * 0.32 + 1.2)
    fig = plt.figure(figsize=(fig_w, 4.2))

    gs = fig.add_gridspec(
        2, 1,
        height_ratios=[3.5, 1],
        hspace=0.06,
        left=0.06, right=0.94,
        top=0.90, bottom=0.28,
    )
    ax_a = fig.add_subplot(gs[0])
    ax_c = fig.add_subplot(gs[1], sharex=ax_a)

    # ── Panel a: vertical multi-metric dot plot ────────────────────────────────
    metric_offsets = [-0.22, 0.0, 0.22]

    for (key, label, color, marker), x_off in zip(METRICS, metric_offsets):
        xs, ys, cis = [], [], []
        for i, disease in enumerate(disease_order):
            row = stats_df[(stats_df['disease'] == disease) &
                           (stats_df['metric'] == key)].iloc[0]
            xs.append(x_pos[i] + x_off)
            ys.append(row['mean'])
            cis.append(row['ci'])

        ax_a.errorbar(
            xs, ys, yerr=cis,
            fmt='none',
            ecolor=color, elinewidth=0.6, capsize=1.5, zorder=2,
        )
        ax_a.plot(
            xs, ys, marker=marker,
            color=color, markerfacecolor=color,
            markersize=4.5, linewidth=0, zorder=3, label=label,
        )

    ax_a.axhline(0, color='black', linewidth=0.6, linestyle='--', zorder=1)
    ax_a.set_ylabel(f'Δ Score ({model_name} − baseline)')
    ax_a.grid(axis='y', linestyle=':', linewidth=0.4, alpha=0.4)
    ax_a.set_xlim(-0.8, n - 0.2)
    ax_a.tick_params(axis='x', bottom=False, labelbottom=False)
    ax_a.spines['bottom'].set_visible(False)

    legend_elems = [
        mlines.Line2D([], [], color=color, marker=marker,
                      markerfacecolor=color, linestyle='none',
                      markersize=4, label=label)
        for _, label, color, marker in METRICS
    ]
    ax_a.legend(handles=legend_elems, frameon=False,
                loc='upper right', ncol=3, handletextpad=0.3,
                bbox_to_anchor=(1.0, 1.0))

    # ── Panel c: heatmap (metrics × diseases) ─────────────────────────────────
    heat_data = np.zeros((3, n))
    star_data  = [['' ] * n for _ in range(3)]
    for j, disease in enumerate(disease_order):
        for i, key in enumerate(metric_keys):
            row = stats_df[(stats_df['disease'] == disease) &
                           (stats_df['metric'] == key)].iloc[0]
            heat_data[i, j] = row['mean']
            star_data[i][j] = fdr_stars(row['qval'])

    vmax = np.abs(heat_data).max()
    im = ax_c.imshow(
        heat_data,
        aspect='auto',
        cmap='RdBu_r',
        vmin=-vmax, vmax=vmax,
    )

    for i in range(3):
        for j in range(n):
            stars = star_data[i][j]
            val   = heat_data[i, j]
            txt_color = 'white' if abs(val) > vmax * 0.55 else 'black'
            ax_c.text(j, i, stars, ha='center', va='center',
                      fontsize=7, color=txt_color, fontweight='bold')

    ax_c.set_yticks([0, 1, 2])
    ax_c.set_yticklabels(metric_lbls)
    ax_c.set_xticks(x_pos)
    ax_c.set_xticklabels(wrap_labels(xlabels, width=40),
                         rotation=45, ha='right')

    # Horizontal colorbar inset in panel a, just below the legend
    cbar_ax = ax_a.inset_axes([0.77, 0.76, 0.16, 0.05])
    cbar = fig.colorbar(im, cax=cbar_ax, orientation='horizontal')
    cbar.set_ticks([-vmax, 0, vmax])
    cbar.set_ticklabels([f'{-vmax:.2f}', '0', f'{vmax:.2f}'])
    cbar.set_label('Δ Score', labelpad=2)
    cbar.ax.tick_params(labelsize=8, pad=1)
    cbar_ax.xaxis.set_ticks_position('bottom')
    cbar_ax.xaxis.set_label_position('bottom')

    # ── save ───────────────────────────────────────────────────────────────────
    os.makedirs(figures_dir, exist_ok=True)
    slug = model_name.lower().replace(' ', '_')
    out_path = os.path.join(figures_dir, f'hpp_benchmark_{slug}.png')
    fig.savefig(out_path, dpi=300, bbox_inches='tight')
    print(f'Saved → {out_path}')




# ── main ──────────────────────────────────────────────────────────────────────

def main():
    metrics_file  = 'results/data/hpp_scores_5_models.csv'
    counts_file   = 'results/data/hpp_case_counts.csv'
    figures_dir   = 'results/main/hpp_results'

    df = pd.read_csv(metrics_file)
    df = df[~df['disease'].isin(EXCLUDED)].copy()
    df['disease'] = df['disease'].str.strip()

    counts_df  = pd.read_csv(counts_file)
    case_counts = dict(zip(
        counts_df['medical_condition'].str.strip(),
        counts_df['count'],
    ))

    for model_name, model_df in df.groupby('model_name'):
        model_df   = model_df.copy()
        stats_df      = compute_stats(model_df)
        ## save stats
        stats_df.to_csv(f'{figures_dir}/{model_name}_stats.csv', index=False)        
        
        disease_order = sort_diseases(stats_df)
        make_figure(model_df, stats_df, disease_order, case_counts, figures_dir, model_name)


if __name__ == '__main__':
    main()