scripts/claim3_w2.py

"""Claim 3 -- Theorem 4.5 (finite-t W2 recursion) and Corollary 4.6 (W2 <= A lambda / B).

Theorem 4.5 is proved by a *synchronous coupling* argument: it bounds
E||theta_t - psi_t||^2 when both chains see the same minibatch and the same
injected noise, and then uses W2^2 <= E||theta_t-psi_t||^2.  We therefore
reproduce exactly that quantity, which is a strictly stronger test of the
theorem than testing W2 itself:

  (A) finite-t check of Eq. (14).  Both chains start at theta_hat so
      W2(theta_0,psi_0)=0.  The right-hand side needs C_s = E||psi_s-theta_hat||^4,
      which is measured from the same run.  We check LHS <= RHS at every t.
  (B) stationary check of Corollary 4.6, W2(pi_theta,pi_psi) <= A lambda/B, with
      A = sqrt(96 C_0) tau_4^2 / mu_hat and C_0 = (2/mu)(Mbar2/(8L) + Mbar^2/(4mu))
      read off Equation (F.12); evaluated with the closed-form logistic constants.
  (C) empirical scaling exponents in lambda (predicted 1) and in B (predicted -1).
  (D) a 1-D marginal W2 computed from sorted quantiles gives a LOWER bound, so
      the true W2 is bracketed.
"""

import json
import sys
import time
import numpy as np

sys.path.insert(0, "/home/ubuntu/samuel/sgmcmc-uq-repro/scripts")
from setup_problem import make
from common import run_coupled, w2_marginal_lower

OUT = "/home/ubuntu/samuel/sgmcmc-uq-repro/outputs"
SEED = 20260725
t0 = time.time()
m, that, nm, K = make(N=500, D=3, seed=SEED, gamma=2.0)
D = m.D
L, mu, muh, Mb, Mb2, tau4 = (
    K["L"],
    K["mu_global"],
    K["mu_hat"],
    K["Mbar"],
    K["Mbar2"],
    K["tau4"],
)
print("constants", K)


def A_const(mu_use):
    C0 = (2.0 / mu_use) * (Mb2 / (8 * L) + Mb**2 / (4 * mu_use))
    return float(np.sqrt(96 * C0) * tau4**2 / muh), float(C0)


res = {
    "claim": "Theorem 4.5 / Corollary 4.6",
    "seed": SEED,
    "constants": K,
    "N": 500,
    "D": D,
    "beta": "inf (SGD)",
}

# ---------------------------------------------------------------- (A) ------
print("=== (A) finite-t check of Eq. (14) ===")


def transient(lam, B, T, R, seed):
    """E||theta_t - psi_t||^2 and C_t = E||psi_t - theta_hat||^4 for t=0..T."""
    rng = np.random.default_rng(seed)
    X, y, Gam = m.X, m.y, m.Gamma / m.N
    gh = nm.g
    p = 1 / (1 + np.exp(-X @ that))
    wh = p * (1 - p)
    TH = np.repeat(that[None, :], R, 0)
    PS = TH.copy()
    Lam = lam * np.eye(D)
    d2 = [0.0]
    C4 = [0.0]
    for t in range(T):
        IDX = rng.integers(0, m.N, size=(R, B))
        Xb = X[IDX]
        zT = np.einsum("rbd,rd->rb", Xb, TH)
        gT = (
            np.einsum("rb,rbd->rd", 1 / (1 + np.exp(-zT)) - y[IDX], Xb) / B + TH @ Gam.T
        )
        U = PS - that
        sU = np.einsum("rbd,rd->rb", Xb, U)
        gP = (
            gh[IDX].mean(1)
            + np.einsum("rb,rbd->rd", wh[IDX] * sU, Xb) / B
            + U @ Gam.T
            + that @ Gam.T
        )
        TH = TH - gT @ Lam.T
        PS = PS - gP @ Lam.T
        d2.append(float(((TH - PS) ** 2).sum(1).mean()))
        C4.append(float(((((PS - that) ** 2).sum(1)) ** 2).mean()))
    return np.array(d2), np.array(C4)


trans = []
for lam in [0.02, 0.05]:
    B = 16
    T = 2000
    d2, C4 = transient(lam, B, T, R=4000, seed=SEED + 31)
    beta_bar = 1 - lam * mu * (1 - 2 * lam * L)
    pref = lam * (lam * Mb2 / 2 + Mb**2 / (4 * mu))
    rhs = np.zeros(T + 1)
    acc = 0.0
    for t in range(1, T + 1):
        acc = beta_bar * acc + C4[t - 1]  # sum_{s=1..t} bar_beta^{t-s} C_{s-1}
        rhs[t] = pref * acc
    # theta_1 == psi_1 identically (the proxy gradient at psi_0 = theta_hat IS the
    # true minibatch gradient), so both sides are 0 at t=1 up to rounding; only
    # compare where the measured LHS is above the float noise floor.
    live = d2[1:] > 1e-20
    ratio = rhs[1:][live] / d2[1:][live]
    ok = bool(np.all(d2[1:][live] <= rhs[1:][live]))
    n_live = int(live.sum())
    trans.append(
        dict(
            lam=lam,
            B=B,
            T=T,
            beta_bar=float(beta_bar),
            prefactor=float(pref),
            holds_at_every_t=ok,
            min_slack_rhs_over_lhs=float(ratio.min()),
            lhs_final=float(d2[-1]),
            rhs_final=float(rhs[-1]),
            t_of_min_slack=int(np.argmin(ratio)) + 1,
            n_t_compared=n_live,
            lhs_t1=float(d2[1]),
        )
    )
    print(
        f"  lam={lam} B={B}: Eq.(14) holds at every t = {ok}; "
        f"min RHS/LHS = {ratio.min():.3e} (at t={np.argmin(ratio)+1}); "
        f"final LHS={d2[-1]:.3e} RHS={rhs[-1]:.3e}"
    )
