hypertower/scripts/output_analysis/visualizations/plot_early_stopping_sweep.py

#!/usr/bin/env python3
"""
Simulate early stopping at each epoch N and show what val/holdout AUC
you would have gotten if you stopped there.

For each fold and each candidate stopping epoch N:
  - Find the epoch <= N with the highest val AUC  (checkpoint selection)
  - Record the val AUC and holdout AUC at that epoch

Then plot mean ± std across all folds as a function of N.

Usage
-----
    python scripts/output_analysis/visualizations/plot_early_stopping_sweep.py \
        --run-dir analysis_data/pipeline_10x5 \
        --eval-mode binary \
        --tower-mode ensemble \
        --head fused
"""
from __future__ import annotations

import argparse
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd


HEAD_COL = {
    "fused":   ("ensemble_val_auc",     "ensemble_holdout_auc"),
    "img":     ("ensemble_val_auc_img", "ensemble_holdout_auc_img"),
    "md":      ("ensemble_val_auc_md",  "ensemble_holdout_auc_md"),
    "classic": ("classic_val_auc",      "classic_holdout_auc"),
}


def load_fold_logs(run_dir: Path, eval_mode: str, tower_mode: str,
                   val_col: str, hld_col: str):
    logs = []
    for rep_dir in sorted(run_dir.glob("rep*")):
        mode_dir = rep_dir / eval_mode / tower_mode
        if not mode_dir.exists():
            continue
        for fd in sorted(
            [d for d in mode_dir.iterdir() if d.is_dir() and d.name.startswith("fold")],
            key=lambda p: int(p.name.replace("fold", "")),
        ):
            log = fd / "epoch_log.csv"
            if not log.exists():
                continue
            df = pd.read_csv(log)
            if val_col not in df.columns or hld_col not in df.columns:
                continue
            df = df[["epoch", val_col, hld_col]].dropna()
            logs.append(df.reset_index(drop=True))
    return logs


def sweep(logs: list[pd.DataFrame], val_col: str, hld_col: str):
    max_epoch = max(df["epoch"].max() for df in logs)
    epochs = np.arange(1, int(max_epoch) + 1)

    val_mat = np.full((len(logs), len(epochs)), np.nan)
    hld_mat = np.full((len(logs), len(epochs)), np.nan)

    for i, df in enumerate(logs):
        for j, n in enumerate(epochs):
            window = df[df["epoch"] <= n]
            if window.empty:
                continue
            best_idx = window[val_col].idxmax()
            val_mat[i, j] = window.loc[best_idx, val_col]
            hld_mat[i, j] = window.loc[best_idx, hld_col]

    return epochs, val_mat, hld_mat


