feat(state-space): dense joint predictive covariance - #760
Merged
Conversation
… posterior Implements the acceptance criteria from #651: `StateSpacePrior.predict`, `StateSpaceConjugateModel.predict`, and `StateSpacePosterior.__call__` now support `covariance="dense"`, matching the dense `ConjugateModel` predictive (mean and full covariance) to ~1e-8 for Matern-1/2, 3/2, 5/2. - `rts_smoother` gains an opt-in `return_gains=True` that exposes its already-computed per-step smoother gains (no extra numerical work, just not discarding them). Existing 2-tuple callers are unaffected. - New `gpjax.state_space.prediction._dense_smoothed_test_covariance` chains those gains into the M x M joint covariance across test points, following the RTS smoother cross-covariance recursion (Sarkka & Solin 2019 SS12.2): Cov(x_i, x_j | y) = G_i...G_{j-1} P_j^smooth for i < j. It never inverts a gain product (ill-conditioned for widely separated points, since gains shrink with lag) and never forms an N x N gram over the training set -- cost is O(N d^3) for the filter/smoother pass plus O(M^2 d^3) for the cross-covariance chaining, both linear in N. - `StateSpacePrior.predict` (no conditioning data) returns the kernel's own dense gram for `covariance="dense"`, since the SDE is an exact representation of the kernel -- no Kalman machinery needed. Design decision (the issue underspecifies this): `StateSpacePosterior.filtered` / `StateSpaceConjugateModel.predict_filter` (the *causal* predictive) keep raising NotImplementedError for `covariance="dense"`. Each filtered test point 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 -- the issue's acceptance criteria only compares against the dense conjugate predictive, which is the smoothed quantity, and an existing repo test (test_state_space_posterior_predict_filter_dense_raises) already pinned the filtered-raises behaviour. Extending the filtered path is left to a future issue if there's demand. Tests: mean+covariance equivalence to the dense ConjugateModel (parametrized over Matern12/32/52 and jitter), diagonal/dense consistency, caller-order preservation, the M=1 degenerate case, a joint-sampling smoke test, jit and grad cleanliness, a machine-precision numpy-oracle check on the exposed smoother gains, and a larger-N robustness smoke test. Existing tests that pinned the old "dense raises" behaviour for the now-supported paths are updated to positive equivalence/plumbing checks instead. Note for reviewers: gpjax/state_space/inference.py::rts_smoother is implemented against its current covariance-form recursion, not a square-root rework (that's issue #668, tracked separately). If #668 lands a true square-root smoother, this cross-covariance recursion is worth revisiting -- a square-root form may simplify the derivation -- but nothing here blocks on it. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011UmdHyjMN5UdYLMD2JzE6s
|
📖 Docs preview: https://pr-760--endearing-crepe-c2d5fe.netlify.app Smoke render — the expensive notebooks run with reduced budgets, so |
# Conflicts: # gpjax/state_space/inference.py # tests/test_state_space/test_prediction.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Checklist
uv run poe formatbefore committing.Description
The state-space predictive surface only supported diagonal (marginal) covariance;
covariance="dense"raisedNotImplementedError. Implements the dense joint predictive by chaining the RTS smoother's already-computed per-step gains into the M×M test-point covariance (Särkkä & Solin 2019 §12.2), never inverting a gain product and never materialising an N×N gram over the training set — cost is O(N d³) + O(M² d³), independent of a quadratic-in-N blowup.Scoping note:
StateSpacePosterior.filtered/predict_filter(the causal predictive) still raiseNotImplementedErrorforcovariance="dense"— each filtered test point conditions on a different information set, so a "joint filtered covariance" isn't the same kind of object and isn't covered by the issue's stated recursion or acceptance criteria (which specify matching the dense smoothed predictive). An existing pinned-behaviour test corroborates this scoping.Addresses #651.
Test plan
uv run pytest tests/test_state_space/— 306 passed, 1 skipped (pre-existing, unrelated)uv run pytest tests/test_conditioning.py tests/test_gps.py— 276 passed (regression check on shared modules)uv run xdoctest gpjax/state_space— 9/9 passeduv run poe lint— clean