scripts/analyze_456.py

"""Aggregate the three experimental claims into compact summaries + figures."""

import glob
import json
import os
from collections import defaultdict

import matplotlib

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

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
OUT = os.path.join(ROOT, "outputs")
FIGS = os.path.join(ROOT, "figs")
os.makedirs(FIGS, exist_ok=True)
plt.rcParams.update({"font.size": 11, "figure.dpi": 130})
summary = {}

# ================================================================= claim 4
p4 = os.path.join(OUT, "claim4_runs.json")
if os.path.exists(p4):
    d = json.load(open(p4))
    W = d["W_hat"]
    groups = defaultdict(list)
    for r in d["runs"]:
        key = (
            r["mode"],
            r["sigma_sq"],
            r["C"],
            r["rho"],
            r.get("alpha_init", "paper_code"),
        )
        groups[key].append(r)
    s4 = {"W_hat": W, "configs": []}
    fig, ax = plt.subplots(1, 2, figsize=(10.5, 3.8))
    for key, rs in sorted(groups.items(), key=lambda kv: str(kv[0])):
        mode, sig, C, rho, ai = key
        its = rs[0]["iters"]
        L = min(len(r["values"]) for r in rs)
        vals = np.array([r["values"][:L] for r in rs])
        gaps = W - vals
        mean, std = gaps.mean(0), gaps.std(0)
        n_div = sum(1 for r in rs if r["nonfinite_iter"] is not None)
        n_ovf = sum(1 for r in rs if r["overflow_iter"] is not None)
        blown = sum(1 for r in rs if (W - r["values"][-1]) > 1.0)

        # first iteration at which the mean gap goes below thresholds
        def first_below(th):
            for i, g in enumerate(mean):
                if g < th:
                    return its[i]
            return None

        cfg = {
            "mode": mode,
            "sigma_sq": sig,
            "C": C,
            "rho": rho,
            "alpha_init": ai,
            "n_seeds": len(rs),
            "iters": its[:L],
            "mean_gap": mean.tolist(),
            "std_gap": std.tolist(),
            "final_mean_gap": float(mean[-1]),
            "nonfinite_runs": n_div,
            "overflow_runs": n_ovf,
            "runs_with_final_gap_gt_1": blown,
            "max_exp_arg": float(max(r["max_exp_arg"] for r in rs)),
            "iters_to_gap_0.1": first_below(0.1),
            "iters_to_gap_0.05": first_below(0.05),
            "iters_to_gap_0.02": first_below(0.02),
        }
        s4["configs"].append(cfg)
        if (mode == "dual" and C == 1e-3) or (
            mode == "semidual" and ai == "paper_code"
        ):
            lbl = (
                rf"baseline dual $\sigma^2$={sig}, C={C}"
                if mode == "dual"
                else rf"proposed $\rho$={rho}, C={C}"
            )
            ax[0].plot(its[:L], mean, lw=1.5, label=lbl)
            ax[0].fill_between(
                its[:L], np.maximum(mean - std, 1e-6), mean + std, alpha=0.2
            )
    ax[0].set_yscale("log")
    ax[0].set_xscale("log")
    ax[0].set_xlabel("kernel-SGD iteration")
    ax[0].set_ylabel("optimality gap")
    ax[0].set_title(r"eOT, $\varepsilon=0.01$ (20 seeds, mean $\pm$ 1 s.d.)")
    ax[0].legend(fontsize=7)
    ax[0].grid(alpha=0.3, which="both")

    # right panel: baseline with larger stepsizes
    for key, rs in sorted(groups.items(), key=lambda kv: str(kv[0])):
        mode, sig, C, rho, ai = key
        if mode != "dual" or C == 1e-3:
            continue
        its = rs[0]["iters"]
        L = min(len(r["values"]) for r in rs)
        gaps = W - np.array([r["values"][:L] for r in rs])
        ax[1].plot(
            its[:L],
            np.abs(gaps.mean(0)),
            lw=1.2,
            ls="--",
            label=rf"$\sigma^2$={sig}, C={C}",
        )
    ax[1].set_yscale("log")
    ax[1].set_xscale("log")
    ax[1].set_xlabel("kernel-SGD iteration")
    ax[1].set_ylabel("|optimality gap|")
    ax[1].set_title("baseline dual with larger stepsizes (divergence)")
    ax[1].legend(fontsize=6, ncol=2)
    ax[1].grid(alpha=0.3, which="both")
    fig.tight_layout()
    fig.savefig(os.path.join(FIGS, "claim4.png"), bbox_inches="tight")
    plt.close(fig)
    summary["claim4"] = s4

