From e0908a6998cac753e83cf5d3573a28068621e5a5 Mon Sep 17 00:00:00 2001
From: Tzu-Chi Yen <tzuchi.yen@colorado.edu>
Date: Sat, 22 Aug 2026 19:17:18 -0600
Subject: [PATCH 1/4] Fix transition covariance for non-normal drift matrices

The transition covariance used by sample, log_prob and conditional.covariance
was built from the eigendecomposition of (A + A^T)/2, which is exact only when
the transformed drift D^-1/2 A D^1/2 is a normal matrix. For other drifts the
samples and log-probabilities were inexact.

- transition_cov computes int_0^dt exp(-A s) exp(-A^T s) ds exactly in the
  eigenbasis of A (any stable, diagonalizable A)
- transition_cov_eigh branches on a normality flag set in preprocessing: the
  existing O(d^2)-per-step formula for normal drifts, one eigendecomposition
  per step otherwise; sample and log_prob read the covariance only through it
- ProcessedDriftMatrix gains noise_cov_eigbasis and is_normal
- tests against Van Loan / Lyapunov references independent of thermox,
  including gradients of log_prob with respect to A; results for normal
  drifts are unchanged
---
 README.md                 |   2 +-
 tests/test_conditional.py |  27 ++++--
 tests/test_nonnormal.py   | 191 ++++++++++++++++++++++++++++++++++++++
 thermox/conditional.py    |   9 +-
 thermox/prob.py           |  35 +++----
 thermox/sampler.py        |  51 +++++++++-
 thermox/utils.py          |  14 ++-
 7 files changed, 286 insertions(+), 43 deletions(-)
 create mode 100644 tests/test_nonnormal.py

diff --git a/README.md b/README.md
index a847e06..9c60f44 100644
--- a/README.md
+++ b/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.
+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.
 
 You can access log-probabilities of the OU process by running `thermox.log_prob`:
 
diff --git a/tests/test_conditional.py b/tests/test_conditional.py
index 07fba84..f4f9375 100644
--- a/tests/test_conditional.py
+++ b/tests/test_conditional.py
@@ -4,26 +4,41 @@
 import thermox
 
 
+def van_loan_covariance(A, D, t):
+    """int_0^t exp(-A s) D exp(-A^T s) ds via Van Loan (1978)."""
+    d = A.shape[0]
+    M = jnp.block([[-A, D], [jnp.zeros((d, d)), A.T]]) * t
+    F = jax.scipy.linalg.expm(M)
+    return F[:d, d:] @ jax.scipy.linalg.expm(-A.T * t)
+
+
 def test_mean_and_cov():
     jax.config.update("jax_enable_x64", True)
     dim = 2
     t = 1.0
 
-    A = jnp.array([[3, 2.5], [2, 4.0]])
+    A = jnp.array([[3, 2.5], [2, 4.0]])  # not symmetric, not normal
     b = jax.random.normal(jax.random.PRNGKey(1), (dim,))
     x0 = jax.random.normal(jax.random.PRNGKey(2), (dim,))
     D = 2 * jnp.eye(dim)
 
+    # References independent of thermox
+    mean_ref = b + jax.scipy.linalg.expm(-A * t) @ (x0 - b)
+    cov_ref = van_loan_covariance(A, D, t)
+
     mean = thermox.conditional.mean(t, x0, A, b, D)
-    samples = jax.vmap(
-        lambda k: thermox.sample(k, jnp.array([0.0, t]), x0, A, b, D)[-1]
-    )(jax.random.split(jax.random.PRNGKey(0), 1000000))
     assert mean.shape == (dim,)
-    assert jnp.allclose(mean, jnp.mean(samples, axis=0), atol=1e-2)
+    assert jnp.allclose(mean, mean_ref, atol=1e-10)
 
     cov = thermox.conditional.covariance(t, A, D)
     assert cov.shape == (dim, dim)
-    assert jnp.allclose(cov, jnp.cov(samples.T), atol=1e-3)
+    assert jnp.allclose(cov, cov_ref, atol=1e-10)
+
+    samples = jax.vmap(
+        lambda k: thermox.sample(k, jnp.array([0.0, t]), x0, A, b, D)[-1]
+    )(jax.random.split(jax.random.PRNGKey(0), 1000000))
+    assert jnp.allclose(mean_ref, jnp.mean(samples, axis=0), atol=1e-2)
+    assert jnp.allclose(cov_ref, jnp.cov(samples.T), atol=1e-3)
 
     mean_and_cov = thermox.conditional.mean_and_covariance(t, x0, A, b, D)
     assert mean_and_cov[0].shape == (dim,)
