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
• predict_exposure
• 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: 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

Wrapper for Kalman filter used with this solver

RTS(kalman_results) → Any

Wrapper for RTS smoother used with this solver

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

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

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

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:

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: 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

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

Wrapper for Kalman filter used with this solver

RTS(kalman_results) → Any

Wrapper for RTS smoother used with this solver