scripts/claim2_thm43.py

"""Claim 2 -- Theorem 4.3: explicit formula for the expected minibatch noise
covariance of the proxy, with NO constant-noise assumption:

    C_psi = (1/B) [ Ihat - Gamma th th^T Gamma^T / N^2
                    + (1/N) sum_n J_n Sigma_psi J_n - Jbar Sigma_psi Jbar ]

(times (N-B)/(N-1) without replacement).

Four independent checks, in increasing order of strength:
  (T1) fully empirical Monte Carlo: draw psi from the proxy chain itself, draw
       real minibatches, take the empirical covariance of the realised proxy
       gradients.  Nothing from Eq. (12) is used.
  (T2) EXACT finite-domain check.  Eq. (12) only uses E[psi-that]=0 and
       E[(psi-that)(psi-that)^T]=Sigma.  So replace the stationary law by a
       finite 2D-atom distribution with exactly those two moments; then the
       expectation over psi and the batch covariance are both finite sums and
       the identity can be checked to machine precision, with no sampling.
  (T3) EXHAUSTIVE without-replacement check: enumerate every one of the C(N,B)
       minibatches on a tiny problem and verify the (N-B)/(N-1) factor.
  (T4) consequence check: plug C_psi into Eq. (11), solve for Sigma_psi, and
       compare against the covariance of a simulated proxy chain.

Counterexample audit: the constant-noise surrogate C = Ihat/B (used by CT and by
DQ+const) and the Ziyin et al. (2022) Eq. (6) large-sample well-specified
surrogate are scored on the same problems.
"""

import itertools
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 (
    NoiseModel,
    Logistic,
    stationary_cov_proxy,
    relf,
    sym,
    run_coupled,
    cov_from_blocks,
)

OUT = "/home/ubuntu/samuel/sgmcmc-uq-repro/outputs"
SEED = 20260725
t0 = time.time()
res = {"claim": "Theorem 4.3", "seed": SEED}

# ---------------------------------------------------------------- T1 -------
print("=== T1  fully empirical Monte Carlo ===")
t1 = []
for design in ["gaussian", "heavy", "binary"]:
    m, that, nm, K = make(N=500, D=3, seed=SEED, gamma=2.0, design=design)
    B, lam = 16, 0.15
    Lam = lam * np.eye(3)
    S = stationary_cov_proxy(nm, Lam, B)
    rng = np.random.default_rng(SEED + 11)
    # run the proxy chain, and at each recorded state form REAL minibatch
    # gradients and accumulate their empirical covariance about the exact mean
    R, T, burn = 600, 4000, 1000
    U = np.zeros((R, 3))
    gh, wh, X, Gam = nm.g, None, m.X, m.Gamma / m.N
    p = 1 / (1 + np.exp(-X @ that))
    wh = p * (1 - p)
    acc = np.zeros((3, 3))
    cnt = 0
    for t in range(T):
        IDX = rng.integers(0, m.N, size=(R, B))
        Xb = X[IDX]
        sU = np.einsum("rbd,rd->rb", Xb, U)
        G = (
            gh[IDX].mean(1)
            + np.einsum("rb,rbd->rd", wh[IDX] * sU, Xb) / B
            + U @ Gam.T
            + that @ Gam.T
        )
        Gmean = nm.H @ U.T  # E[G|psi] = H (psi - that)
        if t >= burn:
            E = G - Gmean.T
            acc += E.T @ E
            cnt += R
        U = U - G @ Lam.T
    C_mc = acc / cnt
    C_f = nm.C(S, B)
    t1.append(
        dict(
            design=design,
            rel_err_formula_vs_MC=relf(C_mc, C_f),
            n_gradient_samples=int(cnt),
            norm_C=float(np.linalg.norm(C_f, "fro")),
        )
    )
    print(
        f"  design={design:<9} ||C_MC - C_eq12||/||C_eq12|| = {relf(C_mc, C_f):.3e} "
        f"({cnt} minibatch gradients)"
    )
