scripts/smoke.py

import numpy as np
import sys

sys.path.insert(0, "/home/ubuntu/samuel/sgmcmc-uq-repro/scripts")
from common import (
    Logistic,
    NoiseModel,
    stationary_cov_proxy,
    eq11_residual,
    stationary_cov_momentum,
    eqB3_residual,
    relf,
    sgd_paths,
)

rng = np.random.default_rng(0)
N, D = 400, 3
X = rng.standard_normal((N, D))
th_true = rng.standard_normal(D)
y = (rng.random(N) < 1 / (1 + np.exp(-X @ th_true))).astype(float)
Gam = 2.0 * np.eye(D)
m = Logistic(X, y, Gam)
that = m.map_estimate()
print("grad at MAP", np.linalg.norm(m.gradL(that)))
nm = NoiseModel(m.grad_n(that), m.hess_n(that), that, Gam, N)
print("H eig", np.linalg.eigvalsh(nm.H))

B, lam = 16, 0.1
Lam = lam * np.eye(D)
S = stationary_cov_proxy(nm, Lam, B)
print("Sigma_psi\n", S)
print("Eq(11) residual  ", eq11_residual(nm, Lam, S, B))
Cf = nm.C(S, B)
Cb = nm.C_bruteforce(S, B, n_mc=400000, rng=np.random.default_rng(1))
print("Thm4.3 rel err vs brute force", relf(Cb, Cf))

Sm, P, cross = stationary_cov_momentum(nm, Lam, B, 0.0)
print("momentum kappa=0 vs non-momentum", relf(Sm, S))
for k in [0.0, 1e-6, 0.05, 0.2, 0.5]:
    Sk, _, _ = stationary_cov_momentum(nm, Lam, B, k)
    print(
        f"  kappa={k:<8} B.3 residual={eqB3_residual(nm, Lam, Sk, B, k):.3e}  "
        f"||Sk-S0||/||S0||={relf(Sk, S):.4f}"
    )

# simulate the proxy and compare
out = sgd_paths(
    m,
    that,
    Lam,
    B,
    T=4000,
    R=2000,
    rng=np.random.default_rng(2),
    proxy_nm=nm,
    burn=1500,
    record_every=10,
)
psi = out["psi"]
Semp = np.cov(psi - that, rowvar=False, bias=True) + np.outer(
    (psi - that).mean(0), (psi - that).mean(0)
)
print("proxy sim vs exact Sigma_psi", relf(Semp, S))
thet = out["theta"]
Sth = np.cov(thet, rowvar=False, bias=True)
print("true SGD Sigma_theta rel diff from Sigma_psi", relf(Sth, S))
print("coupling E||th-psi||^2", out["d2"].mean())