bundle_v1/scripts/claim2_isotropy.py

"""Claim 2 — Isotropy equalizes contraction and reduces drift amplification.

  (A) Bound checks (paper App. B.1.4, Eqs. 27-33), over random PD covariances
      with a fixed variance budget Tr(Sigma) = d:
        - lambda_min(Sigma) <= d/d = 1, equality iff isotropic
          (worst-case contraction is maximized by isotropy);
        - worst-case drift amplification for fixed ||b'||:
          max_{||e||=||b'||=1} |e^T Sigma^{-1} b'| = 1/lambda_min,
          minimized iff isotropic;
        - the paper's Eq. 33 bound |e^T Sigma^{-1} b'| <= kappa(Sigma) with
          b' = Sigma w*' and ||w*'|| <= 1 holds everywhere (and we also record
          the parameterization-sharp bound |e^T w*'| <= 1 to document that
          Eq. 33 is valid but loose in that exact fixed-Sigma parameterization).

  (B) App. F synthetic drifting-target experiment: online SGD linear regression
      tracking w* that switches at steps 100 and 200 (3 random unit-sphere
      targets), features ~ N(0, Sigma) with kappa in {1, 10, 100}, Tr(Sigma)=d,
      d in {2, 8, 32}, 50 seeds. Metrics: divergence steps (||e_{t+1}||>||e_t||),
      mean recovery time after each switch, final error.

Deterministic (seeded), CPU-only. Exits non-zero if any check fails.
"""
import csv
import json
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", "claim2")
os.makedirs(OUT, exist_ok=True)


def make_sigma(kappa, d, rng, rotate=True):
    eig = np.ones(d) if kappa == 1 else np.geomspace(1.0, kappa, d)
    eig = eig * (d / eig.sum())
    if rotate and d > 1:
        q, _ = np.linalg.qr(rng.normal(size=(d, d)))
        return (q * eig) @ q.T, np.sort(eig)
    return np.diag(eig), np.sort(eig)


def check_A(d=8, n_random=2000, seed=1):
    rng = np.random.default_rng(seed)
    ok = True
    lmins = []
    for _ in range(n_random):
        a = rng.normal(size=(d, 3 * d))
        sig = a @ a.T
        sig *= d / np.trace(sig)
        lmins.append(np.linalg.eigvalsh(sig).min())
    lmins = np.asarray(lmins)
    ok &= bool(np.all(lmins <= 1.0 + 1e-9))
    iso_lmin = np.linalg.eigvalsh(np.eye(d)).min()
    ok &= abs(iso_lmin - 1.0) < 1e-12
    print(f"[A1] lambda_min <= trace/d for {n_random} random Sigma: "
          f"max observed {lmins.max():.4f} (bound 1.0) -> {'PASS' if ok else 'FAIL'}; "
          f"isotropic attains {iso_lmin:.4f}")

    # Drift amplification for fixed ||b'||: sup |e^T Sigma^{-1} b'| = 1/lambda_min
    amp_ok = True
    for kappa in (1, 10, 100):
        sig, eig = make_sigma(kappa, d, rng)
        sig_inv = np.linalg.inv(sig)
        best = 0.0
        for _ in range(20000):
            e = rng.normal(size=d); e /= np.linalg.norm(e)
            bd = rng.normal(size=d); bd /= np.linalg.norm(bd)
            best = max(best, abs(e @ sig_inv @ bd))
        sup = 1.0 / eig[0]
        amp_ok &= best <= sup * (1 + 1e-9)
        # the sup is attained along the min-eigenvector direction
        w = np.linalg.eigh(sig)[1][:, 0]
        attained = abs(w @ sig_inv @ w)
        amp_ok &= abs(attained - sup) / sup < 1e-9
        print(f"[A2] kappa={kappa:>3}: empirical max |e^T Sig^-1 b'| = {best:.3f}, "
              f"analytic sup 1/lambda_min = {sup:.3f} (attained {attained:.3f})")
    ok &= amp_ok

    # Paper Eq. 33: with b' = Sigma w*', ||w*'||<=1: |e^T Sigma^{-1} b'| <= kappa
    eq33_ok = True
    loose = []
    for kappa in (1, 10, 100):
        sig, eig = make_sigma(kappa, d, rng)
        sig_inv = np.linalg.inv(sig)
        worst_val = 0.0
        for _ in range(20000):
            e = rng.normal(size=d); e /= np.linalg.norm(e)
            ws = rng.normal(size=d); ws /= np.linalg.norm(ws)
            val = abs(e @ sig_inv @ (sig @ ws))
            worst_val = max(worst_val, val)
        eq33_ok &= worst_val <= kappa * (1 + 1e-9)
        loose.append((kappa, worst_val))
        print(f"[A3] Eq.33 kappa={kappa:>3}: max |e^T Sig^-1 (Sig w*')| = {worst_val:.4f} "
              f"<= kappa = {kappa} (exact value of the term is |e^T w*'| <= 1: bound holds, "
              f"loose by {kappa / max(worst_val, 1e-12):.1f}x)")
    ok &= eq33_ok
    return ok, lmins, loose


