bundle_v1/scripts/claim3_gaussianity.py

"""Claim 3 — Gaussianity controls higher-order moments / tail behavior.

  (A) Isserlis' theorem: for zero-mean Gaussian, E[x_i x_j x_k x_l] =
      S_ij S_kl + S_ik S_jl + S_il S_jk. Verified by Monte Carlo; a Laplace
      distribution with the *identical covariance* violates it (heavier 4th
      moments), i.e. covariance does not pin down higher-order moments once
      Gaussianity is dropped.

  (B) Stein residual (paper App. B.1.5, Eqs. 45-50): E[phi f(phi)] =
      Sigma E[grad f] + E[f(phi) r(phi)] with r == 0 iff Gaussian. We check
      ||E[phi f(phi)] - Sigma E[grad f]|| ~ 0 for Gaussian and > 0 for
      Laplace / Student-t at the same covariance, for two test functions.

  (C) App. F tail experiment: online SGD tracking with isotropic Gaussian vs
      isotropic Laplace features at identical covariance; Laplace should show
      more frequent transient error increases (spikes) and larger worst-case
      single-step updates.

Deterministic (seeded), CPU-only. Exits non-zero if any check fails.
"""
import csv
import os
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(__file__))
from plot_style import SERIES, apply_style

import matplotlib.pyplot as plt

OUT = os.path.join(os.path.dirname(__file__), "..", "outputs", "claim3")
os.makedirs(OUT, exist_ok=True)


def sample_features(dist, n, chol, rng):
    d = chol.shape[0]
    if dist == "gaussian":
        z = rng.normal(size=(n, d))
    elif dist == "laplace":
        z = rng.laplace(size=(n, d)) / np.sqrt(2)  # unit variance components
    elif dist == "student5":
        z = rng.standard_t(df=5, size=(n, d)) / np.sqrt(5 / 3)  # unit variance
    else:
        raise ValueError(dist)
    return z @ chol.T


def check_A(d=4, n=4_000_000, seed=3):
    rng = np.random.default_rng(seed)
    a = rng.normal(size=(d, d))
    sig = a @ a.T + d * np.eye(d)
    sig *= d / np.trace(sig)
    chol = np.linalg.cholesky(sig)
    idxs = [(0, 0, 0, 0), (0, 1, 0, 1), (0, 1, 2, 3), (1, 1, 2, 2), (0, 0, 0, 1)]
    ok = True
    rows = []
    for dist in ("gaussian", "laplace"):
        x = sample_features(dist, n, chol, rng)
        emp_cov = x.T @ x / n
        cov_err = np.abs(emp_cov - sig).max()
        for (i, j, k, l) in idxs:
            emp = float(np.mean(x[:, i] * x[:, j] * x[:, k] * x[:, l]))
            isr = sig[i, j] * sig[k, l] + sig[i, k] * sig[j, l] + sig[i, l] * sig[j, k]
            rel = abs(emp - isr) / max(abs(isr), 1e-9)
            rows.append(dict(dist=dist, idx=str((i, j, k, l)), empirical=emp,
                             isserlis=float(isr), rel_dev=rel))
        devs = [r["rel_dev"] for r in rows if r["dist"] == dist]
        if dist == "gaussian":
            ok &= max(devs) < 0.02
            print(f"[A] gaussian: max |cov err| {cov_err:.4f}; max rel dev from "
                  f"Isserlis {max(devs):.4f} -> {'PASS' if max(devs) < 0.02 else 'FAIL'}")
        else:
            ok &= max(devs) > 0.10
            print(f"[A] laplace (same covariance, max |cov err| {cov_err:.4f}): max rel dev "
                  f"{max(devs):.4f} (expected >> 0, Isserlis violated) -> "
                  f"{'PASS' if max(devs) > 0.10 else 'FAIL'}")
    return ok, rows


def check_B(d=6, n=4_000_000, seed=4):
    rng = np.random.default_rng(seed)
    a = rng.normal(size=(d, d))
    sig = a @ a.T + d * np.eye(d)
    sig *= d / np.trace(sig)
    chol = np.linalg.cholesky(sig)
    v = rng.normal(size=d); v /= np.linalg.norm(v)

    fns = {
        "tanh(v.phi)": (lambda x: np.tanh(x @ v),
                        lambda x: (1 - np.tanh(x @ v) ** 2)[:, None] * v[None, :]),
        "(v.phi)^3": (lambda x: (x @ v) ** 3,
                      lambda x: (3 * (x @ v) ** 2)[:, None] * v[None, :]),
    }
    ok = True
    rows = []
    resid = {}
    for dist in ("gaussian", "laplace", "student5"):
        x = sample_features(dist, n, chol, rng)
        for name, (f, gradf) in fns.items():
            lhs = (x * f(x)[:, None]).mean(axis=0)
            rhs = sig @ gradf(x).mean(axis=0)
            # scale-free residual
            rnorm = float(np.linalg.norm(lhs - rhs) / max(np.linalg.norm(rhs), 1e-9))
            rows.append(dict(dist=dist, f=name, stein_residual_rel=rnorm))
            resid[(dist, name)] = rnorm
            print(f"[B] {dist:>9} f={name:<11} rel Stein residual = {rnorm:.4f}")
    for name in fns:
        g = resid[("gaussian", name)]
        ok &= g < 0.02
        ok &= resid[("laplace", name)] > 3 * max(g, 1e-4)
        ok &= resid[("student5", name)] > 3 * max(g, 1e-4)
    print(f"[B] gaussian residual ~ 0; laplace/student-t residual > 0 -> "
          f"{'PASS' if ok else 'FAIL'}")
    return ok, rows