def plot(epochs, val_mat, hld_mat, out_path: Path, title: str, inverted: bool = False, ci_tol: float = 0.0):
    val_mean = np.nanmean(val_mat, axis=0)
    val_std  = np.nanstd(val_mat,  axis=0)
    hld_mean = np.nanmean(hld_mat, axis=0)
    hld_std  = np.nanstd(hld_mat,  axis=0)

    if inverted:
        # compute cost per fold, then aggregate — avoids max-of-mean bias
        best_val_per_fold = np.nanmax(val_mat, axis=1, keepdims=True)  # (n_folds, 1)
        best_hld_per_fold = np.nanmax(hld_mat, axis=1, keepdims=True)
        delta_val = best_val_per_fold - val_mat   # (n_folds, n_epochs)
        delta_hld = best_hld_per_fold - hld_mat
        y_val  = np.nanmean(delta_val, axis=0)
        y_hld  = np.nanmean(delta_hld, axis=0)
        sy_val = np.nanstd(delta_val,  axis=0)
        sy_hld = np.nanstd(delta_hld,  axis=0)
    else:
        y_val, y_hld   = val_mean, hld_mean
        sy_val, sy_hld = val_std,  hld_std

    fig, ax = plt.subplots(figsize=(11, 5))

    if inverted:
        # faint per-fold lines
        for i in range(delta_val.shape[0]):
            ax.plot(epochs, delta_val[i], color="steelblue", linewidth=0.6, alpha=0.18)
            ax.plot(epochs, delta_hld[i], color="firebrick", linewidth=0.6, alpha=0.18)

    ax.plot(epochs, y_val, color="steelblue", linewidth=2.0,
            label="Best val − val@N  (val cost of stopping early)" if inverted
                  else "Val AUC (best ckpt up to N)")
    ax.fill_between(epochs, y_val - sy_val, y_val + sy_val, color="steelblue", alpha=0.15)

    ax.plot(epochs, y_hld, color="firebrick", linewidth=2.0,
            label="Best hld − hld@N  (holdout cost of stopping early)" if inverted
                  else "Holdout AUC (at best val ckpt)")
    ax.fill_between(epochs, y_hld - sy_hld, y_hld + sy_hld, color="firebrick", alpha=0.15)

    if inverted:
        ax.axhline(0, color="black", linewidth=1.0, linestyle="--", alpha=0.4)

        # CI-crosses-zero regions (with optional tolerance)
        val_ci_zero  = (y_val - sy_val) <= ci_tol
        hld_ci_zero  = (y_hld - sy_hld) <= ci_tol
        both_ci_zero = val_ci_zero & hld_ci_zero

        ymax = max(np.nanmax(y_val), np.nanmax(y_hld)) * 1.15
        ax.fill_between(epochs, 0, ymax, where=val_ci_zero,
                        color="steelblue", alpha=0.12, label="val CI ≤ 0")
        ax.fill_between(epochs, 0, ymax, where=hld_ci_zero,
                        color="firebrick", alpha=0.12, label="holdout CI ≤ 0")
        ax.fill_between(epochs, 0, ymax, where=both_ci_zero,
                        color="purple", alpha=0.20, label="both CI ≤ 0")

        ax.set_ylabel("AUC lost vs best achievable")
        ax.set_ylim(-0.05, ymax)
        legend_loc = "upper right"
    else:
        # gap curve on twin axis
        gap_mean = val_mean - hld_mean
        ax2 = ax.twinx()
        ax2.plot(epochs, gap_mean, color="darkorange", linewidth=1.5,
                 linestyle="--", alpha=0.7, label="Val−Holdout gap")
        ax2.set_ylabel("Val − Holdout gap", color="darkorange", fontsize=9)
        ax2.tick_params(axis="y", labelcolor="darkorange")
        ax2.set_ylim(-0.1, 0.4)
        lines2, labels2 = ax2.get_legend_handles_labels()

        best_hld_ep  = epochs[np.nanargmax(hld_mean)]
        best_hld_val = hld_mean[np.nanargmax(hld_mean)]
        ax.axvline(best_hld_ep, color="firebrick", linewidth=1.2, linestyle=":",
                   alpha=0.8, label=f"peak holdout @ epoch {best_hld_ep} ({best_hld_val:.3f})")
        stable = epochs >= 3
        min_gap_ep = epochs[stable][np.nanargmin(gap_mean[stable])]
        ax.axvline(min_gap_ep, color="darkorange", linewidth=1.2, linestyle=":",
                   alpha=0.8, label=f"min gap @ epoch {min_gap_ep}")
        ax.set_ylabel("AUC")
        ax.set_ylim(0.5, 1.05)
        legend_loc = "lower right"

    ax.set_xlabel("Stopping epoch N")
    ax.set_title(title)
    ax.grid(axis="y", linestyle="--", alpha=0.35)

    lines1, labels1 = ax.get_legend_handles_labels()
    if not inverted:
        lines1 += lines2; labels1 += labels2
    ax.legend(lines1, labels1, fontsize=8, loc=legend_loc)

    out_path.parent.mkdir(parents=True, exist_ok=True)
    fig.tight_layout()
    fig.savefig(out_path, dpi=150)
    plt.close(fig)
    print(f"Saved → {out_path}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--run-dir",    default="analysis_data/pipeline_10x5")
    ap.add_argument("--eval-mode",  default="binary")
    ap.add_argument("--tower-mode", default="ensemble")
    ap.add_argument("--head",       default="fused", choices=list(HEAD_COL))
    ap.add_argument("--out",        default=None)
    ap.add_argument("--inverted",   action="store_true",
                    help="Plot best-achievable minus current (cost of stopping early)")
    ap.add_argument("--ci-tol",     type=float, default=0.0,
                    help="Tolerance for CI-crosses-zero shading (default 0.0)")
    args = ap.parse_args()

    run_dir = Path(args.run_dir)
    val_col, hld_col = HEAD_COL[args.head]

    logs = load_fold_logs(run_dir, args.eval_mode, args.tower_mode, val_col, hld_col)
    if not logs:
        print("No epoch_log.csv files found.")
        return
    print(f"Loaded {len(logs)} fold logs")

    epochs, val_mat, hld_mat = sweep(logs, val_col, hld_col)

    tol_tag = f"_tol{args.ci_tol}" if args.ci_tol else ""
    suffix = f"_inverted{tol_tag}" if args.inverted else ""
    out = Path(args.out) if args.out else (
        run_dir / args.eval_mode / args.tower_mode / "plots" /
        f"early_stopping_sweep_{args.head}{suffix}.png"
    )
    title = ("Simulated early stopping — cost of stopping at epoch N\n"
             if args.inverted else
             "Simulated early stopping sweep\n")
    title += f"{run_dir.name} · {args.eval_mode}/{args.tower_mode} · head={args.head}"
    if args.inverted and args.ci_tol:
        title += f"  (CI tol={args.ci_tol})"
    plot(epochs, val_mat, hld_mat, out, title, inverted=args.inverted, ci_tol=args.ci_tol)


if __name__ == "__main__":
    main()