diff --git a/cuthbert/gaussian/kalman.py b/cuthbert/gaussian/kalman.py index 7ee2fd9c..0598fba5 100644 --- a/cuthbert/gaussian/kalman.py +++ b/cuthbert/gaussian/kalman.py @@ -17,8 +17,16 @@ from cuthbert.inference import Filter, Smoother from cuthbert.utils import dummy_tree_like from cuthbertlib.kalman import filtering, smoothing +from cuthbertlib.kalman.filtering import ( + SteadyStateFilterParams, + compute_steady_state_filter_params, +) from cuthbertlib.types import Array, ArrayTree, ArrayTreeLike, KeyArray +# Expose this function at the top level of the module, +# as it can be useful for users to precompute steady-state parameters. +_ = compute_steady_state_filter_params + class KalmanFilterState(NamedTuple): """Kalman filter state.""" @@ -65,6 +73,7 @@ def build_filter( get_init_params: GetInitParams, get_dynamics_params: GetDynamicsParams, get_observation_params: GetObservationParams, + steady_state_params: SteadyStateFilterParams | None = None, ) -> Filter: """Builds an exact Kalman filter object for linear Gaussian SSMs. @@ -77,6 +86,12 @@ def build_filter( get_observation_params: Function to get observation parameters, H, d, chol_R, y given model inputs sufficient to define p(y_t | x_t) = N(H @ x_t + d, chol_R @ chol_R^T). + steady_state_params: Optional precomputed steady-state filter parameters + (see + [`compute_steady_state_filter_params`][cuthbert.gaussian.kalman.compute_steady_state_filter_params]). + When provided the per-step QR decomposition is skipped and only the + observation-dependent quantities are recomputed, giving a significant + speed-up for long time series with time-invariant models. Returns: Filter object for exact Kalman filter. Suitable for associative scan. @@ -90,6 +105,7 @@ def build_filter( filter_prepare, get_dynamics_params=get_dynamics_params, get_observation_params=get_observation_params, + steady_state_params=steady_state_params, ), filter_combine=filter_combine, associative=True, @@ -164,6 +180,7 @@ def filter_prepare( model_inputs: ArrayTreeLike, get_dynamics_params: GetDynamicsParams, get_observation_params: GetObservationParams, + steady_state_params: SteadyStateFilterParams | None = None, key: KeyArray | None = None, ) -> KalmanFilterState: """Prepare a state for an exact Kalman filter step. @@ -172,6 +189,11 @@ def filter_prepare( model_inputs: Model inputs. get_dynamics_params: Function to get dynamics parameters, F, c, chol_Q. get_observation_params: Function to get observation parameters, H, d, chol_R, y. + steady_state_params: Optional precomputed steady-state filter parameters + (see + [`compute_steady_state_filter_params`][cuthbert.gaussian.kalman.compute_steady_state_filter_params]). + When provided the per-step QR decomposition is skipped and only the + observation-dependent quantities are recomputed. key: JAX random key - not used. Returns: @@ -180,7 +202,9 @@ def filter_prepare( model_inputs = tree.map(lambda x: jnp.asarray(x), model_inputs) F, c, chol_Q = get_dynamics_params(model_inputs) H, d, chol_R, y = get_observation_params(model_inputs) - elem = filtering.associative_params_single(F, c, chol_Q, H, d, chol_R, y) + elem = filtering.associative_params_single( + F, c, chol_Q, H, d, chol_R, y, steady_state_params + ) return KalmanFilterState(elem=elem, model_inputs=model_inputs) diff --git a/cuthbertlib/kalman/__init__.py b/cuthbertlib/kalman/__init__.py index e1b2c6cf..ef45585f 100644 --- a/cuthbertlib/kalman/__init__.py +++ b/cuthbertlib/kalman/__init__.py @@ -1,4 +1,8 @@ from cuthbertlib.kalman import generate -from cuthbertlib.kalman.filtering import predict +from cuthbertlib.kalman.filtering import ( + SteadyStateFilterParams, + compute_steady_state_filter_params, + predict, +) from cuthbertlib.kalman.filtering import update as filter_update from cuthbertlib.kalman.smoothing import update as smoother_update diff --git a/cuthbertlib/kalman/filtering.py b/cuthbertlib/kalman/filtering.py index 8c125d1c..0c6deb10 100644 --- a/cuthbertlib/kalman/filtering.py +++ b/cuthbertlib/kalman/filtering.py @@ -62,6 +62,7 @@ def update( chol_R: ArrayLike, y: ArrayLike, log_normalizing_constant: ArrayLike = 0.0, + steady_state_params: "SteadyStateFilterParams | None" = None, ) -> tuple[tuple[Array, Array], Array]: """Update the mean and square root covariance with a linear Gaussian observation. @@ -74,6 +75,16 @@ def update( y: Observation. log_normalizing_constant: Optional input of log normalizing constant to be added to log normalizing constant of the Bayesian update. + steady_state_params: Optional precomputed steady-state parameters + (see + [`SteadyStateFilterParams`][cuthbertlib.kalman.filtering.SteadyStateFilterParams]). + When provided the QR + decomposition is skipped; `elem.U` is used as the posterior + Cholesky covariance and `K` as the Kalman gain. For the + sequential filter the caller is responsible for supplying a + `SteadyStateFilterParams` whose `K` and `elem.U` reflect the + Riccati steady state, not the parallel-scan values produced by + [`compute_steady_state_filter_params`][cuthbert.gaussian.kalman.compute_steady_state_filter_params]. Returns: Updated mean and square root covariance as well as the log marginal likelihood. @@ -95,26 +106,36 @@ def update( y_hat = H @ m + d y_diff = y - y_hat - M = jnp.block( - [ - [H @ chol_P, chol_R], - [chol_P, jnp.zeros((n_x, n_y), dtype=chol_P.dtype)], - ] - ) - chol_S = tria(M) - chol_Py = chol_S[n_y:, n_y:] - - Gmat = chol_S[n_y:, :n_y] - Imat = chol_S[:n_y, :n_y] - - my = m + Gmat @ solve_triangular(Imat, y_diff, lower=True) + if steady_state_params is None: + M = jnp.block( + [ + [H @ chol_P, chol_R], + [chol_P, jnp.zeros((n_x, n_y), dtype=chol_P.dtype)], + ] + ) + chol_S = tria(M) + chol_Py = chol_S[n_y:, n_y:] + Gmat = chol_S[n_y:, :n_y] + Imat = chol_S[:n_y, :n_y] + my = m + Gmat @ solve_triangular(Imat, y_diff, lower=True) + else: + Imat = steady_state_params.chol_S + chol_Py = steady_state_params.elem.U + my = m + steady_state_params.K @ y_diff ell = multivariate_normal.logpdf(y, y_hat, Imat, nan_support=False) return (my, chol_Py), jnp.asarray(ell + log_normalizing_constant) def associative_params_single( - F: Array, c: Array, chol_Q: Array, H: Array, d: Array, chol_R: Array, y: Array + F: Array, + c: Array, + chol_Q: Array, + H: Array, + d: Array, + chol_R: Array, + y: Array, + steady_state_params: "SteadyStateFilterParams | None" = None, ) -> FilterScanElement: """Single time step for scan element for square root parallel Kalman filter. @@ -126,6 +147,14 @@ def associative_params_single( d: Observation shift. chol_R: Generalized Cholesky factor of the observation noise covariance. y: Observation. + steady_state_params: Optional precomputed steady-state parameters + (see + [`SteadyStateFilterParams`][cuthbertlib.kalman.filtering.SteadyStateFilterParams]). + When provided the QR + decomposition is skipped; `A`, `U`, and `Z` are reused from + the stored element and only the observation-dependent `b`, + `eta`, and `ell` are evaluated, replacing the expensive + per-step QR with cheap matrix–vector products. Returns: Prepared scan element for the square root parallel Kalman filter. @@ -136,41 +165,170 @@ def associative_params_single( ny, nx = H.shape - # joint over the predictive and the observation - Psi_ = jnp.block([[H @ chol_Q, chol_R], [chol_Q, jnp.zeros((nx, ny))]]) + inn = y - H @ c - d # innovation relative to the prior prediction - Tria_Psi_ = tria(Psi_) + if steady_state_params is None: + # joint over the predictive and the observation + Psi_ = jnp.block([[H @ chol_Q, chol_R], [chol_Q, jnp.zeros((nx, ny))]]) + Tria_Psi_ = tria(Psi_) - Psi11 = Tria_Psi_[:ny, :ny] - Psi21 = Tria_Psi_[ny : ny + nx, :ny] - U = Tria_Psi_[ny : ny + nx, ny:] + Psi11 = Tria_Psi_[:ny, :ny] + Psi21 = Tria_Psi_[ny : ny + nx, :ny] + U = Tria_Psi_[ny : ny + nx, ny:] + + # pre-compute inverse of Psi11: we apply it to matrices and vectors alike. + Psi11_inv = solve_triangular(Psi11, jnp.eye(ny), lower=True) + + K = Psi21 @ Psi11_inv # local Kalman gain + HF = H @ F + A = F - K @ HF # corrected transition matrix - # pre-compute inverse of Psi11: we apply it to matrices and vectors alike. - Psi11_inv = solve_triangular(Psi11, jnp.eye(ny), lower=True) + # information filter (Z_filter is the pre-padded (nx, ny) version) + Z_filter = HF.T @ Psi11_inv.T + if nx > ny: + Z = jnp.concatenate([Z_filter, jnp.zeros((nx, nx - ny))], axis=1) + else: + Z = tria(Z_filter) + else: + A, _, U, _, Z, _ = steady_state_params.elem + K = steady_state_params.K + Psi11 = steady_state_params.chol_S + Psi11_inv = steady_state_params.Psi11_inv + Z_filter = steady_state_params.Z_filter + + b = c + K @ inn # corrected transition offset + eta = Z_filter @ (Psi11_inv @ inn) + ell = jnp.asarray( + multivariate_normal.logpdf(y, H @ c + d, Psi11, nan_support=False) + ) - # predictive model given one observation - K = Psi21 @ Psi11_inv # local Kalman gain - HF = H @ F # temporary variable - A = F - K @ HF # corrected transition matrix + return FilterScanElement(A, b, U, eta, Z, ell) - b = c + K @ (y - H @ c - d) # corrected transition offset - # information filter - Z = HF.T @ Psi11_inv.T - eta = Psi11_inv @ (y - H @ c - d) - eta = Z @ eta +class SteadyStateFilterParams(NamedTuple): + """Precomputed, time-invariant quantities for a steady-state Kalman filter. + + In a time-invariant linear Gaussian SSM the filter gain and posterior + covariance converge to constants. Passing an instance of this class as the + `steady_state_params` argument to + [`associative_params_single`][cuthbertlib.kalman.filtering.associative_params_single] + or [`update`][cuthbertlib.kalman.filtering.update] skips the expensive + per-step QR decomposition and reuses the + constant `A`, `U`, and `Z` blocks. + + Attributes: + elem: A `FilterScanElement` whose `A`, `U`, and `Z` fields hold + the steady-state values. The `b`, `eta`, and `ell` fields are + unused (they depend on the observation). + K: Steady-state Kalman gain matrix, shape `(nx, ny)`. + chol_S: Lower-triangular Cholesky factor of the steady-state innovation + covariance `S = H P_pred H^T + R`, shape `(ny, ny)`. + Psi11_inv: Inverse of `chol_S`, shape `(ny, ny)`. Precomputed so + that the per-step cost is pure matrix–vector products with no + triangular solve. + Z_filter: Pre-padded information gain matrix `(H F)^T S^{-T}`, + shape `(nx, ny)`. Used to form `eta` without the QR + decomposition; distinct from the square `Z` in `elem` which is + used by the associative combine operator. + """ + + elem: FilterScanElement + K: Array + chol_S: Array + Psi11_inv: Array + Z_filter: Array + + +def compute_steady_state_filter_params( + F: Array, + chol_Q: Array, + H: Array, + chol_R: Array, +) -> SteadyStateFilterParams: + """Compute steady-state filter parameters for a time-invariant linear Gaussian SSM. + + For a time-invariant model the `A`, `U`, and `Z` fields of each + associative scan element are identical at every time step — they depend + only on `F`, `chol_Q`, `H`, and `chol_R`, not on the observation. + This function extracts those constant fields once (along with the Kalman + gain `K` and the innovation Cholesky `chol_S`) by performing the same + block-triangularization as + [`associative_params_single`][cuthbertlib.kalman.filtering.associative_params_single] + but without + evaluating the observation-dependent terms. + + The returned + [`SteadyStateFilterParams`][cuthbertlib.kalman.filtering.SteadyStateFilterParams] + can be passed to + [`associative_params_single`][cuthbertlib.kalman.filtering.associative_params_single] + or [`update`][cuthbertlib.kalman.filtering.update] to skip the per-step QR + decomposition and only recompute the cheap observation-dependent + quantities `b`, `eta`, and `ell` at every step. + + This function is intended to be called **outside of JIT** as a one-off + pre-computation. The returned params can then be passed into a JIT-compiled + filter call for all subsequent runs. + + Note: + The `K` stored here is the *parallel-scan* gain, derived from + `chol_Q` alone (see + [`associative_params_single`][cuthbertlib.kalman.filtering.associative_params_single]). + When using [`update`][cuthbertlib.kalman.filtering.update] in a + sequential filter, supply a + [`SteadyStateFilterParams`][cuthbertlib.kalman.filtering.SteadyStateFilterParams] + whose `K` and `elem.U` reflect the Riccati steady state instead. + + Args: + F: State transition matrix, shape `(nx, nx)`. + chol_Q: Generalized Cholesky factor of the transition noise covariance, + shape `(nx, nx)`. + H: Observation matrix, shape `(ny, nx)`. + chol_R: Generalized Cholesky factor of the observation noise covariance, + shape `(ny, ny)`. + Returns: + [`SteadyStateFilterParams`][cuthbertlib.kalman.filtering.SteadyStateFilterParams] + containing the constant `A`, `U`, + `Z`, gain `K`, innovation Cholesky `chol_S`, precomputed + `Psi11_inv`, and pre-padded information gain `Z_filter`. + """ + F = jnp.asarray(F) + chol_Q = jnp.asarray(chol_Q) + H = jnp.asarray(H) + chol_R = jnp.asarray(chol_R) + + ny, nx = H.shape + + # Mirrors the block-triangularization in associative_params_single, + # stopping before the observation-dependent b, eta, ell. + Psi_ = jnp.block([[H @ chol_Q, chol_R], [chol_Q, jnp.zeros((nx, ny))]]) + Tria_Psi_ = tria(Psi_) + + chol_S = Tria_Psi_[:ny, :ny] # innovation Cholesky + Psi21 = Tria_Psi_[ny : ny + nx, :ny] + U = Tria_Psi_[ny : ny + nx, ny:] # posterior sqrt covariance + + Psi11_inv = solve_triangular(chol_S, jnp.eye(ny), lower=True) + K = Psi21 @ Psi11_inv # Kalman gain + + HF = H @ F + A = F - K @ HF + + Z_filter = HF.T @ Psi11_inv.T # (nx, ny); used for eta, before padding + Z = Z_filter if nx > ny: Z = jnp.concatenate([Z, jnp.zeros((nx, nx - ny))], axis=1) else: Z = tria(Z) - # local log marginal likelihood - ell = jnp.asarray( - multivariate_normal.logpdf(y, H @ c + d, Psi11, nan_support=False) - ) + dummy_b = jnp.zeros(nx) + dummy_eta = jnp.zeros(nx) + dummy_ell = jnp.array(0.0) - return FilterScanElement(A, b, U, eta, Z, ell) + elem = FilterScanElement(A=A, b=dummy_b, U=U, eta=dummy_eta, Z=Z, ell=dummy_ell) + return SteadyStateFilterParams( + elem=elem, K=K, chol_S=chol_S, Psi11_inv=Psi11_inv, Z_filter=Z_filter + ) def filtering_operator( diff --git a/cuthbertlib/quadrature/common.py b/cuthbertlib/quadrature/common.py index a1c975f4..0dc05a6a 100644 --- a/cuthbertlib/quadrature/common.py +++ b/cuthbertlib/quadrature/common.py @@ -1,6 +1,6 @@ """Common types and protocols for quadrature.""" -from typing import NamedTuple, Protocol, Self, runtime_checkable +from typing import NamedTuple, Protocol, runtime_checkable import jax.numpy as jnp @@ -45,7 +45,7 @@ def mean(self) -> Array: return jnp.dot(self.wm, self.points) # Should this be property too? - def covariance(self, other: Self | None = None) -> Array: + def covariance(self, other: "SigmaPoints | None" = None) -> Array: """Computes the covariance between the sigma points and the other sigma points. Args: diff --git a/docs/examples/kalman_tracking.md b/docs/examples/kalman_tracking.md index cf7bee1d..8adaebb5 100644 --- a/docs/examples/kalman_tracking.md +++ b/docs/examples/kalman_tracking.md @@ -238,7 +238,7 @@ car's position. transformations like `jit`, `vmap`, and automatic differentiation. ## Next Steps - +- **Steady-state Kalman filtering**: Explore the [steady-state Kalman filter](steady_state_kalman.md) for time-invariant systems, which can be more efficient for long time series. - **Smoothing**: Use [`cuthbert.smoother`](../api_cuthbert/smoothing.md) for backward pass smoothing. - **Parameter Learning**: Combine with optimization libraries like diff --git a/docs/examples/steady_state_kalman.md b/docs/examples/steady_state_kalman.md new file mode 100644 index 00000000..84dbe614 --- /dev/null +++ b/docs/examples/steady_state_kalman.md @@ -0,0 +1,233 @@ +# Steady-state Kalman filter + +For **time-invariant** linear Gaussian SSMs the Kalman gain converges to a constant after a few steps. +`cuthbert` lets you pre-compute that constant gain once — outside +of JIT — and reuse it at every time step, replacing the per-step QR decomposition +inside [`filter_prepare`][cuthbert.gaussian.kalman.filter_prepare] with cheap +matrix–vector products. + +This example shows how to use [`compute_steady_state_filter_params`][cuthbertlib.kalman.filtering.compute_steady_state_filter_params] +together with [`build_filter`][cuthbert.gaussian.kalman.build_filter], and benchmarks +the result against the standard filter. + +### Setup and imports + +```{.python #steady-state-kalman-imports} +import timeit + +import jax +import jax.numpy as jnp +from jax import jit, random + +from cuthbert import filter as run_filter +from cuthbert.gaussian import kalman +``` + +### Build a synthetic LG-SSM + +We use a random time-invariant system with state dimension $n_x = 20$ and +observation dimension $n_y = 4$. The larger dimensions make the QR savings more +pronounced than in the small car-tracking example. + +```{.python #steady-state-kalman-model} +KEY = random.key(42) +NUM_STEPS = 5_000 +NX, NY = 20, 4 + +key, sk = random.split(KEY) +F = jnp.eye(NX) * 0.99 + 0.01 * random.normal(sk, (NX, NX)) / NX +key, sk = random.split(key) +_noise = random.normal(sk, (NX, NX)) / NX**0.5 +chol_Q = jnp.linalg.cholesky(jnp.eye(NX) + 0.1 * _noise @ _noise.T) +key, sk = random.split(key) +H = random.normal(sk, (NY, NX)) / NX**0.5 +chol_R = 0.5 * jnp.eye(NY) +c = jnp.zeros(NX) +d = jnp.zeros(NY) +m0 = jnp.zeros(NX) +chol_P0 = jnp.eye(NX) + +# Simulate trajectory +key, sk = random.split(key) +x = m0 + chol_P0 @ random.normal(sk, (NX,)) +ys = [] +for _ in range(NUM_STEPS): + key, sk1, sk2 = random.split(key, 3) + x = F @ x + chol_Q @ random.normal(sk1, (NX,)) + y = H @ x + chol_R @ random.normal(sk2, (NY,)) + ys.append(y) +ys = jnp.stack(ys) +``` + +### Build both filters + +The standard filter is built in the usual way. The steady-state filter requires +only one extra call — [`compute_steady_state_filter_params`][cuthbertlib.kalman.filtering.compute_steady_state_filter_params] +— which performs the block-triangularization of `associative_params_single` **once** +and packages the constant results into a +[`SteadyStateFilterParams`][cuthbertlib.kalman.filtering.SteadyStateFilterParams] struct. + +```{.python #steady-state-kalman-build} +def get_init_params(model_inputs): + return m0, chol_P0 + + +def get_dynamics_params(model_inputs): + return F, c, chol_Q + + +def get_observation_params(model_inputs): + return H, d, chol_R, ys[model_inputs - 1] + + +model_inputs = jnp.arange(len(ys) + 1) + +# Standard filter +standard_filter = kalman.build_filter( + get_init_params, get_dynamics_params, get_observation_params +) + +# Steady-state filter: the only required change is passing ss_params. +ss_params = kalman.compute_steady_state_filter_params(F, chol_Q, H, chol_R) +steady_state_filter = kalman.build_filter( + get_init_params, get_dynamics_params, get_observation_params, + steady_state_params=ss_params, +) +``` + +!!! note "What does `compute_steady_state_filter_params` compute?" + In the parallel associative scan each `FilterScanElement` carries constant + matrices `A`, `U`, `Z` (and the gain `K`) that are identical at every time step + for a time-invariant model. `compute_steady_state_filter_params` extracts these + by running the same block-QR as `associative_params_single` **once** (outside JIT), + so that per-step `filter_prepare` calls need only evaluate the + observation-dependent `b`, `eta`, and `ell` — purely matrix–vector products. + + This is distinct from the *Riccati* steady-state gain used in sequential + filters; see [`SteadyStateFilterParams`][cuthbertlib.kalman.filtering.SteadyStateFilterParams] + for details on the distinction. + +### JIT-compile and verify correctness + +```{.python #steady-state-kalman-run} +jitted_filter = jit(run_filter, static_argnames=["filter_obj", "parallel"]) + +# Warm up (triggers XLA compilation) +standard_result = jitted_filter(standard_filter, model_inputs, parallel=True) +jax.block_until_ready(standard_result.mean) + +ss_result = jitted_filter(steady_state_filter, model_inputs, parallel=True) +jax.block_until_ready(ss_result.mean) + +# Both filters compute mathematically identical operations; any difference is +# purely float32 rounding from evaluating the same QR inside vs. outside vmap. +max_abs_err = float(jnp.max(jnp.abs(standard_result.mean - ss_result.mean))) +mean_magnitude = float(jnp.mean(jnp.abs(standard_result.mean))) +rel_err = max_abs_err / (mean_magnitude + 1e-8) +print(f"Max absolute difference in posterior means: {max_abs_err:.2e} (float32 rounding)") +print(f"Relative error (max abs / mean magnitude): {rel_err:.2e}") +assert rel_err < 5e-2, ( + f"Steady-state and standard filters disagree beyond float32 rounding: " + f"relative error = {rel_err:.2e}" +) +``` + +### Benchmark + +We compare wall-clock time for both **parallel** (associative scan, +$\mathcal{O}(\log T)$ depth) and **sequential** ($\mathcal{O}(T)$) modes. + +In the parallel case the bottleneck is the `filtering_operator` combine step — +which also contains `tria` calls — so the gain from skipping `filter_prepare`'s +QR is modest. In the sequential case every `filter_prepare` call is on the +critical path, so the savings are larger. + +```{.python #steady-state-kalman-benchmark} +REPS = 30 + + +def time_filter(filter_obj, parallel): + def run(): + result = jitted_filter(filter_obj, model_inputs, parallel=parallel) + jax.block_until_ready(result.mean) + return timeit.timeit(run, number=REPS) / REPS + + +# Warm up sequential variants +_ = jitted_filter(standard_filter, model_inputs, parallel=False) +jax.block_until_ready(_.mean) +_ = jitted_filter(steady_state_filter, model_inputs, parallel=False) +jax.block_until_ready(_.mean) + +t_par_std = time_filter(standard_filter, parallel=True) +t_par_ss = time_filter(steady_state_filter, parallel=True) +t_seq_std = time_filter(standard_filter, parallel=False) +t_seq_ss = time_filter(steady_state_filter, parallel=False) + +print(f"\nBenchmark over {NUM_STEPS} time steps, nx={NX}, ny={NY} ({REPS} reps each)") +print(f" {'':30s} {'standard':>12s} {'steady-state':>12s} {'speed-up':>9s}") +print(f" {'parallel (associative scan)':30s} {t_par_std*1e3:12.2f} {t_par_ss*1e3:12.2f} {t_par_std/t_par_ss:9.2f}x") +print(f" {'sequential (scan)':30s} {t_seq_std*1e3:12.2f} {t_seq_ss*1e3:12.2f} {t_seq_std/t_seq_ss:9.2f}x") +``` + +### Observed output + +#### CPU (float32) + +```text +Max absolute difference in posterior means: 2.81e-03 (float32 rounding) +Relative error (max abs / mean magnitude): 1.05e-03 + +Benchmark over 5000 time steps, nx=20, ny=4 (30 reps each) + standard steady-state speed-up + parallel (associative scan) 668.16 621.28 1.08x + sequential (scan) 344.74 278.22 1.24x +``` + +#### GPU (float64) + +```text +Max absolute difference in posterior means: 3.69e-12 +Relative error (max abs / mean magnitude): 1.47e-12 + +Benchmark over 5000 time steps, nx=20, ny=4 (30 reps each) + standard steady-state speed-up + parallel (associative scan) 85.27 80.65 1.06x + sequential (scan) 5452.67 4225.71 1.29x +``` + +## Key Takeaways + +- **One extra line**: the only change from a standard filter is calling + `compute_steady_state_filter_params` once and passing the result to + `build_filter`. +- **No approximation**: results are mathematically identical to the standard + filter (differences are float32 rounding only; vanish with float64). +- **Where the speedup is**: the QR decomposition is eliminated from every + `filter_prepare` call. The gain is most visible in sequential mode and + grows with state / observation dimension and on GPU. +- **Parallel-scan vs Riccati**: `compute_steady_state_filter_params` computes + the *parallel-scan* gain (prior = $Q$, single block-QR), not the Riccati + steady-state gain. For the sequential `update` function, supply a + `SteadyStateFilterParams` built from the DARE solution instead. + +## Next Steps + +- **Kalman tracking**: See [the car tracking example](kalman_tracking.md) for a + gentle introduction to `cuthbert`'s filter API. +- **Temporal parallelization**: The [temporal parallelization example](temporal_parallelization_kalman.md) + explains the parallel associative scan in more depth. +- **Parameter estimation**: Combine with gradient-based optimization — see the + [EM example](parameter_estimation_em.md). + + + + diff --git a/tests/cuthbertlib/kalman/test_filtering.py b/tests/cuthbertlib/kalman/test_filtering.py index f26c2fab..5e0618d9 100644 --- a/tests/cuthbertlib/kalman/test_filtering.py +++ b/tests/cuthbertlib/kalman/test_filtering.py @@ -1,9 +1,17 @@ import chex import jax import jax.numpy as jnp +import numpy as np import pytest +from scipy.linalg import solve_discrete_are -from cuthbertlib.kalman.filtering import predict, update +from cuthbertlib.kalman.filtering import ( + FilterScanElement, + SteadyStateFilterParams, + compute_steady_state_filter_params, + predict, + update, +) from cuthbertlib.kalman.generate import generate_lgssm @@ -85,3 +93,206 @@ def test_update(seed, x_dim, y_dim): chex.assert_trees_all_close((m, P), (des_m, des_P), rtol=1e-10) chex.assert_trees_all_close(ell, des_ell, rtol=1e-10) + + +# --------------------------------------------------------------------------- +# Helpers shared by the steady-state tests +# --------------------------------------------------------------------------- + + +def _make_lgssm(seed, nx, ny): + """Return (F, c, chol_Q, H, d, chol_R, Q, R) for a stable random LG-SSM.""" + rng = np.random.default_rng(seed) + # Stable transition: random orthogonal matrix scaled below 1 + _A = rng.standard_normal((nx, nx)) + U, _, Vt = np.linalg.svd(_A) + F = 0.8 * (U @ Vt) + c = rng.standard_normal(nx) * 0.1 + _L = rng.standard_normal((nx, nx)) + Q = _L @ _L.T / nx + np.eye(nx) + chol_Q = np.linalg.cholesky(Q) + H = rng.standard_normal((ny, nx)) / np.sqrt(nx) + _M = rng.standard_normal((ny, ny)) + R = _M @ _M.T / ny + np.eye(ny) + chol_R = np.linalg.cholesky(R) + d = rng.standard_normal(ny) * 0.1 + return F, c, chol_Q, H, d, chol_R, Q, R + + +# --------------------------------------------------------------------------- +# Test 1 – compute_steady_state_filter_params against closed-form expressions +# +# The parallel-scan gain is derived from a *single* block-QR where the prior +# covariance is Q (process noise), NOT from a Riccati iteration. It satisfies +# the closed-form: +# +# K = Q H^T (H Q H^T + R)^{-1} (Kalman gain, prior = Q) +# U = chol(Q − K H Q) (posterior sqrt covariance) +# A = F − K H F (corrected transition) +# Z_filter = (H F)^T S^{−T} (pre-padded information gain) +# Psi11_inv = chol_S^{−1} (precomputed inverse) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("seed", [0, 42, 99]) +@pytest.mark.parametrize("nx, ny", [(3, 2), (4, 4), (5, 3)]) +def test_compute_steady_state_filter_params_closed_form(seed, nx, ny): + F, c, chol_Q, H, d, chol_R, Q, R = _make_lgssm(seed, nx, ny) + + ss = compute_steady_state_filter_params( + jnp.array(F), + jnp.array(chol_Q), + jnp.array(H), + jnp.array(chol_R), + ) + + # closed-form quantities + S = H @ Q @ H.T + R # innovation covariance + K_ref = Q @ H.T @ np.linalg.inv(S) + P_post_ref = Q - K_ref @ H @ Q # posterior covariance + HF = H @ F + + # gain + chex.assert_trees_all_close(ss.K, jnp.array(K_ref), atol=1e-10) + + # innovation Cholesky: chol_S @ chol_S.T == S + S_reconstructed = ss.chol_S @ ss.chol_S.T + chex.assert_trees_all_close(S_reconstructed, jnp.array(S), atol=1e-10) + + # Psi11_inv is the inverse of chol_S + I_check = ss.chol_S @ ss.Psi11_inv + chex.assert_trees_all_close(I_check, jnp.eye(ny), atol=1e-10) + + # Z_filter = (H F)^T chol_S^{-T} + Z_filter_ref = HF.T @ np.linalg.inv(np.array(ss.chol_S)).T + chex.assert_trees_all_close(ss.Z_filter, jnp.array(Z_filter_ref), atol=1e-10) + + # posterior covariance: U @ U.T == P_post + P_post_reconstructed = ss.elem.U @ ss.elem.U.T + chex.assert_trees_all_close(P_post_reconstructed, jnp.array(P_post_ref), atol=1e-10) + + # corrected transition: A = F - K H F + A_ref = F - K_ref @ HF + chex.assert_trees_all_close(ss.elem.A, jnp.array(A_ref), atol=1e-10) + + +# --------------------------------------------------------------------------- +# Test 2 – Riccati-based SteadyStateFilterParams for the sequential update +# +# The TRUE steady-state gain for the sequential filter is the Riccati K, +# obtained by solving the discrete algebraic Riccati equation (DARE): +# +# solve_discrete_are(F, H.T, Q, R) → P_pred_ss +# K_riccati = P_pred_ss H^T (H P_pred_ss H^T + R)^{-1} +# +# We construct a SteadyStateFilterParams from Riccati quantities and verify +# that update(..., steady_state_params=ss_riccati) produces the same posterior +# mean as the standard update once the running filter has converged. +# --------------------------------------------------------------------------- + + +def _riccati_steady_state_params(F, Q, H, R): + """Build SteadyStateFilterParams from the DARE solution (Riccati K). + + `scipy.linalg.solve_discrete_are(a, b, q, r)` solves + `a.T X a - X - (a.T X b) inv(r + b.T X b) (b.T X a) + q = 0`, + so to obtain the Kalman predictive-covariance DARE + `P = F P_post F.T + Q - F P_post H.T inv(S) H P_post F.T` + the first argument must be `F.T`, not `F`. + """ + # DARE: solve_discrete_are(F.T, H.T, Q, R) returns P_pred_ss + P_pred = solve_discrete_are(F.T, H.T, Q, R) + S = H @ P_pred @ H.T + R + chol_S = np.linalg.cholesky(S) + Psi11_inv = np.linalg.inv(chol_S) + K = P_pred @ H.T @ Psi11_inv.T @ Psi11_inv # K = P_pred H^T S^{-1} + P_post = P_pred - K @ H @ P_pred + chol_P_post = np.linalg.cholesky(P_post) + # Z_filter and A are needed only by associative_params_single, not update; + # populate them correctly for completeness. + HF = H @ F + Z_filter = HF.T @ Psi11_inv.T + nx = F.shape[0] + ny = H.shape[0] + if nx > ny: + Z = np.concatenate([Z_filter, np.zeros((nx, nx - ny))], axis=1) + else: + from scipy.linalg import qr + + _, Z = qr(Z_filter.T, mode="economic") + Z = Z.T + A = F - K @ HF + dummy = np.zeros(nx) + elem = FilterScanElement( + A=jnp.array(A), + b=jnp.array(dummy), + U=jnp.array(chol_P_post), + eta=jnp.array(dummy), + Z=jnp.array(Z), + ell=jnp.array(0.0), + ) + return SteadyStateFilterParams( + elem=elem, + K=jnp.array(K), + chol_S=jnp.array(chol_S), + Psi11_inv=jnp.array(Psi11_inv), + Z_filter=jnp.array(Z_filter), + ) + + +@pytest.mark.parametrize("seed", [0, 42, 99]) +@pytest.mark.parametrize("nx, ny", [(3, 2), (4, 4), (5, 3)]) +def test_update_steady_state_riccati(seed, nx, ny): + """Riccati steady-state update matches the standard update initialized at the Riccati covariance. + + When the standard filter is started at the Riccati posterior covariance, its + running K is exactly the Riccati K at every subsequent step. The two filters + must therefore produce identical posterior means. + """ + F, c, chol_Q, H, d, chol_R, Q, R = _make_lgssm(seed, nx, ny) + ss_riccati = _riccati_steady_state_params(F, Q, H, R) + + rng = np.random.default_rng(seed + 1) + + # Both filters start from the same arbitrary mean and the Riccati posterior covariance. + # With chol_P = chol_P_riccati, one predict+update cycle in the standard filter + # uses exactly K_riccati, so both means must agree at every step. + m0 = jnp.array(rng.standard_normal(nx)) + chol_P_riccati = ss_riccati.elem.U + + m_std = m0 + chol_P_std = chol_P_riccati + m_ss = m0 + + for _ in range(20): + # standard: full predict + update + m_std, chol_P_std = predict( + m_std, + chol_P_std, + jnp.array(F), + jnp.array(c), + jnp.array(chol_Q), + ) + y = jnp.array(rng.standard_normal(ny)) + (m_std, chol_P_std), _ = update( + m_std, + chol_P_std, + jnp.array(H), + jnp.array(d), + jnp.array(chol_R), + y, + ) + + # steady-state: same predict step, then update with fixed Riccati params + m_pred_ss = jnp.array(F @ np.array(m_ss) + c) + (m_ss, _), _ = update( + m_pred_ss, + chol_P_riccati, + jnp.array(H), + jnp.array(d), + jnp.array(chol_R), + y, + steady_state_params=ss_riccati, + ) + + chex.assert_trees_all_close(m_ss, m_std, atol=1e-10)