solvers

solvers#

In smolgp, “solvers” provide a swappable low-level interface for the Bayesian filtering and smoothing algorithms required for GP conditioning. New solvers can be contributed as external packages or pull requests to the smolgp GitHub repository.

The four built-in solvers are:

StateSpaceSolver: Standard Kalman filter and RTS smoother for instantaneous kernels (see smolgp.kernels.base). This is the default solver.
IntegratedStateSpaceSolver: Kalman filter and RTS smoother for integrated (time-averaged) measurement kernels (see smolgp.kernels.integrated).
ParallelStateSpaceSolver: GPU-parallelised version of StateSpaceSolver with \(O(\log N)\) complexity on compatible hardware.
ParallelIntegratedStateSpaceSolver: GPU-parallelised version of IntegratedStateSpaceSolver.

All four inherit from Solver, which fixes the interface (Kalman, RTS, condition, predict) and supplies the shared state-order bookkeeping and a default marginal likelihood built from any filter’s innovations. The parallel solvers subclass their sequential counterparts, implement associative scans for their Kalman and RTS methods, and fix log_probability to the default behavior rather than inherit a sequential scan.

All solvers are exact up to numerical precision.

Users generally do not need to instantiate solvers directly; GaussianProcess selects the appropriate solver automatically based on the kernel type.

Submodules#

Classes#

`Solver`	Base class for a smolgp solver.
`StateSpaceSolver`	A solver that implements Kalman filtering and RTS smoothing for state space GPs.
`ParallelStateSpaceSolver`	A solver that uses `jax.lax.associative_scan` to implement
`IntegratedStateSpaceSolver`	A solver that uses `jax.lax.scan` to implement Kalman filtering
`ParallelIntegratedStateSpaceSolver`	A solver that uses `jax.lax.associative_scan` to implement

Package Contents#

class smolgp.solvers.Solver[source]#

Bases: equinox.Module

Base class for a smolgp solver.

Subclasses must implement Kalman(), RTS(), smoothing_gains() and predict(). condition() is provided here too, being only a sequencing of Kalman() and RTS(), but a solver whose smoother takes different arguments overrides it.

The likelihood is implemented here in terms of the filter’s innovations, which every Kalman filter produces, but a subclass may override log_probability() with a more efficient scan that only accumulates the likelihood contributions (rather than the full filter outputs).

kernel#

The kernel defining the state space model.

Type:: StateSpaceModel

X#

The observed input coordinates.

Type:: JAXArray

noise#

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

Type:: JAXArray

state_coords#

State-level bookkeeping. One state per observation for an instantaneous kernel, two (exposure start and end) for an integrated kernel.

Type:: StateCoords

X: tinygp.helpers.JAXArray#

kernel: smolgp.kernels.base.StateSpaceModel#

noise: tinygp.helpers.JAXArray#

state_coords: smolgp.solvers.state_coords.StateCoords#

property t_states: tinygp.helpers.JAXArray#: The chronologically sorted time coordinate of every state.

_to_state_order(*arrays: tinygp.helpers.JAXArray) → tuple[tinygp.helpers.JAXArray, Ellipsis][source]#

Gather per-observation arrays into the solver’s (chronologically sorted) state order.

self.X and everything derived from it (y, noise) are kept in the caller’s input order; the filter/smoother step chronologically, so they need the sorted order instead. state_coords.obsid is exactly that permutation. Results come back in state order and are mapped back by the usual sort-by-obsid machinery.

abstractmethod Kalman(y, return_v_S: bool = False) → Any[source]#

Run this solver’s Kalman filter.

Returns (m_filtered, P_filtered, m_predicted, P_predicted),: plus (v, S) when return_v_S is True.

abstractmethod RTS(kalman_results) → Any[source]#: Run this solver’s RTS smoother over Kalman()’s output.

abstractmethod smoothing_gains(P_filtered, P_predicted) → tinygp.helpers.JAXArray[source]#: The y-independent RTS smoothing gains for this state timeline.

condition(y, return_v_S: bool = False) → Any[source]#

Filter then smooth, giving the posterior at the data.

Implemented here rather than per solver: it is only a sequencing of Kalman() and RTS() plus packaging, so every solver whose RTS takes the filter’s four outputs shares it verbatim. A solver whose smoother has a different signature overrides it – see ParallelStateSpaceSolver, whose parallel smoother consumes only the filtered pair.

