kalman

kalman#

Functions#

KalmanFilter(kernel, X, y, R[, return_v_S])

Wrapper for jitted kalman_filter function

kalman_filter(A, Q, H_all, R, t, y, m0, P0)

Jax implementation of the Kalman filter algorithm

KalmanLoglike(kernel, X, y, R)

Wrapper for the jitted kalman_loglike function

kalman_loglike(A, Q, H_all, R, t, y, m0, P0)

Optimized function to return just the marginal log likelihood.

kalman_gains(A, H_all, R, t, P_predicted)

The y-independent per-step quantities needed to replay kalman_filter's

kalman_filter_batched_mean(A_all, H_all, K_all, ...)

Batched-mean-path replay of kalman_filter's forward recursion, given

Module Contents#

smolgp.solvers.kalman.KalmanFilter(kernel, X, y, R, return_v_S=False)[source]#

Wrapper for jitted kalman_filter function

Parameters:

  • kernel – StateSpaceModel kernel

  • X – data coordinates, e.g. time or (time, texp, instid)

  • y – observations, shape (N, D)

  • R – observation noise covariance, shape (N, D, D)

Returns:

filtered means P_filtered: filtered covariances m_predicted: predicted means P_predicted: predicted covariances

Return type:

m_filtered

smolgp.solvers.kalman.kalman_filter(A, Q, H_all, R, t, y, m0, P0)[source]#

Jax implementation of the Kalman filter algorithm

See Theorem 4.2 (pdf page 77) in “Bayesian Filtering and Smoothing” by Simo S{“a}rkk{“a} for detailed description of the algorithm and notation.

e.g. _prev is _{k-1}

_pred is _k^{-}

Total runtime complexity is O(N*d^3) where N is the number of time steps and d is the dimension of the state vector.

smolgp.solvers.kalman.KalmanLoglike(kernel, X, y, R)[source]#

Wrapper for the jitted kalman_loglike function

Same arguments as KalmanFilter(), but returns just the marginal log likelihood.

smolgp.solvers.kalman.kalman_loglike(A, Q, H_all, R, t, y, m0, P0)[source]#

Optimized function to return just the marginal log likelihood. The performance boost here comes from three optimisations:

  1. Replace solve with a simple division in

    kalman_gain() when D==1.

  2. Prebuilds A(Δ) and Q(Δ) with vmap instead of computing them

    on-the-fly inside the scan.

  3. Splitting: separate scans for the y-independant covariance and the

    y-dependent mean, so each carries one array instead of (m, P).

#3 is only relevant for d == 2 kernels (with nontrivial Q). At that size, the covariance half of the scan fits just under a code-generation threshold, so splitting it from the mean scan (which is always small) is a ~10x speedup as both scans compile into short instructions whereas carrying both is over the threshold and compiles to ~7x more instructions per step. d == 1 fits as a single scan, so splitting is actually marginally slower (10%), and d > 2 is over the threshold regardless, so splitting is irrelevant.

smolgp.solvers.kalman.kalman_gains(A, H_all, R, t, P_predicted)[source]#

The y-independent per-step quantities needed to replay kalman_filter’s mean-path recursion for many different observation vectors, given P_predicted from ONE prior call to kalman_filter (any y – P_predicted never depends on it).

Pointwise in k (jax.vmap, no scan needed): A_k, H_k, and the Kalman gain K_k (given P_predicted[k]) don’t depend on any other step.

Parameters:

H_all – shape (N, D, dim) – the observation model already evaluated over the full coordinates (jax.vmap(kernel.observation_model)(X)), NOT the bare function; see KalmanFilter’s own note on why.

Returns:

shape (N, dim, dim) – A(0, Delta_k) H_all: shape (N, D, dim) – passed through unchanged K_all: shape (N, dim, D) – Kalman gain at step k

Return type:

A_all

smolgp.solvers.kalman.kalman_filter_batched_mean(A_all, H_all, K_all, y_batch, m0)[source]#

Batched-mean-path replay of kalman_filter’s forward recursion, given PRECOMPUTED, y-independent gains (kalman_gains). Processes M observation/residual batches in a single jax.lax.scan.

Runtime: O(N * M * dim * D), vs. O(M * N * dim^3) for M independent calls to kalman_filter.

Parameters:

  • A_all – (N, dim, dim), H_all: (N, D, dim), K_all: (N, dim, D) – from kalman_gains

  • y_batch – (M, N, D) – batch of M observation/residual arrays

  • m0 – (dim,) – prior mean, shared/broadcast across the batch

Returns:

(M, N, dim) m_predicted_batch: (M, N, dim)

Return type:

m_filtered_batch