repro-bundle/claim3_spectral.py

"""
Claim 3 — Each self-training iteration is a direction-dependent spectral filter:
  * signal along a strong spike survives with factor  (s/(s+tau))^(t+1) ~ 1,
  * the isotropic/noise subspace is suppressed at rate  (1+tau)^(-t),
  * multi-spike: strong directions preserved, weak ones forgotten (soft selection).

We measure the empirical signal-alignment coefficient  a_j(t) = <bhat_t, u_j>/r_j
and compare it to the predicted survival factor  kappa_j^(t+1), kappa_j=s_j/(s_j+tau).
"""
import sys, numpy as np, pandas as pd
sys.path.insert(0, "src")
from linear_selftrain import build_spiked_covariance, iterative_self_train, spiked_theory, SpikedCov
from plotting import new_fig, add_theory_sim, save, PALETTE
import plotly.graph_objects as go

OUT = "outputs/claim3"
p, trials, T = 1200, 12, 12
rho = 2.0
n = int(round(p / rho))
tau = rho - 1
rows = []

# ---------- (1) survival factor of the signal for varying spike strength ------
figS = new_fig("Claim 3: signal survival <bhat_t,u1>/r vs (s/(s+tau))^(t+1), rho=2",
               "iteration t", "signal survival coefficient")
for i, s in enumerate([2.0, 5.0, 25.0, 100.0]):
    _, U = build_spiked_covariance(p, [s]); u1 = U[:, 0]
    cov = SpikedCov(p, [s], U)
    beta = u1.copy()  # r=1
    align = np.zeros(T + 1)
    for tr in range(trials):
        rng = np.random.default_rng(400 + 10 * i + tr)
        _, betas = iterative_self_train(None, beta, n, 1.0, T, rng=rng, return_betas=True, cov=cov)
        align += np.array([b @ u1 for b in betas])
    align /= trials
    surv = (s / (s + tau)) ** (np.arange(T + 1) + 1)
    add_theory_sim(figS, np.arange(T + 1), surv, align, color=PALETTE[i], name=f"s={s:g}")
    print(f"survival s={s:g}: t=1 emp={align[1]:.3f} thy={surv[1]:.3f} | "
          f"t={T} emp={align[T]:.3f} thy={surv[T]:.3f} (kappa={s/(s+tau):.3f})")
    for t in range(T + 1):
        rows.append(dict(part="survival", s=s, t=t, emp_align=align[t], theory_survival=surv[t]))
save(figS, pd.DataFrame(rows), OUT, "claim3_survival_factor")

# ---------- (2) noise-only suppression at rate (1+tau)^-t ----------------------
figN = new_fig("Claim 3: signal-free noise suppression V_t = (sigma^2/tau)(1+tau)^-t",
               "iteration t", "stochastic error (log scale)")
figN.update_yaxes(type="log")
rowsN = []
for i, rho2 in enumerate([1.5, 2.0, 2.5]):
    n2 = int(round(p / rho2)); tau2 = rho2 - 1
    s = 5.0
    _, U = build_spiked_covariance(p, [s]); cov = SpikedCov(p, [s], U)
    beta0 = np.zeros(p)  # signal-free: r=0
    acc = np.zeros(T + 1)
    for tr in range(trials):
        rng = np.random.default_rng(500 + 10 * i + tr)
        r = iterative_self_train(None, beta0, n2, 1.0, T, rng=rng, cov=cov)
        acc += r
    acc /= trials
    theoryV = (1.0 / tau2) * (1.0 + tau2) ** (-np.arange(T + 1))  # sigma^2=1
    add_theory_sim(figN, np.arange(T + 1), theoryV, acc, color=PALETTE[i], name=f"ρ={rho2}")
    ratio = np.mean(acc[1:6] / acc[0:5])
    print(f"noise-free rho={rho2}: emp decay ratio~{ratio:.3f}  1/(1+tau)={1/(1+tau2):.3f}")
    for t in range(T + 1):
        rowsN.append(dict(part="noise_decay", rho=rho2, t=t, emp=acc[t], theory=theoryV[t]))
save(figN, pd.DataFrame(rowsN), OUT, "claim3_noise_suppression")

# ---------- (3) multi-spike soft feature selection ----------------------------
spikes = [80.0, 8.0, 1.5]           # strong, medium, weak
_, U = build_spiked_covariance(p, spikes)
cov = SpikedCov(p, spikes, U)
beta = U[:, 0] + U[:, 1] + U[:, 2]  # equal signal power r_j=1 in each direction
aligns = np.zeros((3, T + 1))
for tr in range(trials):
    rng = np.random.default_rng(600 + tr)
    _, betas = iterative_self_train(None, beta, n, 1.0, T, rng=rng, return_betas=True, cov=cov)
    for j in range(3):
        aligns[j] += np.array([b @ U[:, j] for b in betas])
aligns /= trials
figM = new_fig("Claim 3: multi-spike soft feature selection (rho=2)",
               "iteration t", "signal survival per direction")
rowsM = []
for j, s in enumerate(spikes):
    surv = (s / (s + tau)) ** (np.arange(T + 1) + 1)
    add_theory_sim(figM, np.arange(T + 1), surv, aligns[j], color=PALETTE[j],
                   name=f"s_{j+1}={s:g} (κ={s/(s+tau):.2f})")
    print(f"multispike s={s:g}: survival t={T} emp={aligns[j][T]:.3f} thy={surv[T]:.3f}")
    for t in range(T + 1):
        rowsM.append(dict(part="multispike", spike=s, t=t, emp=aligns[j][t], theory=surv[t]))
save(figM, pd.DataFrame(rowsM), OUT, "claim3_multispike_selection")
print("\nClaim 3: strong spikes survive (~1), weak spikes forgotten; noise ~ (1+tau)^-t.")