Fix transition covariance for non-normal drift matrices

[junipertcy]

Edited to add: In the second commit, the covariance is now built by a block matrix exponential and doubling, not in the eigenbasis; see my comment below.

When teaching myself how thermox simulates SPU trajectories, I noticed that thermox is exact only when the transformed drift matrix $A_y = L^{-1} A L$ (with $D = LL^\intercal$, as built in thermox.preprocess) is normal. This is because thermox uses the eigendecomposition of $A_y$'s symmetric part

$$A_s = \frac{1}{2}(A_y + A_y^\intercal) = V_s \Lambda_s V_s^\intercal$$

and evaluates, with $\lambda_s$ the diagonal of $\Lambda_s$,

$$\tilde{\Sigma}_t = V_s \mathrm{diag}\left(\frac{1 - e^{-2\lambda_s t}}{2\lambda_s}\right) V_s^\intercal = \int_0^t e^{-(A_y + A_y^\intercal)s} ds$$

Whereas in principle, we should evaluate

$$\Sigma_t = \int_0^t e^{-A_y s} e^{-A_y^\intercal s} , ds$$

The two equations are equal if and only if $A_y$ is normal, that is, $A_y$ and $A_y^\intercal$ commute. But normality is a property of the transformed drift, so a symmetric $A$ is fine with $D = \sigma^2 I$ (this is why thermox.linalg.solve/inv are not affected; expnegm of a non-normal matrix is), but not with a general $D$ (unless $A$ and $D$ commute); a non-symmetric $A$ is wrong even with $D = I$.

The mean is not affected: thermox computes it from the eigendecomposition of $A_y$ itself (expm_vp). The problem is the covariance, and three functions inherit it: sample, conditional.covariance and log_prob (and its gradients).

This PR evaluates $\Sigma_t$ exactly in the eigenbasis of $A_y$ that thermox already computes for the mean:

$$\Sigma_t = V \left[ Q_{ij} \frac{1 - e^{-(\lambda_i + \bar{\lambda}_j) t}}{\lambda_i + \bar{\lambda}_j} \right]_{ij} V^{H}, \qquad A_y = V \Lambda V^{-1}, \quad Q = V^{-1} V^{-H},$$

which holds for any stable, diagonalizable $A_y$ and reduces to the old formula when $A_y$ is normal ($V$ unitary, $Q = I$). In code: transition_cov is this formula; preprocess_drift_matrix additionally stores $Q$ (noise_cov_eigbasis) and a normality flag (is_normal); transition_cov_eigh is the single place that branches on the flag, and sample, log_prob and conditional.covariance read the covariance only through it. A normal $A_y$ keeps the existing $O(d^2)$-per-step path, so nothing changes for current users. Otherwise, $\Sigma_t$ is eigendecomposed at each step, $O(d^3)$.

The new tests check covariance, log-probability and its gradient against Van Loan and Lyapunov references that do not go through thermox, and that the normal cases are unchanged. I tried to keep the edits as small as possible.

[junipertcy]

I also looked at why the old tests did not catch this. test_log_prob_numeric in tests/test_log_prob.py is the only upstream test whose reference does not come from thermox itself (it integrates the covariance numerically), and its case is non-normal: A symmetric $A$ with a dense $D$, where the covariance is 3.5% off. It asserts rtol=1e-2, though, and the actual discrepancy on main is below $2 \times 10^{-3}$. The reason is that the samples come from thermox's own model; the expected log-likelihood is stationary at the data-generating distribution, so the first-order error averages out over the transitions and what remains (a KL divergence) is second order in the covariance error. With this PR it is about $10^{-6}$, so rtol=1e-4 would fail on main but pass here.

test_mean_and_cov in tests/test_conditional.py does use a non-symmetric $A$ (covariance 3 to 4% off), but it compares conditional.covariance with samples drawn by thermox from the same formula, so the two sides move together and the test cannot see the error. I changed it to compare against expm and Van Loan references instead. Lastly, examples/matrix_exponentials is asymmetric by intent, but it uses orthogonal matrices, which are normal, so it happened to be exact.

