scripts/exp_prescription.py

"""Claims 2 & 3, sharp form: does the theorem's OWN (h, N) prescription work?

Theorem 4.3 prescribes, for the strongly convex case,
    h = Theta~( eps / (kappa [tr H]^{1/2}) ),   N = Theta~( kappa^{3/2} beta^{-1/2} [tr H]^{1/2} / eps )
and asserts KL(mu (P')^N || pi) <= eps^2.  Note the product
    N h = Theta~( kappa^{1/2} beta^{-1/2} ),
i.e. the prescribed *simulated time* grows only like sqrt(kappa).  ULD with the
mandated friction gamma = sqrt(32 beta) contracts at rate alpha/gamma, so the
burn-in time alone is Theta(gamma/alpha) = Theta(kappa beta^{-1/2}).  If that is
right, the prescription must break for large kappa.

Theorem 5.2 prescribes h = Theta~(beta^{-1/6}[tr H]^{-1/3} eps^{2/3}),
N = Theta~(kappa [beta^{-1} tr H]^{1/3} eps^{-2/3}), whose product is
N h = Theta~(kappa beta^{-1/2}) -- consistent with the burn-in.

Protocol (no unknown constants needed): calibrate the hidden Theta~ constants at
a base condition number kappa0 by taking the *empirically optimal* step size
h0 there and the exact number of steps N0 that reaches eps^2 at h0.  Then follow
the theorem's scaling for larger kappa and measure the resulting KL.
"""

import numpy as np

import common as C
import ulmc_core as U

NGRID = U.n_grid(int(8e9), 1.03)
res = {
    "protocol": "calibrate Theta~ constants at kappa0, then follow the theorem's "
    "h(kappa), N(kappa) scaling and measure the achieved KL"
}

BETA, TRH, D, EPS = 1.0, 20.0, 22, 1e-2
KAPPAS = [2.0, 4.0, 8.0, 16.0, 32.0, 64.0, 128.0, 256.0, 512.0]


def family(kappa):
    alpha = BETA / kappa
    a = np.array([BETA, alpha])
    m = np.array([20.0, 2.0])
    return a, m


for scheme, hexp, Nexp in (("ulmc", -1.0, 1.5), ("rmd", 0.0, 1.0)):
    a, m = family(KAPPAS[0])
    S0, m0 = C.init_cold(a, BETA)
    base = C.sweep_neps(
        scheme, a, m, BETA, EPS, S0, m0, hs=C.hgrid(BETA, n=26, ratio=1.32), ngrid=NGRID
    )
    h0, N0, k0 = base["h_star"], base["N_eps"], KAPPAS[0]
    rows = []
    for kappa in KAPPAS:
        a, m = family(kappa)
        S0, m0 = C.init_cold(a, BETA)
        h = h0 * (kappa / k0) ** hexp
        N = int(round(N0 * (kappa / k0) ** Nexp))
        if scheme == "ulmc":
            mp = U.ulmc_maps(a, h, C.gamma_of(BETA))
            kl = U.kl_at(N, mp, S0, m0, a, m)
        else:
            mp = U.rmd_maps(a, h, C.gamma_of(BETA), nq=48)
            fin = U.ulmc_maps(a, h, C.gamma_of(BETA))
            kl = U.kl_at(N, mp, S0, m0, a, m, final_ulmc=fin)
        # how many steps are ACTUALLY needed at that h?
        vals = U.kl_curve(
            mp,
            S0,
            m0,
            a,
            m,
            NGRID,
            final_ulmc=(
                U.ulmc_maps(a, h, C.gamma_of(BETA)) if scheme == "rmd" else None
            ),
        )
        Nreq = U.first_below(NGRID, vals, EPS * EPS)
        rows.append(
            {
                "kappa": kappa,
                "h_prescribed": float(h),
                "N_prescribed": N,
                "KL_at_prescribed_N": float(kl),
                "eps2": EPS * EPS,
                "prescription_holds": bool(kl <= EPS * EPS),
                "N_actually_required_at_that_h": Nreq,
                "shortfall_factor": (float(Nreq / N) if Nreq else None),
            }
        )
    ks = [r["kappa"] for r in rows]
    sreq, r2 = C.fit_exponent(ks, [r["N_actually_required_at_that_h"] for r in rows])
    res[scheme] = {
        "calibration": {
            "kappa0": k0,
            "h0": h0,
            "N0": N0,
            "eps": EPS,
            "trH": float((a * m).sum()),
            "beta": BETA,
            "d": D,
        },
        "rows": rows,
        "theorem_N_exponent_in_kappa": Nexp,
        "measured_N_exponent_in_kappa_at_prescribed_h": sreq,
        "r2": r2,
        "n_kappa_where_prescription_fails": int(
            sum(1 for r in rows if not r["prescription_holds"])
        ),
        "max_shortfall_factor": float(max(r["shortfall_factor"] or 0 for r in rows)),
    }
    print(
        scheme,
        "measured N exponent at prescribed h:",
        round(sreq, 3),
        "theorem:",
        Nexp,
        "failures:",
        res[scheme]["n_kappa_where_prescription_fails"],
        "/",
        len(rows),
        "max shortfall x",
        round(res[scheme]["max_shortfall_factor"], 2),
        flush=True,
    )

# ------------------------------------------------------------------
# The same statement expressed as simulated time N*h (constant-free).
# Thm 4.3 => N h ~ kappa^{1/2};  Thm 5.2 => N h ~ kappa^{1}.
# ------------------------------------------------------------------
for scheme in ("ulmc", "rmd"):
    rows = []
    for kappa in KAPPAS:
        a, m = family(kappa)
        S0, m0 = C.init_cold(a, BETA)
        r = C.sweep_neps(
            scheme,
            a,
            m,
            BETA,
            EPS,
            S0,
            m0,
            hs=C.hgrid(BETA, n=26, ratio=1.32),
            ngrid=NGRID,
        )
        rows.append(
            {
                "kappa": kappa,
                "N_eps": r["N_eps"],
                "h_star": r["h_star"],
                "Nh": r["N_eps"] * r["h_star"],
            }
        )
    s, r2 = C.fit_exponent([r["kappa"] for r in rows], [r["Nh"] for r in rows])
    sN, _ = C.fit_exponent([r["kappa"] for r in rows], [r["N_eps"] for r in rows])
    res[scheme + "_simulated_time"] = {
        "rows": rows,
        "Nh_exponent_in_kappa": s,
        "r2": r2,
        "N_exponent_in_kappa": sN,
        "theorem_Nh_exponent": 0.5 if scheme == "ulmc" else 1.0,
    }
    print(
        scheme,
        "min simulated time N*h exponent in kappa:",
        round(s, 3),
        "theorem:",
        res[scheme + "_simulated_time"]["theorem_Nh_exponent"],
        flush=True,
    )

C.dump("prescription", res)