scripts/claim1_thm41.py

"""Claim 1 -- Theorem 4.1:  ||Sigma_theta - Sigma_psi|| / ||Sigma_theta|| <= C_v lambda^{1/2}
                            |sigma_{theta,d} - sigma_{psi,d}| / sigma_{theta,d} <= C_s lambda^{1/2}

Independent test.
  * Sigma_psi is computed EXACTLY (no simulation): the one-step second-moment
    identity of the proxy recursion is affine in Sigma, so it is solved as a
    D^2 linear system.  (This also re-derives Prop. 4.2 / Eq. (11).)
  * Sigma_theta is estimated from SGD run on the TRUE logistic loss.  The proxy
    is run on the SAME minibatches, so Sigma_theta_hat - Sigma_psi_hat is a
    control variate with a variance orders of magnitude smaller than either
    covariance alone; the reported Sigma_theta = Sigma_psi_exact + (hat diff).
  * The empirical scaling exponent of the relative error in lambda is fitted and
    compared with the exponent 1/2 that the theorem asserts as an upper bound.
  * Boundary audit: step size pushed past lambda < 1/(2L), and Assumption (C)
    (strong convexity) removed by a rank-deficient design with no prior.

CPU only.  Seeds are explicit.
"""

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,
    stationary_cov_proxy,
    eq11_residual,
    cov_from_blocks,
    relf,
    w2_marginal_lower,
)

OUT = "/home/ubuntu/samuel/sgmcmc-uq-repro/outputs"
SEED = 20260725
B = 16
LAMBDAS = [0.02, 0.05, 0.10, 0.20, 0.35, 0.50]

t_start = time.time()
m, that, nm, K = make(N=500, D=3, seed=SEED, gamma=2.0)
D = m.D
print("constants", K)
print("1/(2L) =", 1.0 / (2 * K["L"]), " mu_hat", K["mu_hat"], " L_hat", K["L_hat"])

rows = []
for lam in LAMBDAS:
    Lam = lam * np.eye(D)
    S_psi = stationary_cov_proxy(nm, Lam, B)
    r11 = eq11_residual(nm, Lam, S_psi, B)
    T = int(4000 + 400 / lam)
    burn = T // 5
    rng = np.random.default_rng(SEED + int(1e5 * lam))
    out = run_coupled(
        m, that, Lam, B, T=T, R=1000, rng=rng, burn=burn, thin=40, nblock=10
    )
    nb = out["accT"].shape[0]

    def stat(mask):
        St = cov_from_blocks(out["accT"], out["sT"], out["cnt"], mask)
        Sp = cov_from_blocks(out["accP"], out["sP"], out["cnt"], mask)
        Sth = S_psi + (St - Sp)  # control-variate estimate
        return float(np.linalg.norm(Sth - S_psi, 'fro') / np.linalg.norm(Sth, 'fro'))

    from common import jackknife

    rel_cov, se_cov = jackknife(stat, nb)
    St = cov_from_blocks(out["accT"], out["sT"], out["cnt"])
    Sp = cov_from_blocks(out["accP"], out["sP"], out["cnt"])
    S_theta = S_psi + (St - Sp)
    sd_t, sd_p = np.sqrt(np.diag(S_theta)), np.sqrt(np.diag(S_psi))
    rel_sd = float(np.max(np.abs(sd_t - sd_p) / sd_t))
    # coupling / W2 quantities (used again by claim 3)
    d2 = out["accD2"].sum() / out["cnt"].sum()
    w2_up = float(np.sqrt(d2))
    w2_lo = w2_marginal_lower(out["thetaS"], out["psiS"])
    rows.append(
        dict(
            lam=lam,
            B=B,
            T=T,
            R=1000,
            eq11_residual=r11,
            norm_Sigma_psi=float(np.linalg.norm(S_psi, "fro")),
            norm_Sigma_theta=float(np.linalg.norm(S_theta, "fro")),
            rel_cov_err=rel_cov,
            rel_cov_err_se=se_cov,
            rel_sd_err=rel_sd,
            Cv_emp=rel_cov / np.sqrt(lam),
            Cs_emp=rel_sd / np.sqrt(lam),
            w2_coupling_upper=w2_up,
            w2_marginal_lower=w2_lo,
            raw_sim_proxy_err=relf(Sp, S_psi),
        )
    )
    print(
        f"lam={lam:<6} relcov={rel_cov:.5f}+-{se_cov:.5f}  relsd={rel_sd:.5f}  "
        f"Cv={rel_cov/np.sqrt(lam):.4f}  W2up={w2_up:.3e}  eq11res={r11:.1e}"
    )

lg = np.log(np.array([r["lam"] for r in rows]))
lc = np.log(np.array([r["rel_cov_err"] for r in rows]))
ls = np.log(np.array([r["rel_sd_err"] for r in rows]))
slope_cov, ic = np.polyfit(lg, lc, 1)
slope_sd, _ = np.polyfit(lg, ls, 1)
r2 = 1 - np.sum((lc - (slope_cov * lg + ic)) ** 2) / np.sum((lc - lc.mean()) ** 2)
print(f"\nfitted exponent  cov: {slope_cov:.3f} (R^2={r2:.4f})   sd: {slope_sd:.3f}")
print(
    "theorem asserts the error is bounded by C_v lambda^0.5 -> any fitted "
    "exponent >= 0.5 is consistent; < 0.5 would falsify."
)