res["A_finite_t_eq14"] = trans

# ---------------------------------------------------------------- (B,C) ----
print("=== (B,C) stationary Cor 4.6 and scaling ===")
lam_cond = min(16 * muh / (200 * L**2), 1 / (4 * L))
print(
    f"  Cor 4.6 requires lambda < min(B mu_hat/(200 L^2), 1/(4L)) = {lam_cond:.4f} at B=16"
)
A_mu, C0_mu = A_const(mu)
A_muh, C0_muh = A_const(muh)
print(
    f"  explicit constant A: with global mu={mu:.4g} -> A={A_mu:.4g}; "
    f"with plug-in mu_hat={muh:.4g} -> A={A_muh:.4g}"
)

sweep_lam = np.load(
    f"{OUT}/claim1_lambda_sweep.npy"
)  # lam, relcov, se, relsd, w2up, w2lo
lam_rows = []
for r in sweep_lam:
    lam, w2u, w2l = float(r[0]), float(r[4]), float(r[5])
    lam_rows.append(
        dict(
            lam=lam,
            B=16,
            w2_coupling_upper=w2u,
            w2_marginal_lower=w2l,
            bound_A_mu=A_mu * lam / 16,
            bound_A_muhat=A_muh * lam / 16,
            within_corollary_condition=bool(lam < lam_cond),
            bound_holds_mu=bool(w2u <= A_mu * lam / 16),
            bound_holds_muhat=bool(w2u <= A_muh * lam / 16),
        )
    )
    print(
        f"  lam={lam:<6} W2 in [{w2l:.3e}, {w2u:.3e}]  bound(mu)={A_mu*lam/16:.3e} "
        f"bound(mu_hat)={A_muh*lam/16:.3e}  inCond={lam < lam_cond}"
    )

B_rows = []
lam_fixed = 0.05
for B in [4, 8, 16, 32, 64, 128]:
    R = max(100, 16000 // B)
    T = 12000
    o = run_coupled(
        m,
        that,
        lam_fixed * np.eye(D),
        B,
        T=T,
        R=R,
        rng=np.random.default_rng(SEED + B),
        burn=T // 4,
        thin=100,
        nblock=8,
    )
    w2u = float(np.sqrt(o["accD2"].sum() / o["cnt"].sum()))
    w2l = w2_marginal_lower(o["thetaS"], o["psiS"])
    cond = min(B * muh / (200 * L**2), 1 / (4 * L))
    B_rows.append(
        dict(
            B=B,
            R=R,
            T=T,
            lam=lam_fixed,
            w2_coupling_upper=w2u,
            w2_marginal_lower=w2l,
            bound_A_muhat=A_muh * lam_fixed / B,
            within_corollary_condition=bool(lam_fixed < cond),
            bound_holds_muhat=bool(w2u <= A_muh * lam_fixed / B),
        )
    )
    print(
        f"  B={B:<4} W2 in [{w2l:.3e},{w2u:.3e}] bound={A_muh*lam_fixed/B:.3e} "
        f"inCond={lam_fixed < cond}"
    )

la = np.log([r["lam"] for r in lam_rows])
wa = np.log([r["w2_coupling_upper"] for r in lam_rows])
s_lam, i_lam = np.polyfit(la, wa, 1)
r2_lam = 1 - np.sum((wa - (s_lam * la + i_lam)) ** 2) / np.sum((wa - wa.mean()) ** 2)
lb = np.log([r["B"] for r in B_rows])
wb = np.log([r["w2_coupling_upper"] for r in B_rows])
s_B, i_B = np.polyfit(lb, wb, 1)
r2_B = 1 - np.sum((wb - (s_B * lb + i_B)) ** 2) / np.sum((wb - wb.mean()) ** 2)
print(f"\n  fitted exponent in lambda: {s_lam:.3f} (predicted 1, R^2={r2_lam:.4f})")
print(f"  fitted exponent in B     : {s_B:.3f} (predicted -1, R^2={r2_B:.4f})")

res.update(
    lambda_condition_at_B16=float(lam_cond),
    A_with_global_mu=A_mu,
    A_with_plugin_muhat=A_muh,
    C0_global=C0_mu,
    C0_plugin=C0_muh,
    lambda_sweep=lam_rows,
    batch_sweep=B_rows,
    fitted_exponent_lambda=float(s_lam),
    fit_r2_lambda=float(r2_lam),
    fitted_exponent_B=float(s_B),
    fit_r2_B=float(r2_B),
    runtime_s=time.time() - t0,
)
with open(f"{OUT}/claim3_w2.json", "w") as f:
    json.dump(res, f, indent=1)
print("\nwrote outputs/claim3_w2.json  runtime %.1f s" % res["runtime_s"])