base

Solver helps organize the common elements of all solvers, including

1. the required methods (Kalman(), RTS(), smoothing_gains(), predict()),

2. the shared state-order bookkeeping definitions,

3. and a default implementation of the marginal likelihood in terms of the filter’s innovations, which every Kalman filter produces. A subclass may override log_probability() with a more efficient scan that only accumulates the likelihood contributions (rather than the full filter outputs).

Subclasses of Solver only inherit when the parent’s method would be still be correct (though perhaps slower). As such, the inheritance tree looks like:

Solver |– StateSpaceSolver instantaneous, sequential Kalman/RTS | `– ParallelStateSpaceSolver associative-scan Kalman/RTS `– IntegratedStateSpaceSolver exposure-aware, sequential Kalman/RTS (over K = 2N states)

`– ParallelIntegratedStateSpaceSolver associative-scan exposure-aware Kalman/RTS (over K = 2N states)

Classes

Solver

Base class for a smolgp solver.

Functions

log_prob_from_v_S(→ tinygp.helpers.JAXArray)

The Gaussian log probability from a Kalman filter's innovations.

Module Contents

smolgp.solvers.base.log_prob_from_v_S(v: tinygp.helpers.JAXArray, S: tinygp.helpers.JAXArray) → tinygp.helpers.JAXArray

The Gaussian log probability from a Kalman filter’s innovations.

Every Kalman filter produces the innovation \(v_k = y_k - H_k m_k^-\) and its covariance \(S_k\), from which the marginal log likelihood is

\[\log p(y) = -\tfrac{1}{2} \sum_k \left( v_k^T S_k^{-1} v_k + \log\det S_k + D \log 2\pi \right),\]

Parameters:

• v – shape (N, D), the innovations.

• S – shape (N, D, D), the innovation covariances.

class smolgp.solvers.base.Solver

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]

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

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

Run this solver’s RTS smoother over Kalman()’s output.

abstractmethod smoothing_gains(P_filtered, P_predicted) → tinygp.helpers.JAXArray

The y-independent RTS smoothing gains for this state timeline.

condition(y, return_v_S: bool = False) → Any

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

The posterior at arbitrary test coordinates.

_log_probability_from_filter(y) → tinygp.helpers.JAXArray

The likelihood via the Kalman filter outputs.

log_probability(y) → tinygp.helpers.JAXArray

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