run_analysis.py

"""Reproduction of the six anchored claims of arXiv:2602.15327 (OpenReview IkjsHRpuYY),
"Prescriptive Scaling Reveals the Evolution of Language Model Capabilities".

Everything here is computed with the independent re-implementation in psl_core.py
(written from the paper's equations) applied to the authors' RELEASED evaluation
tables, vendored under authors/.  No language model is trained or evaluated.

Writes outputs/*.json.  Run:  python run_analysis.py
"""

from __future__ import annotations

import json
import os
import warnings

import numpy as np
import pandas as pd
from scipy.interpolate import BSpline
from scipy.optimize import minimize
from scipy.special import expit, logit
from scipy.stats import binomtest

import psl_core as P

warnings.filterwarnings("ignore", category=RuntimeWarning)

HERE = os.path.dirname(os.path.abspath(__file__))
AUTH = os.path.join(HERE, "authors")
OUT = os.environ.get("PSL_OUTPUT_DIR", os.path.join(HERE, "outputs"))
os.makedirs(OUT, exist_ok=True)

Z24 = np.log10(1000.0)  # 10^24 FLOPs in units of 1e21
PAPER_T1 = {
    "IFEval Raw": 0.828,
    "BBH Raw": 0.700,
    "MATH Lvl 5 Raw": 0.539,
    "GPQA Raw": 0.424,
    "MUSR Raw": 0.535,
    "MMLU-PRO Raw": 0.563,
}
PAPER_T2 = {  # (pinball_ID, pinball_OOD, calib_ID, calib_OOD)
    "Constant": (5.35e-3, 6.23e-3, 4.12e-2, 3.60e-2),
    "Binwise": (4.01e-3, 5.00e-3, 1.66e-2, 2.81e-2),
    "I-spline": (4.00e-3, 4.92e-3, 1.83e-2, 2.41e-2),
    "Sigmoid": (4.08e-3, 4.93e-3, 1.84e-2, 2.21e-2),
}


def dump(name, obj):
    with open(os.path.join(OUT, name), "w") as f:
        json.dump(obj, f, indent=2, default=float)
    print(f"  -> outputs/{name}")


# ===========================================================================
# I-spline (paper Appendix B.3), independently implemented
# ===========================================================================
def ispline_basis(z, edges, degree=3):
    """Clamped I-spline basis: I_j(z) = int_{k0}^{z} M_j, each nondecreasing 0->1."""
    z = np.asarray(z, float)
    edges = np.unique(np.asarray(edges, float))
    if edges.size < 2:
        return np.zeros((z.size, 0))
    p = int(degree)
    z0, z1 = float(edges[0]), float(edges[-1])
    t = np.r_[np.full(p + 1, z0), edges[1:-1], np.full(p + 1, z1)]
    nb = len(t) - p - 1
    if nb <= 0:
        return np.zeros((z.size, 0))
    zc = np.clip(z, z0, z1)
    X = np.zeros((z.size, nb))
    k = p + 1
    for i in range(nb):
        den = float(t[i + k] - t[i])
        if den <= 0:
            continue
        c = np.zeros(nb)
        c[i] = k / den  # normalised M-spline
        I = BSpline(t, c, p, extrapolate=False).antiderivative()
        v = I(zc) - float(I(z0))
        X[:, i] = np.clip(np.nan_to_num(v), 0.0, 1.0)
    return X


def fit_ispline(z, y, edges, tau=P.TAU, kappa=P.KAPPA):
    """q(z) = sigma(a0 + sum_j w_j I_j(z)), w_j >= 0, minimising smoothed pinball."""
    X = ispline_basis(z, edges)
    if X.shape[1] == 0:
        return None
    n = X.shape[1]

    def f(v):
        g = v[0] + X @ v[1:]
        return float(np.mean(P.smooth_pinball(y - expit(g), tau, kappa)))

    q_lo = float(np.clip(np.quantile(y, 0.02), 1e-4, 1 - 1e-4))
    q_hi = float(np.clip(np.quantile(y, tau), 1e-4, 1 - 1e-4))
    best, bo = None, np.inf
    for spread in (0.0, 1.0):
        v0 = np.r_[logit(q_lo), np.full(n, max(logit(q_hi) - logit(q_lo), 0.1) / n * (1 + spread))]
        r = minimize(f, v0, method="L-BFGS-B",
                     bounds=[(None, None)] + [(0.0, None)] * n,
                     options={"maxiter": 2000, "ftol": 1e-14})
        if r.fun < bo:
            bo, best = float(r.fun), r.x
    return best


def ispline_pred(v, z, edges):
    return expit(v[0] + ispline_basis(z, edges) @ v[1:])


def best_constant(y, tau=P.TAU, kappa=P.KAPPA):
    """Scalar c in [0,1] minimising the mean smoothed pinball loss on y.

    d/dc of the mean smoothed pinball loss is (1-tau) - mean(sigma(kappa*(y-c))),
    so the optimum solves mean(sigma(kappa*(y - c))) = 1 - tau.  The left side is
    strictly decreasing in c, so a bisection is exact and needs no optimiser.
    """
    y = np.asarray(y, float)
    y = y[np.isfinite(y)]
    if y.size == 0:
        return np.nan
    lo, hi = 0.0, 1.0
    for _ in range(200):
        mid = 0.5 * (lo + hi)
        if np.mean(expit(kappa * (y - mid))) > 1.0 - tau:
            lo = mid
        else:
            hi = mid
    return 0.5 * (lo + hi)


def binwise_fit(z, y, edges, tau=P.TAU):
    """Per-bin constant minimising the smoothed pinball loss on the training y."""
    c = np.full(len(edges) - 1, best_constant(y, tau))
    for b, m in P.bin_index_masks(z, edges):
        if m.sum() > 0:
            c[b] = best_constant(y[m], tau)
    return c


