repro-bundle-v2/repro/run_experiments_v2.py

#!/usr/bin/env python3
"""Scaled empirical reproduction (v2) for ICML 2026 paper LyhBIrNBXv / arXiv 2602.19241:
"Scaling Laws for Precision in High-Dimensional Linear Regression".

Unlike the v1 script (which evaluated the paper's *analytic* risk formula and fit it
back to itself), this version runs genuine Monte-Carlo experiments on the paper's data
model and compares *measured* excess risk against the theoretical scaling laws.

Data model (Gaussian sequence / high-dimensional linear regression):
    x_i ~ N(0, lambda_i),  lambda_i = i^{-a}  (power-law spectrum, Assumption 3.5)
    theta_i = 1  (so the approximation error of a size-M model is
                  sum_{i>M} lambda_i = zeta(a, M+1) ~ M^{-(a-1)}/(a-1))
    y = sum_i theta_i x_i + sigma * z
The ambient dimension is infinite: features beyond the largest fitted model size
enter y exactly through their (Gaussian) aggregate, so no truncation bias exists.

Estimator: OLS on the first M features of the (possibly quantized) design matrix.
Population excess risk is computed in closed form given the estimate:
    R(theta_hat) = sum_{i<=M} lambda_i (theta_hat_i - 1)^2 + zeta(a, M+1)

Quantizers applied to the stored training design matrix X:
    multiplicative (floating-point-like): X_q = X * (1 + eps*G),  G ~ N(0,1) iid
    additive (integer-like, b bits):      X_q = delta_b * round(clip(X, +-A)/delta_b),
                                          A = 4 (covers +-4 sd of the largest feature),
                                          delta_b = 2A / 2^b, error variance ~ delta_b^2/12

Experiments
    E1  risk vs model size M (M in [8,1024], N=16384)  -> alpha = -(a-1); claims 1,2,3,4
    E2  risk vs sample size N (N in [512,32768], compute-optimal M=N^{1/a})
        -> beta = -(a-1)/a, N_eff shrinkage under multiplicative quant; claims 1,3,4
    E3  crossover: risk vs N under additive b-bit quantization -> precision-limited
        floor vs sample-limited regime; floor scaling ~2^{-2b(a-1)/a}; claim 2
    E4  spectrum of quantized covariance -> tail flattening at delta_b^2/12; claim 2
    E5  trace amplification E[tr(Xq'Xq)]/E[tr(X'X)] = 1 + eps^2; claim 5

Hardware: CPU only, thread-capped to 4 cores (shared box).
Run:  python repro/run_experiments_v2.py            (from PrecisionScalingLaws/)
Outputs: results/results_v2.json + results/fig_*.png
"""

import os

# Thread caps BEFORE importing numpy (4 cores of a shared 20-core CPU box).
for _v in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS",
           "NUMEXPR_NUM_THREADS", "VECLIB_MAXIMUM_THREADS"):
    os.environ[_v] = "4"

import json
import platform
import sys
import time

import numpy as np
from scipy import linalg as sla
from scipy import stats
from scipy.special import zeta

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

HERE = os.path.dirname(os.path.abspath(__file__))
RESULTS_DIR = os.path.normpath(os.path.join(HERE, "..", "results"))
os.makedirs(RESULTS_DIR, exist_ok=True)

SEED_BASE = 20260724

# Categorical palette (validated, fixed order): blue, orange, aqua, yellow, magenta, green, violet, red
PAL = ["#2a78d6", "#eb6834", "#1baf7a", "#eda100", "#e87ba4", "#008300", "#4a3aa7", "#e34948"]
plt.rcParams.update({
    "figure.dpi": 130, "savefig.dpi": 130, "font.size": 9.5,
    "axes.grid": True, "grid.alpha": 0.25, "grid.linewidth": 0.6,
    "axes.spines.top": False, "axes.spines.right": False,
    "lines.linewidth": 2.0, "legend.frameon": False,
})

# ----------------------------------------------------------------------------- helpers

