repro-bundle/v1/scripts/test_ste_repro.py

"""Unit tests for the STE-quantization reproduction.

Verifies:
  1. Quantizer moments against hand-computed values for b=2, omega=1.
  2. STE simulation matches ODE for an unquantized linear regression reference
     (kappa=sigma2=1) — this is exactly solvable and the high-dim ODE is exact
     for any d (no quantization, no isotropy assumption needed).
  3. Quantizer psi matches its differentiable relaxation psi_T as T -> 0.
  4. Input-only fixed point formula reproduces the unquantized ridge limit.
  5. Macroscopic state functions reduce to psi(m)^2 etc. when s -> 0.
"""
import math
import sys
from pathlib import Path

import numpy as np
from scipy.stats import norm

sys.path.insert(0, str(Path(__file__).parent))
from ste_repro import (
    Quantizer, macro_m_psi, macro_q_psi, macro_r_psi, eps_g,
    ode_rhs, input_only_fixed_point, input_only_stability_bound,
    small_eta_fixed_point_prediction, identity_quantizer_moments,
)
from ste_sim_torch import run_ste, STEConfig


def approx(a, b, tol=1e-6):
    return abs(a - b) <= tol * max(1.0, abs(b))


def test_quantizer_moments_b2_omega1():
    q = Quantizer(b=2, omega=1.0)
    # L=2, Delta=1, levels=[-1,0,1], theta=[-0.5, 0.5]
    assert q.L == 2
    assert approx(q.Delta, 1.0)
    assert np.allclose(q.levels, [-1.0, 0.0, 1.0])
    assert np.allclose(q.theta, [-0.5, 0.5])
    # sigma2_psi = 2 (1 - Phi(0.5)) = 2 (1 - 0.6914624612838983)
    expected_s2 = 2.0 * (1.0 - norm.cdf(0.5))
    assert approx(q.sigma2_psi(), expected_s2, 1e-10), q.sigma2_psi()
    # kappa_psi = 2 * phi(0.5)
    expected_k = 2.0 * norm.pdf(0.5)
    assert approx(q.kappa_psi(), expected_k, 1e-10), q.kappa_psi()
    print("  test_quantizer_moments_b2_omega1 OK  (sigma2_psi=%.6f, kappa_psi=%.6f)"
          % (q.sigma2_psi(), q.kappa_psi()))


def test_psi_matches_psi_T_small_T():
    q = Quantizer(b=3, omega=1.5)
    x = np.linspace(-2.0, 2.0, 51)
    hard = q.psi(x)
    soft = q.psi_t(x, T=1e-4)
    assert np.allclose(hard, soft, atol=2e-3), (hard[:5], soft[:5])
    print("  test_psi_matches_psi_T_small_T OK")


def test_macro_s_zero_reduces_to_psi_of_m():
    q = Quantizer(b=3, omega=1.0)
    m = 0.37
    expected_m_psi = float(q.psi(np.array([m]))[0])
    expected_q_psi = expected_m_psi ** 2
    assert approx(macro_m_psi(m, 1e-9, q), expected_m_psi, 1e-3)
    assert approx(macro_q_psi(m, 1e-9, q), expected_q_psi, 1e-3)
    print("  test_macro_s_zero_reduces_to_psi_of_m OK")


