feat(state-space): dense joint predictive covariance

CHANGELOG.md

@@ -10,6 +10,31 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- **Dense joint predictive covariance for state-space GPs.**
+  `StateSpacePrior.predict`, `StateSpaceConjugateModel.predict`, and
+  `StateSpacePosterior.__call__` (`gpjax.state_space`) now accept
+  `covariance="dense"`, returning the full joint covariance across test
+  points rather than marginal variances only
+  ([#651](https://github.com/thomaspinder/GPJax/issues/651)). For the
+  unconditioned prior this is just the kernel's own dense gram — the
+  state-space SDE is an exact representation with no training data to
+  marginalise out. For the conditioned (smoothed) posterior it is built from
+  the RTS smoother's cross-covariance recursion (Särkkä & Solin 2019 §12.2):
+  `rts_smoother` gains an opt-in `return_gains=True` that exposes its
+  already-computed per-step smoother gains, and a new
+  `gpjax.state_space.prediction._dense_smoothed_test_covariance` chains them
+  into the `M x M` test-point covariance. This keeps the state-space
+  formulation's linear-in-`N` cost — no `N x N` gram over the training set is
+  ever formed — with the cross-covariance work landing at `O(M^2 d^3)`,
+  independent of `N` and no larger than the `O(M^2)` already required to
+  store the dense output. `StateSpacePosterior.filtered` /
+  `StateSpaceConjugateModel.predict_filter` (the *causal* predictive) keep
+  raising `NotImplementedError` for `covariance="dense"`: each of their test
+  points conditions on a different information set (observations up to its
+  own timestamp), so a "joint" filtered covariance is not the dense
+  conjugate-predictive-shaped object the smoothed path now matches, and
+  extending it is left to a future issue if there is demand for it.
+
 - **`gpjax.fit_natgrads` and the `gpjax.natural_gradients` module.** Trains a variational
   family by alternating one natural-gradient step on the variational distribution with
   one step of a supplied Optax optimiser on everything else — kernel and likelihood

gpjax/state_space/conditioning.py

@@ -66,11 +66,19 @@ class StateSpacePosterior(Posterior):
     gives the causal (filter-only) alternative, and
     :attr:`log_marginal_likelihood` gives the evidence from the same filter.
 
-    Predictive contract (v1): every query returns diagonal (marginal)
-    covariance only. The marginals are exact; a dense joint predictive is not
-    implemented in v1. This predictive is therefore not Liskov-substitutable
-    for a dense :class:`gpjax.conditioning.ExactPosterior` predictive, and
-    ``covariance="dense"`` raises :class:`NotImplementedError`.
+    Predictive contract: :meth:`__call__` (the smoothed predictive) supports
+    both ``covariance="diagonal"`` (marginal variances, the default) and
+    ``covariance="dense"`` (the full joint covariance across test points, via
+    the RTS smoother's cross-covariance recursion — see
+    :func:`gpjax.state_space.prediction._dense_smoothed_test_covariance`), so
+    it is Liskov-substitutable for a dense
+    :class:`gpjax.conditioning.ExactPosterior` predictive. :meth:`filtered`
+    (the causal predictive) has no dense joint form in v1 — each test point
+    conditions on a *different* information set (data up to its own
+    timestamp), so a "joint" filtered covariance is not the same kind of
+    object as the dense conjugate predictive it would otherwise be compared
+    against — and ``covariance="dense"`` there still raises
+    :class:`NotImplementedError`.
 
     Time ordering: the predictive queries merge and sort the train and test
     grids internally, so they are order-insensitive.
@@ -150,25 +158,26 @@ def __call__(
 
         Args:
             test_inputs: Test timestamps of shape ``(M, 1)``.
-            covariance: Must be ``"diagonal"``; the v1 state-space predictive
-                has no dense joint form.
+            covariance: ``"diagonal"`` returns the marginal variances only;
+                ``"dense"`` returns the full ``M x M`` joint covariance, built
+                from the RTS smoother's cross-covariance recursion (Särkkä &
+                Solin 2019 §12.2) rather than a dense ``N x N`` gram over the
+                training set.
 
         Returns:
             GaussianDistribution: The smoothed predictive, carrying an
-                ``lx.DiagonalLinearOperator`` scale.
-
-        Raises:
-            NotImplementedError: If ``covariance="dense"``.
+                ``lx.DiagonalLinearOperator`` scale when ``covariance`` is
+                ``"diagonal"`` or an ``lx.MatrixLinearOperator`` when it is
+                ``"dense"``.
         """
-        if covariance != "diagonal":
-            raise _dense_not_implemented("prediction")
         from gpjax.state_space.prediction import predict_smoothed
 
         return predict_smoothed(
             self.model,
             self.train_data,
             test_inputs,
             observation_mask=self.observation_mask,
+            covariance=covariance,
         )
 
     def predict(
@@ -187,7 +196,8 @@ def predict(
         Args:
             test_inputs: Test timestamps of shape ``(M, 1)``.
             train_data: Ignored.
-            covariance: Must be ``"diagonal"``.
+            covariance: ``"diagonal"`` for marginal variances or ``"dense"``
+                for the full joint covariance.
 
         Returns:
             GaussianDistribution: The smoothed predictive.

gpjax/state_space/gps.py

@@ -20,22 +20,22 @@
 from gpjax.distributions import GaussianDistribution
 from gpjax.gps import ConjugateModel, Prior
 from gpjax.likelihoods import Gaussian, MultiOutputGaussian
+from gpjax.linalg.utils import add_jitter
 from gpjax.state_space.conditioning import StateSpacePosterior
 from gpjax.typing import Array
 
 
 class StateSpacePrior(Prior):
     """Prior for a state-space (Markovian) GP.
 
-    Identical to ``gpjax.gps.Prior`` except predictions are diagonal-only
-    (the prior is stationary in time, so off-diagonal covariance carries no
-    extra information for v1's diagonal-only predictive contract).
-
-    **Predictive contract (v1):** prediction returns diagonal (marginal)
-    covariance only; the marginals are exact. A dense joint predictive is not
-    implemented in v1 and is tracked as a follow-up. This predictive is
-    therefore not Liskov-substitutable for a dense
-    dense ``gpjax.gps.ConjugateModel`` predictive.
+    Predictions match ``gpjax.gps.Prior`` for both covariance modes.
+    ``covariance="diagonal"`` (the default) uses the SDE's stationary state
+    covariance directly rather than routing through the kernel, since the
+    prior is stationary in time; ``covariance="dense"`` returns the kernel's
+    own dense gram over the test inputs, since the state-space SDE is an
+    *exact* representation of the kernel with no training data to
+    marginalise out. Either way, this predictive is Liskov-substitutable for
+    the dense ``gpjax.gps.Prior`` predictive.
 
     Example:
         >>> import gpjax as gpx
@@ -52,13 +52,16 @@ def __call__(self, test_inputs, *, covariance="diagonal"):
         return self.predict(test_inputs, covariance=covariance)
 
     def predict(self, test_inputs, *, covariance="diagonal"):
-        if covariance != "diagonal":
-            raise NotImplementedError(
-                "State-space prior prediction returns diagonal (marginal) covariance "
-                "only; a dense joint predictive is not implemented in v1 and is "
-                "tracked as a follow-up. The marginal variances returned are exact, "
-                "so for diagonal-only use pass covariance='diagonal'."
+        mean_at_test = self.mean_function(test_inputs)
+        loc = jnp.atleast_1d(mean_at_test.squeeze())
+
+        if covariance == "dense":
+            gram_dense = add_jitter(
+                self.kernel.gram(test_inputs).as_matrix(), self.jitter
             )
+            scale = lx.MatrixLinearOperator(gram_dense)
+            return GaussianDistribution(loc=loc, scale=scale)
+
         from gpjax.state_space.kernels import to_sde
 
         sde = to_sde(self.kernel)
@@ -68,8 +71,6 @@ def predict(self, test_inputs, *, covariance="diagonal"):
         marginal_variance = (H @ P_inf @ H.T).squeeze() + self.jitter
 
         n_test = test_inputs.shape[0]
-        mean_at_test = self.mean_function(test_inputs)
-        loc = jnp.atleast_1d(mean_at_test.squeeze())
         scale = lx.DiagonalLinearOperator(jnp.full(n_test, marginal_variance))
         return GaussianDistribution(loc=loc, scale=scale)
 
@@ -94,11 +95,15 @@ class StateSpaceConjugateModel(ConjugateModel):
       - ``predict_filter``         : ``condition(D).filtered(t)``, the causal
         (filter-only) predictive
 
-    **Predictive contract (v1):** prediction returns diagonal (marginal)
-    covariance only; the marginals are exact. A dense joint predictive is not
-    implemented in v1 and is tracked as a follow-up. This predictive is
-    therefore not Liskov-substitutable for a dense
-    dense ``gpjax.gps.ConjugateModel`` predictive.
+    **Predictive contract:** ``predict``/``__call__`` (the smoothed
+    predictive) supports both ``covariance="diagonal"`` (marginal variances,
+    the default) and ``covariance="dense"`` (the full joint covariance across
+    test points, via the RTS smoother's cross-covariance recursion), so it is
+    Liskov-substitutable for the dense ``gpjax.gps.ConjugateModel``
+    predictive. ``predict_filter`` (the causal predictive) has no dense joint
+    form: each test point conditions on a different information set, so
+    ``covariance="dense"`` there still raises ``NotImplementedError`` — see
+    :class:`~gpjax.state_space.conditioning.StateSpacePosterior`.
 
     Example:
         >>> import gpjax as gpx
@@ -166,8 +171,8 @@ def predict(
         Args:
             test_inputs: Test timestamps of shape ``(M, 1)``.
             train_data: The observations to condition on.
-            covariance: Must be ``"diagonal"``; the v1 state-space predictive
-                has no dense joint form.
+            covariance: ``"diagonal"`` for marginal variances or ``"dense"``
+                for the full joint covariance.
             observation_mask: Optional boolean mask over the training points.
 
         Returns:

gpjax/state_space/inference.py

@@ -321,7 +321,7 @@ def skip_update(args):
     )
 
 
-def rts_smoother(sde, forward_outputs, time_steps):
+def rts_smoother(sde, forward_outputs, time_steps, *, return_gains: bool = False):
     r"""Square-root RTS smoother.
 
     Runs the backward recursion of Rauch, Tung & Striebel (1965) entirely in
@@ -375,6 +375,17 @@ def rts_smoother(sde, forward_outputs, time_steps):
     implementation already recomputed ``transition_matrix_next`` rather than
     threading it through from the forward pass.
 
+    The per-step gains :math:`G_i` are already computed inside the backward
+    scan to produce :math:`P^{\text{smooth}}_i`; ``return_gains`` just adds
+    them to the scan's output rather than discarding them, at no extra
+    numerical cost. They are the building block for the smoother
+    cross-covariance recursion
+    (:math:`\mathrm{Cov}(x_i, x_j \mid y_{1:T}) = G_i \cdots G_{j-1}
+    P^{\text{smooth}}_j` for :math:`i < j`, Särkkä & Solin 2019 §12.2) used by
+    :func:`gpjax.state_space.prediction._dense_test_covariance` to build the
+    joint predictive covariance without materialising an
+    :math:`N \times N` gram.
+
     Args:
         sde (LinearSDE): State-space SDE used in the forward pass;
             ``sde.discretise(dt)`` is called once per backward step.
@@ -384,13 +395,20 @@ def rts_smoother(sde, forward_outputs, time_steps):
         time_steps (Float[Array, "num_train"]): Same ``time_steps`` that drove
             the forward pass; ``time_steps[i+1]`` is the inter-step ``dt``
             between filtered index ``i`` and predicted index ``i + 1``.
+        return_gains (bool, keyword-only): If ``True``, also return the
+            per-step smoother gains :math:`G_i` for :math:`i = 0, \ldots,
+            \text{num\_train} - 2`.
 
     Returns:
         tuple: ``smoothed_means`` of shape ``Float[Array, "num_train state_dim"]``
             and ``smoothed_Ls`` of shape
             ``Float[Array, "num_train state_dim state_dim"]``, lower-triangular
             square roots (``smoothed_Ls[i] @ smoothed_Ls[i].T`` is the smoothed
-            covariance at step ``i``).
+            covariance at step ``i``). If ``return_gains`` is ``True``, a third
+            element ``smoother_gains`` of shape
+            ``Float[Array, "num_train-1 state_dim state_dim"]`` is appended,
+            with ``smoother_gains[i]`` the gain :math:`G_i` connecting smoothed
+            index ``i`` to smoothed index ``i + 1``.
 
     See plans/2026-04-21-state-space-gps-design.md §Stage 3.
     """
@@ -431,7 +449,7 @@ def backward_step(carry, scan_input):
         )
         L_smoothed = _qr_sqrt_sum(smoother_gain @ L_smoothed_next, L_virtual)
 
-        return (mean_smoothed, L_smoothed), (mean_smoothed, L_smoothed)
+        return (mean_smoothed, L_smoothed), (mean_smoothed, L_smoothed, smoother_gain)
 
     # Initial smoother carry = filtered state at the last step (no future).
     init_carry = (means_updated[-1], Ls_updated[-1])
@@ -450,7 +468,9 @@ def backward_step(carry, scan_input):
     _, smoothed_outputs_reversed = jax.lax.scan(
         backward_step, init_carry, backward_inputs_reversed
     )
-    smoothed_means_prefix, smoothed_Ls_prefix = smoothed_outputs_reversed
+    smoothed_means_prefix, smoothed_Ls_prefix, smoother_gains_prefix_reversed = (
+        smoothed_outputs_reversed
+    )
 
     # Un-reverse and append the last step (which equals the filtered state).
     smoothed_means = jnp.concatenate(
@@ -459,4 +479,10 @@ def backward_step(carry, scan_input):
     smoothed_Ls = jnp.concatenate(
         [jnp.flip(smoothed_Ls_prefix, axis=0), Ls_updated[-1:]], axis=0
     )
-    return smoothed_means, smoothed_Ls
+    if not return_gains:
+        return smoothed_means, smoothed_Ls
+
+    # smoother_gains[i] = G_i, connecting smoothed index i to i + 1; indices
+    # 0 .. num_train - 2, matching backward_inputs before the reverse.
+    smoother_gains = jnp.flip(smoother_gains_prefix_reversed, axis=0)
+    return smoothed_means, smoothed_Ls, smoother_gains