scripts/ulmc_core.py

"""
Exact Gaussian-moment machinery for ULMC / RMD / LMC / composite-LMC on
diagonal quadratic targets.

Target:  V(x) = 1/2 sum_i a_i x_i^2 ,  a_i in [alpha, beta].
         => grad^2 V = A = diag(a),  tightest Hessian upper bound H = A,
            tr(H) = sum_i a_i,  kappa = beta/alpha.
         pi(x,p) propto exp(-V(x) - ||p||^2/2)  =>  per coordinate
            pi_i = N(0, diag(1/a_i, 1)).

Because V is quadratic and diagonal, EVERY discretization below is a linear
Gaussian recursion that factorises across coordinates.  So if mu_0 factorises
and is Gaussian, mu_N is Gaussian (ULMC/LMC/composite) and
KL(mu_N || pi) = sum_i KL_2d(mu_N^i || pi_i) is available in CLOSED FORM --
no Monte-Carlo noise anywhere.

For RMD the per-step midpoints u_n, v_n are random, so mu_N is a *mixture* of
Gaussians.  We propagate the exact first and second moments (exact, by
integrating the linear recursion over the laws of u,v with Gauss-Legendre
quadrature) and report
    KL_G := KL( N(m_N, Sigma_N) || pi ),
which by the maximum-entropy property of the Gaussian is a CERTIFIED LOWER
BOUND on the true KL(mu_N||pi).  `rmd_frozen_upper` gives the matching
convexity upper bound  E_{u,v-seq}[ KL(component || pi) ] >= KL(mu_N||pi).

n-step propagation uses repeated squaring of the augmented affine map, so N up
to 1e10 costs O(log N).

All quantities are exact to float64; no seeds are needed for the exact paths.
"""

import numpy as np

# ----------------------------------------------------------------------------
# Brownian functionals of one ULD step  (per coordinate, scalar BM on [0,h])
#   xi1(s) = sqrt(2g) int_0^s (1-e^{-g(s-u)})/g dB_u
#   xi2(h) = sqrt(2g) int_0^h e^{-g(h-u)} dB_u
# ----------------------------------------------------------------------------


def _em1(x):
    """1 - e^{-x}, stable for small x."""
    return -np.expm1(-np.asarray(x, dtype=float))


def var_xi1(s, g):
    """Var(xi1(s)) = (2/g^2) * G(g s),  G(z) = z - 3/2 + 2e^{-z} - e^{-2z}/2."""
    z = np.asarray(g * s, dtype=float)
    big = z > 5e-2
    G = np.empty_like(z)
    zb = z[big] if z.ndim else z
    if z.ndim == 0:
        if big:
            G = z - 1.5 + 2 * np.exp(-z) - 0.5 * np.exp(-2 * z)
        else:
            G = z**3 / 3 - z**4 / 4 + 7 * z**5 / 60 - z**6 / 24
        return 2.0 * G / g**2
    G[big] = z[big] - 1.5 + 2 * np.exp(-z[big]) - 0.5 * np.exp(-2 * z[big])
    zs = z[~big]
    G[~big] = zs**3 / 3 - zs**4 / 4 + 7 * zs**5 / 60 - zs**6 / 24
    return 2.0 * G / g**2


def cov_xi1_xi1(s, t, g):
    """Cov(xi1(s), xi1(t)) for arbitrary s,t>0 (symmetric)."""
    s = np.asarray(s, dtype=float)
    t = np.asarray(t, dtype=float)
    lo = np.minimum(s, t)
    hi = np.maximum(s, t)
    tau = hi - lo
    e = np.exp(-g * tau)
    return (2.0 / g) * (
        (g * lo - _em1(g * lo)) / g
        - e * _em1(g * lo) / g
        + e * _em1(2 * g * lo) / (2 * g)
    )


