gpu/single-gpu/home_benjamin_hdbasin/code/scripts/build_paper_figures.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import math
from pathlib import Path

import numpy as np
import pandas as pd


METHOD_ORDER = [
    "HD-BasinFlow",
    "ASHA Successive Halving",
    "Optuna TPE",
    "Random Search",
    "Sobol Quasi-Random",
]

BASELINE_ORDER = [
    "ASHA Successive Halving",
    "Optuna TPE",
    "Random Search",
    "Sobol Quasi-Random",
]

WORKLOAD_LABELS = {
    "fashion_mnist_tiny_cnn": "Fashion-MNIST CNN",
    "cifar10_tiny_cnn": "CIFAR-10 CNN",
    "sst2_distilbert": "SST-2 DistilBERT",
    "imdb_distilbert": "IMDB DistilBERT",
    "ag_news_tfidf_logreg": "AG News TF-IDF",
    "tabular_credit_hgb": "Tabular Credit HGB",
}

METHOD_COLORS = {
    "HD-BasinFlow": "#1b7f5a",
    "ASHA Successive Halving": "#6d5dfc",
    "Optuna TPE": "#d55e00",
    "Random Search": "#4c78a8",
    "Sobol Quasi-Random": "#7f7f7f",
}


def pct(value: float) -> str:
    if pd.isna(value):
        return "n/a"
    return f"{100 * float(value):.1f}%"


def clean_label(value: str) -> str:
    return WORKLOAD_LABELS.get(value, value.replace("_", " "))


def setup_matplotlib():
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    plt.rcParams.update(
        {
            "figure.dpi": 140,
            "savefig.dpi": 300,
            "font.size": 10,
            "axes.titlesize": 12,
            "axes.labelsize": 10,
            "legend.fontsize": 8,
            "xtick.labelsize": 8,
            "ytick.labelsize": 8,
            "axes.spines.top": False,
            "axes.spines.right": False,
            "pdf.fonttype": 42,
            "ps.fonttype": 42,
        }
    )
    return plt


def savefig(plt, outdir: Path, name: str) -> None:
    plt.tight_layout()
    plt.savefig(outdir / f"{name}.png")
    plt.savefig(outdir / f"{name}.pdf")
    plt.close()


def load_inputs(input_dir: Path) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
    evals = pd.read_csv(input_dir / "combined" / "raw" / "evaluations.csv")
    metrics = pd.read_csv(input_dir / "combined" / "processed" / "metrics.csv")
    decisions = pd.read_csv(input_dir / "target_savings" / "target_savings_decisions.csv")
    return evals, metrics, decisions


def normalize_evals(evals: pd.DataFrame) -> pd.DataFrame:
    df = evals.copy()
    for col in ["cumulative_cost", "best_loss_so_far", "loss", "walltime", "iteration", "seed"]:
        if col in df:
            df[col] = pd.to_numeric(df[col], errors="coerce")
    df["method_label"] = np.where(df["method"].eq("HD-BasinFlow"), "HD-BasinFlow", df["method"])
    df = df[df["method_label"].isin(METHOD_ORDER)].copy()
    df["workload_label"] = df["experiment"].map(clean_label)
    return df


def normalize_metrics(metrics: pd.DataFrame) -> pd.DataFrame:
    df = metrics.copy()
    for col in ["median_final_best_loss", "median_walltime", "median_cost", "runs"]:
        if col in df:
            df[col] = pd.to_numeric(df[col], errors="coerce")
    df["method_label"] = np.where(df["method"].eq("HD-BasinFlow"), "HD-BasinFlow", df["method"])
    df = df[df["method_label"].isin(METHOD_ORDER)].copy()
    df["workload_label"] = df["experiment"].map(clean_label)
    return df


