Fix transition covariance for non-normal drift matrices

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`:
 

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,)

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)

thermox/conditional.py

@@ -1,12 +1,11 @@
-from jax import numpy as jnp
 from jax import Array
 
 from thermox.utils import (
     ProcessedDriftMatrix,
     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
 
 

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,