repro-bundle-v0/repro/run_experiments.py

"""Experiments for all 6 claims of arXiv:2602.13960.

Claim 1: W1(Y^(a), N(0,S_Y)) <= U sqrt(a) log(1/a)  (iid Thm 3.1 +
         Markov Thm 4.1) -- combined evidence from the instantiations below.
Claim 2: SGD, smooth strongly convex (Prop 3.1).
Claim 3: linear SA with Hurwitz matrix + contractive nonlinear SA
         (Props 3.2, 3.3).
Claim 4: non-uniform Berry-Esseen tail bound ~ C a^{1/4} log^{1/2}(1/a)/a.
Claim 5: Markovian noise extension (Prop 4.1) -- including the falsification
         check that the *naive* marginal-variance Gaussian is wrong.
Claim 6: general convex (f = |x|^h/h): Gibbs limit at scaling a^{1/h}
         (Prop 5.1).

Usage: python repro/run_experiments.py       (writes results/claims.json)
"""

from __future__ import annotations

import json
import os
import time

import numpy as np
from scipy import stats

import sa

HERE = os.path.dirname(os.path.abspath(__file__))
RES = os.path.join(os.path.dirname(HERE), "results")
os.makedirs(RES, exist_ok=True)

ALPHAS = [0.2, 0.1, 0.05, 0.025, 0.0125, 0.00625, 0.003125]
N = 400_000
N_TAIL = 1_600_000


def claim2_sgd():
    print("== Claim 2: SGD smooth strongly convex (Prop 3.1) ==")
    var_y = sa.sgd1d_limit_var()
    w1 = []
    for a in ALPHAS:
        t0 = time.time()
        x = sa.sgd1d_stationary(a, N, seed=int(1e6 * a))
        y = x / np.sqrt(a)
        w1.append(sa.w1_to_gaussian(y, var_y))
        print(f"  a={a:<8g} W1={w1[-1]:.5f}  [{time.time()-t0:.0f}s]")
    m, b, r2 = sa.fit_powerlaw(ALPHAS, w1)
    U, r2u = sa.fit_sqrtlog(ALPHAS, w1)
    print(f"  power-law slope={m:.3f} (theory 0.5 + log) R2={r2:.4f}; "
          f"U-model U={U:.3f} R2={r2u:.4f}")
    return {"alphas": ALPHAS, "w1": w1, "var_y": var_y, "slope": m,
            "slope_r2": r2, "U": U, "U_r2": r2u, "N": N}


def claim3_linear_contractive():
    print("== Claim 3: linear SA + contractive nonlinear SA ==")
    out = {}
    S_lin = sa.lyapunov_cov(sa.A_LIN, sa.SIG_LIN)
    S_con = sa.lyapunov_cov(sa.C_CONTRACT * sa._R - np.eye(2), sa.SIG_LIN)
    for name, sampler, S in [("linear", sa.linear_sa_stationary, S_lin),
                             ("contractive", sa.contractive_sa_stationary,
                              S_con)]:
        w1 = []
        for a in ALPHAS:
            t0 = time.time()
            x = sampler(a, N, seed=int(2e6 * a) + hash(name) % 1000)
            y = x / np.sqrt(a)
            v, _ = sa.sliced_w1_to_gaussian(y, S)
            w1.append(v)
            print(f"  {name:11s} a={a:<8g} sliced-W1={v:.5f} "
                  f"[{time.time()-t0:.0f}s]")
        m, b, r2 = sa.fit_powerlaw(ALPHAS, w1)
        U, r2u = sa.fit_sqrtlog(ALPHAS, w1)
        print(f"  {name}: slope={m:.3f} R2={r2:.4f}; U={U:.3f} R2={r2u:.4f}")
        out[name] = {"alphas": ALPHAS, "w1": w1, "Sigma_Y": S.tolist(),
                     "slope": m, "slope_r2": r2, "U": U, "U_r2": r2u,
                     "N": N}
    return out


def claim4_tails():
    print("== Claim 4: Berry-Esseen-type tail bounds ==")
    zeta = np.array([1.0, 1.0]) / np.sqrt(2.0)
    S = sa.lyapunov_cov(sa.A_LIN, sa.SIG_LIN)
    sd = float(np.sqrt(zeta @ S @ zeta))
    a_levels = [0.5, 1.0, 1.5, 2.0, 2.5, 3.0]
    alphas = [0.1, 0.05, 0.025, 0.0125, 0.00625]
    rows = []
    sup_scaled = []
    for al in alphas:
        t0 = time.time()
        x = sa.linear_sa_stationary(al, N_TAIL, seed=int(3e6 * al))
        proj = (x / np.sqrt(al)) @ zeta
        deltas = []
        for a in a_levels:
            p_emp = float((proj > a * sd).mean())
            p_gauss = float(stats.norm.sf(a))
            deltas.append(abs(p_emp - p_gauss))
        # the bound is C a^{1/4} log^{1/2}/a  => a*Delta should be ~flat in a
        scaled = [a * d for a, d in zip(a_levels, deltas)]
        sup_scaled.append(max(scaled))
        rows.append({"alpha": al, "deltas": deltas, "scaled": scaled})
        print(f"  a={al:<8g} sup_a a*|Delta|={max(scaled):.5f} "
              f"Delta(0.5)={deltas[0]:.5f} Delta(3)={deltas[-1]:.6f} "
              f"[{time.time()-t0:.0f}s]")
    m, b, r2 = sa.fit_powerlaw(alphas, sup_scaled)
    # decay in deviation level: mean over alphas of Delta(a)*a / Delta(a)...
    print(f"  sup_a a|Delta| power-law slope={m:.3f} "
          f"(theory 0.25 + log^0.5) R2={r2:.4f}")
    return {"alphas": alphas, "a_levels": a_levels, "rows": rows,
            "sup_scaled": sup_scaled, "slope": m, "slope_r2": r2,
            "sd": sd, "N": N_TAIL}


