scripts/test_sanity.py

"""Self-tests for the exact moment machinery (scripts/ulmc_core.py).

1. repeated-squaring propagation == naive step-by-step iteration
2. exact ULMC moments == Monte-Carlo moments of a direct particle simulation
3. exact RMD moments   == Monte-Carlo moments of a direct particle simulation
4. discretisation bias -> 0 as h -> 0, and is exactly 0 for a = 0
5. LMC / composite-LMC moments == Monte-Carlo
"""

import json
import sys

import numpy as np

sys.path.insert(0, "/home/ubuntu/samuel/ulmc-kl-repro/scripts")
import ulmc_core as U

SEED = 20260725
rng = np.random.default_rng(SEED)
res = {}
BETA = 1.0
G = np.sqrt(32 * BETA)


def naive(T, q, A, S0, m0, n):
    S, m = S0.copy(), m0.copy()
    for _ in range(n):
        S, m = U.compose_step(T, q, A, S, m)
    return S, m


# ---- 1. repeated squaring vs naive -----------------------------------------
a = np.array([1.0, 0.3, 0.01])
h = 0.05
Am, T, q = U.ulmc_maps(a, h, G)
S0 = np.tile(np.array([1.0, 0.0, 1.0]), (3, 1))
m0 = np.tile(np.array([0.7, 0.0]), (3, 1))
for n in (1, 7, 233, 4096):
    S1, m1 = U.propagate(T, q, Am, S0, m0, n)
    S2, m2 = naive(T, q, Am, S0, m0, n)
    err = max(np.abs(S1 - S2).max(), np.abs(m1 - m2).max())
    res[f"squaring_vs_naive_n{n}_maxabs"] = float(err)
    assert err < 1e-8, (n, err)


# ---- 2. ULMC exact moments vs particle Monte-Carlo --------------------------
def simulate_ulmc(a, h, g, n, npart, rng):
    x = np.full((npart,), 0.7)
    p = rng.standard_normal(npart)
    e = np.exp(-g * h)
    c1 = (1 - e) / g
    Gc = (h - c1) / g
    q11 = U.var_xi1(h, g)
    q12 = U.cov_xi1_xi2(h, h, g)
    q22 = U.var_xi2(h, g)
    L = np.linalg.cholesky(np.array([[q11, q12], [q12, q22]]))
    for _ in range(n):
        z = rng.standard_normal((npart, 2)) @ L.T
        xn = x + c1 * p - Gc * a * x + z[:, 0]
        pn = e * p - c1 * a * x + z[:, 1]
        x, p = xn, pn
    return x, p


npart, nstep = 400_000, 40
for aval in (1.0, 0.2):
    x, p = simulate_ulmc(aval, 0.08, G, nstep, npart, rng)
    S1, m1 = U.propagate(
        *U.ulmc_maps(np.array([aval]), 0.08, G)[1:],
        U.ulmc_maps(np.array([aval]), 0.08, G)[0],
        np.array([[0.7 * 0.7, 0.0, 1.0]]),
        np.array([[0.7, 0.0]]),
        nstep,
    )
    prod = np.stack([x * x, x * p, p * p], 1)
    mc = prod.mean(0)
    se = prod.std(0) / np.sqrt(npart)
    z = np.abs(S1[0] - mc) / se
    res[f"ulmc_mc_zscore_a{aval}"] = float(z.max())
    res[f"ulmc_mc_relerr_a{aval}"] = float((np.abs(S1[0] - mc) / np.abs(mc)).max())
    assert z.max() < 4.0, (aval, z, S1[0], mc)


# ---- 3. RMD exact moments vs particle Monte-Carlo ---------------------------
def sample_uv(h, g, n, rng):
    """Inverse-CDF sampling from the laws (3.3)."""
    t = np.linspace(0, 1, 20001)
    c1 = (1 - np.exp(-g * h)) / g
    pu = h * (1 - np.exp(-g * (1 - t) * h)) / (h - c1)
    pv = h * g * np.exp(-g * (1 - t) * h) / (1 - np.exp(-g * h))
    out = []
    for dens in (pu, pv):
        cdf = np.concatenate(
            [[0], np.cumsum(0.5 * (dens[1:] + dens[:-1]) * np.diff(t))]
        )
        cdf /= cdf[-1]
        out.append(np.interp(rng.random(n), cdf, t))
    return out