def lambdas(a, M):
    return np.arange(1, M + 1, dtype=np.float64) ** (-a)


def tail_risk(a, M):
    """sum_{i>M} i^{-a} (exact Hurwitz zeta)."""
    return float(zeta(a, M + 1))


def quant_mult(X, eps, rng):
    return X * (1.0 + eps * rng.standard_normal(X.shape))


def quant_add(X, bits, A=4.0, rng=None):
    """b-bit uniform quantizer with subtractive dither (standard dithered rounding).

    Plain round-to-nearest ZEROES features whose scale is below delta/2 (dead
    columns, rank-deficient covariance) -- a truncation effect, not the paper's
    additive model. Subtractive dither makes the quantization error exactly
    uniform(-delta/2, delta/2), independent of the signal: constant error
    variance delta^2/12 per entry, i.e. precisely the additive (integer-like)
    quantization channel analyzed by Theorems 4.2/4.4.
    """
    delta = 2.0 * A / (2 ** bits)
    if rng is None:
        rng = np.random.default_rng(SEED_BASE + 424242 + bits)
    u = rng.uniform(-delta / 2.0, delta / 2.0, size=X.shape)
    Xq = np.clip(X + u, -A, A)
    return delta * np.round(Xq / delta) - u, delta


def make_design(rng, N, M, a):
    """N x M design with independent N(0, i^{-a}) columns."""
    Z = rng.standard_normal((N, M))
    Z *= np.sqrt(lambdas(a, M))[None, :]
    return Z


def make_labels(rng, X, a, sigma):
    """y = X @ 1 + tail-signal + sigma*noise (tail beyond X's columns, exact)."""
    N, M = X.shape
    tv = tail_risk(a, M)
    y = X.sum(axis=1)
    y += np.sqrt(tv) * rng.standard_normal(N)
    y += sigma * rng.standard_normal(N)
    return y


def risks_over_M_grid(Xq, y, a, M_grid):
    """Excess population risk of OLS on the first M columns of Xq, per M in M_grid."""
    Mmax = M_grid[-1]
    G = Xq[:, :Mmax].T @ Xq[:, :Mmax]
    c = Xq[:, :Mmax].T @ y
    lam = lambdas(a, Mmax)
    out = []
    for M in M_grid:
        try:
            cf = sla.cho_factor(G[:M, :M], lower=True, check_finite=False)
            th = sla.cho_solve(cf, c[:M], check_finite=False)
        except np.linalg.LinAlgError:
            # Aggressive additive quantization rounds small tail features to exactly
            # zero -> singular Gram. Min-norm LS assigns those dead columns theta=0,
            # which is the honest estimate (the feature carries no information).
            th = np.linalg.lstsq(G[:M, :M], c[:M], rcond=None)[0]
        r = float(np.sum(lam[:M] * (th - 1.0) ** 2)) + tail_risk(a, M)
        out.append(r)
    return np.asarray(out)


def fit_loglog(x, y):
    s, b, r, _, se = stats.linregress(np.log(x), np.log(y))
    return float(s), float(np.exp(b)), float(r ** 2), float(se)


def ci95(vals, axis=0):
    vals = np.asarray(vals)
    m = vals.mean(axis=axis)
    n = vals.shape[axis]
    half = 1.96 * vals.std(axis=axis, ddof=1) / np.sqrt(n) if n > 1 else np.zeros_like(m)
    return m, half

# ----------------------------------------------------------------------------- E1: risk vs M

