bundle_v1/scripts/claim1_theorem31.py

"""Claim 1 — Independent numerical check of Theorem 3.1 (tracking-error dynamics).

Verifies, for a linear critic w under gradient flow  w' = -2*Sigma*w + 2*b_t  with
drifting target statistics b_t = Sigma @ w*(t):

  (A) The theorem identity (paper Eq. 1 / Eq. 15, fixed Sigma):
        Gamma' = -4 e^T Sigma e - 2 e^T Sigma^{-1} b_t'
      matches the finite-difference derivative of Gamma = ||e||^2 along an
      RK4-integrated trajectory, for isotropic and ill-conditioned Sigma.

  (B) The full decomposition with time-varying Sigma(t) (paper App. B Eq. 26):
        Gamma' = -4 e^T Sigma e - 2 e^T Sigma^{-1} b' + 2 e^T Sigma^{-1} Sigma' Sigma^{-1} b
      also matches finite differences (checks the dropped third term is correct).

  (C) Qualitative regimes of Fig. 1: with zero drift the error contracts
      monotonically at rate >= 4*lambda_min; with drift, poorly conditioned
      Sigma yields weaker worst-case contraction.

Deterministic (fixed seed), 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

RNG = np.random.default_rng(0)
D = 8
OUT = os.path.join(os.path.dirname(__file__), "..", "outputs", "claim1")
os.makedirs(OUT, exist_ok=True)


def make_sigma(kappa, d=D, rng=RNG):
    """PD covariance with condition number kappa and trace d (fixed variance budget)."""
    if kappa == 1:
        eig = np.ones(d)
    else:
        eig = np.geomspace(1.0, kappa, d)
    eig = eig * (d / eig.sum())
    q, _ = np.linalg.qr(rng.normal(size=(d, d)))
    return (q * eig) @ q.T


def wstar(t, u, v, omega=1.5):
    return np.cos(omega * t) * u + np.sin(omega * t) * v


def dwstar(t, u, v, omega=1.5):
    return omega * (-np.sin(omega * t) * u + np.cos(omega * t) * v)


def rk4(f, y, t, dt):
    k1 = f(t, y)
    k2 = f(t + dt / 2, y + dt / 2 * k1)
    k3 = f(t + dt / 2, y + dt / 2 * k2)
    k4 = f(t + dt, y + dt * k3)
    return y + dt / 6 * (k1 + 2 * k2 + 2 * k3 + k4)


def check_A(kappa):
    """Fixed-Sigma theorem identity vs finite differences."""
    sigma = make_sigma(kappa)
    sigma_inv = np.linalg.inv(sigma)
    q, _ = np.linalg.qr(RNG.normal(size=(D, 2)))
    u, v = q[:, 0], q[:, 1]

    def f(t, w):
        return -2 * sigma @ w + 2 * sigma @ wstar(t, u, v)

    dt, T = 1e-3, 6.0
    n = int(T / dt)
    w = np.zeros(D)
    ts = np.arange(n) * dt
    gam = np.empty(n)
    analytic = np.empty(n)
    for i, t in enumerate(ts):
        e = w - wstar(t, u, v)
        gam[i] = e @ e
        bdot = sigma @ dwstar(t, u, v)
        analytic[i] = -4 * e @ sigma @ e - 2 * e @ sigma_inv @ bdot
        w = rk4(f, w, t, dt)
    fd = np.gradient(gam, dt)
    # compare away from the first steps (transient) using a robust scale
    sl = slice(5, n - 5)
    scale = np.maximum(np.abs(analytic[sl]), np.abs(fd[sl]).max() * 1e-3)
    rel = np.abs(fd[sl] - analytic[sl]) / scale
    return ts, gam, analytic, fd, float(np.median(rel)), float(np.percentile(rel, 99.9))


def check_B():
    """Time-varying Sigma(t): verify full Eq. 26 including the third term."""
    eig = np.geomspace(1.0, 20.0, D)
    eig = eig * (D / eig.sum())
    q0, _ = np.linalg.qr(RNG.normal(size=(D, D)))
    # rotate the eigenbasis slowly in the plane of the first two eigenvectors
    def sigma_t(t, w_rot=0.7):
        c, s = np.cos(w_rot * t), np.sin(w_rot * t)
        rot = np.eye(D)
        rot[0, 0], rot[0, 1], rot[1, 0], rot[1, 1] = c, -s, s, c
        qt = q0 @ rot
        # np.errstate: Apple Accelerate BLAS emits spurious overflow/divide
        # warnings on this finite matmul; outputs are verified finite below.
        with np.errstate(all="ignore"):
            out = (qt * eig) @ qt.T
        assert np.all(np.isfinite(out))
        return out

    def dsigma_t(t, h=1e-6):
        return (sigma_t(t + h) - sigma_t(t - h)) / (2 * h)

    qw, _ = np.linalg.qr(RNG.normal(size=(D, 2)))
    u, v = qw[:, 0], qw[:, 1]

    def b_t(t):
        return sigma_t(t) @ wstar(t, u, v)

    def db_t(t, h=1e-6):
        return (b_t(t + h) - b_t(t - h)) / (2 * h)

    def f(t, w):
        return -2 * sigma_t(t) @ w + 2 * b_t(t)

    dt, T = 5e-4, 4.0
    n = int(T / dt)
    w = RNG.normal(size=D) * 0.5
    ts = np.arange(n) * dt
    gam = np.empty(n)
    analytic = np.empty(n)
    for i, t in enumerate(ts):
        sig = sigma_t(t)
        sig_inv = np.linalg.inv(sig)
        e = w - sig_inv @ b_t(t)
        gam[i] = e @ e
        term1 = -4 * e @ sig @ e
        term2 = -2 * e @ sig_inv @ db_t(t)
        term3 = 2 * e @ sig_inv @ dsigma_t(t) @ sig_inv @ b_t(t)
        analytic[i] = term1 + term2 + term3
        w = rk4(f, w, t, dt)
    fd = np.gradient(gam, dt)
    sl = slice(5, n - 5)
    scale = np.maximum(np.abs(analytic[sl]), np.abs(fd[sl]).max() * 1e-3)
    rel = np.abs(fd[sl] - analytic[sl]) / scale
    return float(np.median(rel)), float(np.percentile(rel, 99.9))


def check_C():
    """Zero drift: Gamma decays at least as fast as exp(-4*lambda_min*t)."""
    results = {}
    for kappa in (1, 10, 100):
        sigma = make_sigma(kappa)
        lmin = np.linalg.eigvalsh(sigma).min()
        wfix = RNG.normal(size=D)
        b = sigma @ wfix

        def f(t, w):
            return -2 * sigma @ w + 2 * b

        dt, T = 1e-3, 2.0
        n = int(T / dt)
        w = np.zeros(D)
        gam = np.empty(n)
        for i in range(n):
            e = w - wfix
            gam[i] = e @ e
            w = rk4(f, w, i * dt, dt)
        ts = np.arange(n) * dt
        bound = gam[0] * np.exp(-4 * lmin * ts)
        monotone = bool(np.all(np.diff(gam) <= 1e-12))
        within = bool(np.all(gam <= bound * (1 + 1e-6)))
        results[kappa] = dict(lambda_min=float(lmin), monotone=monotone,
                              bounded_by_exp=within, final_gamma=float(gam[-1]))
    return results


def main():
    apply_style()
    rows, ok = [], True

    fig, axes = plt.subplots(1, 3, figsize=(12, 3.4))
    for ax, kappa, color in zip(axes, (1, 10, 100), SERIES):
        ts, gam, analytic, fd, med, p999 = check_A(kappa)
        passed = p999 < 1e-2
        ok &= passed
        rows.append(dict(check="A_fixed_sigma", kappa=kappa,
                         median_rel_err=med, p999_rel_err=p999,
                         passed=passed))
        print(f"[A] kappa={kappa:>3}: median rel err {med:.2e}, "
              f"p99.9 {p999:.2e} -> {'PASS' if passed else 'FAIL'}")
        step = 40
        ax.plot(ts[5:-5], fd[5:-5], color=SERIES[1], linewidth=3.5, alpha=0.35,
                label="finite-difference dΓ/dt")
        ax.plot(ts[5:-5:step], analytic[5:-5:step], color=color, linestyle="none",
                marker="o", markersize=3, label="Theorem 3.1 RHS")
        ax.set_title(f"κ(Σ) = {kappa}")
        ax.set_xlabel("t")
        if kappa == 1:
            ax.set_ylabel("dΓ/dt")
        ax.legend(fontsize=8)
    fig.suptitle("Claim 1(A): analytic tracking-error identity vs simulation (d=8)", y=1.04)
    fig.tight_layout()
    fig.savefig(os.path.join(OUT, "claim1_identity.png"), bbox_inches="tight")

    medB, p999B = check_B()
    passedB = p999B < 1e-2
    ok &= passedB
    rows.append(dict(check="B_timevarying_sigma_eq26", kappa="20(rotating)",
                     median_rel_err=medB, p999_rel_err=p999B, passed=passedB))
    print(f"[B] full Eq.26 (time-varying Sigma): median rel err {medB:.2e}, "
          f"p99.9 {p999B:.2e} -> {'PASS' if passedB else 'FAIL'}")

    resC = check_C()
    for kappa, r in resC.items():
        ok &= r["monotone"] and r["bounded_by_exp"]
        rows.append(dict(check="C_zero_drift_contraction", kappa=kappa,
                         median_rel_err="", p999_rel_err="",
                         passed=r["monotone"] and r["bounded_by_exp"],
                         **{k: v for k, v in r.items()}))
        print(f"[C] kappa={kappa:>3}: monotone={r['monotone']}, "
              f"Gamma <= Gamma0*exp(-4*lmin*t)={r['bounded_by_exp']} "
              f"(lambda_min={r['lambda_min']:.3f})")

    keys = sorted({k for row in rows for k in row})
    with open(os.path.join(OUT, "claim1_results.csv"), "w", newline="") as f:
        wtr = csv.DictWriter(f, fieldnames=keys)
        wtr.writeheader()
        wtr.writerows(rows)

    print("CLAIM 1 NUMERICAL CHECK:", "PASS" if ok else "FAIL")
    sys.exit(0 if ok else 1)


if __name__ == "__main__":
    main()