scripts/exp_bias_law.py

"""Claim 1 & 2 & 3 foundation: the exact stationary KL bias of ULMC and RMD.

We measure the *exact* asymptotic (n -> inf) KL divergence between the law of
the discretization and pi on a 1-coordinate Gaussian target V = a x^2/2, as a
function of (a, h, beta), and fit the power law.  This is the object that the
paper's Lemma 4.1 / 5.1 + change-of-measure controls; it is what fixes the
admissible step size h and hence the whole complexity.

Dimension-freeness is a direct corollary: the total bias is
    sum_i F(a_i, h, beta),
so if F(a,h,beta) = C a^s h^b / beta^r with s >= 1, the total is bounded by a
function of tr(H) = sum_i a_i with NO explicit d.
"""

import numpy as np

import common as C
import ulmc_core as U

res = {"description": "exact stationary KL bias of ULMC / RMD on Gaussian targets"}

NBIG = int(2e9)


def floor_of(scheme, a, h, beta):
    g = C.gamma_of(beta)
    av = np.array([a])
    mult = np.array([1.0])
    S0 = np.array([[1.0 / a, 0.0, 1.0]])
    m0 = np.zeros((1, 2))
    if scheme == "ulmc":
        mp = U.ulmc_maps(av, h, g)
        return U.kl_at(NBIG, mp, S0, m0, av, mult)
    mp = U.rmd_maps(av, h, g, nq=64)
    fin = U.ulmc_maps(av, h, g)
    return U.kl_at(NBIG, mp, S0, m0, av, mult, final_ulmc=fin)


for scheme in ("ulmc", "rmd"):
    blk = {}
    # exponent in a (h, beta fixed)
    for beta in (1.0, 4.0):
        avs = np.geomspace(0.01 * beta, 1.0 * beta, 9)
        for h in (0.02 / np.sqrt(beta), 0.01 / np.sqrt(beta)):
            F = np.array([floor_of(scheme, a, h, beta) for a in avs])
            s, r2 = C.fit_exponent(avs, F)
            blk[f"a_exponent_beta{beta}_h{h:.5f}"] = {"exponent": s, "r2": r2}
    # exponent in h (a = beta, i.e. the stiffest direction)
    for beta in (0.5, 1.0, 4.0):
        hs = np.geomspace(0.002, 0.04, 9) / np.sqrt(beta)
        F = np.array([floor_of(scheme, beta, h, beta) for h in hs])
        s, r2 = C.fit_exponent(hs, F)
        blk[f"h_exponent_beta{beta}"] = {"exponent": s, "r2": r2}
        # asymptotic slope from the two smallest h
        blk[f"h_exponent_beta{beta}_asymptotic"] = float(
            np.log(F[1] / F[0]) / np.log(hs[1] / hs[0])
        )
    # exponent in beta at fixed a/beta = 1 and fixed gamma*h
    betas = np.geomspace(0.25, 16.0, 7)
    F = np.array([floor_of(scheme, b, 0.02 / np.sqrt(b), b) for b in betas])
    s, r2 = C.fit_exponent(betas, F)
    blk["beta_exponent_at_fixed_gamma_h_and_a_over_beta"] = {"exponent": s, "r2": r2}
    # calibrated constant
    if scheme == "ulmc":
        blk["constant_C_in_F=C*a^2*h^2/beta"] = [
            float(floor_of("ulmc", a, h, b) * b / (a * a * h * h))
            for (a, h, b) in [(1.0, 0.01, 1.0), (0.3, 0.005, 1.0), (2.0, 0.004, 4.0)]
        ]
        blk["one_over_256"] = 1.0 / 256
    else:
        blk["constant_C_in_F=C*a^2*h^4"] = [
            float(floor_of("rmd", a, h, b) / (a * a * h**4))
            for (a, h, b) in [(1.0, 0.0025, 1.0), (0.3, 0.002, 1.0), (2.0, 0.001, 4.0)]
        ]
    res[scheme] = blk

# ---------------------------------------------------------------------------
# The dimension-free inequality:  total bias <= h^2 tr(H) / 256   for ULMC.
# Sample 200 random spectra with wildly different d at matched tr(H).
# ---------------------------------------------------------------------------
rng = np.random.default_rng(C.SEED)
rows = []
for _ in range(200):
    d = int(rng.integers(3, 400))
    beta = float(np.exp(rng.uniform(-1, 1.5)))
    a = np.sort(np.exp(rng.uniform(np.log(1e-3 * beta), 0.0, d)) * beta)[::-1]
    a[0] = beta
    trH = float(a.sum())
    h = float(np.exp(rng.uniform(np.log(1e-3), np.log(0.05))) / np.sqrt(beta))
    g = C.gamma_of(beta)
    mp = U.ulmc_maps(a, h, g)
    S0 = np.stack([1.0 / a, np.zeros(d), np.ones(d)], 1)
    tot = U.kl_at(NBIG, mp, S0, np.zeros((d, 2)), a, np.ones(d))
    rows.append(
        {
            "d": d,
            "beta": beta,
            "trH": trH,
            "h": h,
            "bias": float(tot),
            "bound_h2trH_over_256": float(h * h * trH / 256),
            "ratio": float(tot / (h * h * trH / 256)),
            "d_dependent_surrogate_h2_beta_d_over_256": float(h * h * beta * d / 256),
        }
    )
ratios = np.array([r["ratio"] for r in rows])
res["dimension_free_bias_bound"] = {
    "n_random_spectra": len(rows),
    "max_ratio_bias_over_h2trH_div256": float(ratios.max()),
    "min_ratio": float(ratios.min()),
    "bound_never_violated": bool(ratios.max() <= 1.0 + 1e-9),
    "seed": C.SEED,
    "examples": rows[:10],
}

C.dump("bias_law", res)
print(json.dumps(res, indent=1)[:4000] if False else "")
for k in ("ulmc", "rmd"):
    print(k, {kk: vv for kk, vv in res[k].items()})
print(
    "dim-free bound holds:",
    res["dimension_free_bias_bound"]["bound_never_violated"],
    "max ratio",
    res["dimension_free_bias_bound"]["max_ratio_bias_over_h2trH_div256"],
)