def run_E1(seeds=12):
    """Risk vs model size M under full precision, multiplicative, and additive quant."""
    t0 = time.time()
    N, Mmax, sigma = 16384, 1024, 0.1
    M_grid = np.unique(np.round(np.geomspace(8, Mmax, 12)).astype(int))
    arms = [("fp64", None), ("mult_eps0.1", 0.1), ("mult_eps0.2", 0.2),
            ("add_b4", 4), ("add_b6", 6), ("add_b8", 8), ("add_b10", 10)]
    a_values = [1.5, 2.0]
    fit_max = {1.5: 512, 2.0: 90}   # keep OLS variance term <~3% of bias term
    res = {"N": N, "sigma": sigma, "M_grid": M_grid.tolist(), "seeds": seeds,
           "arms": {}, "a_values": a_values, "fit_max": {str(k): v for k, v in fit_max.items()}}
    risks = {(a, arm): [] for a in a_values for arm, _ in arms}
    for s in range(seeds):
        for a in a_values:
            rng = np.random.default_rng(SEED_BASE + 1000 * s + int(10 * a))
            X = make_design(rng, N, Mmax, a)
            y = make_labels(rng, X, a, sigma)
            for arm, p in arms:
                if arm == "fp64":
                    Xq = X
                elif arm.startswith("mult"):
                    Xq = quant_mult(X, p, np.random.default_rng(SEED_BASE + 7 + 1000 * s + int(10 * a)))
                else:
                    Xq, _ = quant_add(X, p)
                risks[(a, arm)].append(risks_over_M_grid(Xq, y, a, M_grid))
    for a in a_values:
        for arm, p in arms:
            mean, half = ci95(np.stack(risks[(a, arm)]))
            entry = {"risk_mean": mean.tolist(), "risk_ci95": half.tolist()}
            sel = M_grid <= fit_max[a]
            if arm.startswith(("fp", "mult")):
                # per-seed slope fits -> CI on the exponent
                slopes = [fit_loglog(M_grid[sel], r[sel])[0] for r in risks[(a, arm)]]
                sm, sh = ci95(np.asarray(slopes))
                _, amp, r2, _ = fit_loglog(M_grid[sel], mean[sel])
                entry.update({"alpha_fit": float(sm), "alpha_ci95": float(sh),
                              "alpha_r2": r2, "alpha_theory": -(a - 1.0)})
            else:
                # additive: measured saturation size (argmin of mean risk) vs theory
                M_sat = int(M_grid[int(np.argmin(mean))])
                delta = 2.0 * 4.0 / (2 ** p)
                M_eff_theory = float((delta ** 2 / 12.0) ** (-1.0 / a))
                entry.update({"M_sat_empirical": M_sat, "M_eff_theory": M_eff_theory,
                              "risk_floor": float(mean.min()), "delta": delta})
            res["arms"][f"a{a}_{arm}"] = entry
    res["runtime_s"] = round(time.time() - t0, 1)
    return res, {k: np.stack(v) for k, v in risks.items()}, M_grid


# ----------------------------------------------------------------------------- E2: risk vs N

