scripts/exp_scaling.py

"""Claims 1, 2, 3: empirical iteration complexity of ULMC and RMD.

N_eps := min over admissible step sizes h <= 1/gamma of the smallest number of
steps with KL(mu (P')^N || pi) <= eps^2, computed EXACTLY (Gaussian target,
closed-form KL, no Monte-Carlo).  mu_0 = N(0, beta^{-1} I) x N(0, I).

Every sweep varies ONE of {eps, tr(H), kappa, beta, d} and holds the rest fixed,
then fits the empirical power law and compares against
   Thm 4.3 (ULMC): N = Otilde( kappa^{3/2} beta^{-1/2} tr(H)^{1/2} / eps )
   Thm 5.2 (RMD) : N = Otilde( kappa [beta^{-1} tr(H)]^{1/3} eps^{-2/3} )
Spectra for the eps / tr(H) / kappa / beta sweeps put most mass at a_i = beta so
that tr(H) is a *tight* proxy for the spectrum (sum a_i^2 = beta tr(H)); this is
the regime where the paper's bound is supposed to be sharp.
"""

import sys

import numpy as np

import common as C
import ulmc_core as U

NGRID = U.n_grid(int(4e9), 1.03)
HS = lambda beta: C.hgrid(beta, n=26, ratio=1.32)

PRED = {
    "ulmc": {"eps": -1.0, "trH": 0.5, "kappa": 1.5, "beta": -0.5, "d": 0.0},
    "rmd": {"eps": -2.0 / 3, "trH": 1.0 / 3, "kappa": 1.0, "beta": -1.0 / 3, "d": 0.0},
}


def stiff_spectrum(k, d, alpha, beta):
    """k coordinates at a=beta, d-k at a=alpha.  tr(H)=k beta+(d-k)alpha,
    sum a_i^2 = k beta^2 + (d-k) alpha^2 ~= beta tr(H) when alpha<<beta."""
    a = np.array([beta, alpha])
    m = np.array([float(k), float(d - k)])
    return a, m


def run(scheme, a, mult, beta, eps):
    S0, m0 = C.init_cold(a, beta)
    r = C.sweep_neps(scheme, a, mult, beta, eps, S0, m0, hs=HS(beta), ngrid=NGRID)
    hmax = HS(beta)[0]
    r["bias_limited"] = bool(r["h_star"] is not None and r["h_star"] < hmax * 0.999)
    r.pop("table")
    return r


res = {
    "grid_resolution_pct": 3.0,
    "n_grid_max": int(NGRID[-1]),
    "init": "mu_0 = N(0, beta^-1 I) x N(0,I)",
    "gamma": "sqrt(32 beta)",
}

