kalman

kalman#

Functions#

IntegratedKalmanFilter(kernel, X, y, t_states, obsid, ...)

Wrapper for integrated_kalman_filter function

integrated_kalman_filter(A_aug, Q_aug, H_aug, R, ...)

Jax implementation of the integrated Kalman filter algorithm

integrated_kalman_gains(A_aug, H_aug, RESET, R, X, ...)

The y-independent per-state quantities needed to replay

integrated_kalman_filter_batched_mean(A_all, H_all, ...)

Batched-mean-path replay of integrated_kalman_filter, using PRECOMPUTED

Module Contents#

smolgp.solvers.integrated.kalman.IntegratedKalmanFilter(kernel, X, y, t_states, obsid, instid, stateid, R, return_v_S=False)[source]#

Wrapper for integrated_kalman_filter function

Parameters:

  • kernel – IntegratedStateSpaceModel kernel

  • X – Array of size N, data coordinates (e.g. (time, texp, instid))

  • y – Array of size (N, D), measurements at the data coordinates

  • t_states – Array of size K, sorted time coordinate of all states (exposure starts and ends)

  • obsid – Array of size N, which observation (0,…,N-1) is being made at each state k

  • instid – Array of size N, which instrument (0,…,Ninst-1) recorded observation n

  • stateid – Array of size K, 0 for exposure-start, 1 for exposure-end

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

  • return_v_S – Whether to return innovation and its covariance (for likelihood computation)

Returns:

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

Return type:

m_filtered

smolgp.solvers.integrated.kalman.integrated_kalman_filter(A_aug, Q_aug, H_aug, R, RESET, X, y, t_states, obsid, instid, stateid, m0, P0)[source]#

Jax implementation of the integrated Kalman filter algorithm

See Section 3.2.1 in Rubenzahl & Hattori et al. (in prep) for detailed description of the algorithm and notation.

smolgp.solvers.integrated.kalman.integrated_kalman_gains(A_aug, H_aug, RESET, R, X, t_states, obsid, instid, stateid, P_predicted)[source]#

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

Computes BOTH the Kalman-gain ingredients (meaningful at exposure-end states) and the Reset matrix (meaningful at exposure-start states) at every state k unconditionally (cheap; no scan needed) – the batched-mean recursion selects the right one per state via the same stateid==0 jax.lax.cond dispatch integrated_kalman_filter uses.

Returns:

shape (K, dim, dim) – A_aug(0, Delta_k) H_all: shape (N, D, dim) – H_aug(X), i.e. jax.vmap(H_aug)(X) K_all: shape (K, dim, D) – Kalman gain (meaningful at end states) Reset_all: shape (K, dim, dim) – RESET(instid[obsid[k]]) (meaningful at start states)

Return type:

A_all

smolgp.solvers.integrated.kalman.integrated_kalman_filter_batched_mean(A_all, H_all, K_all, Reset_all, obsid, stateid, y_batch, m0)[source]#

Batched-mean-path replay of integrated_kalman_filter, using PRECOMPUTED (integrated_kalman_gains) A_all/H_all/K_all/Reset_all.

The exposure-start “reset” update is a pure per-sample-batched einsum with no y/data dependence at all (the innovation is always zero at a start state) – jnp.einsum(“ij,mj->mi”, Reset_k, m_pred_batch). The stateid==0 dispatch itself (which state gets which update) is data-independent (known from t_states/obsid alone), so the same jax.lax.cond structure as integrated_kalman_filter applies unchanged, just wrapping batched einsums instead of unbatched matrix-vector @.

Parameters:

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

  • Reset_all – (K, dim, dim) – all from integrated_kalman_gains

  • obsid – (K,) bookkeeping arrays (same as integrated_kalman_filter)

  • stateid – (K,) bookkeeping arrays (same as integrated_kalman_filter)

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

  • m0 – (dim,)

Returns:

(M, K, dim)

Return type:

m_filtered_batch, m_predicted_batch