def run_E2(seeds=12):
    """Risk vs sample size N at the compute-optimal allocation M = round(N^{1/a})."""
    t0 = time.time()
    sigma = 0.25
    N_grid = np.array([512, 1024, 2048, 4096, 8192, 16384, 32768])
    arms = [("fp64", None), ("mult_eps0.1", 0.1), ("mult_eps0.2", 0.2)]
    a_values = [1.5, 2.0]
    res = {"sigma": sigma, "N_grid": N_grid.tolist(), "seeds": seeds,
           "allocation": "M = round(N^{1/a})", "arms": {}}
    risks = {(a, arm): [] for a in a_values for arm, _ in arms}
    for s in range(seeds):
        for a in a_values:
            Mmax = int(round(N_grid[-1] ** (1.0 / a)))
            rng = np.random.default_rng(SEED_BASE + 500000 + 1000 * s + int(10 * a))
            Xfull = make_design(rng, int(N_grid[-1]), Mmax, a)
            yfull = make_labels(rng, Xfull, a, sigma)
            for arm, p in arms:
                if arm == "fp64":
                    Xq = Xfull
                else:
                    Xq = quant_mult(Xfull, p, np.random.default_rng(
                        SEED_BASE + 900000 + 1000 * s + int(10 * a) + int(100 * p)))
                row = []
                for N in N_grid:
                    M = int(round(N ** (1.0 / a)))
                    r = risks_over_M_grid(Xq[:N, :M], yfull[:N], a, np.array([M]))[0]
                    row.append(r)
                risks[(a, arm)].append(np.asarray(row))
    fit_sel = N_grid >= 2048   # drop the two smallest N (finite-size N^{-2(a-1)/a} term)
    for a in a_values:
        beta_th = -(a - 1.0) / a
        fp_mean, _ = ci95(np.stack(risks[(a, "fp64")]))
        s_fp, A_fp, _, _ = fit_loglog(N_grid[fit_sel], fp_mean[fit_sel])
        for arm, p in arms:
            R = np.stack(risks[(a, arm)])
            mean, half = ci95(R)
            slopes = [fit_loglog(N_grid[fit_sel], r[fit_sel])[0] for r in R]
            sm, sh = ci95(np.asarray(slopes))
            _, amp, r2, _ = fit_loglog(N_grid[fit_sel], mean[fit_sel])
            entry = {"risk_mean": mean.tolist(), "risk_ci95": half.tolist(),
                     "beta_fit": float(sm), "beta_ci95": float(sh), "beta_r2": r2,
                     "beta_theory": beta_th}
            if arm.startswith("mult"):
                # N_eff(N): solve R_fp(N_eff) = R_arm(N) on the fitted fp power law
                Neff_ratio = (mean[fit_sel] / A_fp) ** (1.0 / s_fp) / N_grid[fit_sel]
                entry["Neff_ratio_mean"] = float(np.mean(Neff_ratio))
                entry["Neff_ratio_per_N"] = Neff_ratio.tolist()
                entry["Neff_ratio_theory"] = float((1.0 + p ** 2) ** (-a / (a - 1.0)))
            res["arms"][f"a{a}_{arm}"] = entry
    res["runtime_s"] = round(time.time() - t0, 1)
    return res, {k: np.stack(v) for k, v in risks.items()}, N_grid, fit_sel


# ----------------------------------------------------------------------------- E3: crossover

def run_E3(seeds=10):
    """Precision-limited vs sample-limited crossover: risk vs N for b-bit additive quant."""
    t0 = time.time()
    a, sigma = 1.5, 0.25
    N_grid = np.array([512, 1024, 2048, 4096, 8192, 16384, 32768])
    bits_list = [4, 5, 6, 7]
    res = {"a": a, "sigma": sigma, "N_grid": N_grid.tolist(), "bits": bits_list,
           "seeds": seeds, "arms": {}}
    risks = {b: [] for b in bits_list}
    risks_fp = []
    for s in range(seeds):
        Mmax = int(round(N_grid[-1] ** (1.0 / a)))
        rng = np.random.default_rng(SEED_BASE + 300000 + 1000 * s)
        Xfull = make_design(rng, int(N_grid[-1]), Mmax, a)
        yfull = make_labels(rng, Xfull, a, sigma)
        arms = [("fp", None)] + [(f"b{b}", b) for b in bits_list]
        for arm, b in arms:
            Xq = Xfull if b is None else quant_add(Xfull, b)[0]
            row = []
            for N in N_grid:
                M = int(round(N ** (1.0 / a)))
                row.append(risks_over_M_grid(Xq[:N, :M], yfull[:N], a, np.array([M]))[0])
            (risks_fp if b is None else risks[b]).append(np.asarray(row))
    fp_mean, fp_half = ci95(np.stack(risks_fp))
    s_fp, A_fp, _, _ = fit_loglog(N_grid[N_grid >= 2048], fp_mean[N_grid >= 2048])
    res["fp"] = {"risk_mean": fp_mean.tolist(), "risk_ci95": fp_half.tolist(),
                 "beta_fit_pooled": s_fp}
    floors = []
    for b in bits_list:
        mean, half = ci95(np.stack(risks[b]))
        floor = float(mean[-2:].mean())          # plateau level at the largest N
        N_star = float((floor / A_fp) ** (1.0 / s_fp))   # where fp risk crosses floor
        floors.append(floor)
        res["arms"][f"b{b}"] = {"risk_mean": mean.tolist(), "risk_ci95": half.tolist(),
                                "floor": floor, "N_star": N_star}
    # floor ~ 2^{-2 b (a-1)/a}: slope of log2(floor) vs b
    sl, ic, r, _, _ = stats.linregress(bits_list, np.log2(floors))
    res["floor_vs_bits"] = {"slope_fit": float(sl), "slope_theory": -2.0 * (a - 1.0) / a,
                            "r2": float(r ** 2), "floors": floors}
    res["runtime_s"] = round(time.time() - t0, 1)
    return res, np.stack(risks_fp), {b: np.stack(risks[b]) for b in bits_list}, N_grid