def simulate_rmd(a, h, g, n, npart, rng):
    x = np.full((npart,), 0.7)
    p = rng.standard_normal(npart)
    e = np.exp(-g * h)
    c1 = (1 - e) / g
    Gc = (h - c1) / g
    for _ in range(n):
        u, v = sample_uv(h, g, npart, rng)
        su, sv = u * h, v * h
        # joint Gaussian (xi1(su), xi1(sv), xi1(h), xi2(h)) per particle
        Vs = np.empty((npart, 4, 4))
        s = np.stack([su, sv, np.full(npart, h)], 1)
        for i in range(3):
            for j in range(3):
                Vs[:, i, j] = U.cov_xi1_xi1(s[:, i], s[:, j], g)
        for i in range(3):
            Vs[:, i, 3] = Vs[:, 3, i] = U.cov_xi1_xi2(s[:, i], h, g)
        Vs[:, 3, 3] = U.var_xi2(h, g)
        Vs += 1e-14 * np.eye(4)
        L = np.linalg.cholesky(Vs)
        z = np.einsum("nij,nj->ni", L, rng.standard_normal((npart, 4)))
        gu = (su - (1 - np.exp(-g * su)) / g) / g
        gv = (sv - (1 - np.exp(-g * sv)) / g) / g
        cu = (1 - np.exp(-g * su)) / g
        cv = (1 - np.exp(-g * sv)) / g
        xhat1 = (1 - a * gu) * x + cu * p + z[:, 0]
        xhat2 = (1 - a * gv) * x + cv * p + z[:, 1]
        xn = x + c1 * p + z[:, 2] - Gc * a * xhat1
        pn = e * p + z[:, 3] - c1 * a * xhat2
        x, p = xn, pn
    return x, p


npart_r, nstep_r = 200_000, 12
for aval in (1.0, 0.2):
    x, p = simulate_rmd(aval, 0.08, G, nstep_r, npart_r, rng)
    A_, T_, q_ = U.rmd_maps(np.array([aval]), 0.08, G)
    S1, m1 = U.propagate(
        T_, q_, A_, np.array([[0.49, 0.0, 1.0]]), np.array([[0.7, 0.0]]), nstep_r
    )
    prod = np.stack([x * x, x * p, p * p], 1)
    mc = prod.mean(0)
    se = prod.std(0) / np.sqrt(npart_r)
    z = np.abs(S1[0] - mc) / se
    res[f"rmd_mc_zscore_a{aval}"] = float(z.max())
    res[f"rmd_mc_relerr_a{aval}"] = float((np.abs(S1[0] - mc) / np.abs(mc)).max())
    assert z.max() < 4.0, (aval, z, S1[0], mc)

# ---- 4. bias -> 0 as h -> 0; exactly 0 for a = 0 ----------------------------
a1 = np.array([1.0])
mult = np.array([1.0])
floors = {}
for h in (0.16, 0.08, 0.04, 0.02, 0.01):
    mp = U.ulmc_maps(a1, h, G)
    fl = U.kl_floor(
        (mp[0], mp[1], mp[2]), np.array([[1.0, 0.0, 1.0]]), np.zeros((1, 2)), a1, mult
    )
    floors[h] = fl
res["ulmc_floor_vs_h"] = {str(k): float(v) for k, v in floors.items()}
hs = np.array(sorted(floors))
fv = np.array([floors[h] for h in hs])
slope = np.polyfit(np.log(hs), np.log(fv), 1)[0]
res["ulmc_floor_h_exponent"] = float(slope)

atiny = np.array([1e-12])
mp = U.ulmc_maps(atiny, 0.1, G)
fl0 = U.kl_at(
    50,
    (mp[0], mp[1], mp[2]),
    np.array([[1.0 / atiny[0], 0.0, 1.0]]),
    np.zeros((1, 2)),
    atiny,
    mult,
)
res["ulmc_floor_a0"] = float(fl0)


# ---- 5. LMC + composite moments vs Monte-Carlo ------------------------------
def simulate_lmc(a, h, n, npart, rng):
    x = np.full((npart,), 0.7)
    for _ in range(n):
        x = (1 - h * a) * x + np.sqrt(2 * h) * rng.standard_normal(npart)
    return x


for aval in (1.0, 0.2):
    x = simulate_lmc(aval, 0.05, 60, 400_000, rng)
    A_, T_, q_ = U.lmc_maps(np.array([aval]), 0.05)
    S1, _ = U.propagate(
        T_, q_, A_, np.array([[0.49, 0.0, 1.0]]), np.array([[0.7, 0.0]]), 60
    )
    rel = abs(S1[0, 0] - np.mean(x * x)) / np.mean(x * x)
    res[f"lmc_mc_relerr_a{aval}"] = float(rel)
    assert rel < 0.02, (aval, rel)


def simulate_comp(a, h, alpha, n, npart, rng):
    x = np.full((npart,), 0.7)
    ea = np.exp(-alpha * h)
    c = (1 - ea) / alpha
    sd = np.sqrt((1 - np.exp(-2 * alpha * h)) / alpha)
    for _ in range(n):
        x = ea * x - c * (a - alpha) * x + sd * rng.standard_normal(npart)
    return x


x = simulate_comp(1.0, 0.05, 0.05, 60, 400_000, rng)
A_, T_, q_ = U.composite_lmc_maps(np.array([1.0]), 0.05, 0.05)
S1, _ = U.propagate(
    T_, q_, A_, np.array([[0.49, 0.0, 1.0]]), np.array([[0.7, 0.0]]), 60
)
rel = abs(S1[0, 0] - np.mean(x * x)) / np.mean(x * x)
res["composite_mc_relerr"] = float(rel)
assert rel < 0.02, rel

res["seed"] = SEED
res["all_passed"] = True
print(json.dumps(res, indent=1))
with open("/home/ubuntu/samuel/ulmc-kl-repro/outputs/test_sanity.json", "w") as f:
    json.dump(res, f, indent=1)