def figure_pass_rate(plt, decisions: pd.DataFrame, outdir: Path) -> str:
    summary = (
        decisions.groupby("baseline", as_index=False)
        .agg(comparisons=("baseline", "size"), passes=("decision", lambda s: int((s == "pass").sum())))
    )
    summary["pass_rate"] = summary["passes"] / summary["comparisons"]
    summary["baseline"] = pd.Categorical(summary["baseline"], BASELINE_ORDER, ordered=True)
    summary = summary.sort_values("baseline")
    fig, ax = plt.subplots(figsize=(6.6, 3.8))
    bars = ax.bar(summary["baseline"].astype(str), summary["pass_rate"] * 100, color="#1b7f5a")
    ax.set_ylim(0, 105)
    ax.set_ylabel("pass rate (%)")
    ax.set_title("Target-savings pass rate by competitor")
    ax.set_xticks(np.arange(len(summary)))
    ax.set_xticklabels(summary["baseline"].astype(str), rotation=25, ha="right")
    for bar, row in zip(bars, summary.itertuples(index=False)):
        ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 2, f"{row.passes}/{row.comparisons}", ha="center", va="bottom")
    savefig(plt, outdir, "fig01_pass_rate_by_baseline")
    return "fig01_pass_rate_by_baseline"


def figure_savings_by_baseline(plt, decisions: pd.DataFrame, outdir: Path) -> str:
    summary = (
        decisions.groupby("baseline", as_index=False)
        .agg(
            eval_savings=("evaluation_savings_fraction", "median"),
            wall_savings=("walltime_savings_fraction", "median"),
        )
    )
    summary["baseline"] = pd.Categorical(summary["baseline"], BASELINE_ORDER, ordered=True)
    summary = summary.sort_values("baseline")
    x = np.arange(len(summary))
    width = 0.36
    fig, ax = plt.subplots(figsize=(7.2, 4.0))
    ax.bar(x - width / 2, summary["eval_savings"] * 100, width, label="evaluations", color="#1b7f5a")
    ax.bar(x + width / 2, summary["wall_savings"] * 100, width, label="wall time", color="#4c78a8")
    ax.axhspan(15, 30, color="#e6f2ec", zorder=0, label="15-30% target band")
    ax.set_ylabel("median savings (%)")
    ax.set_title("Median compute savings at matched quality target")
    ax.set_xticks(x)
    ax.set_xticklabels(summary["baseline"].astype(str), rotation=25, ha="right")
    ax.legend(frameon=False)
    savefig(plt, outdir, "fig02_median_savings_by_baseline")
    return "fig02_median_savings_by_baseline"


def figure_quality_savings_scatter(plt, decisions: pd.DataFrame, outdir: Path) -> str:
    fig, ax = plt.subplots(figsize=(7.0, 4.4))
    for baseline in BASELINE_ORDER:
        g = decisions[decisions["baseline"].eq(baseline)]
        if g.empty:
            continue
        ax.scatter(
            g["evaluation_savings_fraction"] * 100,
            g["loss_delta_pct"] * 100,
            s=70,
            label=baseline,
            alpha=0.86,
            edgecolor="white",
            linewidth=0.7,
        )
    ax.axhline(1.0, color="#8f2f2f", linestyle="--", linewidth=1.2, label="1% quality tolerance")
    ax.axvline(20.0, color="#444444", linestyle=":", linewidth=1.2, label="20% savings target")
    ax.set_xlabel("evaluation savings (%)")
    ax.set_ylabel("HD loss delta vs baseline (%)")
    ax.set_title("Quality-compute tradeoff across workloads")
    ax.legend(frameon=False, ncols=2)
    savefig(plt, outdir, "fig03_quality_vs_compute_savings")
    return "fig03_quality_vs_compute_savings"