[junipertcy]

The script below reproduces the numbers above on main and on this branch.

Standalone check (jax + thermox only): relative Frobenius error of the transition covariance against references that do not go through thermox

"""Relative Frobenius error of thermox's transition covariance against references
that do not go through thermox. Run it on main and on this branch."""
import jax, jax.numpy as jnp, thermox
jax.config.update("jax_enable_x64", True)

def van_loan(A, D, t):  # Sigma_t = int_0^t e^{-As} D e^{-A^T s} ds (block exponential)
    d = A.shape[0]
    F = jax.scipy.linalg.expm(jnp.block([[-A, D], [jnp.zeros((d, d)), A.T]]) * t)
    return F[:d, d:] @ jax.scipy.linalg.expm(-A.T * t)

def lyapunov(A, D):  # stationary covariance: A S + S A^T = D
    d = A.shape[0]; I = jnp.eye(d)
    K = jnp.kron(I, A) + jnp.kron(A, I)
    return jnp.linalg.solve(K, D.reshape(-1, order="F")).reshape(d, d, order="F")

def relerr(X, ref):
    return float(jnp.linalg.norm(X - ref) / jnp.linalg.norm(ref))

def is_normal(A, D):  # of the transformed drift A_y = L^{-1} A L, D = L L^T
    L = jnp.linalg.cholesky(D); Ay = jnp.linalg.solve(L, A @ L)
    return float(jnp.linalg.norm(Ay @ Ay.T - Ay.T @ Ay) / jnp.linalg.norm(Ay) ** 2) < 1e-10

A_sym = jnp.array([[3.0, 2, 1], [2, 4, 2], [1, 2, 5]])
D_dense = jnp.array([[1.0, 0.3, -0.1], [0.3, 1, 0.2], [-0.1, 0.2, 1]])
cases = [
    ("symmetric A, D = I",                          A_sym, jnp.eye(3)),
    ("rotation-like A = [[1,2],[-2,1]], D = I",     jnp.array([[1.0, 2], [-2, 1]]), jnp.eye(2)),
    ("symmetric A, D = diag(1,4,9)",                A_sym, jnp.diag(jnp.array([1.0, 4, 9]))),
    ("symmetric A, dense D   (test_log_prob)",      A_sym, D_dense),
    ("A = [[3,2.5],[2,4]], D = 2I  (test_conditional)", jnp.array([[3.0, 2.5], [2, 4]]), 2 * jnp.eye(2)),
    ("triangular A, D = I",                         jnp.array([[2.0, 1.5, 0], [0, 3, 1.5], [0, 0, 4]]), jnp.eye(3)),
]
print(f"{'case':50s} {'A_y normal':>10s} {'Sigma_t(0.7)':>13s} {'Sigma_inf':>10s} {'sampled Sigma_inf':>18s}")
for name, A, D in cases:
    d = A.shape[0]; S_inf = lyapunov(A, D)
    e_t = relerr(thermox.conditional.covariance(0.7, A, D), van_loan(A, D, 0.7))
    e_inf = relerr(thermox.conditional.covariance(200.0, A, D), S_inf)
    ts = jnp.arange(0.0, 10000.0, 0.5)
    xs = thermox.sample(jax.random.PRNGKey(0), ts, jnp.zeros(d), A, jnp.zeros(d), D)
    e_mc = relerr(jnp.cov(xs[2000:].T), S_inf)
    print(f"{name:50s} {str(is_normal(A, D)):>10s} {e_t:13.1e} {e_inf:10.1e} {e_mc:18.1e}")

On main:

case                                               A_y normal  Sigma_t(0.7)  Sigma_inf  sampled Sigma_inf
symmetric A, D = I                                       True       2.2e-15    5.3e-16            1.6e-02
rotation-like A = [[1,2],[-2,1]], D = I                  True       1.5e-16    0.0e+00            7.2e-03
symmetric A, D = diag(1,4,9)                            False       1.7e-01    2.6e-01            1.2e-01
symmetric A, dense D   (test_log_prob)                  False       3.5e-02    4.7e-02            2.9e-02
A = [[3,2.5],[2,4]], D = 2I  (test_conditional)         False       3.1e-02    4.0e-02            1.4e-02
triangular A, D = I                                     False       6.3e-02    9.6e-02            3.3e-02