abstractmethod predict(X_test, conditioned_results) → Any[source]#: The posterior at arbitrary test coordinates.

_log_probability_from_filter(y) → tinygp.helpers.JAXArray[source]#: The likelihood via the Kalman filter outputs.

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

The marginal log likelihood of the data, y.

By default, runs the Kalman filter and reduces its innovations. However, the Kalman filter computes more than is necessary if one only wants the likelihood. Hence, a Solver can override this function with an optimized method, e.g. log_probability()

class smolgp.solvers.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:

Retrodiction — test point precedes all observations: smoothed backward from the first data point using the stationary prior.
Interpolation — test point falls between two observations: Kalman-predicted forward from the nearest past point, then RTS-smoothed backward from the nearest future point.
Extrapolation — test point follows all observations: Kalman-predicted forward from the final filtered state.

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

Bases: smolgp.solvers.solver.StateSpaceSolver

A solver that uses jax.lax.associative_scan to implement parallel Kalman filtering and RTS smoothing.

Inherits from StateSpaceSolver and overrides the Kalman and RTS methods to use the parallel implementations. Methods which do not benefit from associative scans are inherited from StateSpaceSolver.

Kalman(y, return_v_S=True) → Any[source]#: Wrapper for Kalman filter used with this solver

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

The marginal log likelihood, reduced from this solver’s own filter.

Overrides StateSpaceSolver.log_probability() deliberately, as that is an optimized sequential scan. The generic path to reuse the Kalman filtered v and S is better here, as those are determined via associative scan, hence the likelihood stays log-depth.

RTS(kalman_results) → Any[source]#: Wrapper for RTS smoother used with this solver

condition(y, return_v_S=False) → tinygp.helpers.JAXArray[source]#: Compute the Kalman predicted, filtered, and RTS smoothed means and covariances at each of the input coordinates

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

Bases: smolgp.solvers.base.Solver

A solver that uses jax.lax.scan to implement Kalman filtering and RTS smoothing for integrated measurements

kernel#

X#

noise#

state_coords#

Kalman(y, return_v_S=False) → Any[source]#: Wrapper for Kalman filter used with this solver

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; see smolgp.solvers.solver.StateSpaceSolver.smoothing_gains().

Additionally threads the exposure reset bookkeeping, which the integrated smoother applies at every stateid == 0 state.

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

Batched-mean-path variant of condition() for M residual/observation arrays sharing this solver’s (kernel, X, noise). See StateSpaceSolver.condition_batched_mean for the general idea; this additionally threads the reset-matrix bookkeeping through integrated_kalman_gains/integrated_kalman_filter_batched_mean.

Parameters:: y_batch – shape (M, N) or (M, N, D)
Returns:: shape (M, K, dim)
Return type:: m_smoothed_batch

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

Algorithm for making predictions at arbitrary coordinates X_test

Parameters:

X_test – The test coordinates. If a tuple (t, delta, instid) with any delta > 0 and y is also given, those test points are predicted as exposure-integrated averages (see predict_exposure()) rather than instantaneous values.
conditioned_results – The output of self.condition()
y – The training data used to condition this solver. Only needed for delta>0 test points.

Returns:

Predicted means of the states at X_test pred_var : Predicted variances of the states at X_test

Return type:

pred_mean

There are three cases for each test point:

Retrodictionsmoothing from the first data point
using the prior as the prediction
Interpolationfiltering from most recent data point
and smoothing from next future point
Extrapolation : predicting from final filtered point

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

Bases: smolgp.solvers.integrated.solver.IntegratedStateSpaceSolver

A solver that uses jax.lax.associative_scan to implement parallel Kalman filtering and RTS smoothing for integrated measurements

_instid_per_state: tinygp.helpers.JAXArray#

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

The marginal log likelihood, reduced from this solver’s own filter.

Overrides IntegratedStateSpaceSolver.log_probability() deliberately, as that is an optimized sequential scan. The generic path to reuse the Kalman filtered v and S is better here, as those are determined via associative scan, hence the likelihood stays log-depth.

Kalman(y, return_v_S=True) → Any[source]#: Wrapper for Kalman filter used with this solver

RTS(kalman_results) → Any[source]#: Wrapper for RTS smoother used with this solver