scripts/claim6_propB1.py

"""Claim 6 -- Proposition B.1: the stationary covariance of SGLD-with-momentum,
Eq. (B.3), and its recovery of the non-momentum result as kappa -> 0.

Method (fully independent, no simulation needed for the headline result).
The momentum proxy Eq. (B.2) is a linear recursion in the joint state
z=(psi-that, nu) whose driving noise has a covariance that is LINEAR in
psi psi'.  Hence E[z z'] solves an exact (2D)^2 linear fixed point, which is what
`stationary_cov_momentum` does.  Given that exact Sigma_psi we can test Eq. (B.3)
to machine precision, at any kappa, D, lambda, B and temperature beta.

  (1) kappa -> 0 recovery: Sigma_psi(kappa) -> Sigma_psi(0) and Eq. (B.3) reduces
      algebraically to Eq. (11).
  (2) SGD case (beta = infinity): residual of Eq. (B.3) at machine precision.
  (3) SGLD case (beta < infinity): the temperature term (1+kappa^2) 2 Lambda/beta
      is checked, and against the corrected term derived here.
  (4) symbolic scalar derivation (sympy) that pins down every coefficient.
  (5) direct simulation of Eq. (B.2) as a cross-check of the exact solver.
"""
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 (stationary_cov_proxy, stationary_cov_momentum,
                    eqB3_residual, eq11_residual, relf)

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

# ---------------------------------------------------------------- (4) ------
print("=== (4) symbolic scalar derivation ===")
import sympy as sp
lam, h, k, c, q, Suu, Sum_, Smm = sp.symbols(
    "lam h kappa c q S_uu S_um S_mm", positive=True)
a, b = 1 - lam * h, -k * lam
sol = sp.solve([
    sp.Eq(Suu, a**2 * Suu + 2 * a * b * Sum_ + b**2 * Smm + lam**2 * c + q),
    sp.Eq(Sum_, a * h * Suu + (a * k + b * h) * Sum_ + b * k * Smm - lam * c),
    sp.Eq(Smm, h**2 * Suu + 2 * h * k * Sum_ + k**2 * Smm + c),
], [Suu, Sum_, Smm], dict=True)[0]
S_sym = sp.simplify(sol[Suu])
# Eq. (B.3) in the scalar (commuting) case, exactly as printed in the paper
lhs = (1 - k) * (2 * lam * h * S_sym) + k / (1 - k**2) * (2 * lam**2 * h**2 * S_sym)
rhs = lam**2 * c + (1 + k**2) / (1 - k**2) * lam**2 * h**2 * S_sym + (1 + k**2) * q
gap = sp.simplify(lhs - rhs)
gap_noq = sp.simplify(gap.subs(q, 0))
print("  Eq.(B.3) residual, scalar case, beta = infinity (q=0):", gap_noq)
Cq = sp.simplify(sp.solve(sp.Eq(sp.simplify(
    (1 - k) * (2 * lam * h * S_sym) + k / (1 - k**2) * (2 * lam**2 * h**2 * S_sym)
    - lam**2 * c - (1 + k**2) / (1 - k**2) * lam**2 * h**2 * S_sym
    - sp.Symbol("Cq") * q), 0), sp.Symbol("Cq"))[0])
print("  correct temperature coefficient (paper prints 1+kappa^2):", sp.factor(Cq))
print("  correct coefficient rewritten:", sp.simplify(Cq - ((1 - k)**2 + 2 * k * lam * h / (1 + k))))
res["symbolic"] = {
    "B3_residual_beta_inf": str(gap_noq),
    "paper_temperature_coeff": "1 + kappa^2",
    "derived_temperature_coeff": str(sp.factor(Cq)),
    "derived_minus_(1-k)^2_minus_2*k*lam*h/(1+k)": str(
        sp.simplify(Cq - ((1 - k) ** 2 + 2 * k * lam * h / (1 + k)))),
}

# ---------------------------------------------------------------- (1,2,3) --
print("\n=== (1,2,3) numerical residuals of Eq. (B.3) ===")
rows = []
for D in [1, 2, 3, 5]:
    for design in ["gaussian", "heavy"]:
        m, that, nm, K = make(N=400, D=D, seed=SEED + D, gamma=2.0, design=design)
        for lamv in [0.05, 0.15, 0.4]:
            Lam = lamv * np.eye(D)
            for B in [8, 64]:
                for kap in [0.0, 1e-6, 0.05, 0.2, 0.5, 0.8, 0.95]:
                    row = dict(D=D, design=design, lam=lamv, B=B, kappa=kap)
                    Si, _, cr = stationary_cov_momentum(nm, Lam, B, kap, beta=np.inf)
                    row["B3_resid_SGD"] = eqB3_residual(nm, Lam, Si, B, kap, beta=np.inf)
                    row["cross_moment_asymmetry"] = float(
                        np.linalg.norm(cr - cr.T) / max(np.linalg.norm(cr), 1e-300))
                    Sb, _, _ = stationary_cov_momentum(nm, Lam, B, kap, beta=400.0)
                    row["B3_resid_SGLD_paper"] = eqB3_residual(
                        nm, Lam, Sb, B, kap, beta=400.0, temp="paper")
                    row["B3_resid_SGLD_corrected"] = eqB3_residual(
                        nm, Lam, Sb, B, kap, beta=400.0, temp="corrected")
                    if kap == 0.0:
                        S0 = stationary_cov_proxy(nm, Lam, B, beta=np.inf)
                        row["kappa0_vs_nonmomentum"] = relf(Si, S0)
                        row["kappa0_eq11_resid"] = eq11_residual(nm, Lam, Si, B)
                    rows.append(row)