def figure_hd_vs_baseline_savings(plt, decisions: pd.DataFrame, outdir: Path) -> str:
    data = decisions.copy()
    data["comparison"] = data["experiment"].map(clean_label) + "\nvs " + data["baseline"]
    data = data.sort_values("evaluation_savings_fraction", ascending=True)
    colors = np.where(data["decision"].eq("pass"), "#1b7f5a", "#8f2f2f")
    fig, ax = plt.subplots(figsize=(8.6, 8.2))
    y = np.arange(len(data))
    ax.barh(y, data["evaluation_savings_fraction"] * 100, color=colors)
    ax.axvline(20, color="#222222", linestyle="--", linewidth=1.2, label="20% savings target")
    ax.axvspan(15, 30, color="#e6f2ec", zorder=0, label="15-30% target band")
    ax.set_yticks(y)
    ax.set_yticklabels(data["comparison"], fontsize=7)
    ax.set_xlabel("HD-BasinFlow evaluation savings vs baseline (%)")
    ax.set_title("Direct comparison: HD-BasinFlow compute savings vs each baseline")
    ax.legend(frameon=False, loc="lower right")
    for idx, row in enumerate(data.itertuples(index=False)):
        val = row.evaluation_savings_fraction * 100
        ax.text(val + 1, idx, "pass" if row.decision == "pass" else "fail", va="center", fontsize=7)
    savefig(plt, outdir, "fig09_hd_vs_baseline_savings")
    return "fig09_hd_vs_baseline_savings"


def figure_hd_vs_baseline_quality_delta(plt, decisions: pd.DataFrame, outdir: Path) -> str:
    data = decisions.copy()
    data["comparison"] = data["experiment"].map(clean_label) + "\nvs " + data["baseline"]
    data = data.sort_values("loss_delta_pct", ascending=True)
    colors = np.where(data["quality_pass"], "#1b7f5a", "#8f2f2f")
    fig, ax = plt.subplots(figsize=(8.6, 8.2))
    y = np.arange(len(data))
    ax.barh(y, data["loss_delta_pct"] * 100, color=colors)
    ax.axvline(0, color="#222222", linewidth=1.0)
    ax.axvline(1, color="#8f2f2f", linestyle="--", linewidth=1.2, label="1% quality tolerance")
    ax.set_yticks(y)
    ax.set_yticklabels(data["comparison"], fontsize=7)
    ax.set_xlabel("HD-BasinFlow median loss delta vs baseline (%)")
    ax.set_title("Direct comparison: HD-BasinFlow quality delta vs each baseline")
    ax.legend(frameon=False, loc="lower right")
    savefig(plt, outdir, "fig10_hd_vs_baseline_quality_delta")
    return "fig10_hd_vs_baseline_quality_delta"


def _heatmap(plt, data: pd.DataFrame, value_col: str, title: str, cbar_label: str, outdir: Path, name: str) -> str:
    pivot = data.pivot(index="experiment", columns="baseline", values=value_col)
    pivot = pivot.reindex(sorted(pivot.index, key=lambda v: clean_label(v))).reindex(columns=BASELINE_ORDER)
    values = pivot.to_numpy(dtype=float)
    fig, ax = plt.subplots(figsize=(8.2, 4.6))
    vmax = np.nanmax(np.abs(values)) if np.isfinite(values).any() else 1.0
    if value_col == "loss_delta_pct":
        vmax = max(vmax, 0.03)
        im = ax.imshow(values * 100, cmap="RdYlGn_r", vmin=-vmax * 100, vmax=vmax * 100, aspect="auto")
    else:
        im = ax.imshow(values * 100, cmap="YlGn", vmin=0, vmax=max(80, np.nanmax(values * 100)), aspect="auto")
    ax.set_xticks(np.arange(len(pivot.columns)))
    ax.set_xticklabels(pivot.columns, rotation=25, ha="right")
    ax.set_yticks(np.arange(len(pivot.index)))
    ax.set_yticklabels([clean_label(v) for v in pivot.index])
    ax.set_title(title)
    for i in range(values.shape[0]):
        for j in range(values.shape[1]):
            if np.isfinite(values[i, j]):
                ax.text(j, i, f"{values[i, j] * 100:.1f}", ha="center", va="center", color="black", fontsize=8)
    cbar = fig.colorbar(im, ax=ax, shrink=0.9)
    cbar.set_label(cbar_label)
    savefig(plt, outdir, name)
    return name


