diff --git a/README.md b/README.md index a847e06..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. +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_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..a963b00 --- /dev/null +++ b/tests/test_nonnormal.py @@ -0,0 +1,302 @@ +"""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.sampler import _scan_linear_recurrence, uniform_dt +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)) + ) + + +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, ts): + d = A.shape[0] + b = jnp.arange(1.0, d + 1.0) + 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, ts): + d = A.shape[0] + b = jnp.arange(1.0, d + 1.0) + 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 + + +@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 + # 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) + + +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/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..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 +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). + 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 @@ -57,20 +63,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, @@ -79,18 +71,36 @@ 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) 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, @@ -100,3 +110,25 @@ def logpt(yt, y0, dt): ) 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 7b83610..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). + 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,11 +99,128 @@ 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. Exact for any stable A. + """ + return transition_expm_and_cov(A.val, dt)[1] + + +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 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( diff --git a/thermox/utils.py b/thermox/utils.py index c47e038..ad2d4cd 100644 --- a/thermox/utils.py +++ b/thermox/utils.py @@ -7,7 +7,7 @@ class ProcessedDriftMatrix(NamedTuple): - """Stores eigendecompositions of A, (A+A^T)/2""" + """Stores eigendecompositions of A, (A+A^T)/2 and whether A is normal (is_normal).""" val: Array eigvals: Array @@ -15,6 +15,7 @@ class ProcessedDriftMatrix(NamedTuple): eigvecs_inv: Array sym_eigvals: Array sym_eigvecs: Array + is_normal: Array def preprocess_drift_matrix(A: Array) -> ProcessedDriftMatrix: @@ -33,6 +34,11 @@ def preprocess_drift_matrix(A: Array) -> ProcessedDriftMatrix: symA = 0.5 * (A + A.T) symA_eigvals, symA_eigvecs = jnp.linalg.eigh(symA) + # 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 +46,7 @@ def preprocess_drift_matrix(A: Array) -> ProcessedDriftMatrix: A_eigvecs_inv, symA_eigvals, symA_eigvecs, + is_normal, )