def test_input_only_fixed_point_unquantized_ridge():
    """For kappa_x = sigma2_x = 1 (unquantized input), the input-only fixed
    point reduces to ridge regression with weights converging to w* / (1+lambda).
    eps_g* = rho + sigma2 - rho / (1+lambda)  (small eta)."""
    lam = 1.0
    rho = 1.0
    sigma2 = 0.0
    eta = 1e-3
    kappa_x, sigma2_x = 1.0, 1.0
    m_s, q_s, e_s = input_only_fixed_point(kappa_x, sigma2_x, eta, lam, rho, sigma2)
    # m* = rho * 1 / (1 + lam) = 1/2  (ridge solution: w = w*/(1+lam))
    assert approx(m_s, 0.5, 1e-10), m_s
    # q* = 1/(1+lam)^2 = 1/4 in the small-eta limit. With finite eta: q* = 1/(4 - eta).
    # eps_g* = rho + q* - 2 m* = 1 + 1/(4-eta) - 1 = 1/(4-eta) -> 1/4 as eta->0.
    assert approx(q_s, 1.0 / (4.0 - eta), 1e-10), q_s
    assert abs(e_s - 1.0 / (4.0 - eta)) < 1e-6, e_s
    # Stability bound: 2 * (1 + 1) / 1 = 4
    assert approx(input_only_stability_bound(1.0, 1.0), 4.0, 1e-10)
    print("  test_input_only_fixed_point_unquantized_ridge OK  (m*=%.4f q*=%.4f eps*=%.6f)"
          % (m_s, q_s, e_s))


def test_ste_matches_ode_unquantized():
    """STE with both psi_w = identity and psi_x = identity is plain linear
    regression. The high-dim ODE is exact for any d in this limit. Check that
    STE (d=200) tracks ODE for unquantized SGD.

    The ODE in this case:
        dm/dtau = -eta [(1+lambda) m - 1]            (rho=1, kappa=sigma2=1)
        dq/dtau = -2 eta [(1+lambda) m - m] + eta^2 eps_g
    We just check: STE(eps_g) vs ODE(eps_g) at several taus agree to ~1e-2.
    """
    from scipy.integrate import solve_ivp
    lam = 1.0
    eta = 0.05
    rho = 1.0
    sigma2 = 0.0
    d = 200
    n_steps = 50_000
    log_period = 200

    # ODE rhs with identity quantizers (kappa_x = sigma2_x = 1, qw = identity).
    # m_psi = m, q_psi = q^? -> no, q_psi = q (since psi_w(w) = w -> ||w||^2/d = q).
    # r_psi = m * m_psi + Delta * s * sum phi(...) -> without weight quant,
    # r_psi = w^T w / d = q.   So:
    #   dm/dtau = -eta [(1+lam) m - 1]
    #   dq/dtau = -2 eta [(1+lam) q - m] + eta^2 * eps_g
    # eps_g = 1 + 1 * q - 2 * 1 * m
    def rhs(tau, y):
        m, q = y
        e = 1.0 + q - 2.0 * m
        dmd = -eta * ((1.0 + lam) * m - 1.0)
        dqd = -2.0 * eta * ((1.0 + lam) * q - m) + eta * eta * 1.0 * e
        return [dmd, dqd]

    # initial: w ~ N(0,1), so m0 = 0, q0 = 1
    sol = solve_ivp(rhs, [0, n_steps / d], [0.0, 1.0],
                    t_eval=np.arange(0, n_steps / d + 1, log_period / d),
                    rtol=1e-9, atol=1e-12, method="DOP853")
    m_ode = sol.y[0]
    q_ode = sol.y[1]
    eps_ode = 1.0 + q_ode - 2.0 * m_ode

    # STE: use Quantizer with very large omega and b such that psi_w ≈ identity.
    # We just call run_ste with qx = None (identity input) and qw = a Quantizer
    # that approximates identity by using a large omega range. But that introduces
    # clipping. Instead, we use a "no-quant" mode: pass a special flag.
    # Easier: add a custom "identity weight" path by using Quantizer with omega large enough.
    # Simpler: use very large omega (=10) and b=8 so levels are dense in [-10,10].
    # For d=200, w stays O(1), so clipping at 10 is negligible.
    qw = Quantizer(b=8, omega=10.0)   # near-identity on the relevant range
    cfg = STEConfig(
        d=d, eta=eta, lam=lam, rho=rho, sigma2=sigma2,
        n_steps=n_steps, log_period=log_period, n_seeds=4,
        w_init_std=1.0, w_star_value=1.0,
        sigma2_x=1.0, kappa_x=1.0, device="cuda",
    )
    taus, metrics, wall = run_ste(
        (qw.levels, qw.theta, qw.omega, qw.Delta),
        None,  # identity input
        cfg,
    )
    eps_ste_mean = metrics["eps_g"].mean(axis=1)
    # compare at matching taus
    # ODE has T = n_steps / d + 1 points; STE same length
    assert len(eps_ode) == len(eps_ste_mean), (len(eps_ode), len(eps_ste_mean))
    # exclude very first point (initial condition may differ)
    diff = np.abs(eps_ode[1:] - eps_ste_mean[1:])
    rel = diff / np.maximum(np.abs(eps_ode[1:]), 1e-2)
    # We expect small error (finite-d only).
    max_rel = rel.max()
    # d=200 should give < 5% relative error after initial transient.
    assert max_rel < 0.10, ("STE-ODE mismatch max_rel=%.4f" % max_rel)
    print("  test_ste_matches_ode_unquantized OK  (max_rel=%.4f, wall=%.2fs)"
          % (max_rel, wall))