res["T1_empirical_mc"] = t1

# ---------------------------------------------------------------- T2 -------
print("=== T2  exact finite-domain check (no sampling at all) ===")


def exact_finite_check(nm, Sigma, B, replace=True):
    """psi-that takes 2D values +-sqrt(D) c_i  (columns of the Cholesky factor),
    each with probability 1/(2D).  Then E[u]=0 and E[u u^T]=Sigma exactly."""
    D = nm.D
    Lc = np.linalg.cholesky(Sigma)
    atoms = np.concatenate([np.sqrt(D) * Lc.T, -np.sqrt(D) * Lc.T], 0)  # (2D,D)
    w = np.full(2 * D, 1.0 / (2 * D))
    # verify the atoms reproduce the moments
    mom1 = (w[:, None] * atoms).sum(0)
    mom2 = np.einsum("a,ad,ae->de", w, atoms, atoms)
    acc = np.zeros((D, D))
    for a, wa in zip(atoms, w):
        V = nm.g + np.einsum("nde,e->nd", nm.J, a)  # per-sample proxy gradient
        vbar = V.mean(0)
        Cb = (V.T @ V / nm.N - np.outer(vbar, vbar)) / B
        if not replace:
            Cb = Cb * (nm.N - B) / (nm.N - 1)
        acc += wa * Cb
    return sym(acc), float(np.linalg.norm(mom1)), relf(mom2, Sigma)


t2 = []
rng2 = np.random.default_rng(SEED + 99)
for design in ["gaussian", "heavy", "binary"]:
    for gamma in [0.0, 2.0, 25.0]:
        m, that, nm, K = make(N=400, D=4, seed=SEED + 3, gamma=gamma, design=design)
        for B in [1, 8, 64]:
            for rep in [True, False]:
                # Eq. (12) is an identity for ANY Sigma with mean 0, so test both
                # the actual stationary covariance (when it exists) and random PD ones.
                cands = []
                try:
                    Sst = stationary_cov_proxy(nm, 0.1 * np.eye(4), B, replace=rep)
                    if np.linalg.eigvalsh(Sst)[0] > 0:
                        cands.append(("stationary", Sst))
                except Exception:
                    pass
                for r in range(2):
                    Z = rng2.standard_normal((4, 4))
                    cands.append(
                        (
                            "random%d" % r,
                            (Z @ Z.T) * 10.0 ** (-2 - r) + 1e-4 * np.eye(4),
                        )
                    )
                for tag, S in cands:
                    C_ex, m1, m2e = exact_finite_check(nm, S, B, rep)
                    C_f = nm.C(S, B, rep)
                    t2.append(
                        dict(
                            design=design,
                            gamma=gamma,
                            B=B,
                            replace=rep,
                            sigma=tag,
                            rel_err=relf(C_ex, C_f),
                            atom_mean_norm=m1,
                            atom_second_moment_err=m2e,
                        )
                    )
worst = max(r["rel_err"] for r in t2)
print(
    f"  {len(t2)} exact configurations (design x prior x B x with/without "
    f"replacement); worst relative error = {worst:.3e}"
)
res["T2_exact_finite_domain"] = dict(
    n_configs=len(t2), worst_rel_err=float(worst), detail=t2
)