def track_sgd(kappa, d, seed, steps=300, switch=(100, 200), lr=None):
    """Online SGD on squared loss, one sample per step, tracking switching w*.

    All condition numbers share the same variance budget Tr(Sigma) = d, so the
    fair comparison uses one common learning rate set by that budget:
    lr = 0.5 / d = 0.5 / E[||phi||^2] (single-sample LMS overshoots only when
    lr * ||phi||^2 > 2, i.e. rarely). The paper (App. F) does not state its lr.
    """
    rng = np.random.default_rng(seed)
    sig, eig = make_sigma(kappa, d, rng, rotate=(d > 2))
    chol = np.linalg.cholesky(sig)
    if lr is None:
        lr = 0.5 / d
    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])
    for t in range(steps):
        if phase < len(switch) and t == switch[phase]:
            phase += 1
        wstar = targets[phase]
        phi = chol @ rng.normal(size=d)
        w = w - lr * (phi @ w - phi @ wstar) * phi
        err[t + 1] = np.linalg.norm(w - wstar)
    return err


def recovery_time(err, start, end, thresh=0.1):
    seg = err[start:end]
    idx = np.where(seg < thresh)[0]
    return int(idx[0]) if len(idx) else end - start


def check_B():
    dims = (2, 8, 32)
    kappas = (1, 10, 100)
    n_seeds = 50
    switch = (100, 200)
    steps = 300
    stats = {}
    for d in dims:
        for kappa in kappas:
            div_counts, recs, finals = [], [], []
            for s in range(n_seeds):
                err = track_sgd(kappa, d, seed=1000 * d + s)
                # count real transient increases, not machine-precision jitter
                inc = np.sum((err[1:] > err[:-1] * (1 + 1e-9)) & (err[:-1] > 1e-8))
                div_counts.append(int(inc))
                r = [recovery_time(err, st, st + 100) for st in (0,) + switch]
                recs.append(float(np.mean(r)))
                finals.append(float(err[-1]))
            stats[(d, kappa)] = dict(
                divergence_steps_mean=float(np.mean(div_counts)),
                divergence_steps_std=float(np.std(div_counts)),
                recovery_steps_mean=float(np.mean(recs)),
                recovery_steps_std=float(np.std(recs)),
                final_err_mean=float(np.mean(finals)),
            )
    ok = True
    for d in dims:
        m1 = stats[(d, 1)]; m100 = stats[(d, 100)]
        ordered = (m1["divergence_steps_mean"] <= m100["divergence_steps_mean"]
                   and m1["recovery_steps_mean"] <= m100["recovery_steps_mean"])
        ok &= ordered
        print(f"[B] d={d:>2}: divergence steps k1={m1['divergence_steps_mean']:.1f} "
              f"k10={stats[(d,10)]['divergence_steps_mean']:.1f} "
              f"k100={m100['divergence_steps_mean']:.1f} | recovery k1="
              f"{m1['recovery_steps_mean']:.1f} k100={m100['recovery_steps_mean']:.1f} "
              f"-> {'PASS' if ordered else 'FAIL'}")
    return ok, stats


def main():
    apply_style()
    okA, lmins, loose = check_A()
    okB, stats = check_B()

    # Figure: example trajectories (d=2, 3 seeds like paper Figs. 18-20) + summary
    fig, axes = plt.subplots(2, 3, figsize=(12, 6))
    for col, seed in enumerate((0, 1, 2)):
        ax = axes[0][col]
        for kappa, color in zip((1, 10, 100), (SERIES[0], SERIES[3], SERIES[5])):
            err = track_sgd(kappa, 2, seed=2000 + seed)
            ax.plot(err, color=color, label=f"κ={kappa}", linewidth=1.6)
        for st in (100, 200):
            ax.axvline(st, color="#999", linestyle=":", linewidth=1)
        ax.set_title(f"d=2, seed {seed}")
        ax.set_xlabel("SGD step")
        if col == 0:
            ax.set_ylabel("‖w−w*‖₂")
            ax.legend(fontsize=8)
    for col, d in enumerate((2, 8, 32)):
        ax = axes[1][col]
        kap = (1, 10, 100)
        xs = np.arange(len(kap))
        dv = [stats[(d, k)]["divergence_steps_mean"] for k in kap]
        dvs = [stats[(d, k)]["divergence_steps_std"] for k in kap]
        rc = [stats[(d, k)]["recovery_steps_mean"] for k in kap]
        ax.bar(xs - 0.18, dv, 0.32, yerr=dvs, color=SERIES[0], label="divergence steps")
        ax.bar(xs + 0.18, rc, 0.32, color=SERIES[5], label="recovery steps")
        ax.set_xticks(xs, [f"κ={k}" for k in kap])
        ax.set_title(f"d={d} (50 seeds)")
        if col == 0:
            ax.set_ylabel("steps")
            ax.legend(fontsize=8)
    fig.suptitle("Claim 2: isotropic Gaussian features track drifting targets most stably", y=1.02)
    fig.tight_layout()
    fig.savefig(os.path.join(OUT, "claim2_tracking.png"), bbox_inches="tight")

    with open(os.path.join(OUT, "claim2_tracking_stats.csv"), "w", newline="") as f:
        wtr = csv.writer(f)
        wtr.writerow(["d", "kappa", "divergence_steps_mean", "divergence_steps_std",
                      "recovery_steps_mean", "recovery_steps_std", "final_err_mean"])
        for (d, k), v in sorted(stats.items()):
            wtr.writerow([d, k, v["divergence_steps_mean"], v["divergence_steps_std"],
                          v["recovery_steps_mean"], v["recovery_steps_std"],
                          v["final_err_mean"]])
    with open(os.path.join(OUT, "claim2_bounds.json"), "w") as f:
        json.dump({"eq33_worst_values": loose,
                   "lambda_min_max_over_random_sigma": float(np.max(lmins))}, f, indent=2)

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


if __name__ == "__main__":
    main()