scripts/exp_genconvex.py

"""Claim 4 -- Theorems 4.4 / 5.4: dimension-free guarantees in the general
(non-strongly) convex regime alpha = 0.

A quadratic target with a genuinely zero eigenvalue has no normalisable pi, so
the operational meaning of "alpha = 0" is: the guarantee must not use strong
convexity, i.e. it must hold on a horizon short compared with the strongly
convex mixing time gamma/alpha.  We therefore

  * take a target whose small eigenvalues are so tiny that alpha plays no role
    on the horizon reached (verified: N h << gamma/alpha, reported per run),
  * initialise mu_0 = N(m0, Sigma_pi^x) x N(0,I) with the mean displaced by W in
    the stiffest direction, so that W_2(mu_0, pi) = W is an explicit, controlled
    parameter and KL(mu_0||pi) = beta W^2 / 2 does NOT grow with d,
  * measure N_eps and fit its exponents in eps, W, tr(H) and d.

Theorem 4.4 (ULMC, alpha=0):  N = Otilde( beta [tr H]^{1/2} W^3 / eps^4 )
Theorem 5.4 (RMD,  alpha=0):  N = Otilde( beta [tr H]^{1/4} W^{5/2} / eps^3 )
(the eps exponents -4 and -3 are unambiguous in both the theorem statements and
Table 1, and are the discriminating quantity here).
"""

import numpy as np

import common as C
import ulmc_core as U

NGRID = U.n_grid(int(8e9), 1.03)
HS = lambda b: C.hgrid(b, n=26, ratio=1.32)
res = {
    "setting": "alpha -> 0 (general convex); mu_0 mean-shifted by W in the "
    "stiff direction, covariance already at pi"
}

ALPHA = 1e-6  # effectively zero: gamma/alpha = 5.7e6 time units
BETA = 1.0


def family(d, nstiff, trH_target):
    """1 stiff coordinate at beta, nstiff-1 more at beta, rest at a tiny value
    chosen to hit tr(H)."""
    rest = d - nstiff
    c = (trH_target - nstiff * BETA) / rest
    assert ALPHA <= c <= BETA, c
    a = np.array([BETA, c])
    m = np.array([float(nstiff), float(rest)])
    return a, m


def run(scheme, a, m, eps, W, beta=BETA):
    S0, m0 = C.init_shifted(a, W, idx=0)
    r = C.sweep_neps(scheme, a, m, beta, eps, S0, m0, hs=HS(beta), ngrid=NGRID)
    r.pop("table")
    alpha_eff = float(np.min(a))
    r["Nh"] = (r["N_eps"] * r["h_star"]) if r["N_eps"] else None
    r["gamma_over_alpha_eff"] = float(C.gamma_of(beta) / alpha_eff)
    r["strong_convexity_irrelevant"] = bool(
        r["Nh"] is not None and r["Nh"] < 0.05 * r["gamma_over_alpha_eff"]
    )
    r["KL0"] = float(0.5 * beta * W * W)
    return r


for scheme, pe, pw, pt in (("ulmc", -4.0, 3.0, 0.5), ("rmd", -3.0, 2.5, 0.25)):
    blk = {}
    a, m = family(400, 4, 6.0)  # trH = 6, d = 400
    # ---- eps ----
    W = 3.0
    epss = np.geomspace(0.05, 0.6, 8)
    rows = [dict(eps=float(e), **run(scheme, a, m, e, W)) for e in epss]
    s, r2 = C.fit_exponent([r["eps"] for r in rows], [r["N_eps"] for r in rows])
    blk["eps"] = {
        "rows": rows,
        "fit_exponent": s,
        "r2": r2,
        "predicted": pe,
        "W": W,
        "trH": 6.0,
        "d": 400,
    }
    # ---- W ----
    eps = 0.15
    Ws = np.geomspace(1.0, 16.0, 7)
    rows = [dict(W=float(w), **run(scheme, a, m, eps, float(w))) for w in Ws]
    s, r2 = C.fit_exponent([r["W"] for r in rows], [r["N_eps"] for r in rows])
    blk["W"] = {
        "rows": rows,
        "fit_exponent": s,
        "r2": r2,
        "predicted": pw,
        "eps": eps,
        "trH": 6.0,
        "d": 400,
    }
    # ---- tr(H) at fixed d ----
    rows = []
    for nst in (2, 4, 8, 16, 32, 64):
        a2, m2 = family(400, nst, float(nst) + 0.5)
        rows.append(
            dict(
                trH=float((a2 * m2).sum()), nstiff=nst, **run(scheme, a2, m2, 0.15, 3.0)
            )
        )
    s, r2 = C.fit_exponent([r["trH"] for r in rows], [r["N_eps"] for r in rows])
    blk["trH"] = {
        "rows": rows,
        "fit_exponent": s,
        "r2": r2,
        "predicted": pt,
        "eps": 0.15,
        "W": 3.0,
        "d": 400,
    }
    # ---- d at fixed tr(H): the dimension-free test in the alpha=0 regime ----
    rows = []
    for d in (20, 50, 100, 200, 400, 800, 1600, 3200):
        a2, m2 = family(d, 4, 6.0)
        rows.append(dict(d=d, trH=6.0, **run(scheme, a2, m2, 0.15, 3.0)))
    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": 0.0,
        "growth_factor_over_160x_d": float(rows[-1]["N_eps"] / rows[0]["N_eps"]),
        "sqrt_d_would_predict": float(np.sqrt(3200 / 20)),
        "eps": 0.15,
        "W": 3.0,
    }
    res[scheme] = blk
    print(
        scheme,
        {k: round(v["fit_exponent"], 3) for k, v in blk.items()},
        "predicted",
        {"eps": pe, "W": pw, "trH": pt, "d": 0.0},
        flush=True,
    )

C.dump("genconvex", res)