Factor the transition operator once on uniform grids

Author: junipertcy

README.md

@@ -9,7 +9,7 @@ To collect samples from this process, define sampling times `ts`, initial state
 ```python
 thermox.sample(key, ts, x0, A, b, D) 
 ```
-Samples are then collected by exact diagonalization (therefore there is no discretization error) and JAX scans. This holds for any stable drift matrix `A` (not necessarily symmetric or normal) and any positive definite diffusion matrix `D`; when `D^{-1/2} A D^{1/2}` is a normal matrix a cheaper O(d^2)-per-step formula is used automatically.
+Samples are then collected by exact diagonalization (therefore there is no discretization error) and JAX scans. This holds for any stable drift matrix `A` (not necessarily symmetric or normal) and any positive definite diffusion matrix `D`; when `D^{-1/2} A D^{1/2}` is a normal matrix, or the time grid is uniform, sampling stays O(d^2) per step.
 
 You can access log-probabilities of the OU process by running `thermox.log_prob`:
 

tests/test_nonnormal.py

@@ -8,6 +8,7 @@
 import pytest
 
 import thermox
+from thermox.sampler import _scan_linear_recurrence, uniform_dt
 from thermox.utils import preprocess_drift_matrix
 
 jax.config.update("jax_enable_x64", True)
@@ -98,21 +99,29 @@ def transition_logpdf(x1, x0, dt):
     )
 
 
+GRIDS = [
+    pytest.param(
+        jnp.array([0.0, 0.1, 0.5, 0.6, 1.4]), id="non-uniform"
+    ),  # per-step path
+    pytest.param(jnp.arange(0.0, 1.5, 0.1), id="uniform"),  # factor-once path
+]
+
+
+@pytest.mark.parametrize("ts", GRIDS)
 @pytest.mark.parametrize("A,D", NONNORMAL_CASES)
-def test_log_prob_matches_reference_gaussian(A, D):
+def test_log_prob_matches_reference_gaussian(A, D, ts):
     d = A.shape[0]
     b = jnp.arange(1.0, d + 1.0)
-    ts = jnp.array([0.0, 0.1, 0.5, 0.6, 1.4])
     xs = jax.random.normal(jax.random.PRNGKey(3), (len(ts), d))
     ref = reference_log_prob(ts, xs, A, b, D)
     assert jnp.isclose(thermox.log_prob(ts, xs, A, b, D), ref, rtol=1e-8)
 
 
+@pytest.mark.parametrize("ts", GRIDS)
 @pytest.mark.parametrize("A,D", NONNORMAL_CASES)
-def test_log_prob_grad_wrt_drift_matches_reference(A, D):
+def test_log_prob_grad_wrt_drift_matches_reference(A, D, ts):
     d = A.shape[0]
     b = jnp.arange(1.0, d + 1.0)
-    ts = jnp.array([0.0, 0.1, 0.5, 0.6, 1.4])
     xs = jax.random.normal(jax.random.PRNGKey(3), (len(ts), d))
     g = jax.grad(lambda A: thermox.log_prob(ts, xs, A, b, D))(A)
     g_ref = jax.grad(lambda A: reference_log_prob(ts, xs, A, b, D))(A)
@@ -134,9 +143,11 @@ def test_log_prob_grad_symmetric_parametrization_matches_reference():
     assert relerr(g, g_ref) < 1e-8
 
 
-def test_sample_covariance_matches_lyapunov_for_anisotropic_noise():
+@pytest.mark.parametrize("jitter", [0.0, 0.5], ids=["uniform-grid", "jittered-grid"])
+def test_sample_covariance_matches_lyapunov_for_anisotropic_noise(jitter):
     A, D = A_SYM, D_DIAG
     ts = jnp.arange(0.0, 20000.0, 0.5)
+    ts = jnp.sort(ts + jitter * jax.random.uniform(jax.random.PRNGKey(1), ts.shape))
     xs = thermox.sample(jax.random.PRNGKey(0), ts, jnp.zeros(3), A, jnp.zeros(3), D)
     emp = jnp.cov(xs[2000:].T)
     # Monte Carlo error of the sample covariance is ~1e-2 here; the symmetric-part
