solver

solver#

Classes#

StateSpaceSolver

A solver that implements Kalman filtering and RTS smoothing for state space GPs.

Module Contents#

class smolgp.solvers.solver.StateSpaceSolver(kernel: smolgp.kernels.base.StateSpaceModel, X: tinygp.helpers.JAXArray, noise: tinygp.helpers.JAXArray)[source]#

Bases: smolgp.solvers.base.Solver

A solver that implements Kalman filtering and RTS smoothing for state space GPs.

Given a StateSpaceModel kernel and a set of observed coordinates, this solver computes the Kalman filtered and Rauch-Tung-Striebel (RTS) smoothed posterior means and covariances using jax.lax.scan for efficient JIT-compiled sequential computation.

Predictions at arbitrary test coordinates are handled by predict(), which dispatches among retrodiction, interpolation, and extrapolation depending on whether each test point falls before, between, or after the observed data.

Parameters:

  • kernel (StateSpaceModel) – The kernel function; must be a StateSpaceModel instance.

  • X (JAXArray) – The input coordinates with leading dimension of size N.

  • noise (JAXArray) – Observation noise covariance array of shape (N, D, D), where D is the observation dimension.

kernel#

The kernel defining the state space model.

Type:

StateSpaceModel

X#

The observed input coordinates.

Type:

JAXArray

noise#

Per-observation noise covariance matrices, shape (N, D, D).

Type:

JAXArray

state_coords#

State-level bookkeeping. For an instantaneous kernel this is the degenerate one-state-per-observation case (see instantaneous()).

Type:

StateCoords

kernel#
X#
noise#
state_coords#
Kalman(y, return_v_S=False) → Any[source]#

Wrapper for Kalman filter used with this solver

log_probability(y) → tinygp.helpers.JAXArray[source]#

The marginal log likelihood, without running the full filter.

Uses kalman_loglike(), whose scan only computes the parts of the Kalman filter needed for the likelihood

RTS(kalman_results) → Any[source]#

Wrapper for RTS smoother used with this solver

smoothing_gains(P_filtered, P_predicted) → tinygp.helpers.JAXArray[source]#

The RTS smoothing gains G_k for this solver’s state timeline.

These are y-independent, so they can be rebuilt on demand from the covariances a previous condition() already produced

condition_batched_mean(y_batch: tinygp.helpers.JAXArray) → tinygp.helpers.JAXArray[source]#

Batched-mean-path variant of condition(): computes the RTS-smoothed posterior mean for M residual/observation vectors that all share this solver’s (kernel, X, noise), at O(M) cost in the cheap mean-path recursion only – the O(N*dim^3) covariance recursion runs exactly ONCE, via the unmodified Kalman filter (called with a dummy y=zeros; exact, since P_filtered/P_predicted never depend on y), rather than once per sample.

Parameters:

y_batch – shape (M, N) or (M, N, D)

Returns:

shape (M, N, dim)

Return type:

m_smoothed_batch

predict(X_test, conditioned_results) → tinygp.helpers.JAXArray[source]#

Algorithm for making predictions at arbitrary coordinates X_test.

Parameters:

  • X_test (JAXArray) – The test coordinates; same shape as self.X.

  • conditioned_results (tuple) – The output of condition().

Returns:

A pair (pred_mean, pred_var) of arrays with leading dimension M = len(X_test), giving the predicted state means and covariances at each test coordinate.

Return type:

tuple

Each test point is handled by one of three cases depending on its position relative to the observed data:

  1. Retrodiction — test point precedes all observations: smoothed backward from the first data point using the stationary prior.

  2. Interpolation — test point falls between two observations: Kalman-predicted forward from the nearest past point, then RTS-smoothed backward from the nearest future point.

  3. Extrapolation — test point follows all observations: Kalman-predicted forward from the final filtered state.