# ---------------- boundary audits -----------------------------------------
audits = {}
# (a) step size past lambda < 1/(2L): proxy chain must lose stationarity
Lhat = K["L_hat"]
aud = []
for f in [0.5, 1.0, 1.5, 1.9, 2.0, 2.2]:
    lam = f / Lhat
    Lam = lam * np.eye(D)
    rho = float(np.max(np.abs(np.linalg.eigvals(np.eye(D) - Lam @ nm.H))))
    try:
        S = stationary_cov_proxy(nm, Lam, B)
        ev = float(np.linalg.eigvalsh(S)[0])
        ok = ev > 0
    except Exception:
        ev, ok = float("nan"), False
    aud.append(
        dict(
            lam=lam,
            lam_times_Lhat=f,
            spectral_radius=rho,
            min_eig_Sigma_psi=ev,
            positive_definite=bool(ok),
        )
    )
    print(
        f"  audit-a lam*Lhat={f}: rho(I-Lam H)={rho:.4f} minEig(Sigma_psi)={ev:.3e} PD={ok}"
    )
audits["stepsize_boundary"] = aud

# (b) Assumption (C) removed: rank-deficient design, no prior -> mu = mu_hat = 0
m2, that2, nm2, K2 = make(N=500, D=3, seed=SEED, gamma=1e-8, rank_deficient=True)
print(
    f"  audit-b  mu_hat with Assumption (C) violated = {K2['mu_hat']:.3e} "
    f"(vs {K['mu_hat']:.3e} when it holds)"
)
audits["strong_convexity_violation"] = dict(
    mu_hat_violated=K2["mu_hat"],
    mu_hat_ok=K["mu_hat"],
    cond_H_violated=float(np.linalg.cond(nm2.H)),
    cond_H_ok=float(np.linalg.cond(nm.H)),
)
try:
    Sbad = stationary_cov_proxy(nm2, 0.1 * np.eye(3), B)
    audits["strong_convexity_violation"]["norm_Sigma_psi"] = float(np.linalg.norm(Sbad))
    audits["strong_convexity_violation"]["min_eig"] = float(np.linalg.eigvalsh(Sbad)[0])
    print(
        "           ||Sigma_psi|| when mu=0:",
        np.linalg.norm(Sbad),
        " min eig:",
        np.linalg.eigvalsh(Sbad)[0],
    )
except Exception as e:
    audits["strong_convexity_violation"]["error"] = str(e)

# (c) heavy-tailed covariates: assumptions still hold but M, tau4 inflate -> C_v grows
m3, that3, nm3, K3 = make(N=500, D=3, seed=SEED, gamma=2.0, design="heavy")
lam = 0.2
Lam = lam * np.eye(3)
S3 = stationary_cov_proxy(nm3, Lam, B)
o3 = run_coupled(
    m3,
    that3,
    Lam,
    B,
    T=6000,
    R=1000,
    rng=np.random.default_rng(SEED + 7),
    burn=1500,
    nblock=10,
)
St3 = cov_from_blocks(o3["accT"], o3["sT"], o3["cnt"])
Sp3 = cov_from_blocks(o3["accP"], o3["sP"], o3["cnt"])
Sth3 = S3 + (St3 - Sp3)
rel3 = float(np.linalg.norm(Sth3 - S3, 'fro') / np.linalg.norm(Sth3, 'fro'))
base = [r for r in rows if r["lam"] == lam][0]["rel_cov_err"]
audits["heavy_tailed_design"] = dict(
    lam=lam,
    Mbar_gauss=K["Mbar"],
    Mbar_heavy=K3["Mbar"],
    tau4_gauss=K["tau4"],
    tau4_heavy=K3["tau4"],
    rel_cov_err_gauss=base,
    rel_cov_err_heavy=rel3,
    Cv_gauss=base / np.sqrt(lam),
    Cv_heavy=rel3 / np.sqrt(lam),
)
print(
    f'  audit-c heavy-tailed design: Mbar {K["Mbar"]:.3f}->{K3["Mbar"]:.3f}, '
    f"rel cov err {base:.4f}->{rel3:.4f}"
)

res = dict(
    claim="Theorem 4.1",
    model="logistic regression + Gaussian prior",
    N=500,
    D=3,
    B=B,
    beta="inf (SGD)",
    seed=SEED,
    constants=K,
    lambda_upper_bound_1_over_2L=1.0 / (2 * K["L"]),
    sweep=rows,
    fitted_exponent_cov=float(slope_cov),
    fitted_exponent_sd=float(slope_sd),
    fit_r2=float(r2),
    Cv_max=float(max(r["Cv_emp"] for r in rows)),
    Cs_max=float(max(r["Cs_emp"] for r in rows)),
    audits=audits,
    runtime_s=time.time() - t_start,
)
with open(f"{OUT}/claim1_thm41.json", "w") as f:
    json.dump(res, f, indent=1)
np.save(
    f"{OUT}/claim1_lambda_sweep.npy",
    np.array(
        [
            [
                r["lam"],
                r["rel_cov_err"],
                r["rel_cov_err_se"],
                r["rel_sd_err"],
                r["w2_coupling_upper"],
                r["w2_marginal_lower"],
            ]
            for r in rows
        ]
    ),
)
print("\nwrote outputs/claim1_thm41.json   runtime %.1f s" % (time.time() - t_start))