def cov_xi1_xi2(s, h, g):
    """Cov(xi1(s), xi2(h)) for 0 < s <= h."""
    s = np.asarray(s, dtype=float)
    tp = h - s
    return 2.0 * np.exp(-g * tp) * (_em1(g * s) / g - _em1(2 * g * s) / (2 * g))


def var_xi2(h, g):
    return _em1(2 * g * h)


# ----------------------------------------------------------------------------
# vec(S) = (S11, S12, S22) helpers
# ----------------------------------------------------------------------------


def T_from_M(M):
    """3x3 matrix implementing vec(S) -> vec(M S M^T).  M: (...,2,2)."""
    m11, m12 = M[..., 0, 0], M[..., 0, 1]
    m21, m22 = M[..., 1, 0], M[..., 1, 1]
    T = np.empty(M.shape[:-2] + (3, 3))
    T[..., 0, 0] = m11 * m11
    T[..., 0, 1] = 2 * m11 * m12
    T[..., 0, 2] = m12 * m12
    T[..., 1, 0] = m11 * m21
    T[..., 1, 1] = m11 * m22 + m12 * m21
    T[..., 1, 2] = m12 * m22
    T[..., 2, 0] = m21 * m21
    T[..., 2, 1] = 2 * m21 * m22
    T[..., 2, 2] = m22 * m22
    return T


def _augment(T, q):
    """(...,3,3),(...,3) -> (...,4,4) affine map [[T,q],[0,1]]."""
    sh = T.shape[:-2]
    Aug = np.zeros(sh + (4, 4))
    Aug[..., :3, :3] = T
    Aug[..., :3, 3] = q
    Aug[..., 3, 3] = 1.0
    return Aug


def propagate(T, q, A, S0, m0, n):
    """Apply the affine second-moment map n times (repeated squaring).

    T,q : second-moment map  S -> T S + q      (vec form)
    A   : mean map           m -> A m
    returns S_n (…,3), m_n (…,2)
    """
    n = int(n)
    Aug = _augment(T, q)
    Augn = np.linalg.matrix_power(Aug, n)
    S0a = np.concatenate([S0, np.ones(S0.shape[:-1] + (1,))], axis=-1)
    Sn = np.einsum("...ij,...j->...i", Augn, S0a)[..., :3]
    An = np.linalg.matrix_power(A, n)
    mn = np.einsum("...ij,...j->...i", An, m0)
    return Sn, mn


def compose_step(T, q, A, S, m):
    """One extra step of the (T,q,A) map."""
    Sn = np.einsum("...ij,...j->...i", T, S) + q
    mn = np.einsum("...ij,...j->...i", A, m)
    return Sn, mn


# ----------------------------------------------------------------------------
# KL of the phase-space Gaussian against pi
# ----------------------------------------------------------------------------


def kl_per_coord(S, m, a):
    """KL( N(m,Sigma) || N(0, diag(1/a,1)) ) with S = Sigma + m m^T (vec form).

    KL = 1/2 [ a*S11 + S22 - 2 - log(a) - log det Sigma ]
    """
    S11, S12, S22 = S[..., 0], S[..., 1], S[..., 2]
    m1, m2 = m[..., 0], m[..., 1]
    c11 = S11 - m1 * m1
    c12 = S12 - m1 * m2
    c22 = S22 - m2 * m2
    det = c11 * c22 - c12 * c12
    det = np.maximum(det, 1e-300)
    return 0.5 * (a * S11 + S22 - 2.0 - np.log(a) - np.log(det))


# ----------------------------------------------------------------------------
# ULMC  (Eq. 3.2)
# ----------------------------------------------------------------------------


