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

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