On this branch:

case                                               A_y normal  Sigma_t(0.7)  Sigma_inf  sampled Sigma_inf
symmetric A, D = I                                       True       2.8e-15    1.2e-15            1.6e-02
rotation-like A = [[1,2],[-2,1]], D = I                  True       2.9e-16    2.4e-16            7.2e-03
symmetric A, D = diag(1,4,9)                            False       2.4e-15    4.8e-16            1.2e-02
symmetric A, dense D   (test_log_prob)                  False       2.2e-15    1.2e-15            1.3e-02
A = [[3,2.5],[2,4]], D = 2I  (test_conditional)         False       6.6e-16    2.5e-16            2.6e-03
triangular A, D = I                                     False       7.3e-16    5.7e-16            1.7e-02

The last column is the empirical covariance of a thermox.sample trajectory (18k points), so it carries about 1e-2 of Monte Carlo noise on top of any bias; the first two rows show that noise floor.

[SamDuffield]

Hi!! Thanks for investigating and highlighting this! The $O(Td^3)$ cost is scaring me though, do you think there might be a way to support non-normal matrices with cost $O(d^3 + Td^2)$ ?

[junipertcy]

Hey Sam, thanks for the speedy reply. Yes! When the time grid is uniform, the transition operator can be built once, from the block matrix exponential, without using the eigenvectors, and applied to every step by a matmul. So a non-normal matrix costs $O(d^3 + T d^2)$, or $O(d^3 \log T + T d^2)$ with the associative scan. Whether we can improve it for arbitrary grids is open though. Is sampling on non-uniform grids something you expect to need?

This second commit is larger than the first, so here is the map when you review: sample and log_prob now first check whether the time grid is uniform after its first gap (like the burn-in grids in thermox.linalg). If it is, and the transformed drift is not normal, we build the transition operator once in transition_expm_and_cov and apply it to every step: for sample with a scan, either the sequential one or a tree, and for log_prob with a single matmul over all steps, since the density needs no scan. I left everything else as it was, and the normal path still agrees with main to rounding.

Let me know if you find anything.

[SamDuffield]

Yeah I'd say that supporting non-uniform grids is a requirement

[junipertcy]

Good challenge! I might have a idea (solving a smaller optimization problem at every time step). Let me dig around and update.

[junipertcy]

Update: My idea did not work. I wanted to work in the eigenbasis of $A_y$, but then realized that $\Sigma_t$ will span the the whole $S_d(\mathbb{R})$, so full rank, dim = $d(d+1)/2$. Propagating $\Sigma_t$ across time makes it rotate. That's two dense matrices multiplying with each other. Without assumptions on $A_y$, there is no "smaller optimization problem".

And then I learned that we can use established tricks on exact discretization of linear SDEs (because the transition operators compose along the semigroup) and approximation theory, respectively, to reach the bound that you are asking for. Specifically, we used "dyadic ladder" for sampling and "piecewise Chebyshev interpolation" for the probability.

This gives $O( U d^3 + T d^2 )$ for both. For sampling, $U$ is a true constant. ($U$ is the number of binary digits; float64 has $U=53$ and float32 has $U=24$). For probability density, $U \sim \log T$ because it depends on how spread out the time gaps are. I measured that $U$ can go from 170 at $T = 100$ to about 390 at $T = 10^5$. There's a hard cap of $U = 17 \times 53$ though because we will hit the number of significant binary digits (that's $53$) of a float64, and we always use $17$ nodes to fit each time piece for a similar reason.

This PR has been a wild ride for me! Please ask me questions. Not this weekend, but next week I will have time to iron out the missing details. As of now, the PR is backward compatible, and can take non-normal matrices at uniform or non-uniform grids.