repro-bundle/v1/scripts/claim3_fixed_point.py

"""Claim 3: Fixed-point analysis quantifies the asymptotic deviation from the
unquantized linear model.

Reproduces Figure 6 of Ichikawa et al. (arXiv:2510.10693) and verifies the
small-learning-rate closed-form of Theorem V.8:

    eps_g* = eps_g^(0) + sigma2_psi * Delta^2 * p * (1-p) + o(eta)   (interior p)
    eps_g* = eps_g^(0) + o(1/sqrt(log(1/eta)))                       (p in {0,1})
    eps_g* = rho + sigma2 - 2*kappa_psi*omega + sigma2_psi*omega^2 + o(eta)  (|c|>=omega)

where
    c       = kappa_x * rho / (sigma2_x + lambda)
    Delta   = 2*omega / L    (L = 2^b - 2)
    p       = (c - v_{i*}) / Delta   in (0,1),  v_{i*} <= c <= v_{i*+1}
    eps_g0  = rho + sigma2 - 2*kappa_x*c + sigma2_x*c^2

Also reproduces Figure 5 (input-only quantization stability boundary and
steady-state eps_g* vs omega) and the closed-form Proposition VI.1.
"""
import json
import math
import sys
import time
from pathlib import Path

import numpy as np
import matplotlib

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

sys.path.insert(0, str(Path(__file__).parent))
from ste_repro import (
    Quantizer, eps_g, macro_m_psi, macro_q_psi, macro_r_psi,
    ode_rhs, input_only_fixed_point, input_only_stability_bound,
    small_eta_fixed_point_prediction,
)
from ste_sim_batch import run_ste_batch, Setting
from scipy.integrate import solve_ivp


OUT = Path(__file__).resolve().parents[1] / "outputs" / "claim3"
OUT.mkdir(parents=True, exist_ok=True)