res["grid"] = rows
w_sgd = max(r["B3_resid_SGD"] for r in rows)
w_k0 = max(r["kappa0_vs_nonmomentum"] for r in rows if "kappa0_vs_nonmomentum" in r)
paper_pos = [r["B3_resid_SGLD_paper"] for r in rows if r["kappa"] > 0]
corr_pos = [r["B3_resid_SGLD_corrected"] for r in rows if r["kappa"] > 0]
print(f"  {len(rows)} configurations (D x design x lambda x B x kappa)")
print(f"  worst Eq.(B.3) residual, SGD (beta=inf), all kappa : {w_sgd:.3e}")
print(f"  worst ||Sigma_psi(kappa=0) - Sigma_psi(non-momentum)||/||.|| : {w_k0:.3e}")
print(f"  SGLD (beta=400), kappa>0: paper temperature term residual "
      f"median {np.median(paper_pos):.3e}, max {max(paper_pos):.3e}")
print(f"  SGLD (beta=400), kappa>0: corrected term residual "
      f"median {np.median(corr_pos):.3e}, max {max(corr_pos):.3e}")
res["summary"] = dict(n_configs=len(rows), worst_B3_resid_SGD=float(w_sgd),
                      worst_kappa0_recovery=float(w_k0),
                      SGLD_paper_resid_median=float(np.median(paper_pos)),
                      SGLD_paper_resid_max=float(max(paper_pos)),
                      SGLD_corrected_resid_median=float(np.median(corr_pos)),
                      SGLD_corrected_resid_max=float(max(corr_pos)))

# kappa -> 0 continuity
print("\n=== kappa -> 0 continuity ===")
m, that, nm, K = make(N=400, D=3, seed=SEED, gamma=2.0)
Lam = 0.15 * np.eye(3)
S0 = stationary_cov_proxy(nm, Lam, 16)
cont = []
for kap in [0.0, 1e-8, 1e-6, 1e-4, 1e-2, 0.1, 0.3, 0.6]:
    Sk, _, _ = stationary_cov_momentum(nm, Lam, 16, kap)
    cont.append(dict(kappa=kap, rel_diff_to_kappa0=relf(Sk, S0),
                     B3_resid=eqB3_residual(nm, Lam, Sk, 16, kap)))
    print(f"  kappa={kap:<8} ||Sigma(k)-Sigma(0)||/||Sigma(0)||={relf(Sk, S0):.3e} "
          f"B.3 resid={eqB3_residual(nm, Lam, Sk, 16, kap):.2e}")
res["kappa_continuity"] = cont

# ---------------------------------------------------------------- (5) ------
print("\n=== (5) direct simulation of Eq. (B.2) ===")


def simulate(nm, model, that, Lam, B, kap, T, R, rng, burn, beta=np.inf):
    D = nm.D
    U = np.zeros((R, D)); M_ = np.zeros((R, D))
    gh, Jn, Gam_ = nm.g, nm.J, model.Gamma / model.N
    Lch = np.linalg.cholesky(2.0 / beta * Lam) if np.isfinite(beta) else None
    acc = np.zeros((D, D)); cnt = 0
    for t in range(T):
        IDX = rng.integers(0, model.N, size=(R, B))
        G = (gh[IDX].mean(1) + np.einsum("rbde,re->rbd", Jn[IDX], U).mean(1)
             + U @ Gam_.T + that @ Gam_.T)
        M_ = kap * M_ + G
        U = U - M_ @ Lam.T
        if Lch is not None:
            U = U + rng.standard_normal((R, D)) @ Lch.T
        if t >= burn:
            acc += U.T @ U; cnt += R
    return acc / cnt


sim = []
for kap in [0.0, 0.3, 0.6]:
    for beta in [np.inf, 400.0]:
        Sx, _, _ = stationary_cov_momentum(nm, Lam, 16, kap, beta=beta)
        Ss = simulate(nm, m, that, Lam, 16, kap, T=8000, R=3000,
                      rng=np.random.default_rng(SEED + 77), burn=2500, beta=beta)
        from common import solve_B3_for_Sigma
        Sb3 = solve_B3_for_Sigma(nm, Lam, 16, kap, beta=beta)
        sim.append(dict(kappa=kap, beta=("inf" if not np.isfinite(beta) else beta),
                        sim_vs_exact=relf(Ss, Sx), B3solution_vs_sim=relf(Sb3, Ss),
                        B3solution_vs_exact=relf(Sb3, Sx)))
        print(f"  kappa={kap} beta={beta}: sim vs exact solver {relf(Ss, Sx):.4f}; "
              f"Sigma solving Eq.(B.3) vs simulation {relf(Sb3, Ss):.4f}")
res["simulation"] = sim
res["runtime_s"] = time.time() - t0
with open(f"{OUT}/claim6_propB1.json", "w") as f:
    json.dump(res, f, indent=1)
print("\nwrote outputs/claim6_propB1.json  runtime %.1f s" % res["runtime_s"])