scripts/claim4_alg1.py

"""Claim 4 -- Algorithm 1: two-stage tuning.
  Stage 1 (offline): subsample M observations, get theta_hat, estimate the
      sandwich covariance Shat = Jhat^{-1} Ihat Jhat^{-1} (/N).
  Stage 2: solve the coupled equations (11)+(12) for the preconditioned step-size
      matrix Lambda so that the stationary covariance equals Shat, then run
      preconditioned SG(L)D.

Testbed: the misspecified heteroskedastic linear regression of Section 6.1,
    y_n | x_n ~ N(x_n' theta*, 1+||x_n||^2),  x_n ~ N(0,I_D),
with a fraction p of outliers (mean shift b, variance inflation s^2), reduced to
D = 10, N = 2000 so that the D^2 x D^2 linear solves stay on CPU.

Checks
  (i)   Stage 1 is consistent: ||Shat_M - Shat_N||/||Shat_N|| decays like M^{-1/2}.
  (ii)  Stage 2's root solve: residual of Eq. (11), and whether the closed-form
        start Lambda_0 = (S H + H S)(C + H S H)^{-1} (which the authors' released
        code uses instead) already solves it.
  (iii) Internal consistency: the EXACT stationary covariance of the proxy under
        the solved Lambda must equal the target Shat to machine precision.
  (iv)  End-to-end: preconditioned SGD run with each tuning rule; error to the
        full-data sandwich target.
  (v)   Prop 4.4 mixing time 2/mu_min(Lambda H) - 1 versus the measured integrated
        autocorrelation time, and the "O(1) epochs" statement.
  (vi)  Cost: Lambda-computation time / MCMC time (paper Table D.1: 1e-5 - 1e-3).
"""

import json
import sys
import time
import numpy as np

sys.path.insert(0, "/home/ubuntu/samuel/sgmcmc-uq-repro/scripts")
from common import NoiseModel, stationary_cov_proxy, eq11_residual, relf
from linreg import LinReg, sandwich, tune_Lambda

OUT = "/home/ubuntu/samuel/sgmcmc-uq-repro/outputs"
SEED = 20260725
D, N = 10, 2000
P_OUT, B_SHIFT, S_INFL = 0.01, 5.0, 5.0

rng = np.random.default_rng(SEED)
X = rng.standard_normal((N, D))
th_star = rng.standard_normal(D)
sd = np.sqrt(1.0 + (X**2).sum(1))
y = X @ th_star + sd * rng.standard_normal(N)
out_idx = rng.choice(N, size=int(P_OUT * N), replace=False)
y[out_idx] = (
    X[out_idx] @ th_star
    + B_SHIFT
    + S_INFL * sd[out_idx] * rng.standard_normal(len(out_idx))
)
print(
    f"misspecified heteroskedastic linear regression: N={N} D={D} "
    f"{len(out_idx)} outliers"
)

t0 = time.time()
full = LinReg(X, y, sigma=1.0, loss="log")
th_full = full.fit()
g_full = full.grad_n(th_full)
J_full = full.hess_n(th_full)
S_target = sandwich(J_full.mean(0), g_full.T @ g_full / N, N)
res = {
    "claim": "Algorithm 1",
    "seed": SEED,
    "N": N,
    "D": D,
    "n_outliers": int(len(out_idx)),
    "note_alg1_line5": (
        "Algorithm 1 line 5 prints Shat <- Jhat^-1 Ihat Jhat^-1 "
        "with no 1/N; the covariance of an N-sample estimator "
        "requires the 1/N, which the authors' released code does "
        "apply.  The 1/N is used here."
    ),
}

# ------------------------------------------------------------------ (i) ----
print("\n=== (i) Stage 1: sandwich estimate from a subsample of size M ===")
stage1 = []
for M in [100, 200, 400, 800, 1600, 2000]:
    errs = []
    for rep in range(20):
        r2 = np.random.default_rng(SEED + 1000 * rep + M)
        idx = r2.choice(N, size=M, replace=False)
        sub = LinReg(X[idx], y[idx], sigma=1.0, loss="log")
        th_M = sub.fit()
        gM = sub.grad_n(th_M)
        S_M = sandwich(sub.hess_n(th_M).mean(0), gM.T @ gM / M, N)
        errs.append(relf(S_M, S_target))
    stage1.append(
        dict(M=M, mean_rel_err=float(np.mean(errs)), sd=float(np.std(errs)), n_reps=20)
    )
    print(
        f"  M={M:<6} ||Shat_M - Shat_N||/||Shat_N|| = {np.mean(errs):.4f} "
        f"(sd {np.std(errs):.4f})"
    )