# ----------------------------------------------------------------------------- E4: spectrum

def run_E4(seeds=3):
    """Eigen-spectrum of the quantized empirical covariance: additive tail flattening."""
    t0 = time.time()
    a, d, N = 1.5, 2048, 8192
    bits_list = [4, 6, 8]
    spectra = {}
    for arm in ["fp"] + [f"b{b}" for b in bits_list]:
        acc = []
        for s in range(seeds):
            rng = np.random.default_rng(SEED_BASE + 700000 + 1000 * s)
            X = make_design(rng, N, d, a)
            if arm != "fp":
                X = quant_add(X, int(arm[1:]))[0]
            ev = np.linalg.eigvalsh(X.T @ X / N)[::-1]
            acc.append(ev)
        spectra[arm] = np.mean(acc, axis=0)
    res = {"a": a, "d": d, "N": N, "bits": bits_list, "seeds": seeds,
           "floors_theory": {f"b{b}": (2 * 4.0 / 2 ** b) ** 2 / 12.0 for b in bits_list},
           "runtime_s": round(time.time() - t0, 1)}
    # measured tail floor = median of the last 200 eigenvalues
    res["floors_measured"] = {arm: float(np.median(spectra[arm][-200:]))
                              for arm in spectra if arm != "fp"}
    return res, spectra


# ----------------------------------------------------------------------------- E5: trace ratio

def run_E5(seeds=30):
    """Claim 5: E[tr(Xq'Xq)] / E[tr(X'X)] = 1 + eps^2, across eps, with CIs."""
    t0 = time.time()
    a, d, N = 1.5, 3000, 4000
    eps_list = [0.05, 0.10, 0.15, 0.20, 0.30]
    ratios = {e: [] for e in eps_list}
    for s in range(seeds):
        rng = np.random.default_rng(SEED_BASE + 800000 + s)
        X = make_design(rng, N, d, a)
        tr = float(np.sum(X * X))
        for e in eps_list:
            Xq = quant_mult(X, e, np.random.default_rng(SEED_BASE + 810000 + 100 * s + int(1000 * e)))
            ratios[e].append(float(np.sum(Xq * Xq)) / tr)
    res = {"a": a, "d": d, "N": N, "seeds": seeds, "eps": eps_list, "per_eps": {}}
    xs, ys = [], []
    for e in eps_list:
        m, h = ci95(np.asarray(ratios[e]))
        res["per_eps"][str(e)] = {"ratio_mean": float(m), "ratio_ci95": float(h),
                                  "ratio_theory": 1.0 + e ** 2}
        xs.append(e ** 2)
        ys.append(m - 1.0)
    sl, ic, r, _, _ = stats.linregress(xs, ys)
    res["excess_vs_eps2"] = {"slope_fit": float(sl), "slope_theory": 1.0,
                             "intercept": float(ic), "r2": float(r ** 2)}
    res["runtime_s"] = round(time.time() - t0, 1)
    return res, ratios


# ----------------------------------------------------------------------------- figures