def figure_heatmaps(plt, decisions: pd.DataFrame, outdir: Path) -> list[str]:
    names = []
    names.append(
        _heatmap(
            plt,
            decisions,
            "evaluation_savings_fraction",
            "Evaluation savings by workload and competitor",
            "savings (%)",
            outdir,
            "fig04_heatmap_evaluation_savings",
        )
    )
    names.append(
        _heatmap(
            plt,
            decisions,
            "loss_delta_pct",
            "Median loss delta by workload and competitor",
            "loss delta (%)",
            outdir,
            "fig05_heatmap_quality_delta",
        )
    )
    return names


def figure_final_loss_by_workload(plt, metrics: pd.DataFrame, outdir: Path) -> list[str]:
    names = []
    use = metrics[metrics["method_label"].isin(METHOD_ORDER)].copy()
    for experiment, g in use.groupby("experiment"):
        g = g.copy()
        g["method_label"] = pd.Categorical(g["method_label"], METHOD_ORDER, ordered=True)
        g = g.sort_values("method_label")
        fig, ax = plt.subplots(figsize=(6.8, 3.7))
        colors = [METHOD_COLORS.get(m, "#777777") for m in g["method_label"].astype(str)]
        ax.bar(g["method_label"].astype(str), g["median_final_best_loss"], color=colors)
        ax.set_title(f"Final validation loss: {clean_label(experiment)}")
        ax.set_ylabel("median final best loss")
        ax.set_xticks(np.arange(len(g)))
        ax.set_xticklabels(g["method_label"].astype(str), rotation=25, ha="right")
        name = f"fig_loss_{experiment}"
        savefig(plt, outdir, name)
        names.append(name)
    return names


def figure_convergence_curves(plt, evals: pd.DataFrame, outdir: Path) -> list[str]:
    names = []
    trend = (
        evals.groupby(["experiment", "method_label", "cumulative_cost"], as_index=False)["best_loss_so_far"]
        .median()
        .sort_values(["experiment", "method_label", "cumulative_cost"])
    )
    for experiment, exp_df in trend.groupby("experiment"):
        fig, ax = plt.subplots(figsize=(7.0, 4.2))
        for method in METHOD_ORDER:
            g = exp_df[exp_df["method_label"].eq(method)]
            if g.empty:
                continue
            ax.plot(
                g["cumulative_cost"],
                g["best_loss_so_far"],
                label=method,
                linewidth=2.0 if method == "HD-BasinFlow" else 1.4,
                color=METHOD_COLORS.get(method),
            )
        ax.set_title(f"Optimization curve: {clean_label(experiment)}")
        ax.set_xlabel("evaluations")
        ax.set_ylabel("median best loss so far")
        ax.legend(frameon=False)
        name = f"fig_curve_{experiment}"
        savefig(plt, outdir, name)
        names.append(name)
    return names


def figure_active_dimension(plt, evals: pd.DataFrame, outdir: Path) -> str | None:
    hd = evals[(evals["method_label"].eq("HD-BasinFlow")) & evals["active_dim"].notna()].copy()
    if hd.empty:
        return None
    hd["active_dim"] = pd.to_numeric(hd["active_dim"], errors="coerce")
    trend = hd.groupby(["experiment", "cumulative_cost"], as_index=False)["active_dim"].median()
    fig, ax = plt.subplots(figsize=(7.0, 4.1))
    for experiment, g in trend.groupby("experiment"):
        ax.plot(g["cumulative_cost"], g["active_dim"], label=clean_label(experiment), linewidth=1.5)
    ax.set_title("HD-BasinFlow active subspace dimension over search")
    ax.set_xlabel("evaluations")
    ax.set_ylabel("median estimated active dimension")
    ax.legend(frameon=False, ncols=2)
    savefig(plt, outdir, "fig06_active_dimension_over_time")
    return "fig06_active_dimension_over_time"