@@ -189,3 +200,103 @@ def test_preprocess_flags_normality():
     # symmetric A becomes non-normal after transforming with anisotropic D
     A_y, _ = thermox.preprocess(A_SYM, D_DIAG)
     assert not bool(A_y.is_normal)
+
+
+def ill_conditioned_eigenvectors_case():
+    # d = 12, diag(1..3) plus 6 x a strictly upper triangular Gaussian: eigenvector
+    # condition number ~1e9, where a formula in the eigenbasis of A loses everything.
+    d = 12
+    upper = jnp.triu(jax.random.normal(jax.random.PRNGKey(0), (d, d)), 1)
+    return jnp.diag(jnp.linspace(1.0, 3.0, d)) + 6.0 * upper
+
+
+@pytest.mark.parametrize("t", [0.01, 0.3])
+def test_covariance_accurate_for_ill_conditioned_eigenvectors(t):
+    A = ill_conditioned_eigenvectors_case()
+    D = jnp.eye(A.shape[0])
+    assert jnp.linalg.cond(jnp.linalg.eig(A)[1]) > 1e8
+    cov = thermox.conditional.covariance(t, A, D)
+    assert relerr(cov, van_loan_covariance(A, D, t)) < 1e-12
+
+
+def test_covariance_is_zero_at_t_zero():
+    A = ill_conditioned_eigenvectors_case()
+    cov = thermox.conditional.covariance(0.0, A, jnp.eye(A.shape[0]))
+    assert jnp.all(cov == 0.0)
+
+
+def linalg_grid(burnin, num_samples=100, dt=0.1):
+    # The grid thermox.linalg builds: x0 at time 0, one gap of burnin * dt, then dt.
+    ts = jnp.arange(burnin, burnin + num_samples + 1) * dt
+    return jnp.concatenate([jnp.array([0]), ts])
+
+
+@pytest.mark.parametrize(
+    "ts",
+    [
+        pytest.param(jnp.arange(0, 1, 0.01), id="readme-arange"),
+        pytest.param(linalg_grid(0), id="linalg-burnin-0"),
+        pytest.param(linalg_grid(1), id="linalg-burnin-1"),
+        pytest.param(linalg_grid(5), id="linalg-burnin-5"),
+        pytest.param((jnp.arange(0, 10001) * 0.1).astype(jnp.float32), id="float32"),
+        pytest.param(jnp.linspace(0, 100, 300), id="linspace"),
+    ],
+)
+def test_uniform_dt_accepts_grids_uniform_up_to_rounding(ts):
+    is_uniform, dt = uniform_dt(ts)
+    assert bool(is_uniform)
+    assert jnp.isclose(dt, ts[2] - ts[1], rtol=1e-6)
+
+
+def test_uniform_dt_rejects_jittered_grid():
+    ts = jnp.arange(0, 100, 0.1)
+    ts = jnp.sort(ts + jax.random.uniform(jax.random.PRNGKey(0), ts.shape) * 0.1)
+    assert not bool(uniform_dt(ts)[0])
+
+
+def contracting_matrix(key, d=4):
+    return jax.scipy.linalg.expm(
+        -(jax.random.normal(key, (d, d)) / d**0.5 + 3 * jnp.eye(d))
+    )
+
+
+@pytest.mark.parametrize("n", [1, 2, 3, 4, 5, 1000, 1001])
+def test_scan_linear_recurrence_matches_sequential_scan(n):
+    k1, k2, k3 = jax.random.split(jax.random.PRNGKey(n), 3)
+    E = contracting_matrix(k1)
+    y0 = jax.random.normal(k2, (4,))
+    u = jax.random.normal(k3, (n, 4))
+    _, ys = jax.lax.scan(lambda y, u_k: (E @ y + u_k,) * 2, y0, u)
+    assert jnp.allclose(_scan_linear_recurrence(E, y0, u), ys, rtol=1e-12, atol=1e-12)
+
+
+def test_uniform_grid_engines_agree_for_ill_conditioned_eigenvectors():
+    # On a uniform grid both engines apply one exp(-A dt); propagated through the
+    # eigenbasis of A (per-step path) they disagree at ~1e-10 for this family.
+    A = ill_conditioned_eigenvectors_case()
+    d = A.shape[0]
+    ts = jnp.arange(0.0, 1.0, 0.05)
+    key = jax.random.PRNGKey(0)
+    x0, b, D = jnp.ones(d), jnp.zeros(d), jnp.eye(d)
+    xa = thermox.sample(key, ts, x0, A, b, D, associative_scan=True)
+    xs = thermox.sample(key, ts, x0, A, b, D, associative_scan=False)
+    assert relerr(xa, xs) < 1e-12
+
+
+def test_log_prob_on_uniform_grid_exact_for_ill_conditioned_eigenvectors():
+    A = ill_conditioned_eigenvectors_case()
+    d = A.shape[0]
+    ts = jnp.arange(0.0, 1.0, 0.05)
+    x0, b, D = jnp.ones(d), jnp.zeros(d), jnp.eye(d)
+    xs = thermox.sample(jax.random.PRNGKey(0), ts, x0, A, b, D)
+    lp = thermox.log_prob(ts, xs, A, b, D)
+    ref = reference_log_prob(ts, xs, A, b, D)
+    assert jnp.abs(lp - ref) / jnp.abs(ref) < 1e-11
+
+
+def test_linalg_expm_of_nonsymmetric_matrix():
+    # expnegm's whitened drift is non-normal for a non-symmetric input: on
+    # upstream main this estimate is off by 0.6; Monte Carlo noise is ~0.01.
+    M = jnp.array([[-1.0, 3.0], [0.0, -2.0]])
+    est = thermox.linalg.expm(M, num_samples=100000, dt=0.1, burnin=0, alpha=1.0)
+    assert jnp.allclose(est, jax.scipy.linalg.expm(M), atol=1e-1)