def fig_E1(res, risks, M_grid, path):
    fig, axes = plt.subplots(1, 2, figsize=(9.6, 3.9), sharey=False)
    show = [("fp64", "FP64", PAL[0]), ("mult_eps0.1", "mult eps=0.1", PAL[1]),
            ("mult_eps0.2", "mult eps=0.2", PAL[2])]
    for ax, a in zip(axes, [1.5, 2.0]):
        for arm, lab, col in show:
            e = res["arms"][f"a{a}_{arm}"]
            m = np.array(e["risk_mean"]); h = np.array(e["risk_ci95"])
            ax.fill_between(M_grid, m - h, m + h, color=col, alpha=0.18, lw=0)
            ax.plot(M_grid, m, color=col,
                    label=f"{lab}: alpha={e['alpha_fit']:.3f} (R2={e['alpha_r2']:.4f})")
        fm = res["fit_max"][str(a)]
        ref = np.array(res["arms"][f"a{a}_fp64"]["risk_mean"])[0]
        xs = M_grid[M_grid <= fm].astype(float)
        ax.plot(xs, ref * (xs / xs[0]) ** (-(a - 1.0)), ls="--", color="#666666", lw=1.4,
                label=f"theory slope {-(a-1.0):.2f}")
        ax.set_xscale("log"); ax.set_yscale("log")
        ax.set_xlabel("model size M"); ax.set_title(f"a = {a}")
        ax.legend(fontsize=7.5)
    axes[0].set_ylabel("excess risk")
    fig.suptitle("E1: risk vs M — multiplicative quantization preserves the M-exponent (claims 1, 3, 4)", y=1.02)
    fig.tight_layout(); fig.savefig(path, bbox_inches="tight"); plt.close(fig)


def fig_E2(res, N_grid, path):
    fig, axes = plt.subplots(1, 2, figsize=(9.6, 3.9))
    show = [("fp64", "FP64", PAL[0]), ("mult_eps0.1", "mult eps=0.1", PAL[1]),
            ("mult_eps0.2", "mult eps=0.2", PAL[2])]
    for ax, a in zip(axes, [1.5, 2.0]):
        for arm, lab, col in show:
            e = res["arms"][f"a{a}_{arm}"]
            m = np.array(e["risk_mean"]); h = np.array(e["risk_ci95"])
            ax.fill_between(N_grid, m - h, m + h, color=col, alpha=0.18, lw=0)
            extra = f", N_eff/N={e['Neff_ratio_mean']:.2f}" if "Neff_ratio_mean" in e else ""
            ax.plot(N_grid, m, color=col,
                    label=f"{lab}: beta={e['beta_fit']:.3f} (R2={e['beta_r2']:.4f}){extra}")
        ref = np.array(res["arms"][f"a{a}_fp64"]["risk_mean"])[2]
        xs = N_grid[N_grid >= 2048].astype(float)
        ax.plot(xs, ref * (xs / xs[0]) ** (-(a - 1.0) / a), ls="--", color="#666666",
                lw=1.4, label=f"theory slope {-(a-1.0)/a:.3f}")
        ax.set_xscale("log"); ax.set_yscale("log")
        ax.set_xlabel("sample size N"); ax.set_title(f"a = {a},  M = N^(1/a)")
        ax.legend(fontsize=7.5)
    axes[0].set_ylabel("excess risk")
    fig.suptitle("E2: risk vs N at compute-optimal M — data exponent and N_eff shrinkage (claims 1, 3, 4)", y=1.02)
    fig.tight_layout(); fig.savefig(path, bbox_inches="tight"); plt.close(fig)