def track_sgd_dist(dist, d, seed, steps=300, switch=(100, 200), lr=None):
    rng = np.random.default_rng(seed)
    chol = np.eye(d)  # isotropic, unit variance: identical covariance for both dists
    if lr is None:
        lr = 0.5 / d  # same budget-scaled lr convention as the Claim 2 experiment
    targets = []
    for _ in range(len(switch) + 1):
        t = rng.normal(size=d); t /= np.linalg.norm(t)
        targets.append(t)
    w = np.zeros(d)
    phase = 0
    err = np.empty(steps + 1)
    err[0] = np.linalg.norm(w - targets[0])
    max_step = 0.0
    for t in range(steps):
        if phase < len(switch) and t == switch[phase]:
            phase += 1
        wstar = targets[phase]
        phi = sample_features(dist, 1, chol, rng)[0]
        delta = -lr * (phi @ w - phi @ wstar) * phi
        max_step = max(max_step, float(np.linalg.norm(delta)))
        w = w + delta
        err[t + 1] = np.linalg.norm(w - wstar)
    return err, max_step


def check_C(d=8, n_seeds=100):
    stats = {}
    for dist in ("gaussian", "laplace"):
        spikes, max_steps, finals = [], [], []
        for s in range(n_seeds):
            err, mx = track_sgd_dist(dist, d, seed=5000 + s)
            spikes.append(int(np.sum(np.diff(err) > 1e-12)))
            max_steps.append(mx)
            finals.append(float(err[-1]))
        stats[dist] = dict(spikes_mean=float(np.mean(spikes)),
                           spikes_std=float(np.std(spikes)),
                           max_update_mean=float(np.mean(max_steps)),
                           max_update_p99=float(np.percentile(max_steps, 99)),
                           final_err_mean=float(np.mean(finals)))
    g, l = stats["gaussian"], stats["laplace"]
    ok = l["spikes_mean"] > g["spikes_mean"] and l["max_update_p99"] > g["max_update_p99"]
    print(f"[C] spikes/300 steps: gaussian {g['spikes_mean']:.1f}±{g['spikes_std']:.1f} "
          f"vs laplace {l['spikes_mean']:.1f}±{l['spikes_std']:.1f}; "
          f"p99 max |update|: {g['max_update_p99']:.2f} vs {l['max_update_p99']:.2f} "
          f"-> {'PASS' if ok else 'FAIL'}")
    return ok, stats


def main():
    apply_style()
    okA, rowsA = check_A()
    okB, rowsB = check_B()
    okC, stats = check_C()

    fig, axes = plt.subplots(1, 3, figsize=(12, 3.4))
    # (1) Isserlis deviation bar
    ax = axes[0]
    dists = ("gaussian", "laplace")
    labels = [r["idx"] for r in rowsA if r["dist"] == "gaussian"]
    xs = np.arange(len(labels))
    for off, dist, color in zip((-0.18, 0.18), dists, (SERIES[0], SERIES[5])):
        devs = [r["rel_dev"] for r in rowsA if r["dist"] == dist]
        ax.bar(xs + off, devs, 0.32, color=color, label=dist)
    ax.set_xticks(xs, labels, rotation=30, fontsize=7)
    ax.set_ylabel("rel. deviation from Isserlis")
    ax.set_title("4th moments vs covariance prediction")
    ax.legend(fontsize=8)
    # (2) Stein residual
    ax = axes[1]
    fs = sorted({r["f"] for r in rowsB})
    xs = np.arange(len(fs))
    for off, dist, color in zip((-0.25, 0.0, 0.25), ("gaussian", "laplace", "student5"),
                                (SERIES[0], SERIES[5], SERIES[3])):
        vals = [next(r["stein_residual_rel"] for r in rowsB
                     if r["dist"] == dist and r["f"] == f) for f in fs]
        ax.bar(xs + off, vals, 0.22, color=color, label=dist)
    ax.set_xticks(xs, fs, fontsize=8)
    ax.set_ylabel("relative Stein residual")
    ax.set_title("Stein residual (0 iff Gaussian)")
    ax.legend(fontsize=8)
    # (3) example trajectories
    ax = axes[2]
    for dist, color in zip(("gaussian", "laplace"), (SERIES[0], SERIES[5])):
        err, _ = track_sgd_dist(dist, 8, seed=5003)
        ax.plot(err, color=color, label=f"isotropic {dist}", linewidth=1.6)
    for st in (100, 200):
        ax.axvline(st, color="#999", linestyle=":", linewidth=1)
    ax.set_xlabel("SGD step")
    ax.set_ylabel("‖w−w*‖₂")
    ax.set_title("tracking, identical covariance (d=8)")
    ax.legend(fontsize=8)
    fig.suptitle("Claim 3: Gaussian moments are covariance-determined; heavy tails destabilize tracking", y=1.04)
    fig.tight_layout()
    fig.savefig(os.path.join(OUT, "claim3_gaussianity.png"), bbox_inches="tight")

    with open(os.path.join(OUT, "claim3_results.csv"), "w", newline="") as f:
        wtr = csv.writer(f)
        wtr.writerow(["section", "dist", "key", "value"])
        for r in rowsA:
            wtr.writerow(["A_isserlis", r["dist"], r["idx"], r["rel_dev"]])
        for r in rowsB:
            wtr.writerow(["B_stein", r["dist"], r["f"], r["stein_residual_rel"]])
        for dist, v in stats.items():
            for k, val in v.items():
                wtr.writerow(["C_tracking", dist, k, val])

    ok = okA and okB and okC
    print("CLAIM 3 NUMERICAL CHECK:", "PASS" if ok else "FAIL")
    sys.exit(0 if ok else 1)


if __name__ == "__main__":
    main()