def binwise_pred(c, z, edges):
    idx = np.clip(np.searchsorted(edges, z, side="right") - 1, 0, len(c) - 1)
    return c[idx]


# ===========================================================================
# Shared data
# ===========================================================================
print("Loading released Open LLM Leaderboard v2 table ...")
DF = P.load_oll(os.path.join(AUTH, "oll_v2_slim.csv"))
DF_C = DF[np.isfinite(DF["z"])].copy()
print(f"  {len(DF)} rows, {len(DF_C)} with a compute proxy")


def splits():
    """Rolling single-k splits (P_t -> P_{t+1}) for t = 1,2,3."""
    names = [n for n, _, _ in P.PERIOD4]
    return [(k + 1, names[k], names[k + 1]) for k in range(3)]


# ===========================================================================
# Correctness gates
# ===========================================================================
def gates():
    print("\n[GATES] correctness gates")
    g = {}

    # G1 -- compute proxy: x = 6*T*P in 1e21 FLOPs, so 10^24 FLOPs <-> x = 1000
    g["compute_proxy"] = {
        "formula": "x = 6 * tokens[T] * params[B]  (units of 1e21 FLOPs)",
        "x_at_1e24_FLOPs": 1000.0,
        "z_at_1e24_FLOPs": Z24,
        "n_models_with_compute": int(len(DF_C)),
        "z_min": float(DF_C["z"].min()),
        "z_max": float(DF_C["z"].max()),
    }

    # G2 -- the smoothed pinball loss is the paper's formula and kappa=50 controls
    # the sup-norm gap to the sharp check loss: sup_u |l_tau(u) - rho_tau(u)| = log2/kappa
    u = np.linspace(-1, 1, 2_000_001)
    for kap in (10.0, 25.0, 50.0, 200.0):
        gap = float(np.max(np.abs(P.smooth_pinball(u, P.TAU, kap) - P.sharp_pinball(u, P.TAU))))
        g.setdefault("smoothing_gap", {})[str(kap)] = {
            "max_abs_gap": gap, "log2_over_kappa": float(np.log(2) / kap),
            "rel_err": float(abs(gap - np.log(2) / kap) / (np.log(2) / kap)),
        }
    # direct algebraic identity check against the literal paper expression
    uu = np.linspace(-0.4, 0.4, 20001)
    lit = np.log(1 + np.exp(50.0 * uu)) / 50.0 + (P.TAU - 1.0) * uu
    g["loss_identity_max_abs_diff"] = float(np.max(np.abs(lit - P.smooth_pinball(uu, P.TAU, 50.0))))

    # G3 -- Table 1 must come out of an independent fit of the released scores
    t1 = {}
    for t in P.TASKS:
        z, y = P.task_xy(DF_C, t)
        th, _ = P.fit_boundary(z, y)
        t1[t] = {"fitted": float(P.sigmoid_boundary(th, np.array([Z24]))[0]),
                 "paper": PAPER_T1[t], "theta_y0_L_zstar_logb": [float(v) for v in th]}
        t1[t]["abs_diff"] = abs(t1[t]["fitted"] - t1[t]["paper"])
    g["table1_gate"] = t1
    g["table1_max_abs_diff"] = max(v["abs_diff"] for v in t1.values())
    g["table1_gate_passed"] = bool(g["table1_max_abs_diff"] < 1e-3)
    print(f"  Table-1 gate: max |diff| = {g['table1_max_abs_diff']:.5f} "
          f"({'PASS' if g['table1_gate_passed'] else 'FAIL'})")
    print(f"  loss identity max |diff| = {g['loss_identity_max_abs_diff']:.2e}")
    dump("correctness_gates.json", g)
    return g


# ===========================================================================
# Claim 1 -- estimator hyperparameters (tau=0.98, kappa=50, lambda=1e-3)
# ===========================================================================
SOURCE_AUDIT = [
    {"file": "skill_frontier/core/sigmoid.py", "line": 63,
     "text": "DEFAULT_LAMBDA_B: float = 1e-3", "constant": "lambda = 1e-3"},
    {"file": "skill_frontier/core/sigmoid.py", "line": 186,
     "text": "tau: float = 0.98,  # fit_sigmoid_frontier default", "constant": "tau = 0.98"},
    {"file": "skill_frontier/core/sigmoid.py", "line": 196,
     "text": "kappa_final: float = 50.0,", "constant": "kappa = 50"},
    {"file": "skill_frontier/core/sigmoid.py", "line": 989,
     "text": 'p.add_argument("--tau", type=float, default=0.98, ...)', "constant": "tau = 0.98"},
    {"file": "skill_frontier/evaluation/pinball_utils.py", "line": 8,
     "text": "def smooth_pinball_loss(r, tau, k_smooth: float = 50.0)", "constant": "kappa = 50"},
    {"file": "skill_frontier/evaluation/pinball_utils.py", "line": 11,
     "text": "return (np.logaddexp(0.0, k_smooth * r) / k_smooth) + (tau - 1.0) * r",
     "constant": "loss form"},
    {"file": "skill_frontier/evaluation/common.py", "line": 25,
     "text": "kappa_final: float = 50.0,", "constant": "kappa = 50"},
    {"file": "skill_frontier/core/sigmoid_quantile_optimizer.py", "line": 46,
     "text": '"""Mean smooth pinball loss + ridge penalty on b^2 ..."""',
     "constant": "Omega(theta) = beta^2, MEAN not SUM"},
    {"file": "skill_frontier/core/sigmoid_quantile_optimizer.py", "line": 199,
     "text": "tau_sched = [min(0.90, tau), min(0.95, tau), tau]",
     "constant": "continuation in tau, ends at 0.98"},
    {"file": "skill_frontier/core/sigmoid_quantile_optimizer.py", "line": 201,
     "text": "kappa_sched = [10.0, 25.0, float(kappa_final)]",
     "constant": "continuation in kappa, ends at 50"},
]


