Support non-reversible quasiseparable kernels

src/tinygp/kernels/quasisep.py

@@ -41,16 +41,22 @@
 from tinygp.solvers.quasisep.general import GeneralQSM
 
 
+def _matrix_transpose(matrix: JAXArray) -> JAXArray:
+    if isinstance(matrix, Block):
+        return matrix.mT
+    return jnp.swapaxes(matrix, -1, -2)
+
+
 class Quasisep(Kernel):
     """The base class for all quasiseparable kernels
 
     Instead of directly implementing the ``p``, ``q``, and ``a`` elements of the
     :class:`tinygp.solvers.quasisep.core.StrictLowerQSM`, this class implements
     ``h``, ``Pinf``, and ``A``, where:
 
-    - ``q = h``,
-    - ``p = h.T @ Pinf @ A``, and
-    - ``a = A``.
+    - ``q = h.T @ Pinf.T``,
+    - ``p = h.T @ A.T``, and
+    - ``a = A.T``.
 
     This notation follows the notation from state space models for stochastic
     differential equations, and so far it seems like a good way to specify these
@@ -75,7 +81,13 @@ def observation_model(self, X: JAXArray) -> JAXArray:
 
     @abstractmethod
     def transition_matrix(self, X1: JAXArray, X2: JAXArray) -> JAXArray:
-        """The transition matrix between two coordinates"""
+        """The adjoint transition matrix between two coordinates.
+
+        If a column-state mean propagates from ``X1`` to ``X2`` as
+        ``m2 = F @ m1``, this method must return ``F.T``. Equivalently, tinygp's
+        Kalman implementation propagates means using
+        ``transition_matrix(X1, X2).T @ m1``.
+        """
         raise NotImplementedError
 
     def coord_to_sortable(self, X: JAXArray) -> JAXArray:
@@ -94,10 +106,13 @@ def to_symm_qsm(self, X: JAXArray) -> SymmQSM:
             jax.tree_util.tree_map(lambda y: jnp.append(y[0], y[:-1]), X), X
         )
         h = jax.vmap(self.observation_model)(X)
-        q = h
-        p = h @ Pinf
-        d = jnp.sum(p * q, axis=1)
-        p = jax.vmap(lambda x, y: x @ y)(p, a)
+        hP = h @ Pinf
+        d = jnp.sum(hP * h, axis=1)
+        a = _matrix_transpose(a)
+        p = jax.vmap(lambda x, y: x @ y)(h, a)
+        # Pinf is a covariance matrix, so it is symmetric and hP is also the
+        # right generator required by the forward-oriented representation.
+        q = hP
         return SymmQSM(diag=DiagQSM(d=d), lower=StrictLowerTriQSM(p=p, q=q, a=a))
 
     def to_general_qsm(self, X1: JAXArray, X2: JAXArray) -> GeneralQSM:
@@ -107,22 +122,25 @@ def to_general_qsm(self, X1: JAXArray, X2: JAXArray) -> GeneralQSM:
 
         Xs = jax.tree_util.tree_map(lambda x: jnp.append(x[0], x[:-1]), X2)
         Pinf = self.stationary_covariance()
-        a = jax.vmap(self.transition_matrix)(Xs, X2)
+        a_adjoint = jax.vmap(self.transition_matrix)(Xs, X2)
+        a = _matrix_transpose(a_adjoint)
         h1 = jax.vmap(self.observation_model)(X1)
         h2 = jax.vmap(self.observation_model)(X2)
 
-        ql = h2
-        pl = h1 @ Pinf
-        qu = h1
-        pu = h2 @ Pinf
+        ql = h2 @ Pinf.T
+        pl = h1
+        qu = h1 @ Pinf
+        pu = h2
 
         i = jnp.clip(idx, 0, ql.shape[0] - 1)
         Xi = jax.tree_util.tree_map(lambda x: jnp.asarray(x)[i], X2)
-        pl = jax.vmap(lambda x, y: x @ y)(pl, jax.vmap(self.transition_matrix)(Xi, X1))
+        transition = jax.vmap(self.transition_matrix)(Xi, X1)
+        pl = jax.vmap(lambda x, y: x @ y.T)(pl, transition)
 
         i = jnp.clip(idx + 1, 0, pu.shape[0] - 1)
         Xi = jax.tree_util.tree_map(lambda x: jnp.asarray(x)[i], X2)
-        qu = jax.vmap(lambda x, y: x @ y)(jax.vmap(self.transition_matrix)(X1, Xi), qu)
+        transition = jax.vmap(self.transition_matrix)(X1, Xi)
+        qu = jax.vmap(lambda x, y: x @ y)(qu, transition)
 
         return GeneralQSM(pl=pl, ql=ql, pu=pu, qu=qu, a=a, idx=idx)
 
@@ -187,8 +205,8 @@ def evaluate(self, X1: JAXArray, X2: JAXArray) -> JAXArray:
         h2 = self.observation_model(X2)
         return jnp.where(
             self.coord_to_sortable(X1) < self.coord_to_sortable(X2),
-            h2 @ Pinf @ self.transition_matrix(X1, X2) @ h1,
-            h1 @ Pinf @ self.transition_matrix(X2, X1) @ h2,
+            h2 @ self.transition_matrix(X1, X2).T @ Pinf @ h1,
+            h1 @ self.transition_matrix(X2, X1).T @ Pinf @ h2,
         )
 
     def evaluate_diag(self, X: JAXArray) -> JAXArray:

tests/test_kernels/test_quasisep_nonreversible.py

@@ -0,0 +1,106 @@
+import jax
+import jax.numpy as jnp
+import jax.scipy as jsp
+
+from tinygp import GaussianProcess
+from tinygp.kernels import quasisep
+from tinygp.solvers import DirectSolver, QuasisepSolver
+from tinygp.solvers.kalman import KalmanSolver
+from tinygp.test_utils import assert_allclose
+
+
+class CausalFilter(quasisep.Quasisep):
+    """A minimal non-reversible two-state process: driver -> response."""
+
+    def design_matrix(self):
+        return jnp.array([[-1.0, 0.0], [0.8, -2.0]])
+
+    def stationary_covariance(self):
+        return jnp.array([[0.5, 2.0 / 15.0], [2.0 / 15.0, 91.0 / 300.0]])
+
+    def observation_model(self, X):
+        _time, channel = X
+        return jnp.eye(2)[channel]
+
+    def coord_to_sortable(self, X):
+        time, _channel = X
+        return time
+
+    def transition_matrix(self, X1, X2):
+        time1, _channel1 = X1
+        time2, _channel2 = X2
+        return jsp.linalg.expm(self.design_matrix().T * (time2 - time1))
+
+
+def direct_covariance(kernel, X1, X2):
+    def evaluate(x1, x2):
+        t1, _ = x1
+        t2, _ = x2
+        h1 = kernel.observation_model(x1)
+        h2 = kernel.observation_model(x2)
+        Pinf = kernel.stationary_covariance()
+        return jax.lax.cond(
+            t1 < t2,
+            lambda: h2 @ kernel.transition_matrix(x1, x2).T @ Pinf @ h1,
+            lambda: h1 @ kernel.transition_matrix(x2, x1).T @ Pinf @ h2,
+        )
+
+    return jax.vmap(lambda x1: jax.vmap(lambda x2: evaluate(x1, x2))(X2))(X1)
+
+
+def test_nonreversible_covariance_and_cross_matmul():
+    kernel = CausalFilter()
+    time = jnp.array([0.0, 0.3, 0.8, 1.4, 2.2, 3.1])
+    channel = jnp.array([0, 1, 0, 1, 1, 0])
+    X = (time, channel)
+    X_test = (jnp.array([0.1, 0.9, 1.8, 2.8]), jnp.array([1, 0, 1, 0]))
+
+    expected = direct_covariance(kernel, X, X)
+    assert_allclose(kernel(X, X), expected)
+    assert_allclose(kernel.to_symm_qsm(X).to_dense(), expected)
+
+    cross = direct_covariance(kernel, X_test, X)
+    y = jnp.linspace(-0.5, 0.7, time.size)
+    assert_allclose(kernel.matmul(X_test, X, y), cross @ y)
+
+    # The response after a driver impulse differs from the reverse ordering;
+    # reversible test kernels cannot expose this orientation requirement.
+    driver_then_response = direct_covariance(
+        kernel, (jnp.array([1.0]), jnp.array([1])), (jnp.array([0.0]), jnp.array([0]))
+    )[0, 0]
+    response_then_driver = direct_covariance(
+        kernel, (jnp.array([1.0]), jnp.array([0])), (jnp.array([0.0]), jnp.array([1]))
+    )[0, 0]
+    assert not jnp.isclose(driver_then_response, response_then_driver)
+
+
+def test_nonreversible_solvers_and_conditioning_agree():
+    kernel = CausalFilter()
+    time = jnp.array([0.0, 0.3, 0.8, 1.4, 2.2, 3.1])
+    channel = jnp.array([0, 1, 0, 1, 1, 0])
+    X = (time, channel)
+    y = jnp.array([0.2, -0.1, 0.3, 0.15, -0.2, 0.05])
+    diag = jnp.full(time.shape, 0.05)
+
+    gp_direct = GaussianProcess(kernel, X, diag=diag, solver=DirectSolver)
+    gp_quasisep = GaussianProcess(kernel, X, diag=diag, solver=QuasisepSolver)
+    gp_kalman = GaussianProcess(kernel, X, diag=diag, solver=KalmanSolver)
+
+    assert_allclose(gp_quasisep.covariance, gp_direct.covariance)
+    assert_allclose(gp_quasisep.log_probability(y), gp_direct.log_probability(y))
+    assert_allclose(gp_kalman.log_probability(y), gp_direct.log_probability(y))
+
+    conditioned_direct = gp_direct.condition(y)
+    conditioned_quasisep = gp_quasisep.condition(y)
+    assert_allclose(conditioned_quasisep.gp.loc, conditioned_direct.gp.loc)
+    assert_allclose(
+        conditioned_quasisep.gp.covariance, conditioned_direct.gp.covariance
+    )
+
+    X_test = (jnp.array([0.1, 0.9, 1.8, 2.8]), jnp.array([1, 0, 1, 0]))
+    conditioned_direct = gp_direct.condition(y, X_test=X_test)
+    conditioned_quasisep = gp_quasisep.condition(y, X_test=X_test)
+    assert_allclose(conditioned_quasisep.gp.loc, conditioned_direct.gp.loc)
+    assert_allclose(
+        conditioned_quasisep.gp.covariance, conditioned_direct.gp.covariance
+    )