integrated

This subpackage implements the Kalman filter and RTS smoother for an augmented state space model which includes an integral state as a IntegratedStateSpaceSolver. This is intended to be used with smolgp.kernels.integrated state space models to properly account for integrated (e.g. exposure-averaged) measurements.

See Integrated Measurements for a tutorial on using the integrated solvers.

Submodules

• kalman
• parallel
• rts
• solver

Classes

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.integrated.IntegratedStateSpaceSolver(kernel: smolgp.kernels.base.StateSpaceModel, X: tinygp.helpers.JAXArray, noise: tinygp.helpers.JAXArray)

Bases: equinox.Module

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

X: tinygp.helpers.JAXArray

kernel: smolgp.kernels.base.StateSpaceModel

noise: tinygp.helpers.JAXArray

state_coords: tinygp.helpers.JAXArray

normalization() → tinygp.helpers.JAXArray

Kalman(y, return_v_S=False) → Any

Wrapper for Kalman filter used with this solver

RTS(kalman_results) → Any

Wrapper for RTS smoother used with this solver

condition(y, return_v_S=False) → tinygp.helpers.JAXArray

Compute the Kalman predicted, filtered, and RTS smoothed means and covariances at each of the input coordinates

predict(X_test, conditioned_results) → tinygp.helpers.JAXArray

Algorithm for making predictions at arbitrary coordinates X_test

Parameters:

• X_test – The test coordinates.

• conditioned_results – The output of self.condition()

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:

1. Retrodictionsmoothing from the first data point

   using the prior as the prediction

2. Interpolationfiltering from most recent data point

   and smoothing from next future point

3. Extrapolation : predicting from final filtered point

class smolgp.solvers.integrated.ParallelIntegratedStateSpaceSolver(kernel: smolgp.kernels.base.StateSpaceModel, X: tinygp.helpers.JAXArray, noise: tinygp.helpers.JAXArray)

Bases: equinox.Module

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

X: tinygp.helpers.JAXArray

kernel: smolgp.kernels.base.StateSpaceModel

noise: tinygp.helpers.JAXArray

state_coords: tinygp.helpers.JAXArray

_state_coords: tinygp.helpers.JAXArray

normalization() → tinygp.helpers.JAXArray

Kalman(y, return_v_S=True) → Any

Wrapper for Kalman filter used with this solver

RTS(kalman_results) → Any

Wrapper for RTS smoother used with this solver

condition(y, return_v_S=True) → tinygp.helpers.JAXArray

Compute the Kalman predicted, filtered, and RTS smoothed means and covariances at each of the input coordinates

predict(X_test, conditioned_results) → tinygp.helpers.JAXArray

Algorithm for making predictions at arbitrary coordinates X_test

Parameters:

• X_test – The test coordinates.

• conditioned_results – The output of self.condition()

• observation_model – (optional) H for the test points should be a function just like self.kernel.observation_model

There are three cases:

1. Retrodictionsmoothing from the first data point

   using the prior as the prediction

2. Interpolationfiltering from most recent data point

   and smoothing from next future point

3. Extrapolation : predicting from final filtered point