def claim1():
    print("\n[CLAIM 1] tau=0.98, kappa=50, lambda=1e-3")
    res = {"source_audit": SOURCE_AUDIT, "paper_section_2_1":
           "We use tau=0.98, kappa=50, lambda=1e-3."}

    # (a) tau is operative, not a label: refit at several tau and measure coverage.
    tau_probe = {}
    for tau in (0.50, 0.75, 0.90, 0.98):
        cov, pred = [], []
        for t in P.TASKS:
            z, y = P.task_xy(DF_C, t)
            th, _ = P.fit_boundary(z, y, tau=tau)
            yh = P.sigmoid_boundary(th, z)
            cov.append(float(np.mean(y <= yh)))
            pred.append(float(P.sigmoid_boundary(th, np.array([Z24]))[0]))
        tau_probe[f"{tau:.2f}"] = {
            "mean_empirical_coverage": float(np.mean(cov)),
            "target": tau,
            "abs_err": float(abs(np.mean(cov) - tau)),
            "per_task_coverage": dict(zip(P.TASKS, cov)),
            "pred_at_1e24": dict(zip(P.TASKS, pred)),
        }
        print(f"  tau={tau:.2f}: mean empirical coverage = {np.mean(cov):.4f}")
    res["tau_sweep"] = tau_probe

    # (b) kappa controls only a narrow band: sup-gap = log2/kappa, and the fit is
    #     insensitive to kappa in the moderate range.
    kap_probe = {}
    for kap in (10.0, 25.0, 50.0, 200.0, 1000.0):
        pr = {}
        for t in P.TASKS:
            z, y = P.task_xy(DF_C, t)
            th, _ = P.fit_boundary(z, y, kappa=kap)
            pr[t] = float(P.sigmoid_boundary(th, np.array([Z24]))[0])
        kap_probe[str(kap)] = {
            "pred_at_1e24": pr,
            "max_abs_shift_vs_kappa50": None,
            "sup_gap_to_sharp_pinball": float(np.log(2) / kap),
        }
    base = kap_probe["50.0"]["pred_at_1e24"]
    for k, v in kap_probe.items():
        v["max_abs_shift_vs_kappa50"] = float(
            max(abs(v["pred_at_1e24"][t] - base[t]) for t in P.TASKS))
    res["kappa_sweep"] = kap_probe

    # (c) lambda = 1e-3 is a conditioning term, not a capacity knob.
    lam_probe = {}
    for lam in (0.0, 1e-4, 1e-3, 1e-2, 1e-1, 1.0):
        pr, bs = {}, {}
        for t in P.TASKS:
            z, y = P.task_xy(DF_C, t)
            th, _ = P.fit_boundary(z, y, lam=lam)
            pr[t] = float(P.sigmoid_boundary(th, np.array([Z24]))[0])
            bs[t] = float(np.exp(th[3]))
        lam_probe[str(lam)] = {"pred_at_1e24": pr, "fitted_beta": bs}
    b3 = lam_probe["0.001"]["pred_at_1e24"]
    for k, v in lam_probe.items():
        v["max_abs_shift_vs_lambda_1e-3"] = float(
            max(abs(v["pred_at_1e24"][t] - b3[t]) for t in P.TASKS))
    res["lambda_sweep"] = lam_probe

    # (d) mean-vs-sum normalisation of the objective (documented discrepancy)
    z, y = P.task_xy(DF_C, "BBH Raw")
    res["objective_normalisation"] = {
        "paper_form": "sum_i l_tau(...) + lambda*Omega(theta)",
        "code_form": "mean_i l_tau(...) + lambda*beta^2",
        "n_train_full": int(z.size),
        "implied_ratio_lambda_sum_vs_mean": int(z.size),
        "note": ("Multiplying the released MEAN objective by n gives "
                 "sum_i l_tau + n*lambda*beta^2, so lambda=1e-3 in the code is an "
                 "effective n*1e-3 in the paper's SUM convention. We report the "
                 "released convention throughout."),
    }

    # (e) NEGATIVE CONTROL: tau must matter. If tau were cosmetic, refitting at
    #     tau=0.5 would leave the boundary and its coverage unchanged.
    d50 = res["tau_sweep"]["0.50"]
    d98 = res["tau_sweep"]["0.98"]
    res["negative_control_tau"] = {
        "coverage_at_tau_0.50": d50["mean_empirical_coverage"],
        "coverage_at_tau_0.98": d98["mean_empirical_coverage"],
        "coverage_gap": float(d98["mean_empirical_coverage"] - d50["mean_empirical_coverage"]),
        "mean_boundary_drop_at_1e24": float(np.mean(
            [d98["pred_at_1e24"][t] - d50["pred_at_1e24"][t] for t in P.TASKS])),
        "control_broke_as_required": bool(
            d98["mean_empirical_coverage"] - d50["mean_empirical_coverage"] > 0.3),
    }
    dump("claim1_hyperparameters.json", res)
    return res