def test_ode_rhs_zero_at_input_only_fixed_point():
    """At the input-only fixed point (m*, q*), the ODE rhs must vanish."""
    q = Quantizer(b=4, omega=1.0)  # weight quantizer (irrelevant in input-only case)
    # for input-only, m_psi = m, q_psi = q, r_psi = q (since weights not quantized).
    # We can't easily express that with our ode_rhs (which assumes weight quant).
    # Instead, test: for joint weight-input quant with identity weight quant
    # (large omega, large b), the input-only fp should be an approximate fp of joint.
    lam = 1.0
    eta = 1e-3
    rho = 1.0
    sigma2 = 0.0
    kappa_x, sigma2_x = 1.0, 1.0
    m_s, q_s, e_s = input_only_fixed_point(kappa_x, sigma2_x, eta, lam, rho, sigma2)
    # use near-identity weight quantizer
    qw = Quantizer(b=8, omega=10.0)
    s = math.sqrt(max(q_s - m_s * m_s / rho, 1e-12))
    rhs = ode_rhs(0.0, [m_s, q_s], qw, kappa_x, sigma2_x, eta, lam, rho, sigma2)
    # |rhs| should be small (not exactly 0 because weight quantizer is not quite identity,
    # but m_psi(m, s) ≈ m for large omega).
    norm = math.sqrt(rhs[0] ** 2 + rhs[1] ** 2)
    assert norm < 1e-2, ("ode rhs at fp not zero: norm=%.6e" % norm)
    print("  test_ode_rhs_zero_at_input_only_fixed_point OK  (rhs norm=%.2e)" % norm)


def test_small_eta_prediction_monotone_in_omega():
    """For interior p, eps_g* prediction includes Delta^2 p (1-p) * sigma2_x.
    Delta = 2 omega / L, so for fixed b this should change non-trivially with omega.
    """
    qw = Quantizer(b=3, omega=1.0)
    qx = Quantizer(b=4, omega=1.0)
    out = small_eta_fixed_point_prediction(qw, qx, lam=0.0, rho=1.0, sigma2=0.0)
    assert out["regime"] in {"interior", "log", "boundary"}, out
    # With lambda=0, c = kappa_x * rho / sigma2_x = kappa_x / sigma2_x.
    # For symmetric quantizer, kappa_x / sigma2_x < 1 so c is interior.
    # Check the correction term is positive.
    if out["regime"] == "interior":
        assert out["delta2_p1p"] > 0
    print("  test_small_eta_prediction_monotone_in_omega OK  (regime=%s, c=%.4f, eps*=%.4f)"
          % (out["regime"], out["c"], out["eps_star_pred"]))


if __name__ == "__main__":
    print("Running unit tests...")
    test_quantizer_moments_b2_omega1()
    test_psi_matches_psi_T_small_T()
    test_macro_s_zero_reduces_to_psi_of_m()
    test_input_only_fixed_point_unquantized_ridge()
    test_ode_rhs_zero_at_input_only_fixed_point()
    test_small_eta_prediction_monotone_in_omega()
    test_ste_matches_ode_unquantized()
    print("ALL TESTS PASSED")