def ulmc_maps(a, h, g):
    """Return A (mean map), T, q for standard ULMC on coordinate stiffness a."""
    a = np.asarray(a, dtype=float)
    e = np.exp(-g * h)
    c1 = _em1(g * h) / g  # (1-e^{-gh})/g
    G = (h - c1) / g  # (1/g)(h - (1-e^{-gh})/g)
    M = np.empty(a.shape + (2, 2))
    M[..., 0, 0] = 1.0 - a * G
    M[..., 0, 1] = c1
    M[..., 1, 0] = -a * c1
    M[..., 1, 1] = e
    q = np.empty(a.shape + (3,))
    q[..., 0] = var_xi1(h, g)
    q[..., 1] = cov_xi1_xi2(h, h, g)
    q[..., 2] = var_xi2(h, g)
    return M, T_from_M(M), q


# ----------------------------------------------------------------------------
# Randomized midpoint discretization (Eq. 3.4), doubly randomized
# ----------------------------------------------------------------------------


def _rmd_nodes(h, g, nq=48):
    """Gauss-Legendre nodes/weights for the laws (3.3) of u and v on [0,1]."""
    x, w = np.polynomial.legendre.leggauss(nq)
    t = 0.5 * (x + 1.0)
    w = 0.5 * w
    c1 = _em1(g * h) / g
    pu = h * _em1(g * (1 - t) * h) / (h - c1)
    pv = h * g * np.exp(-g * (1 - t) * h) / _em1(g * h)
    wu = w * pu
    wv = w * pv
    return t, wu / wu.sum(), wv / wv.sum()


def rmd_maps(a, h, g, nq=48, literal_typo=False):
    """Return A (mean map), T, q for the doubly-randomized midpoint scheme.

    literal_typo=True reproduces the coefficient (1-e^{-gh})/g printed in front
    of P in the two midpoint lines of Eq. (3.4); the surrounding text says the
    midpoints are obtained by *standard ULMC of step u_n h / v_n h*, which gives
    (1-e^{-g u_n h})/g.  Default False = the text-consistent version.
    """
    a = np.asarray(a, dtype=float)[..., None]  # broadcast over quad nodes
    t, wu, wv = _rmd_nodes(h, g, nq)
    e = np.exp(-g * h)
    c1 = _em1(g * h) / g
    G = (h - c1) / g

    su, sv = t * h, t * h
    gu = (su - _em1(g * su) / g) / g
    gv = (sv - _em1(g * sv) / g) / g
    cu = c1 * np.ones_like(su) if literal_typo else _em1(g * su) / g
    cv = c1 * np.ones_like(sv) if literal_typo else _em1(g * sv) / g

    # X' = [1 - aG(1 - a gu)] X + [c1 - aG cu] P + xi1(h) - aG xi1(u h)
    # P' = [-a c1 (1 - a gv)] X + [e - a c1 cv] P + xi2(h) - a c1 xi1(v h)
    r11 = 1.0 - a * G * (1.0 - a * gu)
    r12 = c1 - a * G * cu
    r21 = -a * c1 * (1.0 - a * gv)
    r22 = e - a * c1 * cv

    E11 = (r11 * wu).sum(-1)
    E12 = (r12 * wu).sum(-1)
    E21 = (r21 * wv).sum(-1)
    E22 = (r22 * wv).sum(-1)
    sh = a.shape[:-1]
    A = np.empty(sh + (2, 2))
    A[..., 0, 0], A[..., 0, 1] = E11, E12
    A[..., 1, 0], A[..., 1, 1] = E21, E22

    # E[M S M^T]: row1 depends on u only, row2 on v only, u indep v.
    T = np.empty(sh + (3, 3))
    T[..., 0, 0] = (r11 * r11 * wu).sum(-1)
    T[..., 0, 1] = 2 * (r11 * r12 * wu).sum(-1)
    T[..., 0, 2] = (r12 * r12 * wu).sum(-1)
    T[..., 1, 0] = E11 * E21
    T[..., 1, 1] = E11 * E22 + E12 * E21
    T[..., 1, 2] = E12 * E22
    T[..., 2, 0] = (r21 * r21 * wv).sum(-1)
    T[..., 2, 1] = 2 * (r21 * r22 * wv).sum(-1)
    T[..., 2, 2] = (r22 * r22 * wv).sum(-1)

    # noise covariance, averaged over u,v
    V1h = var_xi1(h, g)
    V2h = var_xi2(h, g)
    C12h = cov_xi1_xi2(h, h, g)
    Cu = cov_xi1_xi1(su, h, g)  # Cov(xi1(uh), xi1(h))
    Cv = cov_xi1_xi1(sv, h, g)
    V1u = var_xi1(su, g)
    V1v = var_xi1(sv, g)
    C2u = cov_xi1_xi2(su, h, g)  # Cov(xi1(uh), xi2(h))
    C2v = cov_xi1_xi2(sv, h, g)
    aG = a * G
    ac1 = a * c1
    q = np.empty(sh + (3,))
    q[..., 0] = ((V1h - 2 * aG * Cu + aG**2 * V1u) * wu).sum(-1)
    q[..., 2] = ((V2h - 2 * ac1 * C2v + ac1**2 * V1v) * wv).sum(-1)
    # cross term needs E_{u,v}[Cov(xi1(uh), xi1(vh))]
    Cuv = cov_xi1_xi1(su[:, None], sv[None, :], g)  # (nq,nq)
    Ecuv = wu @ Cuv @ wv
    q[..., 1] = (
        C12h
        - (ac1[..., 0]) * (Cv * wv).sum(-1)
        - (aG[..., 0]) * (C2u * wu).sum(-1)
        + (aG[..., 0]) * (ac1[..., 0]) * Ecuv
    )
    return A, T, q