# ===========================================================================
# Claim 2 -- estimator comparison (Table 2)
# ===========================================================================
def claim2():
    print("\n[CLAIM 2] Table 2: sigmoid vs I-spline vs binwise vs constant")
    rows = []
    for k, tr, va in splits():
        for t in P.TASKS:
            a = DF_C[DF_C["period"] == tr]
            b = DF_C[DF_C["period"] == va]
            zt, yt = P.task_xy(a, t)
            zv, yv = P.task_xy(b, t)
            if zt.size < 40 or zv.size < 10:
                continue
            edges = P.equal_mass_bins(zt, 10, 30)
            # restrict OOD evaluation to the train z-range (paper Section 3.1)
            m = (zv >= edges[0]) & (zv <= edges[-1])
            zv, yv = zv[m], yv[m]
            if zv.size < 10:
                continue

            # OOD calibration is measured on bins built from the validation z
            # ("test_fixed"); in-distribution calibration uses the train bins.
            edges_v = P.equal_mass_bins(zv, 10, 30)

            preds = {}
            c = float(np.quantile(yt, P.TAU))  # global tau-quantile (null model)
            preds["Constant"] = (np.full(zt.size, c), np.full(zv.size, c))
            cb = binwise_fit(zt, yt, edges)
            preds["Binwise"] = (binwise_pred(cb, zt, edges), binwise_pred(cb, zv, edges))
            # I-spline knots come from a coarse 3-bin equal-mass partition of z_train
            edges_i = P.equal_mass_bins(zt, 3, 1)
            vi = fit_ispline(zt, yt, edges_i)
            preds["I-spline"] = (ispline_pred(vi, zt, edges_i), ispline_pred(vi, zv, edges_i))
            th, _ = P.fit_boundary(zt, yt)
            preds["Sigmoid"] = (P.sigmoid_boundary(th, zt), P.sigmoid_boundary(th, zv))

            for est, (ht, hv) in preds.items():
                for split, zz, yy, hh, ee in (("ID", zt, yt, ht, edges),
                                              ("OOD", zv, yv, hv, edges_v)):
                    cov = P.coverage_by_bin(zz, yy, hh, ee)
                    rows.append(dict(
                        k=k, task=t, estimator=est, split=split,
                        pinball=float(np.mean(P.smooth_pinball(yy - hh))),
                        calib=float(np.mean(np.abs(cov["signed"]))),
                        n=int(yy.size)))
        print(f"  split t={k} done")
    R = pd.DataFrame(rows)

    tab = {}
    for est in ("Constant", "Binwise", "I-spline", "Sigmoid"):
        e = R[R.estimator == est]
        tab[est] = {
            "pinball_ID": float(e[e.split == "ID"]["pinball"].mean()),
            "pinball_OOD": float(e[e.split == "OOD"]["pinball"].mean()),
            "calib_ID": float(e[e.split == "ID"]["calib"].mean()),
            "calib_OOD": float(e[e.split == "OOD"]["calib"].mean()),
            "paper": {"pinball_ID": PAPER_T2[est][0], "pinball_OOD": PAPER_T2[est][1],
                      "calib_ID": PAPER_T2[est][2], "calib_OOD": PAPER_T2[est][3]},
        }
        print(f"  {est:9s} ID {tab[est]['pinball_ID']:.5f} (paper {PAPER_T2[est][0]:.5f})  "
              f"OOD {tab[est]['pinball_OOD']:.5f} (paper {PAPER_T2[est][1]:.5f})  "
              f"calibOOD {tab[est]['calib_OOD']:.5f} (paper {PAPER_T2[est][3]:.5f})")

    # NEGATIVE CONTROL: destroy the compute-accuracy link by shuffling z. The
    # sigmoid must then be no better than the constant baseline.
    rng = np.random.default_rng(0)
    nc = []
    for k, tr, va in splits():
        for t in P.TASKS:
            a = DF_C[DF_C["period"] == tr]
            b = DF_C[DF_C["period"] == va]
            zt, yt = P.task_xy(a, t)
            zv, yv = P.task_xy(b, t)
            if zt.size < 40 or zv.size < 10:
                continue
            zs = rng.permutation(zt)
            edges = P.equal_mass_bins(zs, 10, 30)
            m = (zv >= edges[0]) & (zv <= edges[-1])
            zv2, yv2 = zv[m], yv[m]
            if zv2.size < 10:
                continue
            th, _ = P.fit_boundary(zs, yt)
            c = float(np.quantile(yt, P.TAU))
            nc.append(dict(
                sig_ID=float(np.mean(P.smooth_pinball(yt - P.sigmoid_boundary(th, zs)))),
                const_ID=float(np.mean(P.smooth_pinball(yt - c)))))
    NC = pd.DataFrame(nc)
    res = {
        "table2": tab,
        "per_run": R.to_dict(orient="records"),
        "sigmoid_minus_ispline_pinball_ID": tab["Sigmoid"]["pinball_ID"] - tab["I-spline"]["pinball_ID"],
        "sigmoid_minus_ispline_pinball_OOD": tab["Sigmoid"]["pinball_OOD"] - tab["I-spline"]["pinball_OOD"],
        "sigmoid_minus_ispline_calib_OOD": tab["Sigmoid"]["calib_OOD"] - tab["I-spline"]["calib_OOD"],
        "negative_control_shuffled_z": {
            "sigmoid_pinball_ID": float(NC["sig_ID"].mean()),
            "constant_pinball_ID": float(NC["const_ID"].mean()),
            "advantage_over_constant": float(NC["const_ID"].mean() - NC["sig_ID"].mean()),
            "real_advantage_over_constant": float(
                tab["Constant"]["pinball_ID"] - tab["Sigmoid"]["pinball_ID"]),
        },
    }
    ncb = res["negative_control_shuffled_z"]
    ncb["advantage_ratio_shuffled_over_real"] = float(
        ncb["advantage_over_constant"] / ncb["real_advantage_over_constant"])
    # DESTRUCTIVE CONTROL: hold the data, temporal splits, quantile, loss, and
    # evaluation metrics fixed while forcing beta=0.  The sigmoid then becomes
    # the intercept-only Constant row already evaluated in the same Table-2
    # protocol, and its OOD calibration advantage must disappear.
    calib_increase = float(tab["Constant"]["calib_OOD"] - tab["Sigmoid"]["calib_OOD"])
    res["destructive_control_beta_zero"] = {
        "mutation": "force sigmoid beta=0 (intercept-only tau-quantile)",
        "registered_sigmoid_calib_OOD": float(tab["Sigmoid"]["calib_OOD"]),
        "beta_zero_calib_OOD": float(tab["Constant"]["calib_OOD"]),
        "calibration_error_increase": calib_increase,
        "calibration_error_ratio": float(tab["Constant"]["calib_OOD"] / tab["Sigmoid"]["calib_OOD"]),
        "control_broke_as_required": bool(calib_increase > 0.01),
    }
    print(f"  neg-control (shuffled z): sigmoid advantage over constant "
          f"{ncb['advantage_over_constant']:.2e} vs real {ncb['real_advantage_over_constant']:.2e}")
    dump("claim2_estimator_comparison.json", res)
    return res