thermox/prob.py

@@ -1,5 +1,5 @@
 import jax.numpy as jnp
-from jax.lax import fori_loop
+from jax.lax import cond, fori_loop
 from jax import Array, vmap
 
 from thermox.utils import (
@@ -8,7 +8,12 @@
     ProcessedDriftMatrix,
     ProcessedDiffusionMatrix,
 )
-from thermox.sampler import expm_vp, transition_cov_eigh
+from thermox.sampler import (
+    expm_vp,
+    transition_cov_eigh,
+    transition_expm_and_cov,
+    uniform_dt,
+)
 
 
 def log_prob(
@@ -28,7 +33,8 @@ def log_prob(
     Assumes x(t_0) is given deterministically.
 
     Preprocessing (diagonalisation) costs O(d^3) and evaluation then costs O(T * d^2),
-    where T=len(ts), or O(T * d^3) when D^-0.5 @ A @ D^0.5 is not a normal matrix.
+    where T=len(ts); when D^-0.5 @ A @ D^0.5 is not a normal matrix,
+    O(d^3 log T + T * d^2) on a uniform time grid and O(T * d^3) otherwise.
 
     By default, this function does the preprocessing on A and D before the evaluation.
     However, the preprocessing can be done externally using thermox.preprocess
@@ -65,7 +71,22 @@ def log_prob_identity_diffusion(
 ) -> float:
     if isinstance(A, Array):
         A = preprocess_drift_matrix(A)
+    if len(ts) < 3:
+        return _log_prob_identity_diffusion_stepwise(ts, xs, A, b)
+    # A non-normal A on a uniform grid: build the transition operator once.
+    is_uniform, dt = uniform_dt(ts)
+    return cond(
+        is_uniform & ~A.is_normal,
+        lambda ts, xs, A, b: _log_prob_identity_diffusion_uniform(ts, xs, A, b, dt),
+        _log_prob_identity_diffusion_stepwise,
+        ts,
+        xs,
+        A,
+        b,
+    )
+
 
+def _log_prob_identity_diffusion_stepwise(ts, xs, A, b):
     def transition_mean(y, dt):
         return b + expm_vp(A, y - b, dt)
 
@@ -89,3 +110,25 @@ def mahalanobis_and_log_det(w, U):
     )
 
     return log_prob_val.real
+
+
+def _log_prob_identity_diffusion_uniform(ts, xs, A, b, dt):
+    """log_prob_identity_diffusion on a uniform grid: the residuals of all
+    steps with gap dt at once, one eigendecomposition of their common
+    covariance, and one term for the first gap."""
+    E1, cov1 = transition_expm_and_cov(A.val, ts[1] - ts[0])
+    E, cov = transition_expm_and_cov(A.val, dt)
+    residuals1 = xs[1] - b - E1 @ (xs[0] - b)
+    residuals = xs[2:] - b - (xs[1:-1] - b) @ E.T
+
+    def log_density(cov, r):
+        w, U = jnp.linalg.eigh(cov)
+        w = jnp.where(w < 1e-20, 1e-20, w)
+        z = (r @ U) / jnp.sqrt(w)
+        n, d = r.shape
+        return (
+            -jnp.sum(z * z) / 2
+            - n * (jnp.sum(jnp.log(w)) + d * jnp.log(2 * jnp.pi)) / 2
+        )
+
+    return log_density(cov1, residuals1[None]) + log_density(cov, residuals)

thermox/sampler.py

@@ -27,7 +27,8 @@ def sample(
     by using exact diagonalization.
 
     Preprocessing (diagonalization) costs O(d^3) and sampling costs O(T * d^2),
-    where T=len(ts), or O(T * d^3) when D^-0.5 @ A @ D^0.5 is not a normal matrix.
+    where T=len(ts); when D^-0.5 @ A @ D^0.5 is not a normal matrix,
+    O(d^3 log T + T * d^2) on a uniform time grid and O(T * d^3) otherwise.
 
     If associative_scan=True then jax.lax.associative_scan is used which will run in
     time O((T/p + log(T)) * d^2) on a GPU/TPU with p cores, still with
@@ -69,10 +70,26 @@ def sample_identity_diffusion(
     b: Array,
     associative_scan: bool = True,
 ) -> Array:
+    if isinstance(A, Array):
+        A = preprocess_drift_matrix(A)
     if associative_scan:
-        return _sample_identity_diffusion_associative_scan(key, ts, x0, A, b)
+        stepwise = _sample_identity_diffusion_associative_scan
     else:
-        return _sample_identity_diffusion_scan(key, ts, x0, A, b)
+        stepwise = _sample_identity_diffusion_scan
+    if len(ts) < 3:
+        return stepwise(key, ts, x0, A, b)
+    # A non-normal A on a uniform grid: build the transition operator once.
+    is_uniform, dt = uniform_dt(ts)
+    return jax.lax.cond(
+        is_uniform & ~A.is_normal,
+        lambda *args: _sample_identity_diffusion_uniform(*args, dt, associative_scan),
+        stepwise,
+        key,
+        ts,
+        x0,
+        A,
+        b,
+    )
 
 
 def expm_vp(A, v, dt):
@@ -82,16 +99,31 @@ def expm_vp(A, v, dt):
     return out.real
 
 
+def transition_expm_and_cov(A, dt, n_doublings=12):
+    """exp(-A dt) and int_0^dt exp(-A s) exp(-A^T s) ds for a d x d matrix A,
+    without an eigendecomposition: Van Loan's block exponential at
+    h = dt / 2**n_doublings, then n_doublings steps of E(2h) = E(h)^2 and
+    cov(2h) = cov(h) + E(h) cov(h) E(h)^T. Exact for any stable A, including
+    dt = 0, for ||A|| dt up to about 1e5 with the default n_doublings.
+    """
+    d = A.shape[0]
+    h = dt / 2**n_doublings
+    zeros, eye = jnp.zeros((d, d), dtype=A.dtype), jnp.eye(d, dtype=A.dtype)
+    # h is small, so expm needs few squarings; its loop always runs max_squarings.
+    F = jax.scipy.linalg.expm(jnp.block([[-A, eye], [zeros, A.T]]) * h, max_squarings=4)
+    E = F[:d, :d]
+    cov = F[:d, d:] @ E.T
+    for _ in range(n_doublings):
+        cov = cov + E @ cov @ E.T
+        E = E @ E
+    return E, 0.5 * (cov + cov.T)
+
+
 def transition_cov(A, dt):
     """Covariance of x_dt given x_0 for dx = -A x dt + dW, i.e.
-    int_0^dt exp(-A s) exp(-A^T s) ds, computed in the eigenbasis of A.
-    Exact for any stable, diagonalizable A.
+    int_0^dt exp(-A s) exp(-A^T s) ds. Exact for any stable A.
     """
-    eigvals_sum = A.eigvals[:, None] + A.eigvals.conj()[None, :]
-    integral = -jnp.expm1(-eigvals_sum * dt) / eigvals_sum
-    cov = A.eigvecs @ (A.noise_cov_eigbasis * integral) @ A.eigvecs.conj().T
-    cov = cov.real
-    return 0.5 * (cov + cov.T)
+    return transition_expm_and_cov(A.val, dt)[1]
 
 
 def transition_cov_eigh(A, dt, apply=lambda w, U: (w, U)):
@@ -130,6 +162,67 @@ def transition_cov_sqrt_vp(A, v, dt):
     )
 
 
+def uniform_dt(ts):
+    """Whether the time grid is uniform after its first gap, up to floating-point
+    rounding of the time stamps, and that step. The first gap is free so that
+    the grids built by thermox.linalg ([0, burnin * dt, dt, ...]) qualify.
+    """
+    n = len(ts) - 2
+    dt = (ts[-1] - ts[1]) / n
+    fitted = ts[1] + dt * jnp.arange(n + 1)
+    eps = jnp.finfo(jnp.result_type(ts, float)).eps
+    is_uniform = jnp.max(jnp.abs(ts[1:] - fitted)) <= 1e3 * eps * jnp.max(jnp.abs(ts))
+    return is_uniform, dt
+
+
+def _scan_linear_recurrence(E, y0, u):
+    """y_k = E y_{k-1} + u_k for k = 1, ..., n, computed like
+    jax.lax.associative_scan but with the level's power of E passed down: a
+    combine at depth j applies E ** (2 ** j), one matmul per level, so the
+    scan costs O(d^3 log n + n d^2) time and O(d^2 log n) memory. Carrying
+    the power inside the scanned elements instead would store one d x d
+    matrix per step, O(n d^2) memory.
+    """
+
+    def scan(elems, M):
+        # The recursion of jax.lax.associative_scan: combine adjacent pairs,
+        # recurse on the pairs, fill in the even positions, interleave.
+        m = elems.shape[0]
+        if m < 2:
+            return elems
+        reduced = elems[0:-1:2] @ M.T + elems[1::2]
+        odd = scan(reduced, M @ M)
+        even = jnp.concatenate(
+            [elems[:1], (odd[:-1] if m % 2 == 0 else odd) @ M.T + elems[2::2]]
+        )
+        # Interleave [even0, odd0, even1, odd1, ...]; len(even) is len(odd) or len(odd) + 1.
+        same = even.shape[0] == odd.shape[0]
+        zero, rest = jnp.zeros((), even.dtype), [(0, 0, 0)] * (even.ndim - 1)
+        return jax.lax.pad(even, zero, [(0, int(same), 1)] + rest) + jax.lax.pad(
+            odd, zero, [(1, int(not same), 1)] + rest
+        )
+
+    return scan(jnp.concatenate([y0[None], u]), E)[1:]
+
+
+def _sample_identity_diffusion_uniform(key, ts, x0, A, b, dt, associative_scan):
+    # One transition operator for the first gap, one for dt, applied to the
+    # same draws as the per-step engines.
+    E1, cov1 = transition_expm_and_cov(A.val, ts[1] - ts[0])
+    E, cov = transition_expm_and_cov(A.val, dt)
+    # The one-sided factor U sqrt(w) of the per-step path, once per operator.
+    w1, U1 = jnp.linalg.eigh(cov1)
+    w, U = jnp.linalg.eigh(cov)
+    z = jax.random.normal(key, (len(ts) - 1,) + x0.shape)
+    y1 = E1 @ (x0 - b) + U1 @ (jnp.sqrt(jnp.maximum(w1, 0.0)) * z[0])
+    u = (z[1:] * jnp.sqrt(jnp.maximum(w, 0.0))) @ U.T
+    if associative_scan:
+        ys = _scan_linear_recurrence(E, y1, u)
+    else:
+        _, ys = jax.lax.scan(lambda y, u_k: (E @ y + u_k,) * 2, y1, u)
+    return jnp.concatenate([x0[None], y1[None] + b, ys + b])
+
+
 def _sample_identity_diffusion_scan(
     key: Array,
     ts: Array,