# ================================================================= claim 5
p5 = os.path.join(OUT, "claim5_kl_dro.json")
if os.path.exists(p5):
    d = json.load(open(p5))
    recs = d["runs"]
    etas = sorted({r["eta"] for r in recs})
    batches = sorted({r["batch"] for r in recs})
    rhos = sorted({r["rho"] for r in recs if r["rho"] is not None})
    nseeds = d["seeds"]
    rows = []
    for b in batches:
        for meth, rho in [("baseline", None)] + [("proposed", r) for r in rhos]:
            per = {}
            for e in etas:
                rs = [
                    x
                    for x in recs
                    if x["method"] == meth
                    and x["batch"] == b
                    and x["eta"] == e
                    and x["rho"] == rho
                ]
                ok = [x for x in rs if x["converged"]]
                per[e] = {
                    "converged": len(ok),
                    "n": len(rs),
                    "mean_final": (
                        float(np.mean([x["final"] for x in ok])) if ok else None
                    ),
                }
            stable = [e for e in etas if per[e]["converged"] == nseeds]
            best = None
            for e in stable:
                if best is None or per[e]["mean_final"] < best[1]:
                    best = (e, per[e]["mean_final"])
            rows.append(
                {
                    "method": meth,
                    "rho": rho,
                    "batch": b,
                    "max_stable_eta": max(stable) if stable else None,
                    "best_eta_allseeds": best[0] if best else None,
                    "best_objective": best[1] if best else None,
                    "per_eta": {str(k): v for k, v in per.items()},
                }
            )
    summary["claim5"] = {
        "L0": d["L0"],
        "n": d["n"],
        "lam": d["lam"],
        "epochs": d["n_epochs"],
        "seeds": nseeds,
        "rows": rows,
    }

    fig, ax = plt.subplots(1, 2, figsize=(10.5, 3.8))
    labels, mat = [], []
    for r in rows:
        labels.append(
            "baseline" if r["method"] == "baseline" else rf"prop. $\rho$={r['rho']:g}"
        )
        mat.append((r["batch"], r["max_stable_eta"]))
    uniq = []
    for lb in labels:
        if lb not in uniq:
            uniq.append(lb)
    M = np.full((len(uniq), len(batches)), np.nan)
    for r in rows:
        lb = "baseline" if r["method"] == "baseline" else rf"prop. $\rho$={r['rho']:g}"
        i, j = uniq.index(lb), batches.index(r["batch"])
        M[i, j] = np.log10(r["max_stable_eta"]) if r["max_stable_eta"] else np.nan
    im = ax[0].imshow(M, cmap="viridis", aspect="auto", vmin=-8, vmax=-4)
    ax[0].set_xticks(range(len(batches)))
    ax[0].set_xticklabels([f"|D|={b}" for b in batches])
    ax[0].set_yticks(range(len(uniq)))
    ax[0].set_yticklabels(uniq, fontsize=8)
    for i in range(M.shape[0]):
        for j in range(M.shape[1]):
            ax[0].text(
                j,
                i,
                "div." if np.isnan(M[i, j]) else f"$10^{{{int(M[i,j])}}}$",
                ha="center",
                va="center",
                color="w",
                fontsize=8,
            )
    ax[0].set_title(r"largest stable stepsize $\eta$ (all 10 seeds)")
    fig.colorbar(im, ax=ax[0], label=r"$\log_{10}\eta$")

    for b, style in zip(batches, ["-", "--", ":"]):
        base = [r for r in rows if r["method"] == "baseline" and r["batch"] == b][0]
        rs = [
            x
            for x in recs
            if x["method"] == "baseline"
            and x["batch"] == b
            and x["eta"] == base["best_eta_allseeds"]
        ]
        if rs:
            ep = sorted(int(k) for k in rs[0]["traj"])
            cur = np.mean([[x["traj"][str(e)] for e in ep] for x in rs], axis=0)
            ax[1].plot(
                ep,
                cur,
                style,
                color="#c0392b",
                lw=1.4,
                label=rf"baseline |D|={b}, $\eta$={base['best_eta_allseeds']:g}",
            )
        prop = sorted(
            [
                r
                for r in rows
                if r["method"] == "proposed"
                and r["batch"] == b
                and r["best_objective"] is not None
            ],
            key=lambda r: r["best_objective"],
        )[0]
        rs = [
            x
            for x in recs
            if x["method"] == "proposed"
            and x["batch"] == b
            and x["eta"] == prop["best_eta_allseeds"]
            and x["rho"] == prop["rho"]
        ]
        if rs:
            ep = sorted(int(k) for k in rs[0]["traj"])
            cur = np.mean([[x["traj"][str(e)] for e in ep] for x in rs], axis=0)
            ax[1].plot(
                ep,
                cur,
                style,
                color="#27ae60",
                lw=1.4,
                label=rf"proposed |D|={b}, $\rho$={prop['rho']:g}, $\eta$={prop['best_eta_allseeds']:g}",
            )
    ax[1].set_xlabel("epoch")
    ax[1].set_ylabel(r"$L(\theta)$")
    ax[1].set_title("California Housing KL-DRO, best stable settings")
    ax[1].legend(fontsize=6)
    ax[1].grid(alpha=0.3)
    ax[1].set_yscale("log")
    fig.tight_layout()
    fig.savefig(os.path.join(FIGS, "claim5.png"), bbox_inches="tight")
    plt.close(fig)