# ===========================================================================
# Claim 3 -- Table 1 at 10^24 FLOPs
# ===========================================================================
def claim3():
    print("\n[CLAIM 3] Table 1 attainable accuracy at 10^24 FLOPs")
    out, rng = {}, np.random.default_rng(0)
    for t in P.TASKS:
        z, y = P.task_xy(DF_C, t)
        th, _ = P.fit_boundary(z, y)
        v = float(P.sigmoid_boundary(th, np.array([Z24]))[0])
        # cold-start (no tau/kappa continuation) as a robustness check
        th_c, _ = P.fit_boundary(z, y, continuation=False)
        vc = float(P.sigmoid_boundary(th_c, np.array([Z24]))[0])
        # DESTRUCTIVE CONTROL: preserve the data, estimator family, compute
        # proxy, optimiser, and evaluation point while destroying the
        # load-bearing registered quantile level.
        th50, _ = P.fit_boundary(z, y, tau=0.50, seed=50)
        v50 = float(P.sigmoid_boundary(th50, np.array([Z24]))[0])
        # NEGATIVE CONTROL: shuffled compute labels destroy the compute ordering,
        # so the "boundary" must collapse to the unconditional tau-quantile.
        sh = []
        for s in range(5):
            ths, _ = P.fit_boundary(rng.permutation(z), y, seed=s)
            sh.append(float(P.sigmoid_boundary(ths, np.array([Z24]))[0]))
        out[t] = {
            "fitted": v, "paper": PAPER_T1[t], "abs_diff": abs(v - PAPER_T1[t]),
            "fitted_cold_start": vc, "cold_start_abs_shift": abs(vc - v),
            "wrong_quantile_tau_0.50": v50,
            "wrong_quantile_abs_diff_from_paper": abs(v50 - PAPER_T1[t]),
            "unconditional_tau_quantile": float(np.quantile(y, P.TAU)),
            "shuffled_z_pred_at_1e24_mean": float(np.mean(sh)),
            "shuffled_z_pred_at_1e24_sd": float(np.std(sh)),
            "n": int(z.size),
        }
        print(f"  {t:16s} {v:.4f} (paper {PAPER_T1[t]:.3f}, diff {v-PAPER_T1[t]:+.4f}) | "
              f"shuffled {np.mean(sh):.4f} vs uncond-q {np.quantile(y, P.TAU):.4f}")
    res = {"per_task": out,
           "max_abs_diff": max(v["abs_diff"] for v in out.values()),
           "all_within_0.001": bool(all(v["abs_diff"] < 1e-3 for v in out.values()))}
    # the control breaks if the shuffled boundary loses its compute dependence
    gaps = [abs(v["shuffled_z_pred_at_1e24_mean"] - v["unconditional_tau_quantile"])
            for v in out.values()]
    real = [abs(v["fitted"] - v["unconditional_tau_quantile"]) for v in out.values()]
    res["negative_control_shuffled_z"] = {
        "mean_gap_shuffled_to_unconditional_quantile": float(np.mean(gaps)),
        "mean_gap_real_to_unconditional_quantile": float(np.mean(real)),
        "collapse_ratio_shuffled_over_real": float(np.mean(gaps) / np.mean(real)),
    }
    wrong = [v["wrong_quantile_abs_diff_from_paper"] for v in out.values()]
    res["destructive_control_wrong_quantile"] = {
        "registered_tau": 0.98,
        "mutated_tau": 0.50,
        "mean_abs_diff_from_paper": float(np.mean(wrong)),
        "max_abs_diff_from_paper": float(np.max(wrong)),
        "min_abs_diff_from_paper": float(np.min(wrong)),
        "all_six_fail_registered_0.001_tolerance": bool(all(x > 1e-3 for x in wrong)),
        "control_broke_as_required": bool(all(x > 1e-3 for x in wrong)),
    }
    dump("claim3_table1.json", res)
    return res


# ===========================================================================
# Claim 4 -- temporal (non-)stationarity
# ===========================================================================
def _temporal_pass(df, label_col="period", seed=None):
    """Fit on P_t, evaluate on P_{t+1} over the train z-range, for t=1,2,3."""
    rows = []
    for k, tr, va in splits():
        a = df[df[label_col] == tr]
        b = df[df[label_col] == va]
        for t in P.TASKS:
            zt, yt = P.task_xy(a, t)
            zv, yv = P.task_xy(b, t)
            if zt.size < 40 or zv.size < 10:
                continue
            edges = P.equal_mass_bins(zt, 10, 30)
            m = (zv >= edges[0]) & (zv <= edges[-1])
            zv, yv = zv[m], yv[m]
            if zv.size < 10:
                continue
            th, _ = P.fit_boundary(zt, yt, seed=0 if seed is None else seed)
            hv = P.sigmoid_boundary(th, zv)
            ht = P.sigmoid_boundary(th, zt)
            covv = P.coverage_by_bin(zv, yv, hv, edges)
            covt = P.coverage_by_bin(zt, yt, ht, edges)
            n_above = int(np.sum(yv > hv))
            bt = binomtest(n_above, int(yv.size), 1.0 - P.TAU, alternative="greater")
            rows.append(dict(
                k=k, task=t, n_ood=int(yv.size), n_above=n_above,
                share_above=float(n_above / yv.size),
                cov_micro=float(np.mean(yv <= hv)),
                signed_micro=float(np.mean(yv <= hv) - P.TAU),
                signed_macro=float(covv["signed"].mean()),
                mae_macro=float(np.mean(np.abs(covv["signed"]))),
                is_signed_micro=float(np.mean(yt <= ht) - P.TAU),
                is_mae_macro=float(np.mean(np.abs(covt["signed"]))),
                pinball_ood=float(np.mean(P.smooth_pinball(yv - hv))),
                p_excess=float(bt.pvalue)))
    return pd.DataFrame(rows)