diff --git a/tests/test_nonnormal.py b/tests/test_nonnormal.py
new file mode 100644
index 0000000..55a1f25
--- /dev/null
+++ b/tests/test_nonnormal.py
@@ -0,0 +1,191 @@
+"""Exactness checks against references (Van Loan, Lyapunov) that do not go
+through thermox: drifts whose transformed form D^{-1/2} A D^{1/2} is not normal,
+plus normal cases whose results must not change.
+"""
+
+import jax
+import jax.numpy as jnp
+import pytest
+
+import thermox
+from thermox.utils import preprocess_drift_matrix
+
+jax.config.update("jax_enable_x64", True)
+
+A_SYM = jnp.array([[3.0, 2.0, 1.0], [2.0, 4.0, 2.0], [1.0, 2.0, 5.0]])
+A_TRI = jnp.array([[2.0, 1.5, 0.0], [0.0, 3.0, 1.5], [0.0, 0.0, 4.0]])
+A_ROT = jnp.array([[1.0, 2.0], [-2.0, 1.0]])
+D_DIAG = jnp.diag(jnp.array([1.0, 4.0, 9.0]))
+D_DENSE = jnp.array([[1.0, 0.3, -0.1], [0.3, 1.0, 0.2], [-0.1, 0.2, 1.0]])
+
+NONNORMAL_CASES = [
+    pytest.param(A_SYM, D_DIAG, id="symmetric-A-diagonal-D"),
+    pytest.param(A_SYM, D_DENSE, id="symmetric-A-dense-D"),
+    pytest.param(A_TRI, jnp.eye(3), id="triangular-A-identity-D"),
+]
+NORMAL_CASES = [
+    pytest.param(A_SYM, jnp.eye(3), id="symmetric-A-identity-D"),
+    pytest.param(A_ROT, jnp.eye(2), id="rotation-A-identity-D"),
+]
+
+
+def van_loan_covariance(A, D, t):
+    """int_0^t exp(-A s) D exp(-A^T s) ds via Van Loan (1978)."""
+    d = A.shape[0]
+    M = jnp.block([[-A, D], [jnp.zeros((d, d)), A.T]]) * t
+    F = jax.scipy.linalg.expm(M)
+    return F[:d, d:] @ jax.scipy.linalg.expm(-A.T * t)
+
+
+def lyapunov_covariance(A, D):
+    """Solve A S + S A^T = D by Kronecker vectorization (small d)."""
+    d = A.shape[0]
+    eye = jnp.eye(d)
+    K = jnp.kron(eye, A) + jnp.kron(A, eye)
+    return jnp.linalg.solve(K, D.reshape(-1, order="F")).reshape(d, d, order="F")
+
+
+def reference_covariance(A, D, t):
+    """Sigma_t = Sigma_inf - exp(-A t) Sigma_inf exp(-A^T t), valid for stable A.
+
+    Numerically benign for large t, unlike the block exponential.
+    """
+    S_inf = lyapunov_covariance(A, D)
+    E = jax.scipy.linalg.expm(-A * t)
+    return S_inf - E @ S_inf @ E.T
+
+
+def relerr(x, y):
+    return jnp.linalg.norm(x - y) / jnp.linalg.norm(y)
+
+
+@pytest.mark.parametrize("A,D", NONNORMAL_CASES + NORMAL_CASES)
+def test_references_agree(A, D):
+    # Guard the references themselves: two independent formulas for Sigma_t.
+    assert (
+        relerr(van_loan_covariance(A, D, 0.7), reference_covariance(A, D, 0.7)) < 1e-10
+    )
+
+
+@pytest.mark.parametrize("A,D", NONNORMAL_CASES + NORMAL_CASES)
+@pytest.mark.parametrize("t", [0.05, 0.7, 5.0])
+def test_conditional_covariance_matches_reference(A, D, t):
+    cov = thermox.conditional.covariance(t, A, D)
+    assert relerr(cov, reference_covariance(A, D, t)) < 1e-8
+
+
+@pytest.mark.parametrize("A,D", NONNORMAL_CASES + NORMAL_CASES)
+def test_stationary_covariance_matches_lyapunov(A, D):
+    lam_min = jnp.min(jnp.linalg.eigvals(A).real)
+    cov = thermox.conditional.covariance(60.0 / lam_min, A, D)
+    assert relerr(cov, lyapunov_covariance(A, D)) < 1e-8
+
+
+def reference_log_prob(ts, xs, A, b, D):
+    """Sum of transition log-densities built from the Van Loan covariance.
+
+    Independent of thermox and differentiable (jax.scipy.linalg.expm has a JVP).
+    """
+
+    def transition_logpdf(x1, x0, dt):
+        mean = b + jax.scipy.linalg.expm(-A * dt) @ (x0 - b)
+        cov = van_loan_covariance(A, D, dt)
+        return jax.scipy.stats.multivariate_normal.logpdf(x1, mean, cov)
+
+    return sum(
+        transition_logpdf(xs[i], xs[i - 1], ts[i] - ts[i - 1])
+        for i in range(1, len(ts))
+    )
+
+
+@pytest.mark.parametrize("A,D", NONNORMAL_CASES)
+def test_log_prob_matches_reference_gaussian(A, D):
+    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("A,D", NONNORMAL_CASES)
+def test_log_prob_grad_wrt_drift_matches_reference(A, D):
+    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)
+    assert relerr(g, g_ref) < 1e-8
+
+
+def test_log_prob_grad_symmetric_parametrization_matches_reference():
+    # A = B B^T with D = I stays on the normal branch; gradients w.r.t. B are exact
+    # there. (At an exactly normal A the derivative w.r.t. A itself in directions
+    # that break normality is that of the symmetric-part formula -- see
+    # thermox.sampler.transition_cov_eigh.)
+    A, D = A_SYM, jnp.eye(3)
+    B = jnp.linalg.cholesky(A)
+    b = jnp.arange(1.0, 4.0)
+    ts = jnp.array([0.0, 0.1, 0.5, 0.6, 1.4])
+    xs = jax.random.normal(jax.random.PRNGKey(3), (len(ts), 3))
+    g = jax.grad(lambda B: thermox.log_prob(ts, xs, B @ B.T, b, D))(B)
+    g_ref = jax.grad(lambda B: reference_log_prob(ts, xs, B @ B.T, b, D))(B)
+    assert relerr(g, g_ref) < 1e-8
+
+
+def test_sample_covariance_matches_lyapunov_for_anisotropic_noise():
+    A, D = A_SYM, D_DIAG
+    ts = jnp.arange(0.0, 20000.0, 0.5)
+    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
+    # formula gives 0.13.
+    assert relerr(emp, lyapunov_covariance(A, D)) < 0.05
+
+
+@pytest.mark.parametrize("A,D", NORMAL_CASES)
+def test_normal_case_equals_symmetric_part_formula(A, D):
+    # For a normal transformed drift the symmetric-part formula is exact; the
+    # general formula must reproduce it, so nothing changes for existing users.
+    t = 0.7
+    A_y, PD = thermox.preprocess(A, D)
+    sym_eigvals, sym_eigvecs = jnp.linalg.eigh(0.5 * (A_y.val + A_y.val.T))
+    old = sym_eigvecs @ jnp.diag(
+        (1 - jnp.exp(-2 * sym_eigvals * t)) / (2 * sym_eigvals)
+    )
+    old = PD.sqrt @ old @ sym_eigvecs.T @ PD.sqrt.T
+    assert relerr(thermox.conditional.covariance(t, A, D), old) < 1e-12
+
+
+@pytest.mark.parametrize("A,D", NONNORMAL_CASES + NORMAL_CASES)
+@pytest.mark.parametrize("t", [0.0, 0.7])
+def test_transition_cov_eigh_factorizes_transition_cov(A, D, t):
+    # (w, U) is the one object sample and log_prob read the covariance through;
+    # both branches must return an orthonormal spectral factorization of
+    # transition_cov, including at t = 0 (the linalg grids contain a zero step).
+    from thermox.sampler import transition_cov, transition_cov_eigh
+
+    A_y, _ = thermox.preprocess(A, D)
+    cov = transition_cov(A_y, t)
+    w, U = transition_cov_eigh(A_y, t)
+    assert jnp.all(w > -1e-12)
+    assert jnp.linalg.norm(U.T @ U - jnp.eye(len(w))) < 1e-10
+    scale = max(1.0, float(jnp.linalg.norm(cov)))
+    assert jnp.linalg.norm(U @ jnp.diag(w) @ U.T - cov) < 1e-10 * scale
+    if bool(A_y.is_normal):
+        # normal branch: precomputed eigenbasis of (A_y + A_y^T)/2, no per-step eigh
+        assert jnp.array_equal(U, A_y.sym_eigvecs)
+    # apply is evaluated inside the branch (so the normal branch never
+    # materializes U per step)
+    trace = transition_cov_eigh(A_y, t, lambda w, U: jnp.sum(w))
+    assert jnp.isclose(trace, jnp.trace(cov), rtol=1e-10, atol=1e-14)
+
+
+def test_preprocess_flags_normality():
+    assert bool(preprocess_drift_matrix(A_SYM).is_normal)
+    assert bool(preprocess_drift_matrix(A_ROT).is_normal)
+    assert not bool(preprocess_drift_matrix(A_TRI).is_normal)
+    # 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)
diff --git a/thermox/conditional.py b/thermox/conditional.py
index 330a2c0..b6e3276 100644
--- a/thermox/conditional.py
+++ b/thermox/conditional.py
@@ -1,4 +1,3 @@
-from jax import numpy as jnp
 from jax import Array
 
 from thermox.utils import (
@@ -6,7 +5,7 @@
     ProcessedDiffusionMatrix,
     handle_matrix_inputs,
 )
