From 90928a8063c0e8de21be9da1a620b5460d5d0794 Mon Sep 17 00:00:00 2001 From: Dan F-M Date: Tue, 8 Sep 2026 15:03:44 -0400 Subject: [PATCH 1/2] Make QuasisepSolver.condition respect parallel=True Fixes #280. The quasiseparable conditional covariance was built with three sequential scans inside qsm_mul, and the conditioned GaussianProcess did not inherit the parallel flag, so its Cholesky was sequential too. This adds associative-scan versions of the two qsm_mul recurrences, threads a parallel kwarg through qsm_mul and SquareQSM.gram, and forwards parallel to the conditioned GP when its covariance is a SymmQSM. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01JTMk4JdA8mBFozHWnxqUqy --- news/280.bugfix | 6 ++ src/tinygp/gp.py | 10 +++ src/tinygp/solvers/quasisep/core.py | 10 ++- src/tinygp/solvers/quasisep/ops.py | 81 ++++++++++++++----- src/tinygp/solvers/quasisep/solver.py | 12 ++- tests/test_solvers/test_quasisep/test_ops.py | 24 +++++- .../test_solvers/test_quasisep/test_solver.py | 1 + 7 files changed, 119 insertions(+), 25 deletions(-) create mode 100644 news/280.bugfix diff --git a/news/280.bugfix b/news/280.bugfix new file mode 100644 index 00000000..d2f83457 --- /dev/null +++ b/news/280.bugfix @@ -0,0 +1,6 @@ +Fixed ``parallel=True`` not reaching the ``QuasisepSolver.condition`` path when +conditioning at the training coordinates: the quasiseparable matrix products +used to build the conditional covariance now use parallel associative scans, and +the conditioned ``GaussianProcess`` inherits the ``parallel`` flag from its +parent so that its Cholesky factorization and all subsequent operations are +parallel as well. diff --git a/src/tinygp/gp.py b/src/tinygp/gp.py index b00203b9..864f0661 100644 --- a/src/tinygp/gp.py +++ b/src/tinygp/gp.py @@ -187,6 +187,15 @@ def condition( if X_test is None: X_test = self.X + # When the conditional covariance is quasiseparable, the new GP will + # use a QuasisepSolver, and it should inherit the parallel flag from + # this GP's solver + solver_kwargs: dict[str, Any] = {} + if isinstance(self.solver, QuasisepSolver) and isinstance( + covariance_value, SymmQSM + ): + solver_kwargs["parallel"] = self.solver.parallel + # The conditional GP will also be a GP with the mean an covariance # specified by a :class:`tinygp.means.Conditioned` and # :class:`tinygp.kernels.Conditioned` respectively. @@ -203,6 +212,7 @@ def condition( ), mean_value=mean_value, covariance_value=covariance_value, + **solver_kwargs, ) return ConditionResult(log_prob, gp) diff --git a/src/tinygp/solvers/quasisep/core.py b/src/tinygp/solvers/quasisep/core.py index c799ea64..fd772ec2 100644 --- a/src/tinygp/solvers/quasisep/core.py +++ b/src/tinygp/solvers/quasisep/core.py @@ -421,16 +421,22 @@ def scale(self, other: JAXArray) -> SquareQSM: upper=self.upper.scale(other), ) - def gram(self) -> SymmQSM: + def gram(self, *, parallel: bool = False) -> SymmQSM: """The inner product of this matrix with itself If this matrix is called ``A``, the Gram matrix is ``A.T @ A``, and that's what this method computes. The result is a :class:`SymmQSM`. + + Args: + parallel: If ``True``, use parallel associative-scan algorithms + for the matrix product. """ + from tinygp.solvers.quasisep.ops import qsm_mul + # We know that this must result in symmetric matrix, but that won't be # enforced; we make it so! It might be possible to make this more # efficient, but perhaps jax is clever enough? - M = self.transpose() @ self + M = qsm_mul(self.transpose(), self, parallel=parallel) return SymmQSM(diag=M.diag, lower=M.lower) @jax.jit diff --git a/src/tinygp/solvers/quasisep/ops.py b/src/tinygp/solvers/quasisep/ops.py index 95c9d26a..16c13d65 100644 --- a/src/tinygp/solvers/quasisep/ops.py +++ b/src/tinygp/solvers/quasisep/ops.py @@ -2,6 +2,7 @@ __all__ = ["elementwise_add", "elementwise_mul", "qsm_mul"] +from functools import partial from typing import TypeVar import jax @@ -49,8 +50,16 @@ def elementwise_mul(a: QSM, b: QSM) -> QSM | None: return construct(diag, lower, upper, is_symm_a and is_symm_b) -@jax.jit -def qsm_mul(a: QSM, b: QSM) -> QSM | None: +@partial(jax.jit, static_argnames=("parallel",)) +def qsm_mul(a: QSM, b: QSM, *, parallel: bool = False) -> QSM | None: + """The product of two quasiseparable matrices + + Args: + a: The left matrix. + b: The right matrix. + parallel: If ``True``, use parallel associative-scan algorithms for the + two recurrences, instead of sequential scans. + """ diag_a, lower_a, upper_a = deconstruct(a) diag_b, lower_b, upper_b = deconstruct(b) @@ -60,28 +69,14 @@ def qsm_mul(a: QSM, b: QSM) -> QSM | None: return DiagQSM(d=diag_a * diag_b) if lower_a is not None and upper_b is not None: - - def calc_phi(phi, data): # type: ignore - a, b, q, g = data - return a @ phi @ b.T + jnp.outer(q, g), phi - - init = jnp.zeros_like(jnp.outer(lower_a.q[0], upper_b.q[0])) - args = (lower_a.a, upper_b.a, lower_a.q, upper_b.q) - _, phi = jax.lax.scan(calc_phi, init, args) - + impl = qsm_mul_phi_parallel if parallel else qsm_mul_phi + phi = impl(lower_a.a, upper_b.a, lower_a.q, upper_b.q) else: phi = None if upper_a is not None and lower_b is not None: - - def calc_psi(psi, data): # type: ignore - a, b, q, g = data - return a.T @ psi @ b + jnp.outer(q, g), psi - - init = jnp.zeros_like(jnp.outer(upper_a.q[-1], lower_b.p[-1])) - args = (upper_a.a, lower_b.a, upper_a.p, lower_b.p) - _, psi = jax.lax.scan(calc_psi, init, args, reverse=True) - + impl = qsm_mul_psi_parallel if parallel else qsm_mul_psi + psi = impl(upper_a.a, lower_b.a, upper_a.p, lower_b.p) else: psi = None @@ -305,6 +300,52 @@ def _shift_bwd(x): return jnp.concatenate((x[1:], jnp.zeros_like(x[-1:])), axis=0) +@jax.jit +def qsm_mul_phi(a, b, q, g): + def impl(phi, data): + a, b, q, g = data + return a @ phi @ b.T + jnp.outer(q, g), phi + + init = jnp.zeros_like(jnp.outer(q[0], g[0])) + _, phi = jax.lax.scan(impl, init, (a, b, q, g)) + return phi + + +@jax.jit +def qsm_mul_phi_parallel(a, b, q, g): + # The recurrence is affine in phi, with a two-sided linear part: + # phi -> a @ phi @ b.T + B. The composition of two such maps is again of + # this form, so we scan over the triples (a, b, B). + B = jnp.einsum("nj,nk->njk", q, g) + _, _, phi = jax.lax.associative_scan(_two_sided_affine_combine, (a, b, B)) + return _shift_fwd(phi) + + +@jax.jit +def qsm_mul_psi(a, b, q, g): + def impl(psi, data): + a, b, q, g = data + return a.T @ psi @ b + jnp.outer(q, g), psi + + init = jnp.zeros_like(jnp.outer(q[-1], g[-1])) + _, psi = jax.lax.scan(impl, init, (a, b, q, g), reverse=True) + return psi + + +@jax.jit +def qsm_mul_psi_parallel(a, b, q, g): + B = jnp.einsum("nj,nk->njk", q, g) + _, _, psi = jax.lax.associative_scan( + _two_sided_affine_combine, (a.mT, b.mT, B), reverse=True + ) + return _shift_bwd(psi) + + +def _two_sided_affine_combine(left, right): + (al, bl, Bl), (ar, br, Br) = left, right + return ar @ al, br @ bl, ar @ Bl @ br.mT + Br + + @jax.jit def lower_matmul(p, q, a, x): def impl(f, data): diff --git a/src/tinygp/solvers/quasisep/solver.py b/src/tinygp/solvers/quasisep/solver.py index bdb63137..8e4eaeba 100644 --- a/src/tinygp/solvers/quasisep/solver.py +++ b/src/tinygp/solvers/quasisep/solver.py @@ -101,6 +101,14 @@ def solve_triangular(self, y: JAXArray, *, transpose: bool = False) -> JAXArray: def dot_triangular(self, y: JAXArray) -> JAXArray: return self.factor.matmul(y, parallel=self.parallel) + def _conditional_delta(self, M: SymmQSM) -> SymmQSM: + # The (QSM) term M @ K^{-1} @ M = (L^{-1} @ M)^T @ (L^{-1} @ M) that + # gets subtracted from M when conditioning at the input coordinates + from tinygp.solvers.quasisep.ops import qsm_mul + + A = qsm_mul(self.factor.inv(), M, parallel=self.parallel) + return A.gram(parallel=self.parallel) + def condition(self, kernel: Kernel, X_test: JAXArray | None, noise: Noise) -> Any: """Compute the covariance matrix for a conditional GP @@ -124,7 +132,7 @@ def condition(self, kernel: Kernel, X_test: JAXArray | None, noise: Noise) -> An # where we are predicting at the input coordinates and a Quasisep kernel if X_test is None and isinstance(kernel, Quasisep): M = kernel.to_symm_qsm(self.X) - delta = (self.factor.inv() @ M).gram() + delta = self._conditional_delta(M) M += noise.to_qsm() return M - delta @@ -154,7 +162,7 @@ def condition_diag( if X_test is None and isinstance(kernel, Quasisep): M = kernel.to_symm_qsm(self.X) - delta = (self.factor.inv() @ M).gram() + delta = self._conditional_delta(M) M += noise.to_qsm() return M.diag.d - delta.diag.d diff --git a/tests/test_solvers/test_quasisep/test_ops.py b/tests/test_solvers/test_quasisep/test_ops.py index 20da5013..2d03b7d3 100644 --- a/tests/test_solvers/test_quasisep/test_ops.py +++ b/tests/test_solvers/test_quasisep/test_ops.py @@ -5,7 +5,12 @@ from numpy import random as np_random from tinygp.kernels.quasisep import Matern32, Matern52 -from tinygp.solvers.quasisep.core import DiagQSM +from tinygp.solvers.quasisep.core import ( + DiagQSM, + LowerTriQSM, + StrictLowerTriQSM, + SymmQSM, +) from tinygp.solvers.quasisep.ops import ( cholesky, cholesky_parallel, @@ -13,6 +18,7 @@ lower_matmul_parallel, lower_solve, lower_solve_parallel, + qsm_mul, symm_inv, symm_inv_parallel, upper_matmul, @@ -74,3 +80,19 @@ def test_symm_inv_parallel(data): assert_allclose(t_p, t_s) assert_allclose(s_p, s_s) assert_allclose(ell_p, ell_s) + + +def test_qsm_mul_parallel(data): + d, p, q, a, x = data + del x + symm = SymmQSM(diag=DiagQSM(d=d), lower=StrictLowerTriQSM(p=p, q=q, a=a)) + lower = LowerTriQSM(diag=symm.diag, lower=symm.lower) + upper = lower.transpose() + square = lower @ symm + + # Exercise: only the phi scan, only the psi scan, and both scans + for left, right in [(lower, symm), (upper, symm), (square.T, square)]: + seq = qsm_mul(left, right) + par = qsm_mul(left, right, parallel=True) + assert_allclose(par.to_dense(), seq.to_dense()) + assert type(par) is type(seq) diff --git a/tests/test_solvers/test_quasisep/test_solver.py b/tests/test_solvers/test_quasisep/test_solver.py index bd4704e6..51e69ccc 100644 --- a/tests/test_solvers/test_quasisep/test_solver.py +++ b/tests/test_solvers/test_quasisep/test_solver.py @@ -81,6 +81,7 @@ def test_consistent_with_direct(kernel_pair, data, parallel): gp1p = gp1.condition(y) gp2p = gp2.condition(y) assert isinstance(gp1p.gp.solver, QuasisepSolver) + assert gp1p.gp.solver.parallel == parallel assert_allclose(gp1p.log_probability, gp2p.log_probability) assert_allclose(gp1p.gp.loc, gp2p.gp.loc) assert_allclose(gp1p.gp.variance, gp2p.gp.variance) From ddc7483634f29e5ef2ea4bc6477a77c7fa6e4244 Mon Sep 17 00:00:00 2001 From: Dan F-M Date: Tue, 8 Sep 2026 16:04:11 -0400 Subject: [PATCH 2/2] Use a rank-J representation for the same-kernel conditional covariance Adversarial review of the previous commit found that forwarding parallel=True to the conditioned GP exposed a numerical weakness: the conditional covariance built as M - (L^{-1} M)^T (L^{-1} M) has 4J generators that encode the difference of two nearly equal matrices, and cholesky_parallel loses many digits (or produces NaN) on that representation. When conditioning with the GP's own kernel, the conditional covariance is N - N K^{-1} N, which has rank J and is far better conditioned for both the sequential and the parallel Cholesky. To use it, kernel=None is now passed through to Solver.condition/condition_diag as the same-kernel signal, and both built-in solvers store their kernel (and QuasisepSolver its noise). Also: QuasisepSolver.covariance respects parallel; qsm_mul no longer mislabels the product of two different SymmQSMs as symmetric; tests cover Block transitions, symmetric operands, Banded noise, chained conditioning, and the conditioned GP's log_probability/sample in both modes. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01JTMk4JdA8mBFozHWnxqUqy --- news/280.bugfix | 19 +++-- src/tinygp/gp.py | 9 ++- src/tinygp/solvers/direct.py | 9 ++- src/tinygp/solvers/kalman.py | 4 +- src/tinygp/solvers/quasisep/ops.py | 8 +- src/tinygp/solvers/quasisep/solver.py | 71 ++++++++++++++-- src/tinygp/solvers/solver.py | 23 +++++- tests/test_solvers/test_quasisep/test_ops.py | 32 +++++++- .../test_solvers/test_quasisep/test_solver.py | 80 +++++++++++++++++++ 9 files changed, 228 insertions(+), 27 deletions(-) diff --git a/news/280.bugfix b/news/280.bugfix index d2f83457..537b7ca6 100644 --- a/news/280.bugfix +++ b/news/280.bugfix @@ -1,6 +1,15 @@ Fixed ``parallel=True`` not reaching the ``QuasisepSolver.condition`` path when -conditioning at the training coordinates: the quasiseparable matrix products -used to build the conditional covariance now use parallel associative scans, and -the conditioned ``GaussianProcess`` inherits the ``parallel`` flag from its -parent so that its Cholesky factorization and all subsequent operations are -parallel as well. +conditioning at the training coordinates. The quasiseparable matrix products +used to build the conditional covariance now support parallel associative +scans, the conditioned ``GaussianProcess`` inherits the ``parallel`` flag from +its parent, and ``QuasisepSolver.covariance`` also respects the flag. As part +of this change, conditioning at the training coordinates with the same kernel +that the GP was built with now uses the identity ``N - N @ K^{-1} @ N`` for the +conditional covariance, which has the same quasiseparable rank as the prior +kernel (instead of four times that rank) and is much better conditioned +numerically. To support this, ``kernel=None`` is now passed through to +``Solver.condition`` and ``Solver.condition_diag`` to signal that the solver's +own kernel should be used, and solvers are expected to store that kernel as +``self.kernel``; third-party ``Solver`` implementations will need to handle +this. Also fixed a bug where the product of two different ``SymmQSM`` matrices +was incorrectly returned as a ``SymmQSM``, dropping its upper triangle. diff --git a/src/tinygp/gp.py b/src/tinygp/gp.py index 864f0661..a87d7173 100644 --- a/src/tinygp/gp.py +++ b/src/tinygp/gp.py @@ -176,14 +176,17 @@ def condition( _check_test_shapes(self.X, X_test) alpha, log_prob, mean_value = self._condition(y, X_test, include_mean, kernel) - if kernel is None: - kernel = self.kernel if noise is None: diag = _default_diag(mean_value) if diag is None else diag noise = Diagonal(diag=jnp.broadcast_to(diag, mean_value.shape)) + # Note: ``kernel`` is passed through unresolved here, since ``None`` + # signals to the solver that we're conditioning using the same kernel + # that it was built with, and it can use a more efficient algorithm covariance_value = self.solver.condition(kernel, X_test, noise) + if kernel is None: + kernel = self.kernel if X_test is None: X_test = self.X @@ -283,7 +286,7 @@ def predict( _, _, mean_value = self._condition(y, None, include_mean, kernel) diag = _default_diag(mean_value) noise = Diagonal(diag=jnp.broadcast_to(diag, mean_value.shape)) - var_value = self.solver.condition_diag(pred_kernel, None, noise) + var_value = self.solver.condition_diag(kernel, None, noise) return mean_value, var_value # Predicting at new test points: build `Ks` once and reuse it for diff --git a/src/tinygp/solvers/direct.py b/src/tinygp/solvers/direct.py index ba4048ac..dcbe176b 100644 --- a/src/tinygp/solvers/direct.py +++ b/src/tinygp/solvers/direct.py @@ -22,6 +22,7 @@ class DirectSolver(Solver): usual constructor. """ + kernel: kernels.Kernel X: JAXArray variance_value: JAXArray covariance_value: JAXArray @@ -45,6 +46,7 @@ def __init__( matrix. This should be equal to the result of calling ``kernel`` and adding ``diag``, but that is not checked. """ + self.kernel = kernel self.X = X self.variance_value = kernel(X) + noise.diagonal() if covariance is None: @@ -73,17 +75,20 @@ def dot_triangular(self, y: JAXArray) -> JAXArray: return jnp.einsum("ij,j...->i...", self.scale_tril, y) def condition( - self, kernel: kernels.Kernel, X_test: JAXArray | None, noise: Noise + self, kernel: kernels.Kernel | None, X_test: JAXArray | None, noise: Noise ) -> Any: """Compute the covariance matrix for a conditional GP Args: kernel: The kernel for the covariance between the observed and - predicted data. + predicted data. If ``None``, the kernel used to construct this + solver is used. X_test: The coordinates of the predicted points. Defaults to the input coordinates. noise: The noise model for the predicted process. """ + if kernel is None: + kernel = self.kernel if X_test is None: Ks = kernel(self.X, self.X) Kss = Ks + noise diff --git a/src/tinygp/solvers/kalman.py b/src/tinygp/solvers/kalman.py index d84f143f..fbd2704c 100644 --- a/src/tinygp/solvers/kalman.py +++ b/src/tinygp/solvers/kalman.py @@ -79,7 +79,9 @@ def dot_triangular(self, y: JAXArray) -> JAXArray: del y raise NotImplementedError - def condition(self, kernel: Kernel, X_test: JAXArray | None, noise: Noise) -> Any: + def condition( + self, kernel: Kernel | None, X_test: JAXArray | None, noise: Noise + ) -> Any: del kernel, X_test, noise raise NotImplementedError diff --git a/src/tinygp/solvers/quasisep/ops.py b/src/tinygp/solvers/quasisep/ops.py index 16c13d65..2cade0ef 100644 --- a/src/tinygp/solvers/quasisep/ops.py +++ b/src/tinygp/solvers/quasisep/ops.py @@ -204,9 +204,11 @@ def impl( diag, lower, upper = impl( diag_a, lower_a, upper_a, diag_b, lower_b, upper_b, phi, psi ) - is_symm_a = isinstance(a, (DiagQSM, SymmQSM)) - is_symm_b = isinstance(b, (DiagQSM, SymmQSM)) - return construct(diag, lower, upper, is_symm_a and is_symm_b) + # Note: the product of two symmetric matrices is not, in general, + # symmetric, so we always return the full (square) result here. Callers + # that know the result must be symmetric (e.g. ``SquareQSM.gram``) can + # re-wrap it. + return construct(diag, lower, upper, False) def deconstruct( diff --git a/src/tinygp/solvers/quasisep/solver.py b/src/tinygp/solvers/quasisep/solver.py index 8e4eaeba..5bb8aa5e 100644 --- a/src/tinygp/solvers/quasisep/solver.py +++ b/src/tinygp/solvers/quasisep/solver.py @@ -12,7 +12,12 @@ from tinygp.helpers import JAXArray from tinygp.kernels.base import Kernel from tinygp.noise import Noise -from tinygp.solvers.quasisep.core import LowerTriQSM, SymmQSM +from tinygp.solvers.quasisep.core import ( + DiagQSM, + LowerTriQSM, + StrictLowerTriQSM, + SymmQSM, +) from tinygp.solvers.solver import Solver @@ -27,7 +32,9 @@ class QuasisepSolver(Solver): usual constructor. """ + kernel: Kernel X: JAXArray + noise: Noise matrix: SymmQSM factor: LowerTriQSM parallel: bool = eqx.field(static=True) @@ -76,7 +83,9 @@ def __init__( if TYPE_CHECKING: assert isinstance(covariance, SymmQSM) matrix = covariance + self.kernel = kernel self.X = X + self.noise = noise self.matrix = matrix self.parallel = parallel self.factor = matrix.cholesky(parallel=parallel) @@ -85,7 +94,8 @@ def variance(self) -> JAXArray: return self.matrix.diag.d def covariance(self) -> JAXArray: - return self.matrix.to_dense() + N = self.matrix.shape[0] + return self.matrix.matmul(jnp.eye(N), parallel=self.parallel) def normalization(self) -> JAXArray: return jnp.sum(jnp.log(self.factor.diag.d)) + 0.5 * self.factor.shape[ @@ -101,15 +111,46 @@ def solve_triangular(self, y: JAXArray, *, transpose: bool = False) -> JAXArray: def dot_triangular(self, y: JAXArray) -> JAXArray: return self.factor.matmul(y, parallel=self.parallel) + def _conditional_at_data(self) -> SymmQSM: + """The conditional covariance at the input coordinates, as a QSM + + When conditioning on the observed data using the same kernel that + this solver was built with, and with ``K = M + N`` the full covariance + matrix, the conditional covariance ``M - M @ K^{-1} @ M`` simplifies + to ``N - N @ K^{-1} @ N``. This form only requires the inverse of + ``K``, and it has the same quasiseparable rank as the original + kernel. It is also much better conditioned than computing + ``M - M @ K^{-1} @ M`` directly, where the difference of two nearly + equal matrices is encoded in the generators. + """ + from tinygp.solvers.quasisep.ops import qsm_mul + + Kinv = self.matrix.inv(parallel=self.parallel) + N = self.noise.to_qsm() + if isinstance(N, DiagQSM): + n = N.d + lam = Kinv.diag.d + t, s, ell = Kinv.lower + return SymmQSM( + diag=DiagQSM(d=n - jnp.square(n) * lam), + lower=StrictLowerTriQSM(p=-n[:, None] * t, q=n[:, None] * s, a=ell), + ) + + P = qsm_mul(N, qsm_mul(Kinv, N, parallel=self.parallel), parallel=self.parallel) + return N - SymmQSM(diag=P.diag, lower=P.lower) + def _conditional_delta(self, M: SymmQSM) -> SymmQSM: # The (QSM) term M @ K^{-1} @ M = (L^{-1} @ M)^T @ (L^{-1} @ M) that # gets subtracted from M when conditioning at the input coordinates + # with a general (cross-)kernel M from tinygp.solvers.quasisep.ops import qsm_mul A = qsm_mul(self.factor.inv(), M, parallel=self.parallel) return A.gram(parallel=self.parallel) - def condition(self, kernel: Kernel, X_test: JAXArray | None, noise: Noise) -> Any: + def condition( + self, kernel: Kernel | None, X_test: JAXArray | None, noise: Noise + ) -> Any: """Compute the covariance matrix for a conditional GP In the case where the prediction is made at the input coordinates with a @@ -121,14 +162,24 @@ def condition(self, kernel: Kernel, X_test: JAXArray | None, noise: Noise) -> An Args: kernel: The kernel for the covariance between the observed and - predicted data. + predicted data. If ``None``, the kernel used to construct this + solver is used, which enables a more efficient and numerically + stable algorithm when ``X_test`` is also ``None``. X_test: The coordinates of the predicted points. Defaults to the input coordinates. noise: The noise model for the predicted process. """ from tinygp.kernels.quasisep import Quasisep - # We can easily compute the conditional as a QSM in the special case + # The most common case: predicting at the input coordinates with the + # kernel that this solver was built with + if X_test is None and kernel is None: + return self._conditional_at_data() + noise.to_qsm() + + if kernel is None: + kernel = self.kernel + + # We can also compute the conditional as a QSM in the special case # where we are predicting at the input coordinates and a Quasisep kernel if X_test is None and isinstance(kernel, Quasisep): M = kernel.to_symm_qsm(self.X) @@ -147,11 +198,11 @@ def condition(self, kernel: Kernel, X_test: JAXArray | None, noise: Noise) -> An return Kss - A.transpose() @ A def condition_diag( - self, kernel: Kernel, X_test: JAXArray | None, noise: Noise + self, kernel: Kernel | None, X_test: JAXArray | None, noise: Noise ) -> JAXArray: """The diagonal of the covariance matrix for a conditional GP - This reuses the same quasiseparable special case as :func:`condition`: + This reuses the same quasiseparable special cases as :func:`condition`: when predicting at the input coordinates with a :class:`tinygp.kernels.quasisep.Quasisep` kernel, the diagonal can be computed in ``O(N)`` (or ``O(N log N)`` with the parallel algorithms) @@ -160,6 +211,12 @@ def condition_diag( """ from tinygp.kernels.quasisep import Quasisep + if X_test is None and kernel is None: + return self._conditional_at_data().diag.d + noise.diagonal() + + if kernel is None: + kernel = self.kernel + if X_test is None and isinstance(kernel, Quasisep): M = kernel.to_symm_qsm(self.X) delta = self._conditional_delta(M) diff --git a/src/tinygp/solvers/solver.py b/src/tinygp/solvers/solver.py index 9c3ad376..c65290c9 100644 --- a/src/tinygp/solvers/solver.py +++ b/src/tinygp/solvers/solver.py @@ -79,11 +79,25 @@ def dot_triangular(self, y: JAXArray) -> JAXArray: raise NotImplementedError @abstractmethod - def condition(self, kernel: Kernel, X_test: JAXArray | None, noise: Noise) -> Any: + def condition( + self, kernel: Kernel | None, X_test: JAXArray | None, noise: Noise + ) -> Any: + """Compute the covariance matrix for a conditional GP + + Args: + kernel: The kernel for the covariance between the observed and + predicted data. If ``None``, the kernel used to construct this + solver is used, and solvers can use this as a signal to enable + specialized algorithms. Solver implementations must store this + kernel as ``self.kernel``. + X_test: The coordinates of the predicted points. Defaults to the + input coordinates. + noise: The noise model for the predicted process. + """ raise NotImplementedError def condition_diag( - self, kernel: Kernel, X_test: JAXArray | None, noise: Noise + self, kernel: Kernel | None, X_test: JAXArray | None, noise: Noise ) -> JAXArray: """The diagonal of the covariance matrix for a conditional GP @@ -98,11 +112,14 @@ def condition_diag( Args: kernel: The kernel for the covariance between the observed and - predicted data. + predicted data. If ``None``, the kernel used to construct this + solver is used. X_test: The coordinates of the predicted points. Defaults to the input coordinates. noise: The noise model for the predicted process. """ + if kernel is None: + kernel = self.kernel # type: ignore if X_test is None: Ks = kernel(self.X, self.X) # type: ignore Kss_diag = kernel(self.X) # type: ignore diff --git a/tests/test_solvers/test_quasisep/test_ops.py b/tests/test_solvers/test_quasisep/test_ops.py index 2d03b7d3..6c4830de 100644 --- a/tests/test_solvers/test_quasisep/test_ops.py +++ b/tests/test_solvers/test_quasisep/test_ops.py @@ -4,7 +4,7 @@ import pytest from numpy import random as np_random -from tinygp.kernels.quasisep import Matern32, Matern52 +from tinygp.kernels.quasisep import SHO, Matern32, Matern52 from tinygp.solvers.quasisep.core import ( DiagQSM, LowerTriQSM, @@ -82,10 +82,19 @@ def test_symm_inv_parallel(data): assert_allclose(ell_p, ell_s) -def test_qsm_mul_parallel(data): +@pytest.mark.parametrize("use_block", [False, True], ids=["dense", "block"]) +def test_qsm_mul_parallel(data, use_block): d, p, q, a, x = data del x - symm = SymmQSM(diag=DiagQSM(d=d), lower=StrictLowerTriQSM(p=p, q=q, a=a)) + if use_block: + # A sum kernel gives block-structured transition matrices + N = len(d) + random = np_random.default_rng(5678) + t = jnp.sort(jnp.asarray(random.uniform(0, 10, N))) + kernel = Matern32(scale=1.3) + SHO(omega=1.5, quality=3.0) + symm = kernel.to_symm_qsm(t) + DiagQSM(jnp.full(N, 0.1)) + else: + symm = SymmQSM(diag=DiagQSM(d=d), lower=StrictLowerTriQSM(p=p, q=q, a=a)) lower = LowerTriQSM(diag=symm.diag, lower=symm.lower) upper = lower.transpose() square = lower @ symm @@ -95,4 +104,21 @@ def test_qsm_mul_parallel(data): seq = qsm_mul(left, right) par = qsm_mul(left, right, parallel=True) assert_allclose(par.to_dense(), seq.to_dense()) + assert_allclose(par.to_dense(), left.to_dense() @ right.to_dense()) assert type(par) is type(seq) + + +def test_qsm_mul_symmetric_operands(data): + # The product of two different symmetric matrices is not symmetric, so + # the full result should be returned + d, p, q, a, x = data + del x + N = len(d) + random = np_random.default_rng(91011) + t = jnp.sort(jnp.asarray(random.uniform(0, 10, N))) + symm1 = SymmQSM(diag=DiagQSM(d=d), lower=StrictLowerTriQSM(p=p, q=q, a=a)) + symm2 = SHO(omega=1.5, quality=3.0).to_symm_qsm(t) + for parallel in [False, True]: + prod = qsm_mul(symm1, symm2, parallel=parallel) + assert not isinstance(prod, SymmQSM) + assert_allclose(prod.to_dense(), symm1.to_dense() @ symm2.to_dense()) diff --git a/tests/test_solvers/test_quasisep/test_solver.py b/tests/test_solvers/test_quasisep/test_solver.py index 51e69ccc..4e634644 100644 --- a/tests/test_solvers/test_quasisep/test_solver.py +++ b/tests/test_solvers/test_quasisep/test_solver.py @@ -7,6 +7,7 @@ from tinygp import GaussianProcess, kernels from tinygp.kernels import quasisep +from tinygp.noise import Banded from tinygp.solvers import DirectSolver, QuasisepSolver from tinygp.test_utils import assert_allclose @@ -142,3 +143,82 @@ def impl(X, y): with pytest.raises(jax.errors.JaxRuntimeError) as exc_info: impl(x_, y_).block_until_ready() assert exc_info.match(r"Input coordinates must be sorted") + + +@pytest.mark.parametrize( + "kernel", + [ + quasisep.Matern32(sigma=1.8, scale=1.5), + quasisep.Cosine(sigma=1.2, scale=0.8), + quasisep.Matern32(sigma=1.8, scale=1.5) + + quasisep.SHO(omega=1.5, quality=3.0, sigma=1.2), + ], + ids=["matern32", "cosine", "matern32+sho"], +) +@pytest.mark.parametrize("parallel", [False, True], ids=["sequential", "parallel"]) +def test_conditioned_gp_operations(kernel, random, parallel): + # Conditioning at the training points with the default jitter gives a + # nearly singular covariance. This checks that the conditioned GP's own + # factorization (which used to be numerically fragile, especially with + # the parallel algorithms) is accurate enough for downstream operations. + with jax.enable_x64(True): + N = 100 + x = jnp.sort(random.uniform(-3, 3, N)) + y1 = random.normal(size=N) + y2 = random.normal(size=N) + diag = 0.1 + 0.05 * random.uniform(size=N) + + gp1 = GaussianProcess( + kernel, x, diag=diag, solver=QuasisepSolver, parallel=parallel + ) + gp2 = GaussianProcess(kernel, x, diag=diag, solver=DirectSolver) + cond1 = gp1.condition(y1) + cond2 = gp2.condition(y1) + + # The conditioned covariance should have the same quasiseparable rank as + # the prior, since we're conditioning with the same kernel + assert isinstance(cond1.gp.solver, QuasisepSolver) + assert cond1.gp.solver.parallel == parallel + assert cond1.gp.solver.matrix.lower.p.shape == gp1.solver.matrix.lower.p.shape + + assert_allclose(cond1.gp.covariance, cond2.gp.covariance) + assert_allclose(cond1.gp.variance, cond2.gp.variance) + assert_allclose(cond1.gp.log_probability(y2), cond2.gp.log_probability(y2)) + assert jnp.isfinite(cond1.gp.sample(jax.random.PRNGKey(0))).all() + + # predict(return_var=True) at the training points uses condition_diag + mu1, var1 = gp1.predict(y1, return_var=True) + mu2, var2 = gp2.predict(y1, return_var=True) + assert_allclose(mu1, mu2) + assert_allclose(var1, var2) + + # Chained conditioning should also be well-behaved + cond1b = cond1.gp.condition(y2) + cond2b = cond2.gp.condition(y2) + assert_allclose(cond1b.log_probability, cond2b.log_probability) + assert_allclose(cond1b.gp.loc, cond2b.gp.loc) + assert_allclose(cond1b.gp.variance, cond2b.gp.variance) + + +@pytest.mark.parametrize("parallel", [False, True], ids=["sequential", "parallel"]) +def test_condition_banded_noise(random, parallel): + N = 50 + x = jnp.sort(random.uniform(-3, 3, N)) + y = random.normal(size=N) + noise = Banded( + diag=0.1 + 0.05 * random.uniform(size=N), + off_diags=0.01 * jnp.ones((N, 1)), + ) + kernel = quasisep.Matern32(sigma=1.8, scale=1.5) + gp1 = GaussianProcess( + kernel, x, noise=noise, solver=QuasisepSolver, parallel=parallel + ) + gp2 = GaussianProcess(kernel, x, noise=noise, solver=DirectSolver) + cond1 = gp1.condition(y) + cond2 = gp2.condition(y) + assert isinstance(cond1.gp.solver, QuasisepSolver) + assert_allclose(cond1.gp.covariance, cond2.gp.covariance) + assert_allclose(cond1.gp.variance, cond2.gp.variance) + assert_allclose( + gp1.predict(y, return_var=True)[1], gp2.predict(y, return_var=True)[1] + )