def claim4():
    print("\n[CLAIM 4] temporal stability / non-stationarity")
    R = _temporal_pass(DF_C)
    per = {}
    STABLE = ["BBH Raw", "GPQA Raw", "MMLU-PRO Raw", "MUSR Raw"]
    DRIFT = ["MATH Lvl 5 Raw", "IFEval Raw"]
    for t in P.TASKS:
        e = R[R.task == t].sort_values("k")
        n_ab, n_tot = int(e["n_above"].sum()), int(e["n_ood"].sum())
        bt = binomtest(n_ab, n_tot, 1.0 - P.TAU, alternative="greater")
        per[t] = {
            "signed_coverage_micro_per_k": e["signed_micro"].round(6).tolist(),
            "signed_coverage_micro_mean": float(e["signed_micro"].mean()),
            "mae_macro_per_k": e["mae_macro"].round(6).tolist(),
            "mae_macro_mean": float(e["mae_macro"].mean()),
            "share_above_per_k": e["share_above"].round(6).tolist(),
            "pooled_share_above": float(n_ab / n_tot),
            "pooled_n_above": n_ab, "pooled_n": n_tot,
            "pooled_binom_p_excess_above_2pct": float(bt.pvalue),
            "in_sample_signed_coverage_mean": float(e["is_signed_micro"].mean()),
            "in_sample_mae_macro_mean": float(e["is_mae_macro"].mean()),
            "group": "drift" if t in DRIFT else "stable",
        }
        print(f"  {t:16s} signed cov {per[t]['signed_coverage_micro_mean']:+.4f}  "
              f"share_above {per[t]['pooled_share_above']:.4f}  "
              f"binom p={per[t]['pooled_binom_p_excess_above_2pct']:.2e}  "
              f"(IS signed {per[t]['in_sample_signed_coverage_mean']:+.4f})")

    # NEGATIVE CONTROL: permute the period labels. If MATH/IFEval under-coverage
    # is genuinely temporal, it must vanish when time is randomised.
    rng = np.random.default_rng(0)
    nc = {t: [] for t in P.TASKS}
    for rep in range(5):
        d = DF_C.copy()
        d["period_shuf"] = rng.permutation(d["period"].to_numpy())
        Rs = _temporal_pass(d, "period_shuf", seed=rep)
        for t in P.TASKS:
            e = Rs[Rs.task == t]
            if len(e):
                nc[t].append(float(e["signed_micro"].mean()))
    ncs = {t: {"mean_signed_coverage": float(np.mean(v)) if v else None,
               "sd": float(np.std(v)) if v else None} for t, v in nc.items()}

    real_drift = float(np.mean([per[t]["signed_coverage_micro_mean"] for t in DRIFT]))
    real_stable = float(np.mean([per[t]["signed_coverage_micro_mean"] for t in STABLE]))
    sh_drift = float(np.mean([ncs[t]["mean_signed_coverage"] for t in DRIFT]))
    sh_stable = float(np.mean([ncs[t]["mean_signed_coverage"] for t in STABLE]))
    res = {
        "per_task": per, "per_run": R.to_dict(orient="records"),
        "group_means": {
            "drift_tasks_signed_coverage": real_drift,
            "stable_tasks_signed_coverage": real_stable,
            "separation": float(real_stable - real_drift)},
        "negative_control_shuffled_periods": {
            "per_task": ncs,
            "drift_tasks_signed_coverage": sh_drift,
            "stable_tasks_signed_coverage": sh_stable,
            "separation": float(sh_stable - sh_drift),
            "control_broke_as_required": bool(
                abs(sh_stable - sh_drift) < 0.3 * abs(real_stable - real_drift))},
    }
    print(f"  group separation real {res['group_means']['separation']:+.4f} vs "
          f"shuffled-time {res['negative_control_shuffled_periods']['separation']:+.4f}")
    dump("claim4_temporal_stability.json", res)
    return res


