scripts/b3_probe.py

"""Probe Eq. (B.3): is the residual an OCR artefact, a symmetry-assumption
artefact, or a genuine error?  Also validate the exact momentum solver against
direct simulation."""

import numpy as np
import sys

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

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()
nm = NoiseModel(m.grad_n(that), m.hess_n(that), that, Gam, N)
B, lam, kap = 16, 0.1, 0.3
Lam = lam * np.eye(D)
H = nm.H
S, P, cross = stationary_cov_momentum(nm, Lam, B, kap)
C = nm.C(S, B)
print("exact Sigma_psi(kappa=%.2f)=\n" % kap, S)
print(
    "lag-1 cross moment asymmetry ||X-X^T||/||X|| =",
    np.linalg.norm(cross - cross.T) / np.linalg.norm(cross),
)

# --- (1) systematic scan over plausible readings of Eq. (B.3) -------------
A1 = Lam @ H @ S + S @ H @ Lam
A2 = Lam @ H @ Lam @ H @ S + S @ H @ Lam @ H @ Lam
R1 = Lam @ C @ Lam
R2 = Lam @ H @ S @ H @ Lam
k = kap
cands_a = {
    "k/(1-k)": k / (1 - k),
    "k^2/(1-k)": k**2 / (1 - k),
    "k/(1-k)^2": k / (1 - k) ** 2,
    "k^2/(1-k)^2": k**2 / (1 - k) ** 2,
    "k": k,
    "k^2": k**2,
    "2k/(1-k)": 2 * k / (1 - k),
}
cands_b = {
    "(1+k)/(1-k)": (1 + k) / (1 - k),
    "(1+k^2)/(1-k)": (1 + k**2) / (1 - k),
    "(1+k)/(1-k)^2": (1 + k) / (1 - k) ** 2,
    "(1+k^2)/(1-k)^2": (1 + k**2) / (1 - k) ** 2,
    "1": 1.0,
    "1+k": 1 + k,
    "1+k^2": 1 + k**2,
    "(1-k)": 1 - k,
}
best = []
for na, a in cands_a.items():
    for nb, b in cands_b.items():
        res = (1 - k) * A1 + a * A2 - R1 - b * R2
        best.append((np.linalg.norm(res) / np.linalg.norm(A1), na, nb))
best.sort()
print("\nbest 6 coefficient readings of Eq. (B.3) (relative residual):")
for r, na, nb in best[:6]:
    print(f"   {r:.4e}   a={na:<12} b={nb}")

# --- (2) is the symmetry assumption in the proof the culprit? -------------
# proof step: P := E[u_{t-1} u_{t-2}^T] solves P = (I-LH)S + kS - k P^T,
# and the paper "solves" it as if P were symmetric: P_sym = ((I-LH)S + kS)/(1+k)
Psym = ((np.eye(D) - Lam @ H) @ S + k * S) / (1 + k)
print("\nexact  E[u_t u_{t-1}^T] =\n", cross)
print("paper-implied symmetric solution =\n", Psym)
print("rel diff exact-vs-symmetric solution:", relf(cross, Psym))
# does B.3 become exact when the exact (asymmetric) cross moment is used?
Mmat = np.eye(D) - Lam @ H
# reconstruct: Sigma = M S M^T + k^2 Lam Mm Lam + Lam C Lam - (Dm + Dm^T) + 2Lam/beta
Mm = 2 * S - cross - cross.T  # Lam E[m m^T] Lam  (beta = inf)
Dm = k * Mmat @ (cross.T - S)  # D = k(I-LH)(E[u_{t-1}u_{t-2}^T] - S)
lhs = Mmat @ S @ Mmat.T + k**2 * Mm + Lam @ C @ Lam - (Dm + Dm.T)
print("\nidentity check with EXACT cross moment: rel resid =", relf(lhs, S))


# --- (3) validate the exact solver by direct simulation -------------------
def simulate_momentum(nm, model, that, Lam, B, kap, T, R, rng, burn):
    D = nm.D
    U = np.zeros((R, D))
    M_ = np.zeros((R, D))
    gh = nm.g
    Jn = nm.J
    acc = np.zeros((D, D))
    cnt = 0
    Gam_ = model.Gamma / model.N
    for t in range(T):
        IDX = rng.integers(0, model.N, size=(R, B))
        gb = gh[IDX].mean(1)
        Ju = np.einsum("rbde,re->rbd", Jn[IDX], U).mean(1)
        G = gb + Ju + U @ Gam_.T + that @ Gam_.T
        M_ = kap * M_ + G
        U = U - M_ @ Lam.T
        if t >= burn:
            acc += U.T @ U
            cnt += R
    return acc / cnt


Ssim = simulate_momentum(
    nm, m, that, Lam, B, kap, T=6000, R=3000, rng=np.random.default_rng(7), burn=2000
)
print("\nsimulated Sigma_psi vs exact solver: rel err =", relf(Ssim, S))
Sb3 = None
try:
    from common import solve_B3_for_Sigma

    Sb3 = solve_B3_for_Sigma(nm, Lam, B, kap)
    print("Sigma solving Eq.(B.3) vs simulation:  rel err =", relf(Sb3, Ssim))
    print("Sigma solving Eq.(B.3) vs exact:       rel err =", relf(Sb3, S))
except Exception as e:
    print("B.3 solve failed", e)