From b74c1d55b44a65a132905b4f25d2d33c0c9ed381 Mon Sep 17 00:00:00 2001 From: Dan F-M Date: Wed, 9 Sep 2026 16:36:52 -0400 Subject: [PATCH 1/2] Refactor the solver conditioning interface around a single hook `Solver.condition` now takes the unresolved kernel (`None` meaning "the kernel this solver was built with", a signal that survives `jax.jit` where object identity does not), the test coordinates, the test noise, and `alpha = K^{-1} (y - mean)`, and returns a `ConditionedComponents` bundle with the conditional kernel, the conditional mean evaluated at the test points, and the conditioned process's solver. The producing solver builds that solver itself, so solver-specific settings such as `parallel` carry over without special cases in `GaussianProcess`, and `GaussianProcess` accepts an already constructed solver instance. `GaussianProcess.predict` is now a thin wrapper around `condition`: under `jax.jit`, XLA eliminates the unused `N_test x N_test` covariance and its Cholesky factorization, so the separate `Solver.condition_diag` hook and the mean-only shortcuts are removed. The generic dense implementation lives in `tinygp.solvers.direct.dense_condition`, which evaluates the cross covariance block once and passes the variance to the resulting `DirectSolver` explicitly (a `variance=` keyword) so that the full matrix stays dead code when only the variance is read. Also: - `kernels.Conditioned` gains a batched `__call__` so that evaluating the full conditional covariance costs two triangular solves rather than one per pair of points. - `means.Conditioned.include_mean` is a static field; previously it was a traced leaf, which broke conditioning a conditioned process under `jax.jit` with the default `include_mean=True`. - The default jitter for the conditioned process is derived from the data dtype rather than the parent's mean dtype. Breaking for third-party `Solver` subclasses: `condition` has a new signature and return type, and `condition_diag` no longer exists. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01HBWVoYn1nqDHDaF1K89Vrd --- news/272.removal | 10 + news/280.bugfix | 6 +- src/tinygp/gp.py | 219 ++++++++---------- src/tinygp/kernels/base.py | 9 + src/tinygp/means.py | 2 +- src/tinygp/solvers/direct.py | 77 +++--- src/tinygp/solvers/kalman.py | 9 +- src/tinygp/solvers/quasisep/solver.py | 131 +++++------ src/tinygp/solvers/solver.py | 92 ++++---- tests/test_gp.py | 11 +- .../test_solvers/test_quasisep/test_solver.py | 2 +- 11 files changed, 292 insertions(+), 276 deletions(-) create mode 100644 news/272.removal diff --git a/news/272.removal b/news/272.removal new file mode 100644 index 00000000..f44f85ba --- /dev/null +++ b/news/272.removal @@ -0,0 +1,10 @@ +The low-level ``Solver.condition`` interface changed: it now receives the +unresolved ``kernel`` (``None`` meaning the kernel the solver was built with) +and the vector ``alpha = K^{-1} (y - mean)``, and returns a +``ConditionedComponents`` bundle holding the conditional kernel, the +conditional mean at the test points, and the conditioned process's solver. +``Solver.condition_diag`` was removed: ``GaussianProcess.predict`` is now a +thin ``jax.jit``-compiled wrapper around ``condition``, which never +materializes the full test covariance unless ``return_cov`` is requested. +Third-party solvers must be updated, and must expose the ``kernel`` and ``X`` +that they were built with as attributes. diff --git a/news/280.bugfix b/news/280.bugfix index 537b7ca6..685558ff 100644 --- a/news/280.bugfix +++ b/news/280.bugfix @@ -8,8 +8,6 @@ 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 +``Solver.condition`` to signal that the solver's own kernel should be used. +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 a87d7173..815ca693 100644 --- a/src/tinygp/gp.py +++ b/src/tinygp/gp.py @@ -49,7 +49,7 @@ class GaussianProcess(eqx.Module): mean (Callable, optional): A callable or constant mean function that will be evaluated with the ``X`` as input: ``mean(X)`` solver: The solver type to be used to execute the required linear - algebra. + algebra, or an already constructed :class:`tinygp.solvers.Solver`. """ num_data: int = eqx.field(static=True) @@ -69,7 +69,7 @@ def __init__( diag: JAXArray | None = None, noise: Noise | None = None, mean: means.MeanBase | Callable[[JAXArray], JAXArray] | JAXArray | None = None, - solver: Any | None = None, + solver: type[Solver] | Solver | None = None, mean_value: JAXArray | None = None, covariance_value: Any | None = None, **solver_kwargs: Any, @@ -98,18 +98,26 @@ def __init__( noise = Diagonal(diag=jnp.broadcast_to(diag, self.mean.shape)) self.noise = noise - if solver is None: - if isinstance(covariance_value, SymmQSM) or isinstance(kernel, Quasisep): - solver = QuasisepSolver - else: - solver = DirectSolver - self.solver = solver( - kernel, - self.X, - self.noise, - covariance=covariance_value, - **solver_kwargs, - ) + if isinstance(solver, Solver): + _check_solver_instance( + solver, self.num_data, covariance_value, solver_kwargs + ) + self.solver = solver + else: + if solver is None: + if isinstance(covariance_value, SymmQSM) or isinstance( + kernel, Quasisep + ): + solver = QuasisepSolver + else: + solver = DirectSolver + self.solver = solver( + kernel, + self.X, + self.noise, + covariance=covariance_value, + **solver_kwargs, + ) @property def loc(self) -> JAXArray: @@ -161,61 +169,75 @@ def condition( diag (JAXArray, optional): Will be passed as the diagonal to the conditioned ``GaussianProcess`` object, so this can be used to introduce, for example, observational noise to predicted data. + noise (Noise, optional): A noise model for the conditioned + ``GaussianProcess``, for more expressive observation noise than + ``diag``. If provided, ``diag`` is ignored. include_mean (bool, optional): If ``True`` (default), the predicted values will include the mean function evaluated at ``X_test``. kernel (Kernel, optional): A kernel to optionally specify the covariance between the observed data and predicted data. See - :ref:`mixture` for an example. + :ref:`mixture` for an example. When this is not provided, the + solver knows that the prediction uses the kernel that the model + was built with, and it can use specialized algorithms; passing + any kernel here, even ``gp.kernel`` itself, disables those. Returns: A named tuple where the first element ``log_probability`` is the log marginal probability of the model, and the second element ``gp`` is the :class:`GaussianProcess` object describing the conditional distribution evaluated at ``X_test``. + + .. note:: + With a :class:`tinygp.solvers.QuasisepSolver` and no ``kernel`` + argument, conditioning stays scalable: at the training inputs + (``X_test=None``) the conditional covariance is quasiseparable, so + the conditioned process is itself a ``QuasisepSolver`` process; at + new test points the conditional mean and variance cost ``O(J^2)`` + per test point, and the conditioned process only builds (and + factorizes) its dense ``N_test x N_test`` covariance if + ``covariance``, ``sample``, or ``log_probability`` are explicitly + requested, at ``O(N_test^3)`` cost on every call. """ _check_test_shapes(self.X, X_test) - alpha, log_prob, mean_value = self._condition(y, X_test, include_mean, kernel) + alpha = self._get_alpha(y) + log_prob = self._compute_log_prob(alpha) + + # Below, we actually want alpha = K^-1 y instead of alpha = L^-1 y + alpha = self.solver.solve_triangular(alpha, transpose=True) + X_eval = self.X if X_test is None else X_test 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)) + n_test = jax.tree_util.tree_leaves(X_eval)[0].shape[0] + diag = _default_diag(alpha) if diag is None else diag + noise = Diagonal(diag=jnp.broadcast_to(jnp.asarray(diag), (n_test,))) # 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 - - # 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 + # that it was built with, and it can use a more efficient algorithm. + # Object identity would not survive ``jax.jit``, but ``None`` does. + comps = self.solver.condition(kernel, X_test, noise, alpha) + mean_value = comps.mean_value + if include_mean: + mean_value = mean_value + jax.vmap(self.mean_function)(X_eval) # 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. + pred_kernel = self.kernel if kernel is None else kernel gp = GaussianProcess( - kernels.Conditioned(self.X, self.solver, kernel), - X_test, + comps.kernel, + X_eval, noise=noise, mean=means.Conditioned( self.X, alpha, - kernel, + pred_kernel, include_mean=include_mean, mean_function=self.mean_function, ), mean_value=mean_value, - covariance_value=covariance_value, - **solver_kwargs, + solver=comps.solver, ) return ConditionResult(log_prob, gp) @@ -245,6 +267,9 @@ def predict( with the ``X`` data provided when instantiating this object. If it is not provided, ``X`` will be used by default, so the predictions will be made. + kernel (Kernel, optional): A kernel to optionally specify the + covariance between the observed data and predicted data; see + :func:`condition`. include_mean (bool, optional): If ``True`` (default), the predicted values will include the mean function evaluated at ``X_test``. return_var (bool, optional): If ``True``, the variance of the @@ -260,56 +285,18 @@ def predict( the variance or covariance of the predicted process will also be returned with shape ``(N_test,)`` or ``(N_test, N_test)`` respectively. - """ - _check_test_shapes(self.X, X_test) - # `return_var` takes priority over `return_cov` (see docstring above), - # so only take the dense `condition()` path — which materializes the - # full N_test x N_test covariance — when the covariance itself was - # requested. - if return_cov and not return_var: - _, cond = self.condition( - y, X_test, kernel=kernel, include_mean=include_mean - ) + This is equivalent to calling :func:`condition` and reading the + ``loc``, ``variance``, or ``covariance`` of the resulting process; since + it is compiled with ``jax.jit``, the full ``N_test x N_test`` conditional + covariance is only ever computed when ``return_cov`` is requested. + """ + _, cond = self.condition(y, X_test, kernel=kernel, include_mean=include_mean) + if return_var: + return cond.loc, cond.variance + if return_cov: return cond.loc, cond.covariance - - if not return_var: - _, _, mean_value = self._condition(y, X_test, include_mean, kernel) - return mean_value - - pred_kernel = self.kernel if kernel is None else kernel - - if X_test is None: - # Predicting at the training points: reuse `_condition`'s O(N) - # mean shortcut, and let the solver pick how to compute the - # variance -- e.g. QuasisepSolver avoids a dense matrix here. - _, _, 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(kernel, None, noise) - return mean_value, var_value - - # Predicting at new test points: build `Ks` once and reuse it for - # both the mean and the variance, instead of evaluating the kernel - # twice via `_condition` -- matters for expensive kernels (e.g. ones - # that differentiate through another kernel). - alpha = self._get_alpha(y) - alpha = self.solver.solve_triangular(alpha, transpose=True) - - Ks = pred_kernel(self.X, X_test) - Kss_diag = pred_kernel(X_test) - mean_offset = ( - jax.vmap(self.mean_function)(X_test) - if include_mean - else jnp.zeros_like(Kss_diag) - ) - - mean_value = jnp.dot(alpha, Ks) + mean_offset - A = self.solver.solve_triangular(Ks) - var_value = ( - Kss_diag - jnp.sum(jnp.square(A), axis=0) + _default_diag(mean_value) - ) - return mean_value, var_value + return cond.loc def sample( self, @@ -360,47 +347,6 @@ def _compute_log_prob(self, alpha: JAXArray) -> JAXArray: def _get_alpha(self, y: JAXArray) -> JAXArray: return self.solver.solve_triangular(y - self.loc) - @partial(jax.jit, static_argnums=(3,)) - def _condition( - self, - y: JAXArray, - X_test: JAXArray | None, - include_mean: bool, - kernel: kernels.Kernel | None = None, - ) -> tuple[JAXArray, JAXArray, JAXArray]: - alpha = self._get_alpha(y) - log_prob = self._compute_log_prob(alpha) - - # Below, we actually want alpha = K^-1 y instead of alpha = L^-1 y - alpha = self.solver.solve_triangular(alpha, transpose=True) - - if X_test is None: - X_test = self.X - - # In this common case (where we're predicting the GP at the data - # points, using the original kernel), the mean is especially fast to - # compute; so let's use that calculation here. - if kernel is None: - delta = self.noise @ alpha - mean_value = y - delta - if not include_mean: - mean_value -= self.loc - - else: - mean_value = kernel.matmul(self.X, y=alpha) - if include_mean: - mean_value += self.loc - - else: - if kernel is None: - kernel = self.kernel - - mean_value = kernel.matmul(X_test, self.X, alpha) - if include_mean: - mean_value += jax.vmap(self.mean_function)(X_test) - - return alpha, log_prob, mean_value - class ConditionResult(NamedTuple): """The result of conditioning a :class:`GaussianProcess` on data @@ -431,7 +377,24 @@ def _default_diag(reference: JAXArray) -> JAXArray: we use sqrt(eps) for the dtype of the mean function because that seems to give sensible results in general. """ - return jnp.sqrt(jnp.finfo(reference).eps) + return jnp.sqrt(jnp.finfo(jnp.result_type(reference)).eps) + + +def _check_solver_instance( + solver: Solver, num_data: int, covariance_value: Any, solver_kwargs: dict +) -> None: + """Check that an already constructed solver is compatible with the process""" + if covariance_value is not None or solver_kwargs: + raise ValueError( + "'covariance_value' and solver keyword arguments cannot be " + "provided alongside an already constructed 'solver'" + ) + n_solver = jax.tree_util.tree_leaves(solver.X)[0].shape[0] + if n_solver != num_data: + raise ValueError( + f"The provided 'solver' was built for {n_solver} data points, " + f"but 'X' has {num_data}" + ) def _check_test_shapes(X: JAXArray, X_test: JAXArray | None) -> None: diff --git a/src/tinygp/kernels/base.py b/src/tinygp/kernels/base.py index 8df5077b..2acdef12 100644 --- a/src/tinygp/kernels/base.py +++ b/src/tinygp/kernels/base.py @@ -152,6 +152,15 @@ def evaluate_diag(self, X: JAXArray) -> JAXArray: K = self.solver.solve_triangular(kernel_vec(self.X, X)) return self.kernel.evaluate_diag(X) - K.transpose() @ K + def __call__(self, X1: JAXArray, X2: JAXArray | None = None) -> JAXArray: + if X2 is None: + return super().__call__(X1) + # Evaluating the full matrix pairwise would require one triangular + # solve per pair; instead, do one batched solve per side + K1 = self.solver.solve_triangular(self.kernel(self.X, X1)) + K2 = K1 if X2 is X1 else self.solver.solve_triangular(self.kernel(self.X, X2)) + return self.kernel(X1, X2) - K1.transpose() @ K2 + class Custom(Kernel): """A custom kernel class implemented as a callable diff --git a/src/tinygp/means.py b/src/tinygp/means.py index 383f9cfe..4415e140 100644 --- a/src/tinygp/means.py +++ b/src/tinygp/means.py @@ -75,7 +75,7 @@ class Conditioned(MeanBase): X: JAXArray alpha: JAXArray kernel: Kernel - include_mean: bool + include_mean: bool = eqx.field(static=True) mean_function: MeanBase | None = None def __call__(self, X: JAXArray) -> JAXArray: diff --git a/src/tinygp/solvers/direct.py b/src/tinygp/solvers/direct.py index dcbe176b..c1e61335 100644 --- a/src/tinygp/solvers/direct.py +++ b/src/tinygp/solvers/direct.py @@ -1,6 +1,6 @@ from __future__ import annotations -__all__ = ["DirectSolver"] +__all__ = ["DirectSolver", "dense_condition"] from typing import Any @@ -11,7 +11,7 @@ from tinygp import kernels from tinygp.helpers import JAXArray from tinygp.noise import Noise -from tinygp.solvers.solver import Solver +from tinygp.solvers.solver import ConditionedComponents, Solver class DirectSolver(Solver): @@ -35,6 +35,7 @@ def __init__( noise: Noise, *, covariance: Any | None = None, + variance: JAXArray | None = None, ): """Build a :class:`DirectSolver` for a given kernel and coordinates @@ -45,10 +46,17 @@ def __init__( covariance: Optionally, a pre-computed array with the covariance matrix. This should be equal to the result of calling ``kernel`` and adding ``diag``, but that is not checked. + variance: Optionally, a pre-computed array with the diagonal of + ``covariance``. If not provided, this is evaluated using + ``kernel`` rather than read off ``covariance``, so that (under + ``jax.jit``) a caller that only needs the variance never forces + the full matrix to be built. """ self.kernel = kernel self.X = X - self.variance_value = kernel(X) + noise.diagonal() + if variance is None: + variance = kernel(X) + noise.diagonal() + self.variance_value = variance if covariance is None: covariance = kernel(X, X) + noise self.covariance_value = covariance @@ -75,26 +83,43 @@ def dot_triangular(self, y: JAXArray) -> JAXArray: return jnp.einsum("ij,j...->i...", self.scale_tril, y) def condition( - 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. 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 - else: - Ks = kernel(self.X, X_test) - Kss = kernel(X_test, X_test) + noise - - A = self.solve_triangular(Ks) - return Kss - A.transpose() @ A + self, + kernel: kernels.Kernel | None, + X_test: JAXArray | None, + noise: Noise, + alpha: JAXArray, + ) -> ConditionedComponents: + return dense_condition(self, kernel, X_test, noise, alpha) + + +def dense_condition( + solver: Solver, + kernel: kernels.Kernel | None, + X_test: JAXArray | None, + noise: Noise, + alpha: JAXArray, +) -> ConditionedComponents: + """The generic implementation of :func:`Solver.condition` + + This computes the conditional covariance as a dense matrix, and it can be + used by any solver since it only requires ``solve_triangular`` and the + ``X`` and ``kernel`` attributes. The cross covariance block ``Ks`` is + evaluated once and shared by the mean, the variance, and the covariance. + The variance is passed to the resulting :class:`DirectSolver` explicitly + so that, under ``jax.jit``, a caller that only reads the conditional mean + and variance never materializes or factorizes the ``N_test x N_test`` + covariance: that computation is dead code and XLA eliminates it. + """ + X_train = solver.X + kernel = solver.kernel if kernel is None else kernel + Xt = X_train if X_test is None else X_test + Ks = kernel(X_train, Xt) + A = solver.solve_triangular(Ks) + var = kernel(Xt) - jnp.sum(jnp.square(A), axis=0) + noise.diagonal() + Kss = (Ks if X_test is None else kernel(Xt, Xt)) - A.transpose() @ A + noise + cond_kernel = kernels.Conditioned(X_train, solver, kernel) + return ConditionedComponents( + kernel=cond_kernel, + mean_value=jnp.dot(alpha, Ks), + solver=DirectSolver(cond_kernel, Xt, noise, covariance=Kss, variance=var), + ) diff --git a/src/tinygp/solvers/kalman.py b/src/tinygp/solvers/kalman.py index fbd2704c..ac593297 100644 --- a/src/tinygp/solvers/kalman.py +++ b/src/tinygp/solvers/kalman.py @@ -54,6 +54,7 @@ def __init__( assert isinstance(noise, Diagonal) assert covariance is None + self.kernel = kernel self.X = X Pinf = kernel.stationary_covariance() self.A = jax.vmap(kernel.transition_matrix)( @@ -80,9 +81,13 @@ def dot_triangular(self, y: JAXArray) -> JAXArray: raise NotImplementedError def condition( - self, kernel: Kernel | None, X_test: JAXArray | None, noise: Noise + self, + kernel: Kernel | None, + X_test: JAXArray | None, + noise: Noise, + alpha: JAXArray, ) -> Any: - del kernel, X_test, noise + del kernel, X_test, noise, alpha raise NotImplementedError diff --git a/src/tinygp/solvers/quasisep/solver.py b/src/tinygp/solvers/quasisep/solver.py index 5bb8aa5e..1bb063ee 100644 --- a/src/tinygp/solvers/quasisep/solver.py +++ b/src/tinygp/solvers/quasisep/solver.py @@ -10,15 +10,16 @@ import numpy as np from tinygp.helpers import JAXArray -from tinygp.kernels.base import Kernel +from tinygp.kernels.base import Conditioned, Kernel from tinygp.noise import Noise +from tinygp.solvers.direct import dense_condition from tinygp.solvers.quasisep.core import ( DiagQSM, LowerTriQSM, StrictLowerTriQSM, SymmQSM, ) -from tinygp.solvers.solver import Solver +from tinygp.solvers.solver import ConditionedComponents, Solver class QuasisepSolver(Solver): @@ -149,81 +150,69 @@ def _conditional_delta(self, M: SymmQSM) -> SymmQSM: return A.gram(parallel=self.parallel) 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 - :class:`tinygp.kernels.quasisep.Quasisep` kernel, this will return the - quasiseparable representation of the conditional matrix. Otherwise, it - will use scalable methods where possible, but return a dense - representation of the covariance, so be careful when predicting at a - large number of test points! - - Args: - kernel: The kernel for the covariance between the observed and - 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. + self, + kernel: Kernel | None, + X_test: JAXArray | None, + noise: Noise, + alpha: JAXArray, + ) -> ConditionedComponents: + """Build the components of the conditioned process + + When predicting at the input coordinates (``X_test=None``) with a + :class:`tinygp.kernels.quasisep.Quasisep` kernel, the conditional + covariance is computed as a quasiseparable matrix, so the conditioned + process is itself a ``QuasisepSolver`` process. When, in addition, the + kernel is the one that this solver was built with (``kernel=None``), a + rank-``J`` representation is used that only requires the inverse of the + training covariance and is numerically much better behaved (see + :func:`_conditional_at_data`). That applies to any ``QuasisepSolver``, + including one that is itself the result of conditioning. + + Otherwise, this falls back on + :func:`tinygp.solvers.direct.dense_condition`, which materializes a + dense conditional covariance, so be careful when predicting at a large + number of test points! """ from tinygp.kernels.quasisep import Quasisep - # 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) - delta = self._conditional_delta(M) - M += noise.to_qsm() - return M - delta - - # Otherwise fall back on the slow method for now :( - if X_test is None: - Kss = Ks = kernel(self.X, self.X) - else: - Kss = kernel(X_test, X_test) - Ks = kernel(self.X, X_test) - - A = self.solve_triangular(Ks) - return Kss - A.transpose() @ A - - def condition_diag( - 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 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) - without ever materializing a dense matrix. Otherwise, this falls back on - :func:`tinygp.solvers.solver.Solver.condition_diag`. - """ - from tinygp.kernels.quasisep import Quasisep + pred_kernel = self.kernel if kernel is None else kernel - if X_test is None and kernel is None: - return self._conditional_at_data().diag.d + noise.diagonal() + if X_test is not None: + comps = dense_condition(self, kernel, X_test, noise, alpha) + if isinstance(pred_kernel, Quasisep): + # The cross-covariance matmul is O((N + M) J^2) for a Quasisep + # kernel, so prefer it to the dense product from ``Ks`` + mean_value = pred_kernel.matmul(X_test, self.X, alpha) + comps = comps._replace(mean_value=mean_value) + return comps 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) - M += noise.to_qsm() - return M.diag.d - delta.diag.d - - return super().condition_diag(kernel, X_test, noise) + covariance = self._conditional_at_data() + noise.to_qsm() + # The conditional mean at the data is M @ alpha = (K - N) @ alpha + # with K = M + N the full training covariance; unlike the kernel's + # matmul, this is cheap for any kernel, including the + # ``Conditioned`` kernel of an already conditioned process + mean_value = self.matrix.matmul(alpha, parallel=self.parallel) + mean_value -= self.noise @ alpha + elif isinstance(pred_kernel, Quasisep): + M = pred_kernel.to_symm_qsm(self.X) + covariance = M + noise.to_qsm() - self._conditional_delta(M) + mean_value = pred_kernel.matmul(self.X, y=alpha) + else: + return dense_condition(self, kernel, X_test, noise, alpha) + + cond_kernel = Conditioned(self.X, self, pred_kernel) + return ConditionedComponents( + kernel=cond_kernel, + mean_value=mean_value, + solver=QuasisepSolver( + cond_kernel, + self.X, + noise, + covariance=covariance, + parallel=self.parallel, + ), + ) def _check_sorted(X: JAXArray) -> None: diff --git a/src/tinygp/solvers/solver.py b/src/tinygp/solvers/solver.py index c65290c9..ba362b53 100644 --- a/src/tinygp/solvers/solver.py +++ b/src/tinygp/solvers/solver.py @@ -1,19 +1,52 @@ from __future__ import annotations -__all__ = ["Solver"] +__all__ = ["Solver", "ConditionedComponents"] from abc import abstractmethod -from typing import Any +from typing import Any, NamedTuple import equinox as eqx -import jax.numpy as jnp from tinygp.helpers import JAXArray from tinygp.kernels.base import Kernel from tinygp.noise import Noise +class ConditionedComponents(NamedTuple): + """The pieces of a conditioned process, as returned by :func:`Solver.condition` + + Everything here describes the conditional process *without* the prior mean + function of the parent process; :class:`tinygp.GaussianProcess` adds that + back itself. + """ + + kernel: Kernel + """The conditional kernel; its diagonal is the conditional variance.""" + + mean_value: JAXArray + """The conditional mean evaluated at the test points.""" + + solver: Solver + """The solver for the conditioned process + + This is built by the solver that produced it, so that solver-specific + settings (and representations) carry over. A solver that can avoid + materializing the full ``N_test x N_test`` conditional covariance should + return a lazy solver here (see, for example, + :class:`tinygp.solvers.direct.LazyDirectSolver`). + """ + + class Solver(eqx.Module): + """The interface for the linear algebra backends used by a GaussianProcess + + Implementations must store the kernel and input coordinates that they were + built with as ``kernel`` and ``X``. + """ + + kernel: Kernel + X: JAXArray + def __init__( self, kernel: Kernel, @@ -80,51 +113,28 @@ def dot_triangular(self, y: JAXArray) -> JAXArray: @abstractmethod def condition( - self, kernel: Kernel | None, X_test: JAXArray | None, noise: Noise - ) -> Any: - """Compute the covariance matrix for a conditional GP + self, + kernel: Kernel | None, + X_test: JAXArray | None, + noise: Noise, + alpha: JAXArray, + ) -> ConditionedComponents: + """Build the components of the process conditioned on observed data 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``. + specialized algorithms. (This signal is used rather than + checking object identity because identity is not preserved + under ``jax.jit``.) 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 | None, X_test: JAXArray | None, noise: Noise - ) -> JAXArray: - """The diagonal of the covariance matrix for a conditional GP + alpha: The vector ``K^{-1} @ (y - mean)`` for the observed data. - Unlike :func:`condition`, which returns the full ``N_test x N_test`` - conditional covariance matrix, this only computes its diagonal, - reusing this solver's existing factorization. The default - implementation below does this in ``O(N_train * N_test)`` time and - memory, without ever materializing a dense ``N_test x N_test`` - matrix. Subclasses can override this to use a more efficient, - solver-specific algorithm where one is available (see e.g. - :class:`tinygp.solvers.QuasisepSolver`). - - Args: - kernel: The kernel for the covariance between the observed and - 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. + Returns: + A :class:`ConditionedComponents` describing the conditional + process, excluding the prior mean function. """ - 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 - else: - Ks = kernel(self.X, X_test) # type: ignore - Kss_diag = kernel(X_test) - A = self.solve_triangular(Ks) - return Kss_diag - jnp.sum(jnp.square(A), axis=0) + noise.diagonal() + raise NotImplementedError diff --git a/tests/test_gp.py b/tests/test_gp.py index 5e65d69e..b8e9d17d 100644 --- a/tests/test_gp.py +++ b/tests/test_gp.py @@ -207,8 +207,8 @@ def test_predict_return_var_takes_priority_over_return_cov(random): def test_predict_quasisep_at_training_points(random): - # Must go through QuasisepSolver.condition_diag's quasiseparable - # shortcut, not fall back to a dense N x N cross-covariance matrix. + # Must go through QuasisepSolver.condition's quasiseparable + # representation, not fall back to a dense N x N cross-covariance matrix. with jax.enable_x64(True): X = jnp.sort(random.uniform(0, 10, 50)) y = jnp.sin(X) + 0.1 * random.normal(size=len(X)) @@ -219,3 +219,10 @@ def test_predict_quasisep_at_training_points(random): mu, var = gp.predict(y, return_var=True) assert_allclose(mu, cond.loc) assert_allclose(var, cond.variance) + + +def test_solver_instance_size_mismatch(random): + X = jnp.sort(random.uniform(0, 10, 20)) + solver = GaussianProcess(kernels.Matern32(1.0), X, diag=0.1).solver + with pytest.raises(ValueError, match="built for 20 data points"): + GaussianProcess(kernels.Matern32(1.0), X[:10], diag=0.1, solver=solver) diff --git a/tests/test_solvers/test_quasisep/test_solver.py b/tests/test_solvers/test_quasisep/test_solver.py index 4e634644..46b615ad 100644 --- a/tests/test_solvers/test_quasisep/test_solver.py +++ b/tests/test_solvers/test_quasisep/test_solver.py @@ -186,7 +186,7 @@ def test_conditioned_gp_operations(kernel, random, parallel): 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 + # predict(return_var=True) at the training points uses the QSM conditional mu1, var1 = gp1.predict(y1, return_var=True) mu2, var2 = gp2.predict(y1, return_var=True) assert_allclose(mu1, mu2) From 5222c16d0e2d8fff28ce007f24d93db9113b0fd3 Mon Sep 17 00:00:00 2001 From: Dan F-M Date: Wed, 9 Sep 2026 16:43:11 -0400 Subject: [PATCH 2/2] Add O(J^2)-per-test-point predictive variance from the QSM Cholesky When a `QuasisepSolver` conditions at new test points with the kernel it was built with, the predictive variance is now computed in O(J^2) per test point (after two O(N J^2) train-only scans) by reusing the quasiseparable Cholesky factorization, instead of a dense O(M^2 N) conditional covariance. The predictive mean was already scalable via `kernel.matmul(X_test, X_train, alpha)` and is unchanged. The conditioned `GaussianProcess` gets a `LazyDirectSolver`, so its dense covariance is only built (and factorized) when `covariance`, `sample`, or `log_probability` are explicitly requested. The variance uses the same per-point cross-covariance row generators as the rectangular quasiseparable product, factored out of `Quasisep.to_general_qsm` into `Quasisep.anchor`, so there is a single anchoring convention. Every propagation runs forward in time, which keeps the result stable across wide training gaps (the naive form, which pulls a test point back across its gap with an inverse transition, overflows for gaps of a few dozen correlation lengths). `anchor` also evaluates masked transitions at a clamped coordinate, so gradients at far extrapolation stay finite; this fixes a latent NaN gradient in `kernel.matmul` for extrapolating test points. `ops.cholesky` and `ops.cholesky_parallel` now also return the inclusive Riccati carry, and the backward congruence recursion shared with `symm_inv_parallel` lives in `ops.congruence_scan`. Noise models with their own quasiseparable states (e.g. `Banded`) enlarge the factorized matrix beyond the kernel's order, so those cases take the dense path, as they do on `DirectSolver`. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01HBWVoYn1nqDHDaF1K89Vrd --- docs/api/solvers.quasisep.rst | 16 + docs/api/solvers.rst | 24 ++ docs/tutorials/quasisep.ipynb | 14 +- news/272.bugfix | 4 + news/272.feature | 11 + src/tinygp/kernels/quasisep.py | 62 ++-- src/tinygp/solvers/direct.py | 66 +++- src/tinygp/solvers/quasisep/core.py | 2 +- src/tinygp/solvers/quasisep/ops.py | 50 ++- src/tinygp/solvers/quasisep/predict.py | 103 ++++++ src/tinygp/solvers/quasisep/solver.py | 68 ++-- tests/test_solvers/test_quasisep/test_ops.py | 9 +- .../test_quasisep/test_predict.py | 315 ++++++++++++++++++ .../test_solvers/test_quasisep/test_solver.py | 1 + 14 files changed, 667 insertions(+), 78 deletions(-) create mode 100644 news/272.bugfix create mode 100644 news/272.feature create mode 100644 src/tinygp/solvers/quasisep/predict.py create mode 100644 tests/test_solvers/test_quasisep/test_predict.py diff --git a/docs/api/solvers.quasisep.rst b/docs/api/solvers.quasisep.rst index 5ce27260..c03eb338 100644 --- a/docs/api/solvers.quasisep.rst +++ b/docs/api/solvers.quasisep.rst @@ -39,3 +39,19 @@ Rectangular Quasiseparable Matrices :toctree: summary GeneralQSM + + +Fast Prediction +--------------- + +.. currentmodule:: tinygp.solvers.quasisep.predict + +.. automodule:: tinygp.solvers.quasisep.predict + +.. autosummary:: + :toctree: summary + + PredictState + precompute + predict_var + ConditionedKernel diff --git a/docs/api/solvers.rst b/docs/api/solvers.rst index 56b8abe2..36f572fd 100644 --- a/docs/api/solvers.rst +++ b/docs/api/solvers.rst @@ -12,6 +12,30 @@ solvers package QuasisepSolver +The solver interface +-------------------- + +.. currentmodule:: tinygp.solvers.solver + +.. autosummary:: + :toctree: summary + + Solver + ConditionedComponents + + +Dense conditioning +------------------ + +.. currentmodule:: tinygp.solvers.direct + +.. autosummary:: + :toctree: summary + + dense_condition + LazyDirectSolver + + Subpackages ----------- diff --git a/docs/tutorials/quasisep.ipynb b/docs/tutorials/quasisep.ipynb index 73fbaa43..a6503f2b 100644 --- a/docs/tutorials/quasisep.ipynb +++ b/docs/tutorials/quasisep.ipynb @@ -225,19 +225,7 @@ "cell_type": "markdown", "id": "11", "metadata": {}, - "source": [ - "This all looks pretty good!\n", - "\n", - "Before closing out this tutorial, here are some technical details to keep in mind when using this solver:\n", - "\n", - "1. This implementation is new, and it hasn't yet been pushed to its limits. If you run into problems, please [open issues or pull requests](https://github.com/dfm/tinygp/issues).\n", - "\n", - "2. The computation of the general conditional model with these kernels is not (yet!) as fast as we might want, and it may be somewhat memory heavy. For very large datasets, it is sometimes sufficient to (a) just compute the conditional at the input points (by omitting the `X_test` parameter in {func}`tinygp.GaussianProcess.condition`), (b) only compute the mean prediction, which should be fast, or (c) only predict at a few test points.\n", - "\n", - "3. For more technical details about these methods, check out the API docs for the {ref}`api-kernels-quasisep`, and the {ref}`api-solvers-quasisep`, as well as the links therein.\n", - "\n", - "4. It should be possible to implement more flexible models using this interface than those supported by `celerite` or `celerite2`, so stay tuned for more tutorials!" - ] + "source": "This all looks pretty good!\n\nBefore closing out this tutorial, here are some technical details to keep in mind when using this solver:\n\n1. This implementation is new, and it hasn't yet been pushed to its limits. If you run into problems, please [open issues or pull requests](https://github.com/dfm/tinygp/issues).\n\n2. Conditioning at test points with the same kernel you fit (the common case) is fast: {func}`tinygp.GaussianProcess.condition` and {func}`tinygp.GaussianProcess.predict` evaluate the predictive mean and variance in `O(J^2)` per test point by reusing the Cholesky factorization, with no dense conditional covariance, and this works under `jax.jit`. Two cases still fall back to dense linear algebra and can be memory heavy for very large test sets: cross-kernel prediction (passing a `kernel` to `condition` — note that this includes explicitly passing the training kernel, so leave `kernel` unset to get the fast path) and requesting the full joint covariance of the prediction (`return_cov=True` or `cond.covariance`). Relatedly, operations on the conditioned GP that need that full covariance, like `sample` and `log_probability`, build and factor a dense matrix on every call, costing `O(M^3)` each time for `M` test points.\n\n3. For more technical details about these methods, check out the API docs for the {ref}`api-kernels-quasisep`, and the {ref}`api-solvers-quasisep`, as well as the links therein.\n\n4. It should be possible to implement more flexible models using this interface than those supported by `celerite` or `celerite2`, so stay tuned for more tutorials!" }, { "cell_type": "code", diff --git a/news/272.bugfix b/news/272.bugfix new file mode 100644 index 00000000..206bec84 --- /dev/null +++ b/news/272.bugfix @@ -0,0 +1,4 @@ +Fixed the conditional covariance returned by a ``QuasisepSolver`` when +conditioning at new test points or with a non-quasiseparable kernel: the +test-point noise was omitted, inconsistently with the ``DirectSolver`` and with +the conditioned process's own ``variance``. diff --git a/news/272.feature b/news/272.feature new file mode 100644 index 00000000..b4186e24 --- /dev/null +++ b/news/272.feature @@ -0,0 +1,11 @@ +Added a fast prediction path for the ``QuasisepSolver``: conditioning at new +test points with the kernel used for fitting (``kernel=None``, the default) +now evaluates the predictive variance in ``O(J^2)`` per test point by reusing +the quasiseparable Cholesky factorization, and the conditioned +``GaussianProcess`` only builds its dense conditional covariance when it is +explicitly requested (``covariance``, ``sample``, or ``log_probability``). +Passing any ``kernel`` argument to ``condition``/``predict``, even the training +kernel itself, uses the dense path instead. The cross-covariance products used +for the predictive mean now handle far extrapolation without ``NaN`` +gradients. ``GaussianProcess`` also accepts an already constructed solver +instance via its ``solver`` argument. diff --git a/src/tinygp/kernels/quasisep.py b/src/tinygp/kernels/quasisep.py index f28ca81b..96d9365a 100644 --- a/src/tinygp/kernels/quasisep.py +++ b/src/tinygp/kernels/quasisep.py @@ -115,34 +115,48 @@ def to_symm_qsm(self, X: JAXArray) -> SymmQSM: 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: - """The generalized quasiseparable representation of this kernel""" - sortable = jax.vmap(self.coord_to_sortable) - idx = jnp.searchsorted(sortable(X2), sortable(X1), side="right") - 1 - - Xs = jax.tree_util.tree_map(lambda x: jnp.append(x[0], x[:-1]), X2) + def anchor(self, x: JAXArray, X: JAXArray) -> tuple[JAXArray, JAXArray, JAXArray]: + """The cross-covariance row generators of one point ``x`` against ``X`` + + Returns ``(idx, pl, qu)`` where ``idx`` is the index of the last + coordinate in (sorted) ``X`` not after ``x``, ``pl`` propagates the + observation model of ``x`` from ``X[idx]`` and ``qu`` propagates it to + ``X[idx + 1]``, so that ``k(x, X[n]) = pl @ a_idx ... a_{n+1} @ q_n`` for + ``n <= idx`` and ``k(x, X[n]) = h_n @ a_n ... a_{idx+2} @ qu`` for + ``n > idx``, with ``(q, a)`` the generators of :func:`to_symm_qsm`. Rows + that extrapolate past the data (``idx < 0`` for ``pl``, ``idx >= N - 1`` + for ``qu``) are zero. The transitions never run backwards in time, so + nothing here amplifies like ``exp(+gap / scale)``. + """ + t = jax.vmap(self.coord_to_sortable)(X) + N = t.shape[0] + idx = jnp.searchsorted(t, self.coord_to_sortable(x), side="right") - 1 + iL = jnp.clip(idx, 0, N - 1) + iR = jnp.clip(idx + 1, 0, N - 1) + xL = jax.tree_util.tree_map(lambda v: jnp.asarray(v)[iL], X) + xR = jax.tree_util.tree_map(lambda v: jnp.asarray(v)[iR], X) + h = self.observation_model(x) Pinf = self.stationary_covariance() - 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 @ Pinf.T - pl = h1 - qu = h1 @ Pinf - pu = h2 + # Masked rows still evaluate the transition -- over a negative, + # unbounded gap, which can overflow -- and a masked ``inf`` poisons + # gradients, so the transition is evaluated at a clamped zero-gap + # coordinate whenever the mask is off (the double-where idiom). + okL = idx >= 0 + xc = jax.tree_util.tree_map(lambda s, v: jnp.where(okL, s, v), x, xL) + pl = jnp.where(okL, self.transition_matrix(xL, xc) @ h, 0.0) - i = jnp.clip(idx, 0, ql.shape[0] - 1) - Xi = jax.tree_util.tree_map(lambda x: jnp.asarray(x)[i], X2) - transition = jax.vmap(self.transition_matrix)(Xi, X1) - pl = jax.vmap(lambda x, y: x @ y.T)(pl, transition) + okR = idx < N - 1 + xc = jax.tree_util.tree_map(lambda s, v: jnp.where(okR, s, v), x, xR) + qu = jnp.where(okR, self.transition_matrix(xc, xR).T @ (Pinf @ h), 0.0) + return idx, pl, qu - i = jnp.clip(idx + 1, 0, pu.shape[0] - 1) - Xi = jax.tree_util.tree_map(lambda x: jnp.asarray(x)[i], X2) - 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) + def to_general_qsm(self, X1: JAXArray, X2: JAXArray) -> GeneralQSM: + """The generalized quasiseparable representation of this kernel""" + idx, pl, qu = jax.vmap(self.anchor, in_axes=(0, None))(X1, X2) + _, q, a = self.to_symm_qsm(X2).lower + h2 = jax.vmap(self.observation_model)(X2) + return GeneralQSM(pl=pl, ql=q, pu=h2, qu=qu, a=a, idx=idx) def matmul( self, diff --git a/src/tinygp/solvers/direct.py b/src/tinygp/solvers/direct.py index c1e61335..90f195f6 100644 --- a/src/tinygp/solvers/direct.py +++ b/src/tinygp/solvers/direct.py @@ -1,6 +1,6 @@ from __future__ import annotations -__all__ = ["DirectSolver", "dense_condition"] +__all__ = ["DirectSolver", "LazyDirectSolver", "dense_condition"] from typing import Any @@ -123,3 +123,67 @@ def dense_condition( mean_value=jnp.dot(alpha, Ks), solver=DirectSolver(cond_kernel, Xt, noise, covariance=Kss, variance=var), ) + + +class LazyDirectSolver(Solver): + """A dense solver whose covariance is built (and factorized) only on demand + + This is meant for conditioned processes whose kernel is cheap to evaluate on + the diagonal but expensive as a full matrix, such as the quasiseparable fast + prediction path: ``variance`` only touches ``kernel(X)``, ``covariance`` + builds the dense ``N x N`` matrix, and ``log_probability``, ``sample``, and + the triangular solves additionally Cholesky factorize it via a + :class:`DirectSolver`. That factorization is rebuilt on each call and not + cached, so those operations cost ``O(N^3)`` every time. (Under ``jax.jit``, + repeated factorizations within one call are de-duplicated by common + subexpression elimination.) + """ + + kernel: kernels.Kernel + X: JAXArray + noise: Noise + + def __init__( + self, + kernel: kernels.Kernel, + X: JAXArray, + noise: Noise, + *, + covariance: Any | None = None, + ): + if covariance is not None: + raise ValueError( + "LazyDirectSolver does not accept a pre-computed covariance" + ) + self.kernel = kernel + self.X = X + self.noise = noise + + def variance(self) -> JAXArray: + return self.kernel(self.X) + self.noise.diagonal() + + def covariance(self) -> JAXArray: + return self.kernel(self.X, self.X) + self.noise + + def _dense(self) -> DirectSolver: + return DirectSolver( + self.kernel, self.X, self.noise, covariance=self.covariance() + ) + + def normalization(self) -> JAXArray: + return self._dense().normalization() + + def solve_triangular(self, y: JAXArray, *, transpose: bool = False) -> JAXArray: + return self._dense().solve_triangular(y, transpose=transpose) + + def dot_triangular(self, y: JAXArray) -> JAXArray: + return self._dense().dot_triangular(y) + + def condition( + self, + kernel: kernels.Kernel | None, + X_test: JAXArray | None, + noise: Noise, + alpha: JAXArray, + ) -> ConditionedComponents: + return dense_condition(self._dense(), kernel, X_test, noise, alpha) diff --git a/src/tinygp/solvers/quasisep/core.py b/src/tinygp/solvers/quasisep/core.py index fd772ec2..cea347e3 100644 --- a/src/tinygp/solvers/quasisep/core.py +++ b/src/tinygp/solvers/quasisep/core.py @@ -541,7 +541,7 @@ def cholesky(self, *, parallel: bool = False) -> LowerTriQSM: (d,) = self.diag p, q, a = self.lower impl = cholesky_parallel if parallel else cholesky - c, w = impl(d, p, q, a) + c, w, _ = impl(d, p, q, a) return LowerTriQSM(diag=DiagQSM(c), lower=StrictLowerTriQSM(p=p, q=w, a=a)) def __neg__(self) -> SymmQSM: diff --git a/src/tinygp/solvers/quasisep/ops.py b/src/tinygp/solvers/quasisep/ops.py index 2cade0ef..3321b198 100644 --- a/src/tinygp/solvers/quasisep/ops.py +++ b/src/tinygp/solvers/quasisep/ops.py @@ -392,8 +392,36 @@ def combine(left, right): return jnp.einsum("nj,njk->nk", q, _shift_bwd(f)) +def congruence_scan(A, B, *, reverse=False, parallel=False): + """The inclusive congruence recursion ``Z_k = A_k @ Z_{k-1} @ A_k^T + B_k`` + + (with ``Z_{k+1}`` in place of ``Z_{k-1}`` when ``reverse``), starting from + zero. + """ + if parallel: + + def combine(left, right): + (Al, Bl), (Ar, Br) = left, right + return Ar @ Al, Ar @ Bl @ Ar.mT + Br + + return jax.lax.associative_scan(combine, (A, B), reverse=reverse)[1] + + def impl(z, data): + Ak, Bk = data + zk = Ak @ z @ Ak.T + Bk + return zk, zk + + return jax.lax.scan(impl, jnp.zeros_like(B[0]), (A, B), reverse=reverse)[1] + + @jax.jit def cholesky(d, p, q, a): + """The Cholesky factor's generators ``(c, w)`` and the inclusive carry ``f`` + + The carry ``f_k = a_k @ f_{k-1} @ a_k^T + w_k @ w_k^T`` is only needed by + the fast prediction path (:mod:`tinygp.solvers.quasisep.predict`). + """ + def impl(carry, data): fp = carry dk, pk, qk, ak = data @@ -401,14 +429,15 @@ def impl(carry, data): tmp = fp @ ak.T wk = (qk - pk @ tmp) / ck fk = ak @ tmp + jnp.outer(wk, wk) - return fk, (ck, wk) + return fk, (ck, wk, fk) init = jnp.zeros_like(jnp.outer(q[0], q[0])) - _, (c, w) = jax.lax.scan(impl, init, (d, p, q, a)) - return c, w + _, (c, w, f) = jax.lax.scan(impl, init, (d, p, q, a)) + return c, w, f def _riccati_scan(d, p, q, a): + """The inclusive Cholesky carry ``f_k`` via a parallel Riccati scan""" J = p.shape[1] I = jnp.eye(J) inv_d = 1.0 / d @@ -426,7 +455,7 @@ def combine(left, right): ) _, f, _ = jax.lax.associative_scan(combine, (A, F, G)) - return _shift_fwd(f) + return f @jax.jit @@ -438,8 +467,8 @@ def emit(f, dk, pk, qk, ak): wk = (qk - pk @ f @ ak.T) / ck return ck, wk - c, w = jax.vmap(emit)(f, d, p, q, a) - return c, w + c, w = jax.vmap(emit)(_shift_fwd(f), d, p, q, a) + return c, w, f @jax.jit @@ -473,7 +502,7 @@ def backward(z, data): @jax.jit def symm_inv_parallel(d, p, q, a): - f = _riccati_scan(d, p, q, a) + f = _shift_fwd(_riccati_scan(d, p, q, a)) def fwd_emit(f, dk, pk, qk, ak): fpk = f @ pk @@ -485,13 +514,8 @@ def fwd_emit(f, dk, pk, qk, ak): ig, s, ell = jax.vmap(fwd_emit)(f, d, p, q, a) - def bwd_combine(left, right): - (Al, Bl), (Ar, Br) = left, right - return Ar @ Al, Ar @ Bl @ Ar.mT + Br - B = jnp.einsum("n,nj,nk->njk", ig, p, p) - _, z = jax.lax.associative_scan(bwd_combine, (ell.mT, B), reverse=True) - z = _shift_bwd(z) + z = _shift_bwd(congruence_scan(ell.mT, B, reverse=True, parallel=True)) def bwd_emit(z, igk, pk, ak, sk): skz = sk @ z diff --git a/src/tinygp/solvers/quasisep/predict.py b/src/tinygp/solvers/quasisep/predict.py new file mode 100644 index 00000000..44223add --- /dev/null +++ b/src/tinygp/solvers/quasisep/predict.py @@ -0,0 +1,103 @@ +"""O(J^2)-per-test-point predictive variance from the QSM Cholesky. + +After the Cholesky factorization at N sorted training points, two train-only +scans (the Cholesky carry itself and one backward congruence) give a state from +which the predictive variance at any test point follows from a binary search +plus an O(J^2) contraction, using the same cross-covariance row generators +(:meth:`~tinygp.kernels.quasisep.Quasisep.anchor`) as the rectangular +quasiseparable product. The predictive *mean* needs none of this: it is that +product, ``kernel.matmul(X_test, X_train, alpha)``. + +With the Cholesky factor ``L = (c; p, w, a)`` and ``v = L^{-1} k_*``, the +variance is ``k** - v^T v``. With ``(idx, pl, qu)`` the anchors of the test +point, ``iL = clip(idx)`` and ``iR = clip(idx + 1)``, the entries split into a +head, ``sum_{n <= iL} v_n^2 = pl^T f_iL pl`` with ``f`` the inclusive Cholesky +carry, and a tail ``sum_{n >= iR} v_n^2 = s^T P_iR s`` where +``s = qu - a_iR f_iL pl`` and ``P`` is the backward congruence +``P_k = A_k^T P_{k+1} A_k + h_k h_k^T / c_k^2`` with +``A_k = a_{k+1} (I - w_k h_k^T / c_k)``. Every propagation in this form runs +forward in time, so nothing here amplifies like ``exp(+gap / scale)``. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import equinox as eqx +import jax +import jax.numpy as jnp + +from tinygp.helpers import JAXArray +from tinygp.kernels.base import Conditioned +from tinygp.solvers.quasisep import ops +from tinygp.solvers.quasisep.block import ensure_dense + +if TYPE_CHECKING: + from tinygp.kernels.quasisep import Quasisep + from tinygp.solvers.quasisep.solver import QuasisepSolver + + +class PredictState(eqx.Module): + """Train-only state for fast prediction; built once by :func:`precompute`.""" + + f: JAXArray # (N, J, J), the inclusive Cholesky carry + P: JAXArray # (N, J, J), the backward congruence accumulator + + +def precompute(solver: QuasisepSolver) -> PredictState: + """Run the two train-only scans and bundle them into a :class:`PredictState`. + + The Cholesky carry is recomputed here (one extra scan) rather than stored + on the solver, so likelihood-only workflows never hold the N x J^2 array. + """ + (d,) = solver.matrix.diag + p, q, a = solver.matrix.lower + c = solver.factor.diag.d + w = solver.factor.lower.q + impl = ops.cholesky_parallel if solver.parallel else ops.cholesky + f = impl(d, p, q, a)[2] + + h = jax.vmap(solver.kernel.observation_model)(solver.X) + a_next = jax.vmap(ensure_dense)( + jax.tree_util.tree_map( + lambda v: jnp.concatenate([v[1:], jnp.eye(v.shape[-1])[None]]), a + ) + ) + inv_c = 1.0 / c + A = a_next @ (jnp.eye(h.shape[1]) - jnp.einsum("n,nj,nk->njk", inv_c, w, h)) + B = jnp.einsum("n,nj,nk->njk", inv_c**2, h, h) + P = ops.congruence_scan( + jnp.swapaxes(A, -1, -2), B, reverse=True, parallel=solver.parallel + ) + return PredictState(f=f, P=P) + + +def predict_var( + kernel: Quasisep, solver: QuasisepSolver, state: PredictState, x_star: JAXArray +) -> JAXArray: + """Predictive (noise-free) variance at one test point.""" + idx, pl, qu = kernel.anchor(x_star, solver.X) + N = state.f.shape[0] + iL = jnp.clip(idx, 0, N - 1) + iR = jnp.clip(idx + 1, 0, N - 1) + f = state.f[iL] # pl == 0 when idx < 0, so no boundary element is needed + a_R = ensure_dense(jax.tree_util.tree_map(lambda v: v[iR], solver.matrix.lower.a)) + s = qu - a_R @ (f @ pl) + tail = jnp.where(idx < N - 1, s @ state.P[iR] @ s, 0.0) + return kernel.evaluate_diag(x_star) - pl @ f @ pl - tail + + +class ConditionedKernel(Conditioned): + """Conditioned kernel with the fast O(J^2) diagonal variance. + + Inherits the dense off-diagonal :meth:`evaluate` and block ``__call__`` + from :class:`tinygp.kernels.Conditioned` (building a full ``M x M`` block + materializes a dense ``N x M`` cross covariance and triangular solve; + prefer :meth:`evaluate_diag` for variances) and overrides only the + diagonal. + """ + + state: PredictState + + def evaluate_diag(self, X: JAXArray) -> JAXArray: + return predict_var(self.kernel, self.solver, self.state, X) diff --git a/src/tinygp/solvers/quasisep/solver.py b/src/tinygp/solvers/quasisep/solver.py index 1bb063ee..5e5ccc8b 100644 --- a/src/tinygp/solvers/quasisep/solver.py +++ b/src/tinygp/solvers/quasisep/solver.py @@ -12,13 +12,15 @@ from tinygp.helpers import JAXArray from tinygp.kernels.base import Conditioned, Kernel from tinygp.noise import Noise -from tinygp.solvers.direct import dense_condition +from tinygp.solvers.direct import LazyDirectSolver, dense_condition +from tinygp.solvers.quasisep import predict as qsp from tinygp.solvers.quasisep.core import ( DiagQSM, LowerTriQSM, StrictLowerTriQSM, SymmQSM, ) +from tinygp.solvers.quasisep.ops import qsm_mul from tinygp.solvers.solver import ConditionedComponents, Solver @@ -124,8 +126,6 @@ def _conditional_at_data(self) -> SymmQSM: ``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): @@ -144,8 +144,6 @@ 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) @@ -165,41 +163,67 @@ def condition( kernel is the one that this solver was built with (``kernel=None``), a rank-``J`` representation is used that only requires the inverse of the training covariance and is numerically much better behaved (see - :func:`_conditional_at_data`). That applies to any ``QuasisepSolver``, + ``_conditional_at_data``). That applies to any ``QuasisepSolver``, including one that is itself the result of conditioning. + When predicting at new test points with the kernel that this solver was + built with (``kernel=None``), the conditional mean and variance are + computed in ``O(J^2)`` per test point by reusing this solver's Cholesky + factorization (see :mod:`tinygp.solvers.quasisep.predict`), and the + conditioned process gets a :class:`tinygp.solvers.direct.LazyDirectSolver` + that only builds the dense conditional covariance on demand. + Otherwise, this falls back on :func:`tinygp.solvers.direct.dense_condition`, which materializes a dense conditional covariance, so be careful when predicting at a large number of test points! """ + # Imported here because ``tinygp.kernels.quasisep`` imports this package from tinygp.kernels.quasisep import Quasisep pred_kernel = self.kernel if kernel is None else kernel - if X_test is not None: - comps = dense_condition(self, kernel, X_test, noise, alpha) - if isinstance(pred_kernel, Quasisep): - # The cross-covariance matmul is O((N + M) J^2) for a Quasisep - # kernel, so prefer it to the dense product from ``Ks`` - mean_value = pred_kernel.matmul(X_test, self.X, alpha) - comps = comps._replace(mean_value=mean_value) - return comps - - if kernel is None: + if X_test is None and kernel is None: + # Conditioning at the data with this solver's kernel: the rank-J + # representation only needs this solver's matrix and noise, so it + # applies to any kernel, including the ``Conditioned`` kernel of an + # already conditioned process. The same goes for the conditional + # mean, M @ alpha = (K - N) @ alpha with K = M + N. covariance = self._conditional_at_data() + noise.to_qsm() - # The conditional mean at the data is M @ alpha = (K - N) @ alpha - # with K = M + N the full training covariance; unlike the kernel's - # matmul, this is cheap for any kernel, including the - # ``Conditioned`` kernel of an already conditioned process mean_value = self.matrix.matmul(alpha, parallel=self.parallel) mean_value -= self.noise @ alpha - elif isinstance(pred_kernel, Quasisep): + + elif not isinstance(pred_kernel, Quasisep): + return dense_condition(self, kernel, X_test, noise, alpha) + + elif X_test is None: M = pred_kernel.to_symm_qsm(self.X) covariance = M + noise.to_qsm() - self._conditional_delta(M) mean_value = pred_kernel.matmul(self.X, y=alpha) + else: - return dense_condition(self, kernel, X_test, noise, alpha) + # The cross-covariance matmul is O((N + M) J^2) for a Quasisep + # kernel, so we always prefer it to a dense product + mean_value = pred_kernel.matmul(X_test, self.X, alpha) + + # When predicting with the kernel that this solver was built with, + # the variance can also be computed in O(J^2) per test point by + # reusing this solver's Cholesky factorization, and the conditioned + # process only needs the dense covariance if explicitly asked for + # it. This requires the factorized matrix to have exactly the + # kernel's quasiseparable generators, which is not the case for a + # noise model with its own states (e.g. ``Banded``). + if kernel is None and isinstance(self.noise.to_qsm(), DiagQSM): + state = qsp.precompute(self) + cond_kernel = qsp.ConditionedKernel(self.X, self, pred_kernel, state) + return ConditionedComponents( + kernel=cond_kernel, + mean_value=mean_value, + solver=LazyDirectSolver(cond_kernel, X_test, noise), + ) + + comps = dense_condition(self, kernel, X_test, noise, alpha) + return comps._replace(mean_value=mean_value) cond_kernel = Conditioned(self.X, self, pred_kernel) return ConditionedComponents( diff --git a/tests/test_solvers/test_quasisep/test_ops.py b/tests/test_solvers/test_quasisep/test_ops.py index 6c4830de..29cf26a4 100644 --- a/tests/test_solvers/test_quasisep/test_ops.py +++ b/tests/test_solvers/test_quasisep/test_ops.py @@ -54,21 +54,22 @@ def test_upper_matmul_parallel(data): def test_cholesky_parallel(data): d, p, q, a, _ = data - c_seq, w_seq = cholesky(d, p, q, a) - c_par, w_par = cholesky_parallel(d, p, q, a) + c_seq, w_seq, f_seq = cholesky(d, p, q, a) + c_par, w_par, f_par = cholesky_parallel(d, p, q, a) assert_allclose(c_par, c_seq) assert_allclose(w_par, w_seq) + assert_allclose(f_par, f_seq) def test_lower_solve_parallel(data): d, p, q, a, x = data - c, w = cholesky(d, p, q, a) + c, w, _ = cholesky(d, p, q, a) assert_allclose(lower_solve_parallel(c, p, w, a, x), lower_solve(c, p, w, a, x)) def test_upper_solve_parallel(data): d, p, q, a, x = data - c, w = cholesky(d, p, q, a) + c, w, _ = cholesky(d, p, q, a) assert_allclose(upper_solve_parallel(c, p, w, a, x), upper_solve(c, p, w, a, x)) diff --git a/tests/test_solvers/test_quasisep/test_predict.py b/tests/test_solvers/test_quasisep/test_predict.py new file mode 100644 index 00000000..fed16e15 --- /dev/null +++ b/tests/test_solvers/test_quasisep/test_predict.py @@ -0,0 +1,315 @@ +# mypy: ignore-errors + +import jax +import jax.numpy as jnp +import numpy as np +import pytest +from numpy import random as np_random + +from tinygp import GaussianProcess +from tinygp.kernels import quasisep as qk +from tinygp.noise import Dense +from tinygp.solvers.direct import LazyDirectSolver +from tinygp.solvers.quasisep import predict +from tinygp.solvers.quasisep.solver import QuasisepSolver +from tinygp.test_utils import assert_allclose + + +@pytest.fixture +def data(): + rng = np_random.default_rng(42) + N, M = 30, 50 + X_train = jnp.sort(jnp.asarray(rng.uniform(0, 10, N))) + y = jnp.asarray(rng.normal(size=N)) + # include extrapolation on both sides + X_test = jnp.sort(jnp.asarray(rng.uniform(-2, 12, M))) + return X_train, y, X_test + + +@pytest.fixture(params=[False, True], ids=["sequential", "parallel"]) +def parallel(request): + return request.param + + +@pytest.fixture( + params=[ + qk.Matern32(scale=1.5), + qk.Matern52(scale=1.5), + qk.SHO(omega=2.0, quality=3.0), + qk.Cosine(scale=1.2), + qk.Matern32(scale=1.0) + qk.Matern52(scale=2.0), + 2.5 * qk.Matern32(scale=1.5), + ], + ids=["Matern32", "Matern52", "SHO", "Cosine", "sum", "scaled"], +) +def kernel(request): + return request.param + + +def test_predict_mean_and_var(data, kernel, parallel): + X_train, y, X_test = data + diag = 0.1 + N = X_train.shape[0] + + gp = GaussianProcess( + kernel, X_train, diag=diag, solver=QuasisepSolver, parallel=parallel + ) + solver = gp.solver + + # dense reference + K = kernel(X_train, X_train) + diag * jnp.eye(N) + Ks = kernel(X_train, X_test) + Kss = jax.vmap(kernel.evaluate_diag)(X_test) + L = jnp.linalg.cholesky(K) + beta_ref = jnp.linalg.solve(K, y) + mu_ref = Ks.T @ beta_ref + A = jax.scipy.linalg.solve_triangular(L, Ks, lower=True) + var_ref = Kss - jnp.sum(A**2, axis=0) + + alpha = solver.solve_triangular(y) + beta = solver.solve_triangular(alpha, transpose=True) + state = predict.precompute(solver) + + mu = kernel.matmul(X_test, X_train, beta) + var = jax.vmap(lambda x: predict.predict_var(kernel, solver, state, x))(X_test) + + assert_allclose(mu, mu_ref) + assert_allclose(var, var_ref) + + +def test_condition_end_to_end(data, kernel, parallel): + X_train, y, X_test = data + diag = 0.1 + N = X_train.shape[0] + + gp = GaussianProcess( + kernel, X_train, diag=diag, solver=QuasisepSolver, parallel=parallel + ) + cond = gp.condition(y, X_test).gp + + K = kernel(X_train, X_train) + diag * jnp.eye(N) + Ks = kernel(X_train, X_test) + Kss = jax.vmap(kernel.evaluate_diag)(X_test) + L = jnp.linalg.cholesky(K) + mu_ref = Ks.T @ jnp.linalg.solve(K, y) + A = jax.scipy.linalg.solve_triangular(L, Ks, lower=True) + var_ref = Kss - jnp.sum(A**2, axis=0) + + assert_allclose(cond.mean, mu_ref) + assert_allclose(cond.variance, var_ref) + + # Re-evaluation at fresh points goes through the Conditioned mean/kernel. + rng = np_random.default_rng(7) + X_new = jnp.asarray(rng.uniform(-1, 11, 5)) + Ks_new = kernel(X_train, X_new) + Kss_new = jax.vmap(kernel.evaluate_diag)(X_new) + mu_new = Ks_new.T @ jnp.linalg.solve(K, y) + A_new = jax.scipy.linalg.solve_triangular(L, Ks_new, lower=True) + var_new = Kss_new - jnp.sum(A_new**2, axis=0) + + assert_allclose(jax.vmap(cond.mean_function)(X_new), mu_new) + assert_allclose(jax.vmap(cond.kernel.evaluate_diag)(X_new), var_new) + + +def _dense_reference(kernel, X_train, y, X_test, diag): + N = X_train.shape[0] + K = kernel(X_train, X_train) + diag * jnp.eye(N) + Ks = kernel(X_train, X_test) + Kss = jax.vmap(kernel.evaluate_diag)(X_test) + mu = Ks.T @ jnp.linalg.solve(K, y) + L = jnp.linalg.cholesky(K) + A = jax.scipy.linalg.solve_triangular(L, Ks, lower=True) + var = Kss - jnp.sum(A**2, axis=0) + return mu, var + + +# Geometries where an anchoring off-by-one would bite: test points coincident +# with training points, duplicated training times, tiny datasets, and +# extrapolation past both ends of the data. +GEOMETRIES = { + "coincident": ( + jnp.array([0.0, 1.0, 2.5, 4.0]), + jnp.array([-1.0, 1.0, 2.5, 4.0, 7.0]), + ), + "duplicates": ( + jnp.array([0.0, 1.0, 1.0, 1.0, 3.0]), + jnp.array([0.5, 1.0, 2.0]), + ), + "single": (jnp.array([1.0]), jnp.array([0.0, 1.0, 2.0])), + "pair": (jnp.array([1.0, 2.0]), jnp.array([0.0, 1.5, 3.0])), + "extrapolation": ( + jnp.array([0.0, 1.0, 2.0]), + jnp.array([-10.0, -0.5, 2.5, 10.0]), + ), +} + + +@pytest.mark.parametrize("geometry", sorted(GEOMETRIES)) +def test_edge_case_geometries(kernel, parallel, geometry): + X_train, X_test = GEOMETRIES[geometry] + rng = np_random.default_rng(99) + y = jnp.asarray(rng.normal(size=X_train.shape[0])) + diag = 0.1 + + gp = GaussianProcess( + kernel, X_train, diag=diag, solver=QuasisepSolver, parallel=parallel + ) + cond = gp.condition(y, X_test).gp + mu_ref, var_ref = _dense_reference(kernel, X_train, y, X_test, diag) + assert_allclose(cond.mean, mu_ref) + assert_allclose(cond.variance, var_ref) + + +def test_far_extrapolation_gradient_finite(parallel): + # Regression test: the anchors were computed unconditionally and only masked + # with a single jnp.where, so for a test point far past the data the + # transition over the (negative) extrapolation gap overflowed to inf in + # Matern-type kernels and the VJP turned the masked inf into a NaN gradient. + rng = np_random.default_rng(0) + X_train = jnp.sort(jnp.asarray(rng.uniform(0, 10, 20))) + y = jnp.asarray(rng.normal(size=20)) + X_test = jnp.array([-5e4, 5.0, 5e4]) + + @jax.grad + def objective(scale): + gp = GaussianProcess( + qk.Matern32(scale=scale), + X_train, + diag=0.1, + solver=QuasisepSolver, + parallel=parallel, + ) + cond = gp.condition(y, X_test).gp + return jnp.sum(cond.mean) + jnp.sum(cond.variance) + + assert jnp.isfinite(objective(1.5)) + + +def test_cross_matmul_gradient_finite_at_far_extrapolation(): + # ``to_general_qsm`` (hence ``kernel.matmul(X_test, X, alpha)``, the + # conditioned mean) evaluated the transition at the clipped neighbour for + # extrapolating rows and only masked the row afterwards, so its gradient + # was NaN once the transition overflowed. + rng = np_random.default_rng(0) + X_train = jnp.sort(jnp.asarray(rng.uniform(0, 10, 20))) + alpha = jnp.asarray(rng.normal(size=20)) + X_test = jnp.array([-5e4, 5.0, 5e4]) + + @jax.grad + def objective(scale): + return jnp.sum(qk.Matern32(scale=scale).matmul(X_test, X_train, alpha)) + + assert jnp.isfinite(objective(1.5)) + + +@pytest.mark.parametrize("gap", [40.0, 500.0]) +def test_wide_training_gap(kernel, parallel, gap): + # Regression test: the anchors used to pull a test point back across its + # whole training gap with an inverse transition, ~exp(+gap / scale), so a + # gap a few dozen correlation lengths wide gave a silently wrong variance + # (underflow) and then NaNs (overflow), in either precision. + X_train = jnp.array([0.0, gap, gap + 1.0]) + y = jnp.array([1.0, -1.0, 0.5]) + X_test = jnp.array([0.5, gap - 1.0, gap + 0.5]) + diag = 0.1 + gp = GaussianProcess( + kernel, X_train, diag=diag, solver=QuasisepSolver, parallel=parallel + ) + mu, var = gp.predict(y, X_test, return_var=True) + mu_ref, var_ref = _dense_reference(kernel, X_train, y, X_test, diag) + assert_allclose(mu, mu_ref) + assert_allclose(var, var_ref) + + +def test_banded_noise_falls_back_to_dense(parallel): + # A noise model with its own quasiseparable states enlarges the factor's + # order beyond the kernel's; the fast path must step aside (it used to + # raise a shape error) and the result must still match the dense one. + from tinygp.noise import Banded + + rng = np_random.default_rng(3) + N, M = 12, 5 + X_train = jnp.sort(jnp.asarray(rng.uniform(0, 10, N))) + y = jnp.asarray(rng.normal(size=N)) + X_test = jnp.sort(jnp.asarray(rng.uniform(-1, 11, M))) + noise = Banded(diag=0.1 * jnp.ones(N), off_diags=0.01 * jnp.ones((N, 1))) + kernel = qk.Matern32(scale=1.5) + + gp = GaussianProcess( + kernel, X_train, noise=noise, solver=QuasisepSolver, parallel=parallel + ) + cond = gp.condition(y, X_test).gp + mu_pred, var_pred = gp.predict(y, X_test, return_var=True) + + K = kernel(X_train, X_train) + noise + Ks = kernel(X_train, X_test) + mu_ref = Ks.T @ jnp.linalg.solve(K, y) + var_ref = jax.vmap(kernel.evaluate_diag)(X_test) - jnp.einsum( + "ij,ij->j", Ks, jnp.linalg.solve(K, Ks) + ) + assert_allclose(cond.mean, mu_ref) + assert_allclose(cond.variance, var_ref) + assert_allclose(mu_pred, mu_ref) + assert_allclose(var_pred, var_ref) + + +def test_conditioned_covariance_includes_full_noise(parallel): + # The lazily built conditional covariance must include the off-diagonal + # part of a Dense test-noise model, like the other conditioning paths do. + rng = np_random.default_rng(5) + N, M = 8, 5 + X_train = jnp.sort(jnp.asarray(rng.uniform(0, 10, N))) + y = jnp.asarray(rng.normal(size=N)) + X_test = jnp.sort(jnp.asarray(rng.uniform(0, 10, M))) + R = rng.normal(size=(M, M)) + R = jnp.asarray(R @ R.T + M * jnp.eye(M)) + + kernel = qk.Matern32(scale=1.5) + diag = 0.1 + gp = GaussianProcess( + kernel, X_train, diag=diag, solver=QuasisepSolver, parallel=parallel + ) + cond = gp.condition(y, X_test, noise=Dense(value=R)).gp + + K = kernel(X_train, X_train) + diag * jnp.eye(N) + Ks = kernel(X_train, X_test) + cov_ref = kernel(X_test, X_test) - Ks.T @ jnp.linalg.solve(K, Ks) + R + assert_allclose(cond.covariance, cov_ref) + + +def test_fast_path_survives_jit(data): + # Regression: the fast path must fire when the GP crosses a jit/pytree + # boundary. Gating it on object identity (``kernel is self.kernel``) silently + # disabled it under jit, because flatten/unflatten produces distinct objects. + X_train, y, X_test = data + gp = GaussianProcess( + qk.Matern32(scale=1.5), X_train, diag=0.1, solver=QuasisepSolver + ) + + # Round-tripping the pytree is exactly what jit does to ``self``. + leaves, treedef = jax.tree_util.tree_flatten(gp) + gp_rt = jax.tree_util.tree_unflatten(treedef, leaves) + cond = gp_rt.condition(y, X_test).gp + assert isinstance(cond.kernel, predict.ConditionedKernel) + assert isinstance(cond.solver, LazyDirectSolver) + + # And the jitted entry point agrees with the eager fast path to precision. + mu_eager, var_eager = gp.predict(y, X_test, return_var=True) + mu_jit, var_jit = jax.jit(lambda g, yy, xx: g.predict(yy, xx, return_var=True))( + gp, y, X_test + ) + assert_allclose(mu_jit, mu_eager) + assert_allclose(var_jit, var_eager) + + +def test_numpy_inputs(data): + # The anchors gather from the training inputs with traced indices, which + # must work when those inputs are plain numpy arrays (as in the tutorials). + X_train, y, X_test = (np.asarray(v) for v in data) + kernel = qk.Matern32(scale=1.5) + gp = GaussianProcess(kernel, X_train, diag=0.1, solver=QuasisepSolver) + cond = gp.condition(y, X_test).gp + mu_ref, var_ref = _dense_reference(kernel, jnp.asarray(X_train), y, X_test, 0.1) + assert_allclose(cond.mean, mu_ref) + assert_allclose(cond.variance, var_ref) + assert_allclose(kernel.matmul(X_test, X_train, y), kernel(X_test, X_train) @ y) diff --git a/tests/test_solvers/test_quasisep/test_solver.py b/tests/test_solvers/test_quasisep/test_solver.py index 46b615ad..f945396b 100644 --- a/tests/test_solvers/test_quasisep/test_solver.py +++ b/tests/test_solvers/test_quasisep/test_solver.py @@ -195,6 +195,7 @@ def test_conditioned_gp_operations(kernel, random, parallel): # Chained conditioning should also be well-behaved cond1b = cond1.gp.condition(y2) cond2b = cond2.gp.condition(y2) + assert isinstance(cond1b.gp.solver, QuasisepSolver) assert_allclose(cond1b.log_probability, cond2b.log_probability) assert_allclose(cond1b.gp.loc, cond2b.gp.loc) assert_allclose(cond1b.gp.variance, cond2b.gp.variance)