-from thermox.sampler import expm_vp
+from thermox.sampler import expm_vp, transition_cov
 
 
 def mean(
@@ -60,11 +59,7 @@ def covariance(
     """
     A_y, D = handle_matrix_inputs(A, D)
 
-    identity_diffusion_cov = (
-        A_y.sym_eigvecs
-        @ jnp.diag((1 - jnp.exp(-2 * A_y.sym_eigvals * t)) / (2 * A_y.sym_eigvals))
-        @ A_y.sym_eigvecs.T
-    )
+    identity_diffusion_cov = transition_cov(A_y, t)
     return D.sqrt @ identity_diffusion_cov @ D.sqrt.T
 
 
diff --git a/thermox/prob.py b/thermox/prob.py
index b0364e9..f4977df 100644
--- a/thermox/prob.py
+++ b/thermox/prob.py
@@ -8,7 +8,7 @@
     ProcessedDriftMatrix,
     ProcessedDiffusionMatrix,
 )
-from thermox.sampler import expm_vp
+from thermox.sampler import expm_vp, transition_cov_eigh
 
 
 def log_prob(
@@ -28,7 +28,7 @@ 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).
+    where T=len(ts), or O(T * d^3) when D^-0.5 @ A @ D^0.5 is not a normal matrix.
 
     By default, this function does the preprocessing on A and D before the evaluation.
     However, the preprocessing can be done externally using thermox.preprocess
@@ -57,20 +57,6 @@ def log_prob(
     return log_prob_ys + D_sqrt_inv_log_det * (len(ts) - 1)
 
 
-def transition_cov_sqrt_inv_vp(A, v, dt):
-    diag = ((1 - jnp.exp(-2 * A.sym_eigvals * dt)) / (2 * A.sym_eigvals)) ** 0.5
-    diag = jnp.where(diag < 1e-20, 1e-20, diag)
-    out = A.sym_eigvecs.T @ v
-    out = out / diag
-    return out.real
-
-
-def transition_cov_log_det(A, dt):
-    diag = (1 - jnp.exp(-2 * A.sym_eigvals * dt)) / (2 * A.sym_eigvals)
-    diag = jnp.where(diag < 1e-20, 1e-20, diag)
-    return jnp.sum(jnp.log(diag))
-
-
 def log_prob_identity_diffusion(
     ts: Array,
     xs: Array,
@@ -84,13 +70,16 @@ def transition_mean(y, dt):
         return b + expm_vp(A, y - b, dt)
 
     def logpt(yt, y0, dt):
-        mean = transition_mean(y0, dt)
-        diff_val = transition_cov_sqrt_inv_vp(A, yt - mean, dt)
-        return (
-            -jnp.dot(diff_val, diff_val) / 2
-            - transition_cov_log_det(A, dt) / 2
-            - jnp.log(2 * jnp.pi) * (yt.shape[0] / 2)
-        )
+        diff = yt - transition_mean(y0, dt)
+
+        def mahalanobis_and_log_det(w, U):
+            w = jnp.where(w < 1e-20, 1e-20, w)
+            diff_val = (U.T @ diff) / jnp.sqrt(w)
+            return jnp.dot(diff_val, diff_val), jnp.sum(jnp.log(w))
+
+        # One factorization of the transition covariance per step serves both terms.
+        quad, log_det = transition_cov_eigh(A, dt, mahalanobis_and_log_det)
+        return -quad / 2 - log_det / 2 - jnp.log(2 * jnp.pi) * (yt.shape[0] / 2)
 
     log_prob_val = fori_loop(
         1,
diff --git a/thermox/sampler.py b/thermox/sampler.py
index 7b83610..2b77fbd 100644
--- a/thermox/sampler.py
+++ b/thermox/sampler.py
@@ -27,7 +27,7 @@ def sample(
     by using exact diagonalization.
 
     Preprocessing (diagonalization) costs O(d^3) and sampling costs O(T * d^2),
-    where T=len(ts).
+    where T=len(ts), or O(T * d^3) when D^-0.5 @ A @ D^0.5 is not a normal matrix.
 
     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
@@ -82,11 +82,52 @@ def expm_vp(A, v, dt):
     return out.real
 
 
+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.
+    """
+    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)
+
+
+def transition_cov_eigh(A, dt, apply=lambda w, U: (w, U)):
+    """Spectral factorization transition_cov(A, dt) = U diag(w) U^T, returned as
+    apply(w, U).
+
+    Branches on A.is_normal. Normal A: U = A.sym_eigvecs is precomputed and w is
+    a closed-form function of dt, O(d^2) per step. Otherwise transition_cov(A, dt)
+    is eigendecomposed at each step, O(d^3). apply is evaluated inside the branch
+    so that only its result, not a d x d matrix per step, leaves the lax.cond.
+
+    Gradients with respect to A follow the branch taken: at an exactly normal A
+    they are those of the normal-branch formula, which depends on A only through
+    (A + A^T)/2.
+    """
+
+    def normal(A, dt):
+        w = (1 - jnp.exp(-2 * A.sym_eigvals * dt)) / (2 * A.sym_eigvals)
+        return apply(w, A.sym_eigvecs)
+
+    def general(A, dt):
+        # eigh rather than Cholesky: stays well defined at dt = 0 (zero covariance).
+        w, U = jnp.linalg.eigh(transition_cov(A, dt))
+        return apply(w, U)
+
+    # apply is evaluated inside the branches so the cond returns a small result;
+    # returning (w, U) and applying it outside made the associative-scan path
+    # measurably slower (vmap's cond batching rule broadcasts the d x d U over all
+    # steps).
+    return jax.lax.cond(A.is_normal, normal, general, A, dt)
+
+
 def transition_cov_sqrt_vp(A, v, dt):
-    diag = ((1 - jnp.exp(-2 * A.sym_eigvals * dt)) / (2 * A.sym_eigvals)) ** 0.5
-    out = diag * v
-    out = A.sym_eigvecs @ out
-    return out.real
+    return transition_cov_eigh(
+        A, dt, lambda w, U: U @ (jnp.sqrt(jnp.maximum(w, 0.0)) * v)
+    )
 
 
 def _sample_identity_diffusion_scan(
diff --git a/thermox/utils.py b/thermox/utils.py
index c47e038..28f5aac 100644
--- a/thermox/utils.py
+++ b/thermox/utils.py
@@ -7,7 +7,8 @@
 
 
 class ProcessedDriftMatrix(NamedTuple):
-    """Stores eigendecompositions of A, (A+A^T)/2"""
+    """Stores eigendecompositions of A, (A+A^T)/2, the noise covariance in the
+    eigenbasis of A (noise_cov_eigbasis) and whether A is normal (is_normal)."""
 
     val: Array
     eigvals: Array
@@ -15,6 +16,8 @@ class ProcessedDriftMatrix(NamedTuple):
     eigvecs_inv: Array
     sym_eigvals: Array
     sym_eigvecs: Array
+    noise_cov_eigbasis: Array
+    is_normal: Array
 
 
 def preprocess_drift_matrix(A: Array) -> ProcessedDriftMatrix:
@@ -33,6 +36,13 @@ def preprocess_drift_matrix(A: Array) -> ProcessedDriftMatrix:
     symA = 0.5 * (A + A.T)
     symA_eigvals, symA_eigvecs = jnp.linalg.eigh(symA)
 
+    noise_cov_eigbasis = A_eigvecs_inv @ A_eigvecs_inv.conj().T
+
+    # A is normal iff A A^T = A^T A; tolerance scales with the working precision.
+    tol = 1e3 * jnp.finfo(A_eigvecs.real.dtype).eps
+    commutator_norm = jnp.linalg.norm(A @ A.T - A.T @ A)
+    is_normal = commutator_norm <= tol * jnp.linalg.norm(A) ** 2
+
     return ProcessedDriftMatrix(
         A,
         A_eigvals,
@@ -40,6 +50,8 @@ def preprocess_drift_matrix(A: Array) -> ProcessedDriftMatrix:
         A_eigvecs_inv,
         symA_eigvals,
         symA_eigvecs,
+        noise_cov_eigbasis,
+        is_normal,
     )
 
 

From d91475d645af7516ff52ce6558e15b4e57a85e7b Mon Sep 17 00:00:00 2001
From: Tzu-Chi Yen <tzuchi.yen@colorado.edu>
Date: Wed, 26 Aug 2026 00:05:05 -0600
Subject: [PATCH 2/4] Factor the transition operator once on uniform grids

---
 README.md               |   2 +-
 tests/test_nonnormal.py | 121 ++++++++++++++++++++++++++++++++++++++--
 thermox/prob.py         |  49 +++++++++++++++-
 thermox/sampler.py      | 113 +++++++++++++++++++++++++++++++++----
 thermox/utils.py        |   7 +--
 5 files changed, 267 insertions(+), 25 deletions(-)

diff --git a/README.md b/README.md
index 9c60f44..0e8e652 100644
--- a/README.md
+++ b/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`:
 
diff --git a/tests/test_nonnormal.py b/tests/test_nonnormal.py
index 55a1f25..a963b00 100644
--- a/tests/test_nonnormal.py
+++ b/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)
diff --git a/thermox/prob.py b/thermox/prob.py
index f4977df..4b328d4 100644
--- a/thermox/prob.py
+++ b/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)
diff --git a/thermox/sampler.py b/thermox/sampler.py
index 2b77fbd..620e9e4 100644
--- a/thermox/sampler.py
+++ b/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,
diff --git a/thermox/utils.py b/thermox/utils.py
index 28f5aac..ad2d4cd 100644
--- a/thermox/utils.py
+++ b/thermox/utils.py
@@ -7,8 +7,7 @@
 
 
 class ProcessedDriftMatrix(NamedTuple):
-    """Stores eigendecompositions of A, (A+A^T)/2, the noise covariance in the
-    eigenbasis of A (noise_cov_eigbasis) and whether A is normal (is_normal)."""
+    """Stores eigendecompositions of A, (A+A^T)/2 and whether A is normal (is_normal)."""
 
     val: Array
     eigvals: Array
@@ -16,7 +15,6 @@ class ProcessedDriftMatrix(NamedTuple):
     eigvecs_inv: Array
     sym_eigvals: Array
     sym_eigvecs: Array
-    noise_cov_eigbasis: Array
     is_normal: Array
 
 
@@ -36,8 +34,6 @@ def preprocess_drift_matrix(A: Array) -> ProcessedDriftMatrix:
     symA = 0.5 * (A + A.T)
     symA_eigvals, symA_eigvecs = jnp.linalg.eigh(symA)
 
-    noise_cov_eigbasis = A_eigvecs_inv @ A_eigvecs_inv.conj().T
-
     # A is normal iff A A^T = A^T A; tolerance scales with the working precision.
     tol = 1e3 * jnp.finfo(A_eigvecs.real.dtype).eps
     commutator_norm = jnp.linalg.norm(A @ A.T - A.T @ A)
@@ -50,7 +46,6 @@ def preprocess_drift_matrix(A: Array) -> ProcessedDriftMatrix:
         A_eigvecs_inv,
         symA_eigvals,
         symA_eigvecs,
-        noise_cov_eigbasis,
         is_normal,
     )
 

From 2cc066a06f45d8d20c7316ae57f71948e861e9c3 Mon Sep 17 00:00:00 2001
From: Tzu-Chi Yen <tzuchi.yen@colorado.edu>
Date: Thu, 27 Aug 2026 14:13:17 -0600
Subject: [PATCH 3/4] Sample non-normal drifts at O(d^2) per step on any time
 grid

On a non-uniform grid, sampling with a non-normal drift factored the
transition covariance at every step, O(T d^3). The noise of every step is
now composed from a fixed set of transition operators, one per binary digit
of the gaps, using cov(a + b) = cov(b) + E(b) cov(a) E(b)^T; the mean is
propagated through the eigenbasis as before.

- _ladder_lattice writes the gaps as integers on a power-of-two lattice
  (2^(e - 52) in float64, 2^(e - 23) in float32), exact to the rounding of
  the time stamps
- _ladder_noise builds M + 1 operator pairs (expm and Cholesky) once and
  applies each to the steps whose gap has that bit set: O(d^3 M + T d^2 M)
- sample_identity_diffusion dispatches: normal drift -> existing path,
  uniform grid -> transition operator once (previous commit), otherwise the
  ladder; both scan engines supported
- tests: composed covariances against the block-exponential reference to
  1e-12, whitened draws over 20 000 irregular steps, engine agreement, vmap
  over keys, zero gaps, operator count independent of T, dispatch, and the
  exactness of the lattice; results for normal drifts and uniform grids are
  unchanged
- d = 64, T = 1000 (CPU, float64): 0.08 s, against 0.75 s for the per-step
  path and 0.006 s for the normal-drift path
---
 README.md               |   2 +-
 tests/test_nonnormal.py | 168 ++++++++++++++++++++++++++++++++++++++++
 thermox/sampler.py      |  87 +++++++++++++++++++--
 3 files changed, 250 insertions(+), 7 deletions(-)

diff --git a/README.md b/README.md
index 0e8e652..d25a48d 100644
--- a/README.md
+++ b/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, or the time grid is uniform, sampling stays O(d^2) per step.
+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`; sampling stays O(d^2) per step on any time grid.
 
 You can access log-probabilities of the OU process by running `thermox.log_prob`:
 
diff --git a/tests/test_nonnormal.py b/tests/test_nonnormal.py
index a963b00..6c26bee 100644
--- a/tests/test_nonnormal.py
+++ b/tests/test_nonnormal.py
@@ -300,3 +300,171 @@ def test_linalg_expm_of_nonsymmetric_matrix():
     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)
+
+
+# --- dyadic ladder: non-normal A on a non-uniform grid ---------------------------
+
+
+def jittered_grid(n, dt=0.1, seed=1):
+    ts = jnp.arange(0.0, n * dt, dt)
+    return jnp.sort(
+        ts + 0.5 * dt * jax.random.uniform(jax.random.PRNGKey(seed), ts.shape)
+    )
+
+
+def random_stable_drift(d, seed=2):
+    # i.i.d. Gaussian / sqrt(d) plus 1.1 I: non-normal, all eigenvalues in the right half-plane
+    return jax.random.normal(jax.random.PRNGKey(seed), (d, d)) / d**0.5 + 1.1 * jnp.eye(
+        d
+    )
+
+
+@pytest.mark.parametrize("A,D", NONNORMAL_CASES)
+def test_ladder_covariance_matches_reference(A, D):
+    # The fold run on covariance matrices instead of draws reproduces Sigma(dt_k)
+    # of the grid as given, for every step, from the ladder's own levels.
+    from thermox.sampler import _ladder_lattice, transition_expm_and_cov
+
+    ts = jittered_grid(301)
+    dts = jnp.diff(ts)
+    A_y, PD = thermox.preprocess(A, D)
+    delta, n, M = _ladder_lattice(ts)
+    assert jnp.max(jnp.abs(n * delta - dts)) <= 1e-12 * jnp.max(ts)
+    d = A.shape[0]
+    C = jnp.zeros((len(dts), d, d))
+    for j in range(M + 1):
+        E, cov = transition_expm_and_cov(A_y.val, delta * 2.0**j)
+        bit = ((n >> j) & 1).astype(bool)
+        C = jnp.where(
+            bit[:, None, None], jnp.einsum("ij,tjk,lk->til", E, C, E) + cov, C
+        )
+    C = jnp.einsum("ij,tjk,lk->til", PD.sqrt, C, PD.sqrt)
+    for k in range(len(dts)):
+        assert relerr(C[k], van_loan_covariance(A, D, dts[k])) < 1e-10
+
+
+@pytest.mark.parametrize("path", ["ladder", "per-step"])
+def test_ladder_whitened_draws_are_standard_normal(path):
+    # 20 000 irregular steps whitened with the reference factor of each step's
+    # covariance: covariance I and mean 0 to Monte Carlo accuracy, on both paths.
+    from thermox.sampler import (
+        _sample_identity_diffusion_ladder,
+        _sample_identity_diffusion_scan,
+        expm_vp,
+    )
+
+    d, T = 8, 20_000
+    A, D = random_stable_drift(d), jnp.eye(d)
+    ts = jnp.sort(jax.random.uniform(jax.random.PRNGKey(3), (T + 1,))) * T * 0.5
+    A_y, _ = thermox.preprocess(A, D)
+    b, x0, key = jnp.zeros(d), jnp.zeros(d), jax.random.PRNGKey(0)
+    if path == "ladder":
+        ys = _sample_identity_diffusion_ladder(key, ts, x0, A_y, b, True)
+    else:
+        ys = _sample_identity_diffusion_scan(key, ts, x0, A_y, b)
+    dts = jnp.diff(ts)
+    means = jax.vmap(lambda y, dt: expm_vp(A_y, y, dt))(ys[:-1], dts)
+    covs = jax.vmap(lambda dt: van_loan_covariance(A, D, dt))(dts)
+    Ls = jnp.linalg.cholesky(covs)
+    z = jax.vmap(lambda L, r: jax.scipy.linalg.solve_triangular(L, r, lower=True))(
+        Ls, ys[1:] - means
+    )
+    assert jnp.max(jnp.abs(z.T @ z / T - jnp.eye(d))) < 4 * (2 / T) ** 0.5
+    assert jnp.max(jnp.abs(jnp.mean(z, axis=0))) < 4 * (1 / T) ** 0.5
+
+
+@pytest.mark.parametrize("dtype,tol", [(jnp.float64, 1e-12), (jnp.float32, 1e-5)])
+def test_ladder_engines_agree(dtype, tol):
+    from thermox.sampler import _sample_identity_diffusion_ladder
+
+    A, D = A_TRI.astype(dtype), jnp.eye(3, dtype=dtype)
+    ts = jittered_grid(200).astype(dtype)
+    A_y, _ = thermox.preprocess(A, D)
+    b, x0, key = jnp.ones(3, dtype), jnp.zeros(3, dtype), jax.random.PRNGKey(0)
+    xa = _sample_identity_diffusion_ladder(key, ts, x0, A_y, b, True)
+    xs = _sample_identity_diffusion_ladder(key, ts, x0, A_y, b, False)
+    # Draws take the default float dtype, as in the other engines: float64 here
+    # because the test module enables x64, float32 in thermox's default setting.
+    assert relerr(xa, xs) < tol
+
+
+def test_ladder_vmap_over_keys_matches_single_draws():
+    A, D = A_TRI, jnp.eye(3)
+    ts = jittered_grid(50)
+    b, x0 = jnp.ones(3), jnp.zeros(3)
+    keys = jax.random.split(jax.random.PRNGKey(0), 4)
+    batched = jax.vmap(lambda k: thermox.sample(k, ts, x0, A, b, D))(keys)
+    single = jnp.stack([thermox.sample(k, ts, x0, A, b, D) for k in keys])
+    assert jnp.array_equal(batched, single)
+
+
+def test_ladder_zero_gap_repeats_the_state():
+    from thermox.sampler import _sample_identity_diffusion_ladder
+
+    ts = jittered_grid(50)
+    ts = jnp.concatenate([ts[:20], ts[19:20], ts[20:]])  # repeated time inside the grid
+    A_y, _ = thermox.preprocess(A_TRI, jnp.eye(3))
+    xs = _sample_identity_diffusion_ladder(
+        jax.random.PRNGKey(0), ts, jnp.zeros(3), A_y, jnp.ones(3), True
+    )
+    assert jnp.all(jnp.isfinite(xs))
+    assert jnp.allclose(
+        xs[20], xs[19], rtol=1e-14, atol=1e-14
+    )  # exp(0) x = x to rounding
+
+
+@pytest.mark.parametrize("dtype,levels", [(jnp.float64, 53), (jnp.float32, 24)])
+def test_ladder_operator_count_is_independent_of_T(dtype, levels):
+    # One transition operator per mantissa bit of ts, built inside one scan; no
+    # factorization per step (the structure does not change with T).
+    from thermox.sampler import _ladder_noise
+
+    A_y, _ = thermox.preprocess(A_TRI.astype(dtype), jnp.eye(3, dtype=dtype))
+    texts = []
+    for n in (100, 200):
+        ts = jittered_grid(n).astype(dtype)
+        texts.append(
+            str(
+                jax.make_jaxpr(lambda k: _ladder_noise(A_y, ts, k))(
+                    jax.random.PRNGKey(0)
+                )
+            )
+        )
+    import re
+
+    for text in texts:
+        assert "eigh" not in text
+        assert text.count(f"length={levels}") == 1  # the one scan over the levels
+    # The loop structure (the level scan and expm's own squaring loop) is the same for both T.
+    assert re.findall(r"length=\d+", texts[0]) == re.findall(r"length=\d+", texts[1])
+
+
+def test_sample_dispatches_to_ladder_on_nonuniform_grid():
+    # thermox.sample takes the ladder for non-normal A on a non-uniform grid.
+    from thermox.sampler import _sample_identity_diffusion_ladder
+
+    A, D = A_SYM, D_DIAG
+    ts = jittered_grid(50)
+    b, x0, key = jnp.ones(3), jnp.zeros(3), jax.random.PRNGKey(0)
+    A_y, PD = thermox.preprocess(A, D)
+    ys = _sample_identity_diffusion_ladder(
+        key, ts, PD.sqrt_inv @ x0, A_y, PD.sqrt_inv @ b, True
+    )
+    direct = jax.vmap(jnp.matmul, in_axes=(None, 0))(PD.sqrt, ys)
+    assert jnp.array_equal(thermox.sample(key, ts, x0, A, b, D), direct)
+
+
+@pytest.mark.parametrize("dtype", [jnp.float64, jnp.float32])
+def test_ladder_lattice_is_exact(dtype):
+    # The unit is exactly a power of two; dyadic gaps have exactly their binary digits.
+    from thermox.sampler import _ladder_lattice
+
+    ts = jnp.array([0.0, 0.5, 0.75, 1.5], dtype)
+    delta, n, M = _ladder_lattice(ts)
+    assert delta.dtype == dtype and M == (52 if dtype == jnp.float64 else 23)
+    assert float(delta) == 2.0 ** (0 - M)  # largest gap 0.75 = 0.75 * 2^0
+    assert [int(v) for v in n] == [
+        2 ** (M - 1),
+        2 ** (M - 2),
+        2 ** (M - 1) + 2 ** (M - 2),
+    ]
diff --git a/thermox/sampler.py b/thermox/sampler.py
index 620e9e4..43c3dda 100644
--- a/thermox/sampler.py
+++ b/thermox/sampler.py
@@ -28,7 +28,7 @@ def sample(
 
     Preprocessing (diagonalization) costs O(d^3) and sampling costs O(T * d^2),
     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.
+    O(d^3 + T * d^2) on any time grid.
 
     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
@@ -78,12 +78,19 @@ def sample_identity_diffusion(
         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.
+    # A non-normal A: build the transition operator once on a uniform grid, or one
+    # per binary digit of the gaps (the ladder) on any other grid.
     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,
+    index = jnp.where(A.is_normal, 0, jnp.where(is_uniform, 1, 2)).astype(jnp.int32)
+    return jax.lax.switch(
+        index,
+        [
+            stepwise,
+            lambda *args: _sample_identity_diffusion_uniform(
+                *args, dt, associative_scan
+            ),
+            lambda *args: _sample_identity_diffusion_ladder(*args, associative_scan),
+        ],
         key,
         ts,
         x0,
@@ -223,6 +230,74 @@ def _sample_identity_diffusion_uniform(key, ts, x0, A, b, dt, associative_scan):
     return jnp.concatenate([x0[None], y1[None] + b, ys + b])
 
 
+def _ladder_lattice(ts):
+    """Gaps of ts as integers on the lattice delta = 2^(e - M), where 2^e is the
+    power of two just above the largest gap and M the number of mantissa bits
+    of ts (52 in float64, 23 in float32): M + 1 binary digits, and every gap is
+    reproduced to the rounding of the time stamps themselves. A power of two
+    keeps dts / delta exact, so the digits do not depend on how it is compiled.
+    """
+    dts = jnp.diff(ts)
+    M = jnp.finfo(ts.dtype).nmant
+    delta = jnp.ldexp(jnp.ones((), ts.dtype), jnp.frexp(jnp.max(dts))[1] - M)
+    n = jnp.round(dts / delta).astype(jnp.int64 if M > 23 else jnp.int32)
+    return delta, n, M
+
+
+def _ladder_noise(A, ts, key):
+    """Noise terms of every step, shape (len(ts) - 1, d), from a dyadic ladder of
+    transition operators: for level j the pair E_j = exp(-A delta 2^j),
+    L_j L_j^T = cov(delta 2^j) is built once and applied to the steps whose gap
+    has bit j set, e <- E_j e + L_j z_j, which composes the covariances exactly
+    (cov(a + b) = cov(b) + E(b) cov(a) E(b)^T). M + 1 operators once, two
+    masked products per level and step.
+    """
+    delta, n, M = _ladder_lattice(ts)
+    d = A.val.shape[0]
+
+    def level(e, j):
+        E, cov = transition_expm_and_cov(A.val, jnp.ldexp(delta, j))
+        # Cholesky, not eigh: every level's gap is positive, and the factor is unique,
+        # so draws at a fixed key do not depend on how the level was compiled.
+        L = jnp.linalg.cholesky(cov)
+        z = jax.random.normal(jax.random.fold_in(key, j), (len(n), d))
+        bit = ((n >> j.astype(n.dtype)) & 1).astype(bool)
+        return jnp.where(bit[:, None], e @ E.T + z @ L.T, e), None
+
+    noise, _ = jax.lax.scan(level, jnp.zeros((len(n), d)), jnp.arange(M + 1))
+    return noise
+
+
+def _sample_identity_diffusion_ladder(key, ts, x0, A, b, associative_scan):
+    # Noise from the ladder, mean through the eigenbasis as on the per-step paths.
+    dts = jnp.diff(ts)
+    noise = _ladder_noise(A, ts, key)
+    if associative_scan:
+
+        @partial(jax.vmap, in_axes=(0, 0))
+        def binary_associative_operator(elem_a, elem_b):
+            t_a, x_a = elem_a
+            t_b, x_b = elem_b
+            return t_a + t_b, expm_vp(A, x_a, t_b) + x_b
+
+        scan_times = jnp.concatenate([ts[:1], dts])
+        scan_values = jnp.concatenate([x0[None] - b, noise])
+        return (
+            jax.lax.associative_scan(
+                binary_associative_operator, (scan_times, scan_values)
+            )[1]
+            + b
+        )
+
+    def step(x, dt_and_u):
+        dt, u = dt_and_u
+        x = b + expm_vp(A, x - b, dt) + u
+        return x, x
+
+    _, xs = jax.lax.scan(step, x0.astype(noise.dtype), (dts, noise))
+    return jnp.concatenate([x0[None], xs])
+
+
 def _sample_identity_diffusion_scan(
     key: Array,
     ts: Array,

From a17e6aa678749bcf7740ecc1021f3f7036bf0515 Mon Sep 17 00:00:00 2001
From: Tzu-Chi Yen <tzuchi.yen@colorado.edu>
Date: Fri, 28 Aug 2026 00:28:20 -0600
Subject: [PATCH 4/4] Evaluate log_prob of non-normal drifts at O(d^2) per step
 on any grid

On a non-uniform grid, log_prob with a non-normal drift factored the
transition covariance at every step, O(T d^3). The two quantities the
density needs, log det cov(dt) and cov(dt)^-1, are smooth functions of the
gap, so they are now interpolated across the gaps: one panel per octave of
the gaps, 17 Chebyshev nodes per panel, one covariance and eigh per node
built once, and a 17-term sum per step. A run-time check on the last two
Chebyshev coefficients of both series (below 1e-10 of the largest in
float64, 3e-4 in float32, above the covariance routine's own accuracy)
guards the interpolation; when it fails, log_prob takes the existing
per-step path.

- _log_prob_panels: octave panels from frexp (exact edges), the cosine
  transform to Chebyshev coefficients, T_n by the three-term recurrence
  (finite gradients with respect to ts), the check, and the fallback
- log_prob_identity_diffusion dispatches: normal drift -> existing path,
  uniform grid -> transition operator once (second commit), otherwise the
  panels with the per-step path as fallback
- _log_prob_identity_diffusion_stepwise is wrapped in jax.checkpoint: cond
  and switch keep every branch's residuals for the backward pass, so
  without it the per-step factorizations were stored even when another
  path ran (grad memory at d = 64, T = 10 000: 9-15 GB -> 1.3 GB on the
  new path, 9.9 GB -> 0.5 GB on the uniform path); the normal-drift
  gradient recomputes its forward pass in exchange, within 15 % of before
  on CPU
- tests: values against the per-step reference to 1e-10 (float64) and
  3e-4 (float32), gradients with respect to the drift and to ts against
  independent references, an oscillatory drift and a zero gap declined
  with the fallback bitwise equal to the per-step path, operator counts
  independent of T, dispatch; results for normal drifts and uniform grids
  are unchanged
- d = 64, T = 1000 (CPU, float64): 0.21 s against 0.85 s for the per-step
  path; T = 10 000: 0.42 s against 8.5 s; gradients 0.90 s and 1.5 s
---
 tests/test_nonnormal.py | 209 ++++++++++++++++++++++++++++++++++++++++
 thermox/prob.py         | 108 +++++++++++++++++++--
 2 files changed, 309 insertions(+), 8 deletions(-)

diff --git a/tests/test_nonnormal.py b/tests/test_nonnormal.py
index 6c26bee..66aace5 100644
--- a/tests/test_nonnormal.py
+++ b/tests/test_nonnormal.py
@@ -468,3 +468,212 @@ def test_ladder_lattice_is_exact(dtype):
         2 ** (M - 2),
         2 ** (M - 1) + 2 ** (M - 2),
     ]
+
+
+# --- Chebyshev panels: log det Sigma(dt) on a non-uniform grid --------------------
+
+
+def random_gaps_grid(n, decades=3.0, seed=4):
+    # Gaps log-uniform over the given number of decades; every gap is different.
+    gaps = 10 ** jax.random.uniform(
+        jax.random.PRNGKey(seed), (n,), minval=-decades / 2, maxval=decades / 2
+    )
+    return jnp.concatenate([jnp.zeros(1), jnp.cumsum(gaps)]) * 0.1
+
+
+def oscillatory_drift():
+    # Eigenvalues 1 +- 10i and 2: |Im lambda| / Re lambda = 10.
+    return jnp.array([[1.0, 10.0, 0.5], [-10.0, 1.0, 0.0], [0.0, 0.0, 2.0]])
+
+
+def test_chebyshev_matrix_is_exact_on_chebyshev_polynomials():
+    from thermox.prob import _chebyshev_matrix
+
+    N = 8
+    x_nodes = jnp.cos(jnp.pi * jnp.arange(N + 1) / N)
+    for k in range(N + 1):
+        # The transform of T_k sampled at the nodes is the unit vector e_k.
+        coeffs = _chebyshev_matrix(N) @ jnp.cos(k * jnp.arccos(x_nodes))
+        assert jnp.allclose(coeffs, jnp.eye(N + 1)[k], atol=1e-12)
+
+
+def reference_terms(A_y, ys, b, dts):
+    """-2 log p of every step from reference_covariance: r^T Sigma^-1 r + log det Sigma
+    + d log(2 pi).
+    """
+    from thermox.sampler import expm_vp
+
+    d = ys.shape[1]
+    r = ys[1:] - b - jax.vmap(lambda y, dt: expm_vp(A_y, y - b, dt))(ys[:-1], dts)
+
+    def term(rk, t):
+        S = reference_covariance(A_y.val, jnp.eye(d), t)
+        return rk @ jnp.linalg.solve(S, rk) + jnp.linalg.slogdet(S)[1]
+
+    return jax.vmap(term)(r, dts) + d * jnp.log(2 * jnp.pi)
+
+
+@pytest.mark.parametrize("A,D", NONNORMAL_CASES)
+def test_panels_match_reference(A, D):
+    # The interpolated value against the per-step reference, relative to the total
+    # magnitude of the terms, to the accuracy of the covariances themselves.
+    from thermox.prob import _log_prob_panels
+
+    ts = random_gaps_grid(200)
+    d = A.shape[0]
+    b = jnp.arange(1.0, d + 1.0)
+    xs = jax.random.normal(jax.random.PRNGKey(3), (len(ts), d))
+    A_y, PD = thermox.preprocess(A, D)
+    ys = jax.vmap(jnp.matmul, in_axes=(None, 0))(PD.sqrt_inv, xs)
+    value, ok = _log_prob_panels(ts, ys, A_y, PD.sqrt_inv @ b)
+    assert bool(ok)
+    terms = reference_terms(A_y, ys, PD.sqrt_inv @ b, jnp.diff(ts))
+    assert jnp.abs(value + 0.5 * jnp.sum(terms)) < 1e-10 * 0.5 * jnp.sum(jnp.abs(terms))
+
+
+def test_panels_match_reference_in_float32():
+    from thermox.prob import _log_prob_panels
+
+    A, D = A_TRI.astype(jnp.float32), jnp.eye(3, dtype=jnp.float32)
+    ts = random_gaps_grid(200).astype(jnp.float32)
+    A_y, _ = thermox.preprocess(A, D)
+    ys = jax.random.normal(jax.random.PRNGKey(3), (len(ts), 3), jnp.float32)
+    value, ok = _log_prob_panels(ts, ys, A_y, jnp.ones(3, jnp.float32))
+    assert bool(ok)
+    terms = reference_terms(
+        preprocess_drift_matrix(A_TRI), ys, jnp.ones(3), jnp.diff(ts)
+    )
+    # float32 accuracy is that of transition_expm_and_cov in float32 (twelve
+    # doublings, about 1e-4).
+    assert jnp.abs(value + 0.5 * jnp.sum(terms)) < 3e-4 * 0.5 * jnp.sum(jnp.abs(terms))
+
+
+def test_panels_grad_wrt_drift_matches_reference():
+    from thermox.prob import _log_prob_panels
+
+    A, D = A_TRI, jnp.eye(3)
+    ts = random_gaps_grid(60)
+    ys = jax.random.normal(jax.random.PRNGKey(3), (len(ts), 3))
+    b = jnp.ones(3)
+
+    def value(A):
+        return _log_prob_panels(ts, ys, preprocess_drift_matrix(A), b)[0]
+
+    g = jax.grad(value)(A)
+    g_ref = jax.grad(lambda A: reference_log_prob(ts, ys, A, b, D))(A)
+    assert relerr(g, g_ref) < 1e-8
+
+
+def test_panels_grad_wrt_ts_matches_reference():
+    # The interpolant is a polynomial in the gap, so the gradient with respect to
+    # the time stamps is finite and equal to the per-step path's.
+    from thermox.prob import _log_prob_identity_diffusion_stepwise, _log_prob_panels
+
+    ts = random_gaps_grid(60)
+    A_y = preprocess_drift_matrix(A_TRI)
+    ys = jax.random.normal(jax.random.PRNGKey(3), (len(ts), 3))
+    g = jax.grad(lambda t: _log_prob_panels(t, ys, A_y, jnp.ones(3))[0])(ts)
+    g_ref = jax.grad(
+        lambda t: _log_prob_identity_diffusion_stepwise(t, ys, A_y, jnp.ones(3))
+    )(ts)
+    assert jnp.all(jnp.isfinite(g))
+    assert relerr(g, g_ref) < 1e-8
+
+
+def test_panels_decline_oscillatory_drift():
+    from thermox.prob import _log_prob_panels
+
+    A_y = preprocess_drift_matrix(oscillatory_drift())
+    ts = random_gaps_grid(200)
+    ys = jax.random.normal(jax.random.PRNGKey(3), (len(ts), 3))
+    assert not bool(_log_prob_panels(ts, ys, A_y, jnp.zeros(3))[1])
+
+
+def test_panels_decline_zero_gap():
+    from thermox.prob import _log_prob_panels
+
+    ts = random_gaps_grid(50)
+    ts = jnp.concatenate([ts[:20], ts[19:20], ts[20:]])  # repeated time inside the grid
+    A_y = preprocess_drift_matrix(A_TRI)
+    ys = jax.random.normal(jax.random.PRNGKey(3), (len(ts), 3))
+    assert not bool(_log_prob_panels(ts, ys, A_y, jnp.zeros(3))[1])
+
+
+@pytest.mark.parametrize("dtype", [jnp.float64, jnp.float32])
+def test_panels_operator_count_is_independent_of_T(dtype):
+    # 17 eigh per used panel, inside one scan over the panels, for T = 100 and
+    # T = 200 alike.
+    from thermox.prob import _log_prob_panels
+
+    import re
+
+    A_y, _ = thermox.preprocess(A_TRI.astype(dtype), jnp.eye(3, dtype=dtype))
+    texts = []
+    for n in (100, 200):
+        ts = random_gaps_grid(n).astype(dtype)
+        ys = jax.random.normal(jax.random.PRNGKey(3), (len(ts), 3), dtype)
+        texts.append(
+            str(
+                jax.make_jaxpr(
+                    lambda ts, ys: _log_prob_panels(ts, ys, A_y, jnp.zeros(3, dtype))
+                )(ts, ys)
+            )
+        )
+    for text in texts:
+        assert (
+            text.count("= eigh[") == 1
+        )  # one call site, batched over the nodes, inside the panel scan
+        assert "[17,3,3]" in text  # the batch of N + 1 node covariances
+    assert re.findall(r"length=\d+", texts[0]) == re.findall(r"length=\d+", texts[1])
+
+
+def test_log_prob_dispatches_to_panels_on_nonuniform_grid():
+    # thermox.log_prob takes the panels on a non-uniform grid with non-normal A, and
+    # their run-time check passes there.
+    from thermox.prob import _log_prob_panels
+
+    A, D = A_SYM, D_DIAG
+    ts = jnp.array([0.0, 0.1, 0.5, 0.6, 1.4])
+    b = jnp.arange(1.0, 4.0)
+    xs = jax.random.normal(jax.random.PRNGKey(3), (len(ts), 3))
+    A_y, PD = thermox.preprocess(A, D)
+    ys = jax.vmap(jnp.matmul, in_axes=(None, 0))(PD.sqrt_inv, xs)
+    value, ok = _log_prob_panels(ts, ys, A_y, PD.sqrt_inv @ b)
+    assert bool(ok)
+    expected = value + jnp.log(jnp.linalg.det(PD.sqrt_inv)) * (len(ts) - 1)
+    assert jnp.array_equal(thermox.log_prob(ts, xs, A, b, D), expected)
+
+
+def test_log_prob_falls_back_to_per_step_path_when_check_fails():
+    # When the check fails, log_prob returns the per-step path's value, bitwise.
+    from thermox.prob import _log_prob_identity_diffusion_stepwise, _log_prob_panels
+
+    A = oscillatory_drift()
+    ts = random_gaps_grid(60)
+    xs = jax.random.normal(jax.random.PRNGKey(3), (len(ts), 3))
+    A_y = preprocess_drift_matrix(A)
+    assert not bool(_log_prob_panels(ts, xs, A_y, jnp.zeros(3))[1])
+    lp = thermox.log_prob(ts, xs, A, jnp.zeros(3), jnp.eye(3))
+    assert jnp.array_equal(
+        lp, _log_prob_identity_diffusion_stepwise(ts, xs, A_y, jnp.zeros(3))
+    )
+
+
+def test_log_prob_operator_count_is_independent_of_T():
+    # log_prob's program on a non-uniform grid carries the batch of 17 node
+    # covariances, and its eigh call sites do not multiply with T (the per-step
+    # fallback branch keeps its own T-length loop, as expected).
+    A, D = A_TRI, jnp.eye(3)
+    texts = []
+    for n in (100, 200):
+        ts = random_gaps_grid(n)
+        xs = jax.random.normal(jax.random.PRNGKey(3), (len(ts), 3))
+        texts.append(
+            str(
+                jax.make_jaxpr(
+                    lambda ts, xs: thermox.log_prob(ts, xs, A, jnp.zeros(3), D)
+                )(ts, xs)
+            )
+        )
+    assert "[17,3,3]" in texts[0]
+    assert texts[0].count("= eigh[") == texts[1].count("= eigh[")
diff --git a/thermox/prob.py b/thermox/prob.py
index 4b328d4..220b1d8 100644
--- a/thermox/prob.py
+++ b/thermox/prob.py
@@ -1,6 +1,6 @@
 import jax.numpy as jnp
-from jax.lax import cond, fori_loop
-from jax import Array, vmap
+from jax.lax import cond, fori_loop, scan, switch
+from jax import Array, checkpoint, vmap
 
 from thermox.utils import (
     handle_matrix_inputs,
@@ -34,7 +34,8 @@ def log_prob(
 
     Preprocessing (diagonalisation) costs O(d^3) and evaluation then costs O(T * d^2),
     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.
+    O(d^3 + T * d^2) on any time grid, or O(T * d^3) when the run-time accuracy
+    check of the interpolation across the gaps fails.
 
     By default, this function does the preprocessing on A and D before the evaluation.
     However, the preprocessing can be done externally using thermox.preprocess
@@ -73,12 +74,19 @@ def log_prob_identity_diffusion(
         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.
+    # A non-normal A: build the transition operator once on a uniform grid, or
+    # interpolate the transition covariance's log-determinant and inverse across
+    # the gaps on any other grid, falling back to the per-step path when that
+    # interpolation fails its run-time check.
     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,
+    index = jnp.where(A.is_normal, 0, jnp.where(is_uniform, 1, 2)).astype(jnp.int32)
+    return switch(
+        index,
+        [
+            _log_prob_identity_diffusion_stepwise,
+            lambda ts, xs, A, b: _log_prob_identity_diffusion_uniform(ts, xs, A, b, dt),
+            _log_prob_identity_diffusion_panels,
+        ],
         ts,
         xs,
         A,
@@ -86,6 +94,9 @@ def log_prob_identity_diffusion(
     )
 
 
+# cond and switch save every branch's residuals for the backward pass; with
+# checkpoint this branch recomputes its per-step intermediates instead of storing them.
+@checkpoint
 def _log_prob_identity_diffusion_stepwise(ts, xs, A, b):
     def transition_mean(y, dt):
         return b + expm_vp(A, y - b, dt)
@@ -132,3 +143,84 @@ def log_density(cov, r):
         )
 
     return log_density(cov1, residuals1[None]) + log_density(cov, residuals)
+
+
+def _chebyshev_matrix(N):
+    """(N + 1) x (N + 1) matrix taking values at the Chebyshev points cos(m pi / N)
+    to the coefficients of the interpolating polynomial in T_0, ..., T_N.
+    """
+    m = jnp.arange(N + 1)
+    h = jnp.where((m == 0) | (m == N), 0.5, 1.0)
+    return (2.0 / N) * h[:, None] * jnp.cos(jnp.pi * jnp.outer(m, m) / N) * h
+
+
+def _log_prob_panels(ts, xs, A, b):
+    """Log-density of the trajectory from Chebyshev panels: for each octave of the
+    gaps that some step falls in, N + 1 = 17 nodes, one eigh per node, the
+    Chebyshev coefficients of log det cov(dt) (scalars) and of cov(dt)^-1
+    (matrices) by one cosine transform, and the interpolants at every step in
+    the panel, O(d^2) per step. Returns (value, ok); ok is false when a panel's
+    last two coefficients are not below the tolerance, a gap is zero, or the
+    gaps span more than P = 53 octaves.
+    """
+    # tol is the accuracy of transition_expm_and_cov in this dtype; N is the
+    # Chebyshev degree per panel and P the number of octaves of gaps handled.
+    dtype = ts.dtype
+    tol = 1e-10 if dtype == jnp.float64 else 3e-4
+    N, P = 16, 53
+    dts = jnp.diff(ts)
+    positive = dts > 0
+    tau_min = jnp.min(jnp.where(positive, dts, jnp.inf))
+    # frexp places every gap in its octave panel p, with edges tau_min 2^p and
+    # tau_min 2^(p + 1) exactly, and x is the gap's coordinate in [-1, 1] there.
+    m, e = jnp.frexp(jnp.where(positive, dts, tau_min) / tau_min)
+    p, x = e - 1, 4.0 * m - 3.0
+    x_nodes = jnp.cos(jnp.pi * jnp.arange(N + 1) / N).astype(dtype)
+    C_N = _chebyshev_matrix(N).astype(dtype)
+    # T_n(x) at every step by the three-term recurrence rather than cos(n arccos x),
+    # whose derivative is infinite at x = -1 (the gradient with respect to ts needs it).
+    W = [jnp.ones_like(x), x]
+    for _ in range(N - 1):
+        W.append(2 * x * W[-1] - W[-2])
+    W = jnp.stack(W)
+    # Residuals of every step, as on the per-step path.
+    residuals = xs[1:] - b - vmap(lambda y, dt: expm_vp(A, y - b, dt))(xs[:-1], dts)
+
+    # Reverse mode recomputes each panel instead of storing its node covariances.
+    @checkpoint
+    def panel(carry, j):
+        terms, ok = carry
+        mask = p == j
+
+        def body(terms, ok):
+            taus = jnp.ldexp(tau_min, j) * (1.5 + 0.5 * x_nodes)
+            covs = vmap(lambda t: transition_expm_and_cov(A.val, t)[1])(taus)
+            w, U = jnp.linalg.eigh(covs)
+            w = jnp.where(w < 1e-20, 1e-20, w)
+            # Chebyshev coefficients of log det cov (c) and of cov^-1 (C_inv).
+            c = C_N @ jnp.sum(jnp.log(w), axis=1)
+            C_inv = jnp.einsum("nm,mij,mj,mkj->nik", C_N, U, 1.0 / w, U)
+            q = jnp.einsum(
+                "tn,nt->t", jnp.einsum("ti,nij,tj->tn", residuals, C_inv, residuals), W
+            )
+            norms = jnp.linalg.norm(C_inv, axis=(1, 2))
+            # Both series have converged: the last two coefficients are below tol.
+            tail_ok = (jnp.max(jnp.abs(c[-2:])) <= tol * jnp.max(jnp.abs(c))) & (
+                jnp.max(norms[-2:]) <= tol * jnp.max(norms)
+            )
+            return jnp.where(mask, q + c @ W, terms), ok & tail_ok
+
+        return cond(jnp.any(mask), body, lambda t, o: (t, o), terms, ok), None
+
+    init = (jnp.zeros(len(dts), dtype), jnp.array(True))
+    (terms, ok), _ = scan(panel, init, jnp.arange(P))
+    value = -0.5 * jnp.sum(terms + xs.shape[1] * jnp.log(2 * jnp.pi))
+    return value, ok & jnp.all(positive) & (jnp.max(e) <= P)
+
+
+def _log_prob_identity_diffusion_panels(ts, xs, A, b):
+    # The panels' value when their run-time check passes, the per-step path's otherwise.
+    value, ok = _log_prob_panels(ts, xs, A, b)
+    return cond(
+        ok, lambda: value, lambda: _log_prob_identity_diffusion_stepwise(ts, xs, A, b)
+    )