# ================================================================= claim 6
p6 = os.path.join(OUT, "claim6_uot_dro_big.json")
if os.path.exists(p6):
    d = json.load(open(p6))
    recs = d["runs"]
    rows = []
    for meth in ["baseline", "proposed", "erm"]:
        for lr in sorted({r["lr"] for r in recs if r["method"] == meth}, reverse=True):
            rs = [r for r in recs if r["method"] == meth and r["lr"] == lr]
            fin = [
                r["final_F"]
                for r in rs
                if not r["diverged"] and r["final_F"] is not None
            ]
            rows.append(
                {
                    "method": meth,
                    "lr": lr,
                    "n": len(rs),
                    "diverged": sum(r["diverged"] for r in rs),
                    "fpe_iters": [r["fpe_iter"] for r in rs],
                    "mean_final_F": float(np.mean(fin)) if fin else None,
                    "std_final_F": float(np.std(fin)) if fin else None,
                    "mean_test_acc": float(
                        np.mean([r["hist"][-1]["test_acc"] for r in rs])
                    ),
                    "mean_val_acc": float(
                        np.mean([r["hist"][-1]["val_acc"] for r in rs])
                    ),
                }
            )
    summary["claim6"] = {
        "config": d["config"],
        "noise_rate": d["noise_rate"],
        "rows": rows,
    }

    fig, ax = plt.subplots(1, 2, figsize=(10.5, 3.8))
    colors = {"baseline": "#c0392b", "proposed": "#27ae60", "erm": "#7f8c8d"}
    styles = {1e-3: ":", 1e-4: "-", 1e-5: "--", 1e-6: "-."}
    for meth in ["baseline", "proposed", "erm"]:
        for lr in sorted({r["lr"] for r in recs if r["method"] == meth}, reverse=True):
            rs = [r for r in recs if r["method"] == meth and r["lr"] == lr]
            L = min(len(r["hist"]) for r in rs)
            steps = [h["step"] for h in rs[0]["hist"][:L]]
            F = np.array(
                [
                    [h["F"] if h["F"] is not None else np.nan for h in r["hist"][:L]]
                    for r in rs
                ]
            )
            ax[0].plot(
                steps,
                np.nanmean(F, 0),
                styles.get(lr, "-"),
                color=colors[meth],
                lw=1.4,
                label=rf"{meth}, $\eta$={lr:g}"
                + (" (FPE)" if any(r["diverged"] for r in rs) else ""),
            )
            acc = np.array([[h["test_acc"] for h in r["hist"][:L]] for r in rs])
            ax[1].plot(
                steps,
                acc.mean(0),
                styles.get(lr, "-"),
                color=colors[meth],
                lw=1.4,
                label=rf"{meth}, $\eta$={lr:g}",
            )
    ax[0].set_xlabel("SGD step (batch size 1)")
    ax[0].set_ylabel(r"$F(\theta)$ of Eq. (20)")
    ax[0].set_title("MNIST UOT-DRO train objective")
    ax[0].legend(fontsize=6)
    ax[0].grid(alpha=0.3)
    ax[1].set_xlabel("SGD step")
    ax[1].set_ylabel("clean test accuracy")
    ax[1].set_title("test accuracy (25% feature-dependent label noise)")
    ax[1].legend(fontsize=6)
    ax[1].grid(alpha=0.3)
    fig.tight_layout()
    fig.savefig(os.path.join(FIGS, "claim6.png"), bbox_inches="tight")
    plt.close(fig)

with open(os.path.join(OUT, "summary_456.json"), "w") as fh:
    json.dump(summary, fh, indent=2)
print(
    json.dumps(
        {
            k: (
                v
                if k != "claim4"
                else {"W_hat": v["W_hat"], "n_configs": len(v["configs"])}
            )
            for k, v in summary.items()
        },
        indent=1,
    )[:3000]
)