scripts/common.py

"""Shared helpers: target families, sweeps, exponent fitting."""

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
OUT = "/home/ubuntu/samuel/ulmc-kl-repro/outputs"


def gamma_of(beta):
    """The friction mandated by Theorems 4.3/4.4/5.2/5.4."""
    return np.sqrt(32.0 * beta)


def spectrum(d, alpha, beta, trH):
    """Diagonal Hessian with exactly d coordinates, min=alpha, max=beta,
    tr(H)=trH.  Returns (distinct a values, multiplicities).  Raises if the
    triple (d, alpha, beta, trH) is infeasible."""
    if d < 2:
        raise ValueError("d>=2")
    c = (trH - alpha - beta) / (d - 2)
    if not (alpha - 1e-12 <= c <= beta + 1e-12):
        raise ValueError(
            f"infeasible: c={c} not in [{alpha},{beta}] (d={d}, trH={trH})"
        )
    a = np.array([beta, alpha, c])
    m = np.array([1.0, 1.0, float(d - 2)])
    return a, m


def init_cold(a, beta):
    """mu_0 = N(0, beta^{-1} I) x N(0, I)  (standard 'start at the mode' warm
    start used throughout the Langevin literature)."""
    K = len(a)
    S0 = np.tile(np.array([1.0 / beta, 0.0, 1.0]), (K, 1))
    m0 = np.zeros((K, 2))
    return S0, m0


def init_shifted(a, w, idx=0):
    """mu_0 = N(m0, Sigma_pi^x) x N(0,I) with the mean displaced by w in
    coordinate `idx`.  KL(mu_0||pi) = 0.5 * a_idx * w^2  (d-independent), and
    W_2(mu_0,pi) = |w|."""
    K = len(a)
    m0 = np.zeros((K, 2))
    m0[idx, 0] = w
    S0 = np.stack([1.0 / a + m0[:, 0] ** 2, np.zeros(K), np.ones(K)], axis=1)
    return S0, m0


def hgrid(beta, n=32, ratio=1.28):
    g = gamma_of(beta)
    hmax = 1.0 / g  # Thm 3.5 requires h <~ 1/gamma ^ gamma/beta
    return hmax / ratio ** np.arange(n)


NGRID = U.n_grid(int(4e9), 1.02)


def sweep_neps(scheme, a, mult, beta, eps, S0, m0, hs=None, ngrid=None):
    """min over h of the number of steps to reach KL <= eps^2."""
    g = gamma_of(beta)
    hs = hgrid(beta) if hs is None else hs
    ngrid = NGRID if ngrid is None else ngrid
    if scheme == "ulmc":
        mk = lambda h: U.ulmc_maps(a, h, g)
        fin = None
    elif scheme == "rmd":
        mk = lambda h: U.rmd_maps(a, h, g, nq=48)
        fin = lambda h: U.ulmc_maps(a, h, g)
    elif scheme == "lmc":
        mk = lambda h: U.lmc_maps(a, h)
        fin = None
    elif scheme == "composite":
        alpha = float(np.min(a))
        mk = lambda h: U.composite_lmc_maps(a, h, alpha)
        fin = None
    else:
        raise ValueError(scheme)
    return U.n_eps_over_h(hs, mk, S0, m0, a, mult, eps * eps, ngrid, final_ulmc_f=fin)


def fit_exponent(xs, ys):
    """log-log slope + R^2."""
    xs, ys = np.asarray(xs, float), np.asarray(ys, float)
    ok = np.isfinite(xs) & np.isfinite(ys) & (ys > 0)
    lx, ly = np.log(xs[ok]), np.log(ys[ok])
    if len(lx) < 2:
        return None, None
    s, b = np.polyfit(lx, ly, 1)
    pred = s * lx + b
    ss = 1 - np.sum((ly - pred) ** 2) / max(np.sum((ly - ly.mean()) ** 2), 1e-300)
    return float(s), float(ss)


def dump(name, obj):
    p = f"{OUT}/{name}.json"
    with open(p, "w") as f:
        json.dump(obj, f, indent=1, default=float)
    print("wrote", p)