def fig_E3(res, N_grid, path):
    fig, axes = plt.subplots(1, 2, figsize=(9.6, 3.9))
    ax = axes[0]
    fp = res["fp"]
    ax.plot(N_grid, fp["risk_mean"], color=PAL[0], label="FP64")
    for i, b in enumerate(res["bits"]):
        e = res["arms"][f"b{b}"]
        col = PAL[(i + 1) % len(PAL)]
        m = np.array(e["risk_mean"]); h = np.array(e["risk_ci95"])
        ax.fill_between(N_grid, m - h, m + h, color=col, alpha=0.18, lw=0)
        ax.plot(N_grid, m, color=col, label=f"int{b}")
        ax.axhline(e["floor"], color=col, ls=":", lw=1.0)
        if N_grid[0] < e["N_star"] < N_grid[-1]:
            ax.axvline(e["N_star"], color=col, ls="--", lw=0.9, alpha=0.6)
    ax.set_xscale("log"); ax.set_yscale("log")
    ax.set_xlabel("sample size N"); ax.set_ylabel("excess risk")
    ax.set_title("precision-limited floors and crossover N*")
    ax.legend(fontsize=8)
    ax2 = axes[1]
    f = res["floor_vs_bits"]
    ax2.plot(res["bits"], np.log2(f["floors"]), "o-", color=PAL[0], ms=7,
             label=f"measured slope {f['slope_fit']:.3f}/bit")
    b0 = res["bits"][0]
    ax2.plot(res["bits"], np.log2(f["floors"][0]) + f["slope_theory"] * (np.array(res["bits"]) - b0),
             ls="--", color="#666666", label=f"theory {f['slope_theory']:.3f}/bit")
    ax2.set_xlabel("bit-width b"); ax2.set_ylabel("log2(risk floor)")
    ax2.set_title(f"floor vs bits (R2={f['r2']:.4f})")
    ax2.legend(fontsize=8)
    fig.suptitle("E3: sample-limited vs precision-limited crossover under additive quantization (claim 2)", y=1.02)
    fig.tight_layout(); fig.savefig(path, bbox_inches="tight"); plt.close(fig)


def fig_E4(res, spectra, res_e1, M_grid, path):
    fig, axes = plt.subplots(1, 2, figsize=(9.6, 3.9))
    ax = axes[0]
    idx = np.arange(1, len(spectra["fp"]) + 1)
    ax.plot(idx, spectra["fp"], color=PAL[0], label="FP64")
    for i, b in enumerate(res["bits"]):
        col = PAL[(i + 1) % len(PAL)]
        ax.plot(idx, spectra[f"b{b}"], color=col, label=f"int{b}")
        ax.axhline(res["floors_theory"][f"b{b}"], color=col, ls=":", lw=1.0)
    ax.set_xscale("log"); ax.set_yscale("log")
    ax.set_xlabel("eigenvalue rank"); ax.set_ylabel("eigenvalue")
    ax.set_title("covariance spectrum: additive noise floors the tail")
    ax.legend(fontsize=8)
    ax2 = axes[1]
    a = 1.5
    for i, arm in enumerate(["fp64", "add_b4", "add_b6", "add_b8", "add_b10"]):
        e = res_e1["arms"][f"a{a}_{arm}"]
        col = PAL[i % len(PAL)]
        lab = "FP64" if arm == "fp64" else f"int{arm[5:]}"
        if "M_sat_empirical" in e:
            lab += f" (M_sat={e['M_sat_empirical']})"
            ax2.axvline(e["M_eff_theory"], color=col, ls="--", lw=0.9, alpha=0.6)
        m = np.array(e["risk_mean"]); h = np.array(e["risk_ci95"])
        ax2.fill_between(M_grid, m - h, m + h, color=col, alpha=0.18, lw=0)
        ax2.plot(M_grid, m, color=col, label=lab)
    ax2.set_xscale("log"); ax2.set_yscale("log")
    ax2.set_xlabel("model size M"); ax2.set_ylabel("excess risk")
    ax2.set_title("risk vs M saturates at M_eff (dashed = theory)")
    ax2.legend(fontsize=7.5)
    fig.suptitle("E4: additive quantization flattens the spectral tail and caps M_eff (claim 2)", y=1.02)
    fig.tight_layout(); fig.savefig(path, bbox_inches="tight"); plt.close(fig)