def run_figure6_small_eta():
    """Figure 6: eps_g* vs omega for weight quantizer b in {2,3} and input
    quantizer b_x in {2,3,4} (plus unquantized input).

    Per paper: STE simulation d=100, eta=1e-4, lambda=0, sigma2=0, tau=8e6.
    The reported curves are eps_g* as omega varies.

    The ODE / theoretical prediction is from Theorem V.8 (small-eta limit).
    """
    print("=" * 60)
    print("[Claim 3 / Figure 6] eps_g* vs omega (small-eta)")
    print("=" * 60)
    lam = 0.0
    rho = 1.0
    sigma2 = 0.0
    # Note: paper uses eta = 1e-4 and tau up to 8e6 (so n_steps = 8e6 * d = 8e8 for d=100).
    # This is too long for our wall-clock budget. We instead:
    #   - Verify the THEORETICAL small-eta prediction against numerical ODE
    #     fixed points (which can be computed at any eta, including eta=1e-4)
    #     by root-finding the ODE right-hand side = 0.
    #   - Run STE at a less aggressive eta=1e-2 and tau=1e5 (still small enough
    #     for the asymptotic prediction to hold) for d=200, n_seeds=3.

    # ---- Theoretical predictions (closed form, Theorem V.8) ----
    omegas = np.linspace(0.1, 2.0, 39)
    # weight quantizer b in {2,3}; input quantizer b_x in {2,3,4} + unquantized
    weight_bs = [2, 3]
    input_configs = [("unquant", None), ("b_x=2", 2), ("b_x=3", 3), ("b_x=4", 4)]
    omega_x = 1.0  # input range (paper Figure 6)

    theory = {}
    for b_w in weight_bs:
        for name, b_x in input_configs:
            qx = Quantizer(b=b_x, omega=omega_x) if b_x is not None else None
            preds = []
            for omega in omegas:
                qw = Quantizer(b=b_w, omega=omega)
                p = small_eta_fixed_point_prediction(qw, qx, lam=lam, rho=rho, sigma2=sigma2)
                preds.append(p)
            theory[f"b_w={b_w},{name}"] = preds

    # ---- ODE fixed-point via root finding ----
    # We solve dm/dtau = 0, dq/dtau = 0 by Newton iteration starting from a
    # reasonable initial guess. For weight-input quant the rhs uses m_psi(m,s)
    # and r_psi(m,s); solving in (m,s) coordinates is easier because the first
    # stationarity condition reduces to m_psi(m,s) = c = kappa_x rho/(sigma2_x+lam).
    from scipy.optimize import brentq

    def find_fp_ode(eta, qw, qx, lam, rho, sigma2, tau_max_ode=2000.0):
        """Find the fixed point by integrating ODE to long tau."""
        kappa_x = 1.0 if qx is None else qx.kappa_psi()
        sigma2_x = 1.0 if qx is None else qx.sigma2_psi()
        m0, q0 = 0.0, 1.0
        sol = solve_ivp(
            ode_rhs, [0.0, tau_max_ode], [m0, q0],
            args=(qw, kappa_x, sigma2_x, eta, lam, rho, sigma2),
            t_eval=np.linspace(0, tau_max_ode, 200),
            rtol=1e-9, atol=1e-12, method="DOP853",
        )
        m_fp = float(sol.y[0, -1])
        q_fp = float(sol.y[1, -1])
        s = math.sqrt(max(q_fp - m_fp ** 2 / rho, 1e-12))
        eps_fp = eps_g(m_fp, s, qw, kappa_x, sigma2_x, rho, sigma2)
        return m_fp, q_fp, eps_fp

    # Compute ODE fixed points at eta = 1e-3 (small but tractable)
    eta_small = 1e-3
    ode_fps = {}
    for b_w in weight_bs:
        for name, b_x in input_configs:
            qx = Quantizer(b=b_x, omega=omega_x) if b_x is not None else None
            fps = []
            for omega in omegas:
                qw = Quantizer(b=b_w, omega=omega)
                m_fp, q_fp, eps_fp = find_fp_ode(eta_small, qw, qx, lam, rho, sigma2)
                fps.append({"omega": float(omega), "m_fp": m_fp, "q_fp": q_fp, "eps_fp": eps_fp})
            ode_fps[f"b_w={b_w},{name}"] = fps
            print(f"  ODE fps for {b_w},{name}: 1 sample: {fps[len(fps)//2]}")

    # ---- STE simulation: at a less aggressive eta, for one curve, to verify ----
    # We pick b_w=2, b_x=unquant (matches paper's left panel of Figure 6)
    # and run STE at d=200, eta=5e-3, n_steps = 2e6 (tau_max = 10000).
    print("\n  STE simulation for one curve (b_w=2, unquant input):")
    d = 200
    n_seeds = 3
    eta_ste = 5e-3
    n_steps_ste = 2_000_000
    log_period_ste = 2000
    # sweep omega values
    omegas_ste = [0.25, 0.5, 0.75, 1.0, 1.25, 1.5]
    settings = []
    for omega in omegas_ste:
        qw = Quantizer(b=2, omega=omega)
        s = Setting(
            name=f"omega={omega}",
            qw_levels=qw.levels.tolist(), qw_theta=qw.theta.tolist(),
            qw_omega=qw.omega, qw_Delta=qw.Delta,
            qx_levels=None,
            eta=eta_ste, lam=lam, sigma2_x=1.0, kappa_x=1.0,
            d=d, rho=rho, sigma2=sigma2,
            w_init_std=1.0, w_star_value=1.0, n_seeds=n_seeds,
        )
        settings.append(s)
    print(f"  n_steps={n_steps_ste}, log_period={log_period_ste}, B={n_seeds * len(settings)}")
    t0 = time.time()
    results, wall = run_ste_batch(settings, n_steps=n_steps_ste, log_period=log_period_ste, device="cuda")
    print(f"  STE wall: {wall:.1f}s ({time.time()-t0:.1f}s incl.)")

    # final eps_g as the "STE fixed point" estimate
    ste_fps = []
    for i, omega in enumerate(omegas_ste):
        eps_final = float(results[i]["eps_g"][-100:].mean())  # average last 100 logs
        eps_final_std = float(results[i]["eps_g"][-100:].std())
        ste_fps.append({
            "omega": omega, "eps_fp_ste": eps_final, "eps_fp_ste_std": eps_final_std,
        })
        print(f"    omega={omega:.2f}: STE eps*={eps_final:.4f} (std {eps_final_std:.4f})")

    # ---- Plot: Figure 6 reproduction ----
    fig, axes = plt.subplots(1, 2, figsize=(12, 5))
    for ax, b_w in zip(axes, weight_bs):
        # theory curves
        for name, b_x in input_configs:
            key = f"b_w={b_w},{name}"
            preds = theory[key]
            eps_star = [p["eps_star_pred"] for p in preds]
            ls = "-" if b_x is None else "--"
            ax.plot(omegas, eps_star, ls, label=f"theory {name}", lw=2)
        # ODE fixed points
        for name, b_x in input_configs:
            key = f"b_w={b_w},{name}"
            fps = ode_fps[key]
            eps_fp = [f["eps_fp"] for f in fps]
            ax.plot(omegas, eps_fp, ":", label=f"ODE fp {name}", lw=1.5)
        # STE points (only for b_w=2, unquant input)
        if b_w == 2:
            ax.errorbar(omegas_ste, [s["eps_fp_ste"] for s in ste_fps],
                        yerr=[s["eps_fp_ste_std"] for s in ste_fps],
                        fmt="o", color="black", ms=6, capsize=3, label="STE (b_w=2, unquant)")
        ax.set_xlabel(r"Weight quantization range $\omega$")
        ax.set_ylabel(r"$\varepsilon_g^*$ (asymptotic)")
        ax.set_title(f"Weight quant b={b_w} (input range ω_x=1)")
        ax.legend(fontsize=8)
        ax.grid(True, alpha=0.3)
    fig.suptitle("Claim 3 / Figure 6: Fixed-point eps_g* vs ω (small-eta limit)")
    fig.tight_layout()
    fig.savefig(OUT / "claim3_figure6_fixed_point.png", dpi=120)
    plt.close(fig)

    # ---- Verify Theorem V.8 prediction against ODE fixed point ----
    # For each (b_w, b_x, omega), compare theoretical eps_g* vs ODE-integrated eps_g*.
    # Theory should match ODE in the small-eta limit.
    verify = []
    for b_w in weight_bs:
        for name, b_x in input_configs:
            for omega in omegas:
                qw = Quantizer(b=b_w, omega=omega)
                qx = Quantizer(b=b_x, omega=omega_x) if b_x is not None else None
                thy = small_eta_fixed_point_prediction(qw, qx, lam=lam, rho=rho, sigma2=sigma2)
                # find ODE fp at this omega
                m_fp, q_fp, eps_fp = find_fp_ode(eta_small, qw, qx, lam, rho, sigma2)
                verify.append({
                    "b_w": b_w, "b_x": name, "omega": float(omega),
                    "theory_eps_star": thy["eps_star_pred"],
                    "theory_regime": thy["regime"],
                    "theory_c": thy["c"],
                    "theory_p": thy["p"],
                    "ode_eps_star": eps_fp,
                    "abs_diff": abs(thy["eps_star_pred"] - eps_fp),
                })
    verify_arr = np.array([v["abs_diff"] for v in verify])
    print(f"\n  Theory vs ODE fixed-point verification:")
    print(f"    n_points = {len(verify)}")
    print(f"    max |theory - ODE| = {verify_arr.max():.6f}")
    print(f"    mean |theory - ODE| = {verify_arr.mean():.6f}")
    # Most should be small (theory holds in small-eta; we use eta=1e-3)

    out = {
        "omegas_theory": omegas.tolist(),
        "omegas_ste": omegas_ste,
        "weight_bs": weight_bs,
        "input_configs": [name for name, _ in input_configs],
        "theory": theory,
        "ode_fps": ode_fps,
        "ste_fps": ste_fps,
        "verification": {
            "n_points": len(verify),
            "max_abs_diff": float(verify_arr.max()),
            "mean_abs_diff": float(verify_arr.mean()),
            "details": verify,
        },
        "_meta": {
            "d": d, "eta_ste": eta_ste, "eta_ode": eta_small,
            "lam": lam, "rho": rho, "sigma2": sigma2, "omega_x": omega_x,
            "n_seeds": n_seeds, "n_steps_ste": n_steps_ste,
        },
    }
    with open(OUT / "claim3_results.json", "w") as f:
        json.dump(out, f, indent=2, default=float)

    print(f"\n[Claim 3] Done. Outputs in {OUT}")
    return out


