scripts/exp_scaling2.py

"""Claims 1, 2, 3 -- corrected scaling sweeps.

Fixes relative to exp_scaling.py:
  * the tr(H) sweep now uses d = k+2 so that the alpha-tail contributes ~0.02 to
    tr(H) instead of ~10 (in the first version tr(H) was inflated at small k,
    which biased the fitted exponent upwards);
  * every sweep is checked to be BIAS-limited (h* strictly inside the grid), so
    that the discretisation error, not the burn-in, is what sets h;
  * each fit is reported both raw and after dividing out the logarithmic burn-in
    factor L = log(KL(mu_0||pi)/eps^2), which Otilde(.) is entitled to hide.
Also measures the W_2 complexity (for the comparison with Liu et al. 2023).
"""

import numpy as np

import common as C
import ulmc_core as U

NGRID = U.n_grid(int(2e10), 1.03)
HS = lambda b: C.hgrid(b, n=30, ratio=1.32)

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


def kl0(a, mult, beta):
    S0, m0 = C.init_cold(a, beta)
    return float(np.sum(mult * U.kl_per_coord(S0, m0, a)))


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)
    r.pop("table")
    r["bias_limited"] = bool(r["h_star"] < HS(beta)[0] * 0.999)
    r["KL0"] = kl0(a, mult, beta)
    r["L"] = float(np.log(max(r["KL0"] / eps**2, 1.0001)))
    r["N_over_L"] = r["N_eps"] / r["L"]
    return r


def fits(rows, xkey, pred):
    x = [r[xkey] for r in rows]
    s, r2 = C.fit_exponent(x, [r["N_eps"] for r in rows])
    sl, r2l = C.fit_exponent(x, [r["N_over_L"] for r in rows])
    sh, _ = C.fit_exponent(x, [r["h_star"] for r in rows])
    return {
        "rows": rows,
        "fit_exponent": s,
        "r2": r2,
        "fit_exponent_log_corrected": sl,
        "r2_log_corrected": r2l,
        "h_star_exponent": sh,
        "predicted": pred,
        "all_bias_limited": bool(all(r["bias_limited"] for r in rows)),
    }


res = {
    "note": "N_eps = min over h<=1/gamma of steps to KL<=eps^2, exact",
    "grid_resolution_pct": 3.0,
}

for scheme in ("ulmc", "rmd"):
    blk = {}
    # ---- eps ----
    a = np.array([1.0, 0.02])
    m = np.array([20.0, 2.0])  # trH=20.04, kappa=50
    epss = (
        np.geomspace(1.5e-3, 2e-2, 9)
        if scheme == "ulmc"
        else np.geomspace(3e-3, 3e-2, 9)
    )
    rows = [dict(eps=float(e), **run(scheme, a, m, 1.0, e)) for e in epss]
    blk["eps"] = fits(rows, "eps", PRED[scheme]["eps"])
    blk["eps"].update(trH=float((a * m).sum()), kappa=50.0, d=22)

    # ---- tr(H): d = k+2 so tr(H) ~= k exactly ----
    eps = 3e-3 if scheme == "ulmc" else 8e-3
    rows = []
    for k in (2, 4, 8, 16, 32, 64, 128, 256):
        a = np.array([1.0, 0.01])
        m = np.array([float(k), 2.0])
        rows.append(
            dict(trH=float(k + 0.02), k=k, d=k + 2, **run(scheme, a, m, 1.0, eps))
        )
    blk["trH"] = fits(rows, "trH", PRED[scheme]["trH"])
    blk["trH"].update(eps=eps, kappa=100.0)

    # ---- kappa ----
    eps = 8e-3 if scheme == "ulmc" else 1.2e-2
    rows = []
    for alpha in np.geomspace(0.5, 1e-3, 10):
        a = np.array([1.0, float(alpha)])
        m = np.array([20.0, 2.0])
        rows.append(
            dict(
                kappa=float(1 / alpha),
                alpha=float(alpha),
                trH=float((a * m).sum()),
                **run(scheme, a, m, 1.0, eps)
            )
        )
    blk["kappa"] = fits(rows, "kappa", PRED[scheme]["kappa"])
    blk["kappa"].update(eps=eps, d=22)

    # ---- beta (kappa, tr(H) fixed) ----
    eps = 5e-3 if scheme == "ulmc" else 1e-2
    rows = []
    for beta in (0.25, 0.5, 1.0, 2.0, 4.0, 8.0):
        k = int(round(20.0 / beta))
        a = np.array([beta, beta / 50.0])
        m = np.array([float(k), 2.0])
        rows.append(
            dict(
                beta=float(beta),
                trH=float((a * m).sum()),
                **run(scheme, a, m, beta, eps)
            )
        )
    blk["beta"] = fits(rows, "beta", PRED[scheme]["beta"])
    blk["beta"].update(eps=eps, kappa=50.0)

    # ---- d at FIXED (alpha, beta, tr(H)) : the dimension-free test ----
    eps = 4e-3 if scheme == "ulmc" else 5e-3
    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)
            )
        )
    blk["d_fixed_trH"] = fits(rows, "d", 0.0)
    blk["d_fixed_trH"].update(
        eps=eps,
        growth_factor_over_82x_d=float(rows[-1]["N_eps"] / rows[0]["N_eps"]),
        sqrt_d_prediction=float(np.sqrt(900 / 11)),
        d_prediction=float(900 / 11),
    )

    # ---- control: H = beta I, tr(H) = beta d ----
    rows = []
    for d in (10, 20, 40, 80, 160, 320, 640):
        a = np.array([1.0, 0.01])
        m = np.array([float(d - 1), 1.0])
        rows.append(dict(d=d, trH=float((a * m).sum()), **run(scheme, a, m, 1.0, eps)))
    blk["d_control_H_eq_betaI"] = fits(rows, "d", PRED[scheme]["trH"])
    res[scheme] = blk
    print(
        scheme,
        {
            k: (
                round(v["fit_exponent"], 3),
                round(v["fit_exponent_log_corrected"], 3),
                v["predicted"],
                v["all_bias_limited"],
            )
            for k, v in blk.items()
        },
        flush=True,
    )


