Add O(J^2)-per-test-point predictive variance from the QSM Cholesky

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

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
 -----------
 

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",

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``.

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.

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.

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.