# ----------------------------------------------------------------------------
# Overdamped baselines
# ----------------------------------------------------------------------------


def lmc_maps(a, h):
    """Euler-Maruyama LMC: x' = (1-h a) x + sqrt(2h) N(0,1). 1-d state, but we
    embed in the same 2-d (x,p) formalism with a dummy p ~ N(0,1) exactly at
    stationarity, so KL is comparable."""
    a = np.asarray(a, dtype=float)
    M = np.zeros(a.shape + (2, 2))
    M[..., 0, 0] = 1.0 - h * a
    M[..., 1, 1] = 0.0  # p redrawn from N(0,1) each step
    q = np.zeros(a.shape + (3,))
    q[..., 0] = 2.0 * h
    q[..., 2] = 1.0
    return M, T_from_M(M), q


def composite_lmc_maps(a, h, alpha):
    """Composite / exponential-integrator LMC in the sense of Freund et al.
    (2022): V = (alpha/2)||x||^2 + R(x) with grad R = (A - alpha I)x; the
    strongly-convex quadratic part is integrated EXACTLY (OU), R explicitly.

        x' = e^{-alpha h} x - ((1-e^{-alpha h})/alpha) (a-alpha) x
             + N(0, (1-e^{-2 alpha h})/alpha)
    """
    a = np.asarray(a, dtype=float)
    ea = np.exp(-alpha * h)
    c = _em1(alpha * h) / alpha
    M = np.zeros(a.shape + (2, 2))
    M[..., 0, 0] = ea - c * (a - alpha)
    M[..., 1, 1] = 0.0
    q = np.zeros(a.shape + (3,))
    q[..., 0] = _em1(2 * alpha * h) / alpha
    q[..., 2] = 1.0
    return M, T_from_M(M), q


# ----------------------------------------------------------------------------
# driver: KL(mu_N || pi) as a function of N
# ----------------------------------------------------------------------------


def kl_at(N, maps, S0, m0, a, mult, final_ulmc=None):
    """KL after N steps of `maps`=(A,T,q) (plus one final ULMC step if given)."""
    A, T, q = maps
    if final_ulmc is not None:
        if N < 1:
            return np.inf
        S, m = propagate(T, q, A, S0, m0, N - 1)
        S, m = compose_step(final_ulmc[1], final_ulmc[2], final_ulmc[0], S, m)
    else:
        S, m = propagate(T, q, A, S0, m0, N)
    return float(np.sum(mult * kl_per_coord(S, m, a)))