# ---------------------------------------------------------------------------
# W2 complexity (Talagrand consequence) -- comparison with Liu et al. (2023),
# who prove Otilde(kappa^{5/3} beta^{-2/3} [tr H]^{1/3} eps^{-2/3}) in W2 for a
# doubly randomized ULD algorithm.  The paper claims its Thm 5.2 implies
# Otilde(kappa^{4/3} beta^{-2/3} [tr H]^{1/2}/eps), i.e. a strictly better kappa.
# ---------------------------------------------------------------------------
def w2_curve(maps, S0, m0, a, mult, grid, final_ulmc=None):
    A, T, q = maps
    Aug = U._augment(T, q)
    K = Aug.shape[0]
    accS = np.broadcast_to(np.eye(4), (K, 4, 4)).copy()
    accM = np.broadcast_to(np.eye(2), (K, 2, 2)).copy()
    S0a = np.concatenate([S0, np.ones((K, 1))], -1)
    out, prev = np.empty(len(grid)), 0
    for i, n in enumerate(grid):
        d = int(n) - prev
        if d > 0:
            accS = accS @ np.linalg.matrix_power(Aug, d)
            accM = accM @ np.linalg.matrix_power(A, d)
            prev = int(n)
        S = np.einsum("kij,kj->ki", accS, S0a)[:, :3]
        mm = np.einsum("kij,kj->ki", accM, m0)
        if final_ulmc is not None:
            S, mm = U.compose_step(final_ulmc[1], final_ulmc[2], final_ulmc[0], S, mm)
        var = np.maximum(S[:, 0] - mm[:, 0] ** 2, 0.0)
        w2sq = mm[:, 0] ** 2 + (np.sqrt(var) - np.sqrt(1.0 / a)) ** 2
        out[i] = float(np.sum(mult * w2sq))
    return np.sqrt(out)


w2 = {}
for scheme in ("ulmc", "rmd"):
    rows = []
    for alpha in np.geomspace(0.5, 2e-3, 8):
        a = np.array([1.0, float(alpha)])
        m = np.array([20.0, 2.0])
        S0, m0 = C.init_cold(a, 1.0)
        g = C.gamma_of(1.0)
        best = None
        for h in HS(1.0):
            mp = U.ulmc_maps(a, h, g) if scheme == "ulmc" else U.rmd_maps(a, h, g, 48)
            fin = U.ulmc_maps(a, h, g) if scheme == "rmd" else None
            v = w2_curve(mp, S0, m0, a, m, NGRID, final_ulmc=fin)
            n = U.first_below(NGRID, v, 0.05)
            if n and (best is None or n < best[0]):
                best = (n, float(h))
        rows.append({"kappa": float(1 / alpha), "N_w2": best[0], "h_star": best[1]})
    s, r2 = C.fit_exponent([r["kappa"] for r in rows], [r["N_w2"] for r in rows])
    w2[scheme] = {"rows": rows, "kappa_exponent": s, "r2": r2, "w2_target": 0.05}
    print("W2", scheme, "kappa exponent", round(s, 3), flush=True)
w2["liu_2023_kappa_exponent"] = 5 / 3
w2["paper_claimed_kappa_exponent_for_RMD_in_W2"] = 4 / 3
res["w2_complexity"] = w2

C.dump("scaling2", res)