# ===========================================================================
# Claim 5 -- balanced I-optimal budgeted design
# ===========================================================================
def claim5():
    print("\n[CLAIM 5] budgeted design: frontier recovery vs budget alpha")
    md = os.path.join(AUTH, "manifests")
    alphas = [5, 10, 20, 50, 100]
    sha = DF_C["Model sha"].astype(str)
    rng = np.random.default_rng(0)
    rows = []
    for k, tr, va in splits():
        train_pool = DF_C[DF_C["period"] == tr]
        val_pool = DF_C[DF_C["period"] == va]
        full = {}
        for t in P.TASKS:
            zt, yt = P.task_xy(train_pool, t)
            if zt.size < 40:
                continue
            th, _ = P.fit_boundary(zt, yt)
            full[t] = (th, P.equal_mass_bins(zt, 10, 30))
        for a in alphas:
            f = os.path.join(md, f"alpha{a}_k{k}_train.txt")
            if not os.path.exists(f):
                continue
            with open(f, encoding="utf-8") as handle:
                ids = set(handle.read().split())
            sub = train_pool[train_pool["Model sha"].astype(str).isin(ids)]
            cost_frac = float(sub["#Params (B)"].sum() / train_pool["#Params (B)"].sum())
            # matched random design: same parameter-count budget, chosen at random
            budget = float(sub["#Params (B)"].sum())
            for t in P.TASKS:
                if t not in full:
                    continue
                th_f, edges = full[t]
                zs, ys = P.task_xy(sub, t)
                zv, yv = P.task_xy(val_pool, t)
                mv = (zv >= edges[0]) & (zv <= edges[-1])
                zv, yv = zv[mv], yv[mv]
                if zs.size < 12 or zv.size < 10:
                    continue
                grid = np.linspace(edges[0], edges[-1], 400)
                qf = P.sigmoid_boundary(th_f, grid)

                th_s, _ = P.fit_boundary(zs, ys)
                qs = P.sigmoid_boundary(th_s, grid)
                hv = P.sigmoid_boundary(th_s, zv)
                cov = P.coverage_by_bin(zv, yv, hv, edges)

                # random-subset control at the same budget
                rnd_dev = []
                for rep in range(3):
                    perm = rng.permutation(len(train_pool))
                    cs = np.cumsum(train_pool["#Params (B)"].to_numpy()[perm])
                    sel = train_pool.iloc[perm[cs <= budget]]
                    zr, yr = P.task_xy(sel, t)
                    if zr.size < 12:
                        continue
                    th_r, _ = P.fit_boundary(zr, yr, seed=rep)
                    rnd_dev.append(float(np.mean(np.abs(P.sigmoid_boundary(th_r, grid) - qf))))

                rows.append(dict(
                    k=k, alpha=a, task=t, n_selected=int(zs.size),
                    cost_fraction=cost_frac,
                    mean_abs_dev_from_full=float(np.mean(np.abs(qs - qf))),
                    max_abs_dev_from_full=float(np.max(np.abs(qs - qf))),
                    random_mean_abs_dev=float(np.mean(rnd_dev)) if rnd_dev else np.nan,
                    oos_mae_macro=float(np.mean(np.abs(cov["signed"]))),
                    oos_pinball=float(np.mean(P.smooth_pinball(yv - hv)))))
        print(f"  split t={k} done")
    R = pd.DataFrame(rows)

    piv = R.pivot_table(index="task", columns="alpha", values="mean_abs_dev_from_full")
    pivr = R.pivot_table(index="task", columns="alpha", values="random_mean_abs_dev")
    pivm = R.pivot_table(index="task", columns="alpha", values="oos_mae_macro")
    print("\n  mean |q_alpha - q_full| over the fitted z-range (design):")
    print(piv.round(4).to_string())
    print("\n  same, matched-budget RANDOM subsets (negative control):")
    print(pivr.round(4).to_string())

    res = {
        "per_run": R.to_dict(orient="records"),
        "frontier_deviation_by_task_alpha": {t: {str(a): float(piv.loc[t, a])
                                                 for a in piv.columns} for t in piv.index},
        "frontier_deviation_random_control": {t: {str(a): float(pivr.loc[t, a])
                                                  for a in pivr.columns} for t in pivr.index},
        "oos_coverage_mae_by_task_alpha": {t: {str(a): float(pivm.loc[t, a])
                                               for a in pivm.columns} for t in pivm.index},
        "macro_mean_frontier_deviation": {str(a): float(piv[a].mean()) for a in piv.columns},
        "macro_mean_random_control": {str(a): float(pivr[a].mean()) for a in pivr.columns},
        "macro_mean_oos_mae": {str(a): float(pivm[a].mean()) for a in pivm.columns},
        "mean_cost_fraction": {str(a): float(R[R.alpha == a]["cost_fraction"].mean())
                               for a in alphas},
    }
    md_ = res["macro_mean_frontier_deviation"]
    res["knee_analysis"] = {
        "rel_improvement_5_to_10": float((md_["5"] - md_["10"]) / md_["5"]),
        "rel_improvement_10_to_20": float((md_["10"] - md_["20"]) / md_["10"]),
        "rel_improvement_20_to_50": float((md_["20"] - md_["50"]) / md_["20"]),
        "rel_improvement_50_to_100": float((md_["50"] - md_["100"]) / md_["50"]),
    }
    res["gpqa_musr_at_5pct"] = {
        t: {"deviation_at_5pct": float(piv.loc[t, 5]),
            "deviation_at_100pct": float(piv.loc[t, 100]),
            "ratio": float(piv.loc[t, 5] / piv.loc[t, 100]),
            "rank_among_6_tasks_at_5pct": int(piv[5].rank().loc[t])}
        for t in ("GPQA Raw", "MUSR Raw")}
    res["all_tasks_ratio_5pct_over_full"] = {
        t: float(piv.loc[t, 5] / piv.loc[t, 100]) for t in piv.index}
    res["negative_control_random_design"] = {
        "design_beats_random_at_alpha": {
            str(a): bool(piv[a].mean() < pivr[a].mean()) for a in piv.columns},
        "mean_ratio_design_over_random": {
            str(a): float(piv[a].mean() / pivr[a].mean()) for a in piv.columns},
    }
    # DESTRUCTIVE CONTROL: hold the released balanced I-optimal procedure,
    # estimator, tasks, splits, and error metric fixed while removing three
    # quarters of the registered 20% budget.  The paper explicitly singles out
    # only GPQA and MUSR as recoverable at 5%; the other four tasks should lose
    # frontier fidelity, and the six-task macro error should increase sharply.
    other = [t for t in piv.index if t not in ("GPQA Raw", "MUSR Raw")]
    ratio_5_over_20 = float(piv[5].mean() / piv[20].mean())
    res["destructive_control_reduce_budget_20_to_5"] = {
        "registered_budget_percent": 20,
        "mutated_budget_percent": 5,
        "macro_deviation_at_20": float(piv[20].mean()),
        "macro_deviation_at_5": float(piv[5].mean()),
        "macro_error_ratio_5_over_20": ratio_5_over_20,
        "non_gpqa_musr_deviation_at_5": {t: float(piv.loc[t, 5]) for t in other},
        "all_four_non_gpqa_musr_exceed_0.04_at_5": bool(all(piv.loc[t, 5] > 0.04 for t in other)),
        "control_broke_as_required": bool(
            ratio_5_over_20 > 1.5 and all(piv.loc[t, 5] > 0.04 for t in other)),
    }
    dump("claim5_budget_design.json", res)
    return res