def n_grid(nmax=int(3e9), ratio=1.03):
    ns, n = [1], 1
    while n < nmax:
        nn = max(n + 1, int(round(n * ratio)))
        ns.append(nn)
        n = nn
    return np.array(ns, dtype=np.int64)


def iters_to_eps(
    maps, S0, m0, a, mult, eps2, nmax=int(3e9), ratio=1.03, final_ulmc=None
):
    """Smallest N on a geometric grid with KL(mu_N||pi) <= eps2; np.inf if none.
    Also returns the KL floor (KL at the largest grid point)."""
    grid = n_grid(nmax, ratio)
    lo, hi = 0, len(grid) - 1
    floor = kl_at(int(grid[hi]), maps, S0, m0, a, mult, final_ulmc)
    if not np.isfinite(floor) or floor > eps2:
        return np.inf, floor
    # KL decreases then plateaus; bisect for first grid index under eps2
    while lo < hi:
        mid = (lo + hi) // 2
        v = kl_at(int(grid[mid]), maps, S0, m0, a, mult, final_ulmc)
        if np.isfinite(v) and v <= eps2:
            hi = mid
        else:
            lo = mid + 1
    return int(grid[lo]), floor


def kl_floor(maps, S0, m0, a, mult, nbig=int(2e9), final_ulmc=None):
    return kl_at(nbig, maps, S0, m0, a, mult, final_ulmc)


def best_over_h(
    hs, make_maps, S0f, m0f, a, mult, eps2, final_ulmc_f=None, nmax=int(3e9)
):
    """min over step sizes h of the number of steps to reach KL <= eps2.
    Returns (Nbest, hbest, table)."""
    table = []
    best = (np.inf, None)
    for h in hs:
        maps = make_maps(h)
        fin = final_ulmc_f(h) if final_ulmc_f is not None else None
        N, fl = iters_to_eps(
            maps, S0f(h), m0f(h), a, mult, eps2, nmax=nmax, final_ulmc=fin
        )
        table.append((float(h), float(N) if np.isfinite(N) else None, float(fl)))
        if N < best[0]:
            best = (N, h)
    return best[0], best[1], table


def kl_curve(maps, S0, m0, a, mult, grid, final_ulmc=None):
    """KL(mu_n||pi) for every n in `grid` (ascending ints). O(len(grid)*log N)."""
    A, T, q = maps
    Aug = _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))], axis=-1)
    out = np.empty(len(grid))
    prev = 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]
        m = np.einsum('kij,kj->ki', accM, m0)
        if final_ulmc is not None:
            S, m = compose_step(final_ulmc[1], final_ulmc[2], final_ulmc[0], S, m)
        out[i] = float(np.sum(mult * kl_per_coord(S, m, a)))
    return out


def first_below(grid, vals, thr):
    idx = np.nonzero(np.isfinite(vals) & (vals <= thr))[0]
    return int(grid[idx[0]]) if len(idx) else None


def n_eps_over_h(hs, make_maps, S0, m0, a, mult, eps2, grid,
                 final_ulmc_f=None):
    """min over h of #steps to reach KL<=eps2.  Returns dict."""
    best_N, best_h, rows = None, None, []
    for h in hs:
        mp = make_maps(h)
        fin = final_ulmc_f(h) if final_ulmc_f is not None else None
        vals = kl_curve(mp, S0, m0, a, mult, grid, final_ulmc=fin)
        N = first_below(grid, vals, eps2)
        rows.append({"h": float(h), "N": N, "floor": float(vals[-1])})
        if N is not None and (best_N is None or N < best_N):
            best_N, best_h = N, float(h)
    return {"N_eps": best_N, "h_star": best_h, "table": rows}