def run_figure5_stability():
    """Figure 5: input-only quantization stability boundary 2/sigma2_psi and
    steady-state eps_g* vs omega. Closed form (Prop. VI.1).
    """
    print("=" * 60)
    print("[Claim 3 / Figure 5] Input-only: stability & eps_g* vs omega")
    print("=" * 60)
    sigma2 = 0.0
    lam = 0.0
    rho = 1.0
    eta = 1e-3  # small eta

    omegas = np.linspace(0.1, 2.0, 39)
    out = {}
    for b in [2, 3, 4, 10]:
        stab = []
        eps_stars = []
        for omega in omegas:
            qx = Quantizer(b=b, omega=omega)
            sigma2_x = qx.sigma2_psi()
            kappa_x = qx.kappa_psi()
            stab.append(input_only_stability_bound(sigma2_x, lam))
            m_s, q_s, e_s = input_only_fixed_point(kappa_x, sigma2_x, eta, lam, rho, sigma2)
            eps_stars.append(e_s)
        out[f"b={b}"] = {
            "omegas": omegas.tolist(),
            "stability_boundary": stab,
            "eps_star": eps_stars,
        }
        # also unquantized baseline: stability 2/(1+lam)=2, eps* = rho - rho/(1+lam) = 0 (with sigma2=0)
        print(f"  b={b}: stability range [{min(stab):.3f}, {max(stab):.3f}], eps* range [{min(eps_stars):.4f}, {max(eps_stars):.4f}]")

    # unquantized baseline
    out["unquant"] = {
        "stability_boundary": 2.0,
        "eps_star": 0.0,  # rho + sigma2 - rho/(sigma2_x+lam) with sigma2_x=1, lam=0 -> 1 - 1 = 0
    }

    fig, axes = plt.subplots(1, 2, figsize=(12, 5))
    ax = axes[0]
    for b in [2, 3, 4, 10]:
        ax.plot(out[f"b={b}"]["omegas"], out[f"b={b}"]["stability_boundary"], label=f"b={b}")
    ax.axhline(2.0, color="k", linestyle="--", label="unquantized")
    ax.set_xlabel(r"Input quantization range $\omega$")
    ax.set_ylabel(r"Stability boundary $2/\sigma_\psi^2$")
    ax.set_title("Stability boundary (input-only quant, λ=0)")
    ax.legend(); ax.grid(True, alpha=0.3)
    ax = axes[1]
    for b in [2, 3, 4, 10]:
        ax.plot(out[f"b={b}"]["omegas"], out[f"b={b}"]["eps_star"], label=f"b={b}")
    ax.axhline(0.0, color="k", linestyle="--", label="unquantized")
    ax.set_xlabel(r"Input quantization range $\omega$")
    ax.set_ylabel(r"$\varepsilon_g^*$")
    ax.set_title("Steady-state generalization error (input-only quant, λ=0)")
    ax.legend(); ax.grid(True, alpha=0.3)
    fig.suptitle("Claim 3 / Figure 5: Input-only quantization stability and eps_g*")
    fig.tight_layout()
    fig.savefig(OUT / "claim3_figure5_stability.png", dpi=120)
    plt.close(fig)

    return out