def fig_E5(res, path):
    fig, ax = plt.subplots(figsize=(5.4, 3.9))
    eps2 = np.array([float(e) ** 2 for e in res["eps"]])
    means = np.array([res["per_eps"][str(e)]["ratio_mean"] for e in res["eps"]]) - 1.0
    halves = np.array([res["per_eps"][str(e)]["ratio_ci95"] for e in res["eps"]])
    ax.errorbar(eps2, means, yerr=halves, fmt="o", color=PAL[0], ms=7, capsize=3,
                label="measured (30 seeds, 95% CI)")
    f = res["excess_vs_eps2"]
    xs = np.linspace(0, eps2.max() * 1.05, 50)
    ax.plot(xs, xs, ls="--", color="#666666", label="theory: ratio - 1 = eps^2")
    ax.plot(xs, f["intercept"] + f["slope_fit"] * xs, color=PAL[1], lw=1.6,
            label=f"fit slope {f['slope_fit']:.4f} (R2={f['r2']:.6f})")
    ax.set_xlabel("eps^2"); ax.set_ylabel("trace ratio - 1")
    ax.set_title("E5: trace amplification tr(Xq'Xq)/tr(X'X) = 1 + eps^2 (claim 5)")
    ax.legend(fontsize=8)
    fig.tight_layout(); fig.savefig(path, bbox_inches="tight"); plt.close(fig)


# ----------------------------------------------------------------------------- main

def main():
    t0 = time.time()
    print("Reproduction v2: Scaling Laws for Precision in High-Dimensional Linear Regression")
    print(f"host={platform.node()}  cpu-only  threads=4  numpy={np.__version__}")

    print("\n[E1] risk vs M (N=16384, M in [8,1024], 7 arms x 2 spectra x 12 seeds)...")
    e1, r1, M_grid = run_E1()
    print(f"     done in {e1['runtime_s']}s")
    print("\n[E2] risk vs N (N in [512,32768], M=N^(1/a), 3 arms x 2 spectra x 12 seeds)...")
    e2, r2, N_grid, _ = run_E2()
    print(f"     done in {e2['runtime_s']}s")
    print("\n[E3] crossover under int-b quantization (b in {4,5,6,7}, 10 seeds)...")
    e3, r3fp, r3, N_grid3 = run_E3()
    print(f"     done in {e3['runtime_s']}s")
    print("\n[E4] covariance spectra (d=2048, N=8192, 3 seeds)...")
    e4, spectra = run_E4()
    print(f"     done in {e4['runtime_s']}s")
    print("\n[E5] trace amplification (d=3000, N=4000, 30 seeds, 5 eps values)...")
    e5, _ = run_E5()
    print(f"     done in {e5['runtime_s']}s")

    print("\nRendering figures...")
    fig_E1(e1, r1, M_grid, os.path.join(RESULTS_DIR, "fig_e1_risk_vs_M.png"))
    fig_E2(e2, N_grid, os.path.join(RESULTS_DIR, "fig_e2_risk_vs_N.png"))
    fig_E3(e3, N_grid3, os.path.join(RESULTS_DIR, "fig_e3_crossover.png"))
    fig_E4(e4, spectra, e1, M_grid, os.path.join(RESULTS_DIR, "fig_e4_spectrum.png"))
    fig_E5(e5, os.path.join(RESULTS_DIR, "fig_e5_trace_ratio.png"))

    summary = {
        "paper": {"openreview": "LyhBIrNBXv", "arxiv": "2602.19241",
                  "title": "Scaling Laws for Precision in High-Dimensional Linear Regression"},
        "hardware": {"host": platform.node(), "platform": platform.platform(),
                     "cpu_only": True, "threads": 4,
                     "numpy": np.__version__},
        "seed_base": SEED_BASE,
        "command": "python repro/run_experiments_v2.py",
        "E1_risk_vs_M": e1, "E2_risk_vs_N": e2, "E3_crossover": e3,
        "E4_spectrum": e4, "E5_trace": e5,
        "total_runtime_s": round(time.time() - t0, 1),
    }
    out = os.path.join(RESULTS_DIR, "results_v2.json")
    with open(out, "w", encoding="utf-8", newline="\n") as f:
        json.dump(summary, f, indent=2)
    print(f"\nTotal runtime: {summary['total_runtime_s']}s")
    print(f"Wrote {out} and 5 figures to {RESULTS_DIR}")
    return 0


if __name__ == "__main__":
    sys.exit(main())