# ===========================================================================
# Claim 6 -- cross-benchmark contamination shift test (Equation 3)
# ===========================================================================
def claim6():
    print("\n[CLAIM 6] AIME-2025 vs MATH-500 cross-benchmark shift test")
    aime = pd.read_csv(os.path.join(AUTH, "aime-2025.csv"))
    math = pd.read_csv(os.path.join(AUTH, "math-500.csv"))
    df = math[["model_id", "math_500_pct"]].merge(
        aime[["model_id", "aime_2025_pct", "release_date"]], on="model_id", how="inner")
    df["math_500_pct"] = pd.to_numeric(df["math_500_pct"], errors="coerce")
    df["aime_2025_pct"] = pd.to_numeric(df["aime_2025_pct"], errors="coerce")
    df["release_date"] = pd.to_datetime(df["release_date"], errors="coerce")
    df = df.dropna(subset=["math_500_pct", "aime_2025_pct", "release_date"])
    n_merged = int(len(df))

    eps = 1e-6
    X = logit(np.clip(df["math_500_pct"].to_numpy() / 100.0, eps, 1 - eps))
    Y = logit(np.clip(df["aime_2025_pct"].to_numpy() / 100.0, eps, 1 - eps))
    cutoff = pd.Timestamp("2025-02-06")  # AIME-2025 release
    g = (df["release_date"] >= cutoff).to_numpy().astype(float)

    # common-support guardrail across the two release groups
    lo = max(X[g == 0].min(), X[g == 1].min())
    hi = min(X[g == 0].max(), X[g == 1].max())
    m = (X >= lo) & (X <= hi)
    Xs, Ys, gs = X[m], Y[m], g[m]

    D = np.column_stack([np.ones_like(Xs), Xs, gs])
    coef, *_ = np.linalg.lstsq(D, Ys, rcond=None)
    resid = Ys - D @ coef
    dof = len(Ys) - D.shape[1]
    s2 = float(resid @ resid) / dof
    se = np.sqrt(np.diag(s2 * np.linalg.inv(D.T @ D)))
    gamma, se_g = float(coef[2]), float(se[2])
    tstat = gamma / se_g

    # stratified permutation of the post-release indicator within X-quantile bins
    rng = np.random.default_rng(0)
    bins = pd.qcut(Xs, 10, labels=False, duplicates="drop")
    perm = np.empty(2000)
    for i in range(2000):
        gp = gs.copy()
        for b in np.unique(bins):
            idx = np.where(bins == b)[0]
            gp[idx] = rng.permutation(gp[idx])
        Dp = np.column_stack([np.ones_like(Xs), Xs, gp])
        perm[i] = np.linalg.lstsq(Dp, Ys, rcond=None)[0][2]
    p_one = float((1 + np.sum(perm >= gamma)) / (len(perm) + 1))

    # NEGATIVE CONTROL: a synthetic +0.8 logit inflation of post-cutoff AIME scores
    # must be detected; if the test cannot see a real shift it is uninformative.
    Yi = Ys + 0.8 * gs
    ci = np.linalg.lstsq(D, Yi, rcond=None)[0]
    permi = np.empty(2000)
    rng2 = np.random.default_rng(1)
    for i in range(2000):
        gp = gs.copy()
        for b in np.unique(bins):
            idx = np.where(bins == b)[0]
            gp[idx] = rng2.permutation(gp[idx])
        Dp = np.column_stack([np.ones_like(Xs), Xs, gp])
        permi[i] = np.linalg.lstsq(Dp, Yi, rcond=None)[0][2]
    p_inj = float((1 + np.sum(permi >= ci[2])) / (len(permi) + 1))

    # placebo: randomise the group label entirely -> gamma should be ~0
    rng3 = np.random.default_rng(2)
    pl = []
    for _ in range(200):
        gp = rng3.permutation(gs)
        Dp = np.column_stack([np.ones_like(Xs), Xs, gp])
        pl.append(float(np.linalg.lstsq(Dp, Ys, rcond=None)[0][2]))

    res = {
        "n_merged": n_merged, "n_in_support": int(m.sum()),
        "n_pre": int((gs == 0).sum()), "n_post": int((gs == 1).sum()),
        "release_cutoff": "2025-02-06",
        "alpha_intercept": float(coef[0]), "beta_slope": float(coef[1]),
        "gamma_group_shift": gamma, "se_gamma": se_g, "t_gamma": float(tstat),
        "p_one_sided_permutation": p_one,
        "paper_p_value": 0.15, "paper_n": 90,
        "gamma_positive": bool(gamma > 0),
        "significant_at_0.05": bool(p_one < 0.05),
        "r2": float(1 - (resid @ resid) / np.sum((Ys - Ys.mean()) ** 2)),
        "odds_ratio": float(np.exp(gamma)),
        "negative_control_injected_shift": {
            "injected_logit_shift": 0.8,
            "recovered_gamma": float(ci[2]),
            "p_one_sided": p_inj,
            "control_broke_as_required": bool(p_inj < 0.05)},
        "placebo_random_group_label": {
            "mean_gamma": float(np.mean(pl)), "sd_gamma": float(np.std(pl)),
            "control_broke_as_required": bool(abs(np.mean(pl)) < 0.5 * abs(gamma))},
    }
    print(f"  n merged={n_merged}, n in support={int(m.sum())} (paper 90)")
    print(f"  gamma={gamma:.4f} (se {se_g:.4f}, t={tstat:.3f}), one-sided p={p_one:.4f} "
          f"(paper 0.15)")
    print(f"  injected-shift control: gamma={ci[2]:.3f}, p={p_inj:.4f}")
    dump("claim6_contamination.json", res)
    return res


# ===========================================================================
if __name__ == "__main__":
    gates()
    claim1()
    claim2()
    claim3()
    claim4()
    claim5()
    claim6()
    print("\nAll analyses complete.")