sl = np.polyfit(
    np.log([r["M"] for r in stage1[:-1]]),
    np.log([r["mean_rel_err"] for r in stage1[:-1]]),
    1,
)[0]
print(f"  fitted decay exponent in M: {sl:.3f} (root-M consistency predicts -0.5)")
res["stage1"] = dict(rows=stage1, fitted_exponent=float(sl))

# ---------------------------------------------------------- (ii,iii,iv,v,vi)
print("\n=== (ii-vi) Stage 2 ===")
nm = NoiseModel(g_full, J_full, th_full, None, N)
H = J_full.mean(0)
Ihat = g_full.T @ g_full / N
A_cov = X.T @ X / N
sig2 = float(np.var(y - X @ th_full))
stage2 = []
for B in [16, 200]:
    n_iters = int(200 * N / B)
    for method in ["CT", "LR+WS", "DQ+const", "DQ+exact"]:
        tl = time.time()
        Lam, rr, Lam0 = tune_Lambda(
            method, H, S_target, Ihat, H, B, nm=nm, A_cov=A_cov, sig2=sig2, sigma=1.0
        )
        tune_t = time.time() - tl
        # (iii) exact stationary covariance of the proxy under this Lambda
        try:
            S_pred = stationary_cov_proxy(nm, Lam, B)
            pred_err = relf(S_pred, S_target)
            r11 = eq11_residual(nm, Lam, S_pred, B)
        except Exception:
            pred_err, r11 = float("nan"), float("nan")
        # residual of the closed-form start
        R0 = nm.C(S_target, B) + H @ S_target @ H
        F0 = Lam0 @ H @ S_target + S_target @ H @ Lam0 - Lam0 @ R0 @ Lam0
        cf_res = float(np.linalg.norm(F0) / np.linalg.norm(Lam0 @ H @ S_target))
        ms = time.time()
        covs, div = full.sgd_paths(
            th_full,
            Lam,
            B,
            n_iters,
            R=30,
            rng=np.random.default_rng(SEED + B),
            burn_frac=0.5,
        )
        mcmc_t = time.time() - ms
        errs = np.array([relf(c, S_target) for c in covs])
        errs[div] = np.inf
        fin = errs[np.isfinite(errs)]
        row = dict(
            B=B,
            method=method,
            tuning_residual=float(rr),
            closedform_residual=cf_res,
            exact_proxy_cov_vs_target=pred_err,
            eq11_residual=r11,
            sgd_cov_err_mean=float(np.mean(errs)) if len(fin) == 30 else float("inf"),
            sgd_cov_err_median=float(np.median(errs)),
            n_diverged=int(div.sum()),
            tune_time_s=tune_t,
            mcmc_time_s=mcmc_t,
            tune_over_mcmc=tune_t / mcmc_t,
            n_iters=n_iters,
        )
        # (v) mixing time for the tuned chain
        if method == "DQ+exact":
            ev = np.linalg.eigvals(Lam @ H).real
            tau_pred = 2.0 / ev.min() - 1.0
            rr2 = np.random.default_rng(SEED + 5)
            covs2, _ = full.sgd_paths(
                th_full, Lam, B, 4000, R=1, rng=rr2, burn_frac=0.0
            )
            # measure tau_int on one long chain
            Xp = []
            TH = th_full.copy()
            path = np.zeros((20000, D))
            for t in range(20000):
                idx = rr2.integers(0, N, size=B)
                gg = full.grad_n(TH, X[idx], y[idx]).mean(0)
                TH = TH - Lam @ gg
                path[t] = TH
            v = np.linalg.eigh(S_target)[1][:, -1]
            f = path[5000:] @ v
            f = f - f.mean()
            ac = np.correlate(f, f, "full")[len(f) - 1 :]
            ac = ac / ac[0]
            cut = np.argmax(ac < 0.05) if np.any(ac < 0.05) else len(ac)
            tau_int = 1 + 2 * ac[1:cut].sum()
            row.update(
                tau_pred_prop44=float(tau_pred),
                tau_int_measured=float(tau_int),
                epochs_to_mix=float(tau_int * B / N),
            )
            print(
                f"    Prop 4.4 predicted tau = {tau_pred:.1f}; measured tau_int = "
                f"{tau_int:.1f}  ({tau_int*B/N:.3f} epochs)"
            )
        stage2.append(row)
        print(
            f"  B={B:<4} {method:<9} root-resid={rr:.2e} closed-form-resid={cf_res:.2e} "
            f"exact-proxy-cov-vs-target={pred_err:.2e}  SGD cov err="
            f"{row['sgd_cov_err_median']:.4g}  tune/mcmc={tune_t/mcmc_t:.2e}"
        )
res["stage2"] = stage2
res["runtime_s"] = time.time() - t0
with open(f"{OUT}/claim4_alg1.json", "w") as f:
    json.dump(res, f, indent=1)
print("\nwrote outputs/claim4_alg1.json  runtime %.1f s" % res["runtime_s"])