diff --git a/src/tinygp/kernels/quasisep.py b/src/tinygp/kernels/quasisep.py index 2ed03f61..f28ca81b 100644 --- a/src/tinygp/kernels/quasisep.py +++ b/src/tinygp/kernels/quasisep.py @@ -41,6 +41,12 @@ 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 @@ -48,9 +54,9 @@ class Quasisep(Kernel): :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: diff --git a/tests/test_kernels/test_quasisep_nonreversible.py b/tests/test_kernels/test_quasisep_nonreversible.py new file mode 100644 index 00000000..b3dffee4 --- /dev/null +++ b/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 + )