def figure_allocation_sources(plt, evals: pd.DataFrame, outdir: Path) -> str | None:
    hd = evals[evals["method_label"].eq("HD-BasinFlow")].copy()
    if hd.empty or "source" not in hd:
        return None
    table = pd.crosstab(hd["experiment"].map(clean_label), hd["source"], normalize="index") * 100
    keep = [c for c in ["initial", "global", "edge_flow", "basin", "repair"] if c in table.columns]
    table = table[keep]
    fig, ax = plt.subplots(figsize=(8.0, 4.5))
    bottom = np.zeros(len(table))
    colors = ["#4c78a8", "#b0b0b0", "#1b7f5a", "#72b7b2", "#d55e00"]
    x = np.arange(len(table))
    for idx, col in enumerate(table.columns):
        vals = table[col].to_numpy()
        ax.bar(x, vals, bottom=bottom, label=col, color=colors[idx % len(colors)])
        bottom += vals
    ax.set_ylabel("share of HD evaluations (%)")
    ax.set_title("HD-BasinFlow allocation source mix")
    ax.set_xticks(x)
    ax.set_xticklabels(table.index, rotation=25, ha="right")
    ax.legend(frameon=False, ncols=len(table.columns))
    savefig(plt, outdir, "fig07_allocation_source_mix")
    return "fig07_allocation_source_mix"


def figure_train_val_curves(plt, evals: pd.DataFrame, outdir: Path) -> str | None:
    rows = []
    epoch_workloads = {"fashion_mnist_tiny_cnn", "cifar10_tiny_cnn", "sst2_distilbert", "imdb_distilbert"}
    for row in evals.itertuples(index=False):
        method = getattr(row, "method_label")
        if method != "HD-BasinFlow":
            continue
        experiment = getattr(row, "experiment")
        if experiment not in epoch_workloads:
            continue
        val_json = getattr(row, "val_curve_json", None)
        train_json = getattr(row, "train_curve_json", None)
        for kind, blob in [("train", train_json), ("validation", val_json)]:
            if not isinstance(blob, str) or not blob.strip():
                continue
            try:
                points = json.loads(blob)
            except json.JSONDecodeError:
                continue
            for point in points:
                rows.append(
                    {
                        "experiment": experiment,
                        "curve": kind,
                        "epoch": point.get("epoch"),
                        "loss": point.get("loss"),
                    }
                )
    curves = pd.DataFrame(rows)
    if curves.empty:
        return None
    curves["epoch"] = pd.to_numeric(curves["epoch"], errors="coerce")
    curves["loss"] = pd.to_numeric(curves["loss"], errors="coerce")
    curves = curves.dropna()
    curves = curves[(curves["epoch"] >= 0) & (curves["epoch"] <= 10)]
    summary = curves.groupby(["experiment", "curve", "epoch"], as_index=False)["loss"].median()
    fig, ax = plt.subplots(figsize=(7.3, 4.3))
    for (experiment, curve), g in summary.groupby(["experiment", "curve"]):
        linestyle = "-" if curve == "validation" else "--"
        ax.plot(g["epoch"], g["loss"], label=f"{clean_label(experiment)} {curve}", linestyle=linestyle, linewidth=1.4)
    ax.set_title("HD-BasinFlow selected training/validation curves")
    ax.set_xlabel("epoch")
    ax.set_ylabel("median loss")
    ax.legend(frameon=False, fontsize=7, ncols=2)
    savefig(plt, outdir, "fig08_hd_train_validation_curves")
    return "fig08_hd_train_validation_curves"