def claim5_markov():
    print("== Claim 5: Markovian noise (Prop 4.1) ==")
    q, c = 0.25, 0.5
    var_correct, var_naive = sa.sgd1d_limit_var(markov=(q, c))
    w1c, w1n = [], []
    for a in ALPHAS:
        t0 = time.time()
        x = sa.sgd1d_stationary(a, N, seed=int(4e6 * a), markov=(q, c))
        y = x / np.sqrt(a)
        w1c.append(sa.w1_to_gaussian(y, var_correct))
        w1n.append(sa.w1_to_gaussian(y, var_naive))
        print(f"  a={a:<8g} W1(correct)={w1c[-1]:.5f} "
              f"W1(naive)={w1n[-1]:.5f} [{time.time()-t0:.0f}s]")
    m, b, r2 = sa.fit_powerlaw(ALPHAS, w1c)
    U, r2u = sa.fit_sqrtlog(ALPHAS, w1c)
    print(f"  correct-S_Y slope={m:.3f} R2={r2:.4f}; U={U:.3f} R2={r2u:.4f}")
    print(f"  naive-S_Y floor: min W1 = {min(w1n):.4f} (does NOT vanish)")
    return {"alphas": ALPHAS, "w1_correct": w1c, "w1_naive": w1n,
            "var_correct": var_correct, "var_naive": var_naive,
            "slope": m, "slope_r2": r2, "U": U, "U_r2": r2u,
            "q": q, "c": c, "N": N}


def claim6_gibbs():
    print("== Claim 6: Gibbs limit for general convex (Prop 5.1) ==")
    out = {}
    alphas6 = [0.2, 0.1, 0.05, 0.025, 0.0125, 0.00625]
    for h in [2, 4]:
        qf = sa.gibbs_target_quantiles(h)
        mean_abs, w1g, w1_bestgauss = [], [], []
        for a in alphas6:
            t0 = time.time()
            x = sa.gibbs_sgd_stationary(a, N, seed=int(5e6 * a) + h, h=h)
            y = x / a ** (1.0 / h)
            mean_abs.append(float(np.abs(x).mean()))
            w1g.append(sa.w1_to_quantiles(y, qf))
            w1_bestgauss.append(sa.w1_to_gaussian(y, float(np.var(y))))
            print(f"  h={h} a={a:<8g} E|X|={mean_abs[-1]:.5f} "
                  f"W1(Gibbs)={w1g[-1]:.5f} "
                  f"W1(best Gauss)={w1_bestgauss[-1]:.5f} "
                  f"[{time.time()-t0:.0f}s]")
        m, b, r2 = sa.fit_powerlaw(alphas6, mean_abs)
        out[f"h{h}"] = {"alphas": alphas6, "mean_abs": mean_abs,
                        "w1_gibbs": w1g, "w1_best_gauss": w1_bestgauss,
                        "scale_slope": m, "scale_r2": r2,
                        "theory_slope": 1.0 / h, "N": N}
        print(f"  h={h}: E|X| ~ alpha^{m:.4f} (theory {1.0/h:.4f}) "
              f"R2={r2:.5f}")
    # kurtosis check at the smallest alpha for h=4 (non-Gaussianity)
    x = sa.gibbs_sgd_stationary(alphas6[-1], N, seed=99, h=4)
    y = x / alphas6[-1] ** 0.25
    kurt = float(stats.kurtosis(y, fisher=False))
    out["h4_kurtosis"] = kurt
    # Gibbs law kurtosis for exp(-y^4/2)-type: ~2.19 (sub-Gaussian, flat top)
    print(f"  h=4 kurtosis={kurt:.3f} (Gaussian=3; Gibbs |y|^4 law ~ 2.19)")
    return out


def main():
    t0 = time.time()
    res = {}
    res["claim2_sgd"] = claim2_sgd()
    res["claim3"] = claim3_linear_contractive()
    res["claim4_tails"] = claim4_tails()
    res["claim5_markov"] = claim5_markov()
    res["claim6_gibbs"] = claim6_gibbs()
    # claim 1 aggregates the W1 rates of claims 2/3/5
    res["claim1_summary"] = {
        "sgd_iid": {"slope": res["claim2_sgd"]["slope"],
                    "U": res["claim2_sgd"]["U"],
                    "U_r2": res["claim2_sgd"]["U_r2"]},
        "linear_iid": {"slope": res["claim3"]["linear"]["slope"],
                       "U": res["claim3"]["linear"]["U"],
                       "U_r2": res["claim3"]["linear"]["U_r2"]},
        "contractive_iid": {"slope": res["claim3"]["contractive"]["slope"],
                            "U": res["claim3"]["contractive"]["U"],
                            "U_r2": res["claim3"]["contractive"]["U_r2"]},
        "sgd_markov": {"slope": res["claim5_markov"]["slope"],
                       "U": res["claim5_markov"]["U"],
                       "U_r2": res["claim5_markov"]["U_r2"]},
    }
    res["wall_seconds"] = time.time() - t0
    with open(os.path.join(RES, "claims.json"), "w", encoding="utf-8") as f:
        json.dump(res, f, indent=1)
    print(f"done in {res['wall_seconds']:.0f}s -> results/claims.json")


if __name__ == "__main__":
    main()