# ---------------------------------------------------------------- T3 -------
print("=== T3  exhaustive enumeration of every minibatch ===")
rng = np.random.default_rng(5)
t3 = []
for N, B in [(6, 2), (8, 3), (9, 4), (10, 5)]:
    D = 3
    X = rng.standard_normal((N, D))
    yv = (rng.random(N) < 0.5).astype(float)
    Gam = 3.0 * np.eye(D)
    mm = Logistic(X, yv, Gam)
    th = mm.map_estimate()
    nmm = NoiseModel(mm.grad_n(th), mm.hess_n(th), th, Gam, N)
    Sig = sym(rng.standard_normal((D, D)))
    Sig = Sig @ Sig.T * 1e-3 + 1e-3 * np.eye(D)
    # exact expectation over psi via the finite atom construction of T2
    Lc = np.linalg.cholesky(Sig)
    atoms = np.concatenate([np.sqrt(D) * Lc.T, -np.sqrt(D) * Lc.T], 0)
    for replace in [True, False]:
        acc = np.zeros((D, D))
        for a in atoms:
            V = nmm.g + np.einsum("nde,e->nd", nmm.J, a)
            if replace:
                combos = itertools.product(range(N), repeat=B)
            else:
                combos = itertools.combinations(range(N), B)
            combos = list(combos)
            G = np.array([V[list(c)].mean(0) for c in combos])
            gm = G.mean(0)
            acc += ((G - gm).T @ (G - gm) / len(combos)) / len(atoms)
        C_f = nmm.C(Sig, B, replace)
        t3.append(
            dict(
                N=N,
                B=B,
                replace=replace,
                n_minibatches=len(combos),
                rel_err=relf(acc, C_f),
            )
        )
        print(
            f"  N={N} B={B} replace={replace}: {len(combos)} minibatches enumerated, "
            f"rel err = {relf(acc, C_f):.3e}"
        )
res["T3_exhaustive"] = t3

# ---------------------------------------------------------------- T4 -------
print("=== T4  consequence: Eq.(11)+Eq.(12) predicts the proxy covariance ===")
m, that, nm, K = make(N=500, D=3, seed=SEED, gamma=2.0)
t4 = []
A_cov = m.X.T @ m.X / m.N
for lam in [0.05, 0.15, 0.35]:
    Lam = lam * np.eye(3)
    S_exact = stationary_cov_proxy(nm, Lam, 16)
    o = run_coupled(
        m,
        that,
        Lam,
        16,
        T=8000,
        R=1200,
        rng=np.random.default_rng(SEED + int(1000 * lam)),
        burn=2000,
        nblock=10,
    )
    Sp = cov_from_blocks(o["accP"], o["sP"], o["cnt"])

    # competitor noise models, solved through the SAME Eq. (11)
    def solve_with(Cfun):
        D = nm.D
        M = np.eye(D) - Lam @ nm.H
        A = np.zeros((D * D, D * D))
        E = np.zeros((D, D))
        for i in range(D):
            for j in range(D):
                E[:] = 0
                E[i, j] = 1
                A[:, i * D + j] = (M @ E @ M.T + Lam @ Cfun(E, lin=True) @ Lam).ravel()
        b = (Lam @ Cfun(np.zeros((D, D)), lin=False) @ Lam).ravel()
        return sym(np.linalg.solve(np.eye(D * D) - A, b).reshape(D, D))

    S_const = solve_with(lambda E, lin: (np.zeros_like(E) if lin else nm.Ihat / 16))
    sig2 = float(np.var(m.y - 1 / (1 + np.exp(-m.X @ that))))

    def C_ziyin(E, lin):
        if lin:
            return (A_cov @ E @ A_cov + np.trace(A_cov @ E) * A_cov) / 16
        return sig2 * A_cov / 16

    S_ziyin = solve_with(C_ziyin)
    t4.append(
        dict(
            lam=lam,
            exact_vs_sim=relf(S_exact, Sp),
            constnoise_vs_sim=relf(S_const, Sp),
            ziyin_vs_sim=relf(S_ziyin, Sp),
        )
    )
    print(
        f"  lam={lam}: Eq.(12) {relf(S_exact, Sp):.4f} | constant-noise "
        f"{relf(S_const, Sp):.4f} | Ziyin Eq.(6) {relf(S_ziyin, Sp):.4f}   "
        f"(vs simulated proxy covariance)"
    )
res["T4_consequence"] = t4
res["runtime_s"] = time.time() - t0
with open(f"{OUT}/claim2_thm43.json", "w") as f:
    json.dump(res, f, indent=1)
print("\nwrote outputs/claim2_thm43.json  runtime %.1f s" % res["runtime_s"])