for scheme in ("ulmc", "rmd"):
    blk = {}

    # ---- eps sweep -------------------------------------------------------
    a, m = stiff_spectrum(20, 22, 0.02, 1.0)  # trH=20.04, kappa=50
    epss = (
        np.geomspace(2e-3, 3e-2, 8) if scheme == "ulmc" else np.geomspace(5e-3, 6e-2, 8)
    )
    rows = [dict(eps=float(e), **run(scheme, a, m, 1.0, e)) for e in epss]
    s, r2 = C.fit_exponent([r["eps"] for r in rows], [r["N_eps"] for r in rows])
    sh, _ = C.fit_exponent([r["eps"] for r in rows], [r["h_star"] for r in rows])
    blk["eps"] = {
        "rows": rows,
        "fit_exponent": s,
        "r2": r2,
        "h_star_exponent": sh,
        "predicted": PRED[scheme]["eps"],
        "trH": float((a * m).sum()),
        "kappa": 50.0,
        "d": 22,
    }

    # ---- tr(H) sweep -----------------------------------------------------
    eps = 4e-3 if scheme == "ulmc" else 1.2e-2
    rows = []
    for k in (2, 4, 8, 16, 32, 64, 128, 256):
        a, m = stiff_spectrum(k, 1000, 0.01, 1.0)
        trH = float((a * m).sum())
        rows.append(dict(trH=trH, k=k, **run(scheme, a, m, 1.0, eps)))
    s, r2 = C.fit_exponent([r["trH"] for r in rows], [r["N_eps"] for r in rows])
    sh, _ = C.fit_exponent([r["trH"] for r in rows], [r["h_star"] for r in rows])
    blk["trH"] = {
        "rows": rows,
        "fit_exponent": s,
        "r2": r2,
        "h_star_exponent": sh,
        "predicted": PRED[scheme]["trH"],
        "eps": eps,
        "kappa": 100.0,
        "d": 1000,
    }

    # ---- kappa sweep -----------------------------------------------------
    eps = 1e-2 if scheme == "ulmc" else 2e-2
    rows = []
    for alpha in np.geomspace(0.5, 2e-3, 9):
        a, m = stiff_spectrum(20, 22, float(alpha), 1.0)
        rows.append(
            dict(
                kappa=float(1.0 / alpha),
                alpha=float(alpha),
                trH=float((a * m).sum()),
                **run(scheme, a, m, 1.0, eps)
            )
        )
    s, r2 = C.fit_exponent([r["kappa"] for r in rows], [r["N_eps"] for r in rows])
    blk["kappa"] = {
        "rows": rows,
        "fit_exponent": s,
        "r2": r2,
        "predicted": PRED[scheme]["kappa"],
        "eps": eps,
        "d": 22,
    }

    # ---- beta sweep (kappa and tr(H) held fixed) --------------------------
    eps = 6e-3 if scheme == "ulmc" else 1.5e-2
    rows = []
    for beta in (0.25, 0.5, 1.0, 2.0, 4.0):
        k = int(round(20.0 / beta))
        a, m = stiff_spectrum(k, k + 2, beta / 50.0, beta)
        rows.append(
            dict(
                beta=float(beta),
                trH=float((a * m).sum()),
                **run(scheme, a, m, beta, eps)
            )
        )
    s, r2 = C.fit_exponent([r["beta"] for r in rows], [r["N_eps"] for r in rows])
    blk["beta"] = {
        "rows": rows,
        "fit_exponent": s,
        "r2": r2,
        "predicted": PRED[scheme]["beta"],
        "eps": eps,
        "kappa": 50.0,
    }

    # ---- d sweep at FIXED (alpha, beta, tr(H)) : the dimension-free test ---
    eps = 5e-3 if scheme == "ulmc" else 1.5e-2
    rows = []
    for d in (11, 20, 40, 80, 160, 320, 640, 900):
        a, m = C.spectrum(d, 0.01, 1.0, 10.0)
        rows.append(
            dict(
                d=d,
                trH=10.0,
                sum_a2=float((a * a * m).sum()),
                **run(scheme, a, m, 1.0, eps)
            )
        )
    s, r2 = C.fit_exponent([r["d"] for r in rows], [r["N_eps"] for r in rows])
    blk["d_fixed_trH"] = {
        "rows": rows,
        "fit_exponent": s,
        "r2": r2,
        "predicted": PRED[scheme]["d"],
        "eps": eps,
        "growth_factor_over_80x_d": float(rows[-1]["N_eps"] / rows[0]["N_eps"]),
        "sqrt_d_would_predict": float(np.sqrt(900 / 11)),
    }

    # ---- control: H = beta I so tr(H) = beta d grows with d ---------------
    rows = []
    for d in (10, 20, 40, 80, 160, 320):
        a, m = np.array([1.0, 0.01]), np.array([float(d - 1), 1.0])
        rows.append(dict(d=d, trH=float((a * m).sum()), **run(scheme, a, m, 1.0, eps)))
    s, r2 = C.fit_exponent([r["d"] for r in rows], [r["N_eps"] for r in rows])
    blk["d_control_H_eq_betaI"] = {
        "rows": rows,
        "fit_exponent": s,
        "r2": r2,
        "predicted": 0.5,
        "eps": eps,
    }
    res[scheme] = blk
    print(scheme, {k: v.get("fit_exponent") for k, v in blk.items()}, flush=True)

C.dump("scaling", res)