def write_index(outdir: Path, generated: list[str], decisions: pd.DataFrame) -> None:
    lines = [
        "# HD-BasinFlow Paper Figure Index",
        "",
        "All figures are generated from `runs/commercial_benchmark_a100`, the fresh A100 benchmark suite.",
        "Each figure is saved as both `.png` and `.pdf` unless noted.",
        "",
        "## Headline Evidence",
        "",
        f"- Reviewable comparisons: {len(decisions)}",
        f"- Operating-point passes: {int((decisions['decision'] == 'pass').sum())}",
        f"- Median evaluation savings: {pct(decisions['evaluation_savings_fraction'].median())}",
        f"- Median walltime savings: {pct(decisions['walltime_savings_fraction'].median())}",
        "",
        "## Figures",
        "",
    ]
    captions = {
        "fig01_pass_rate_by_baseline": "Pass rate by competitor under the 1% quality / 20% savings operating-point rule.",
        "fig02_median_savings_by_baseline": "Median evaluation and walltime savings by competitor; shaded region marks the conservative 15-30% commercial band.",
        "fig03_quality_vs_compute_savings": "Quality-compute tradeoff for every workload-baseline pair.",
        "fig04_heatmap_evaluation_savings": "Evaluation savings heatmap across workloads and competitors.",
        "fig05_heatmap_quality_delta": "Median loss delta heatmap across workloads and competitors.",
        "fig06_active_dimension_over_time": "HD-BasinFlow active subspace estimates during the search.",
        "fig07_allocation_source_mix": "How HD-BasinFlow allocated evaluations across initialization, global search, edge flow, basin continuation, and repairs.",
        "fig08_hd_train_validation_curves": "Median train/validation curves for HD-BasinFlow-selected configurations where epoch curves were logged.",
        "fig09_hd_vs_baseline_savings": "Direct per-workload comparison of HD-BasinFlow evaluation savings against each baseline.",
        "fig10_hd_vs_baseline_quality_delta": "Direct per-workload comparison of HD-BasinFlow median loss delta against each baseline.",
    }
    for name in generated:
        caption = captions.get(name, f"Paper figure generated from benchmark data: `{name}`.")
        lines.append(f"- `{name}.pdf` / `{name}.png`: {caption}")
    lines.extend(
        [
            "",
            "## Suggested Paper Placement",
            "",
            "- Main paper: fig01, fig02, fig03, fig04, fig05.",
            "- Method/results appendix: convergence curves, final-loss bars, active-dimension, allocation-source, and train/validation figures.",
            "- Avoid overclaiming: the evidence supports compute savings at matched validation quality, not global optimality.",
        ]
    )
    (outdir / "FIGURE_INDEX.md").write_text("\n".join(lines) + "\n", encoding="utf-8")


def main() -> int:
    parser = argparse.ArgumentParser(description="Build paper-ready HD-BasinFlow figures.")
    parser.add_argument("--input-dir", type=Path, default=Path("runs/commercial_benchmark_a100"))
    parser.add_argument("--outdir", type=Path, default=Path("runs/paper_figures"))
    args = parser.parse_args()

    args.outdir.mkdir(parents=True, exist_ok=True)
    plt = setup_matplotlib()
    evals, metrics, decisions = load_inputs(args.input_dir)
    evals = normalize_evals(evals)
    metrics = normalize_metrics(metrics)

    generated: list[str] = []
    generated.append(figure_pass_rate(plt, decisions, args.outdir))
    generated.append(figure_savings_by_baseline(plt, decisions, args.outdir))
    generated.append(figure_quality_savings_scatter(plt, decisions, args.outdir))
    generated.append(figure_hd_vs_baseline_savings(plt, decisions, args.outdir))
    generated.append(figure_hd_vs_baseline_quality_delta(plt, decisions, args.outdir))
    generated.extend(figure_heatmaps(plt, decisions, args.outdir))
    active = figure_active_dimension(plt, evals, args.outdir)
    if active:
        generated.append(active)
    allocation = figure_allocation_sources(plt, evals, args.outdir)
    if allocation:
        generated.append(allocation)
    train_val = figure_train_val_curves(plt, evals, args.outdir)
    if train_val:
        generated.append(train_val)
    generated.extend(figure_convergence_curves(plt, evals, args.outdir))
    generated.extend(figure_final_loss_by_workload(plt, metrics, args.outdir))
    decisions.to_csv(args.outdir / "paper_decision_table.csv", index=False)
    metrics.to_csv(args.outdir / "paper_metrics_table.csv", index=False)
    write_index(args.outdir, generated, decisions)
    print(args.outdir)
    print(f"generated_figures={len(generated)}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())