def run_claim3():
    out5 = run_figure5_stability()
    out6 = run_figure6_small_eta()

    # also report key verification number: theory vs ODE fp max diff
    verify = out6["verification"]
    print("\n[Claim 3] Summary:")
    print(f"  Theory vs ODE fixed-point max|diff|: {verify['max_abs_diff']:.6f}")
    print(f"  Theory vs ODE fixed-point mean|diff|: {verify['mean_abs_diff']:.6f}")
    # Note: unquantized baseline eps_g* = rho + sigma2 - kappa^2 rho/(sigma2+lambda)
    # For rho=1, sigma2=0, lam=0: eps_g* = 1 - 1 = 0
    # Quantization deviation = eps_g*(quantized) - eps_g*(unquantized)
    # For weight b=2, input unquant, small eta:
    # c = 1*1/(1+0) = 1, omega >= 1 -> c >= omega case (if omega=1, c=omega exactly -> boundary)
    # For omega=0.5: c=1 > omega=0.5 -> boundary case, eps* = 1 - 2*1*0.5 + 1*0.25 = 0.25
    # For omega=1.5: c=1 < omega=1.5, interior
    #   i*: largest i with v_i <= c=1
    #   For b=2, omega=1.5: levels = [-1.5, -0.5, 0.5, 1.5] (L+1=4, Delta=1)
    #   i* = 2 (v_2=0.5), p = (1 - 0.5)/1 = 0.5
    #   eps_g0 = 1 - 2*1*1 + 1*1 = 0
    #   correction = 1 * 1^2 * 0.5*0.5 = 0.25
    #   eps* = 0 + 0.25 = 0.25
    return {"figure5": out5, "figure6": out6}


if __name__ == "__main__":
    run_claim3()