Filters¶

The consolidated solver API: gaussian_filter (fixed grid), gaussian_filter_adaptive (adaptive, save-at-grid), and rts_smoother for the backward smoothing pass — all returning a single FilterResult. The linearization is a pluggable Correction (EK0 / EK1 / IEKF), and adaptive stepping uses a step-size controller (PController / PIController). Diffusion is calibrated per step by default (calibration="dynamic", a scalar quasi-MLE), so the process noise is rescaled each step; pass calibration="none" for a fixed diffusion. See How to choose.

Dispatch is automatic: a PrecondIWP / PrecondMaternPrior prior selects the preconditioned square-root recursion, and passing obs_model= adds a masked observation update.

The `FilterResult`¶

Both gaussian_filter and gaussian_filter_adaptive return a FilterResult with named fields (recover a covariance with P = P_sqr.T @ P_sqr; see notation):

Field	Meaning
`t`	time grid
`m`, `P_sqr`	filtered posterior mean and square-root covariance
`log_likelihood`	scalar post-calibration ODE-information log-marginal-likelihood (not comparable across `calibration` modes — each defines a different generative model)
`log_likelihood_obs`	observation log-likelihood (`None` without `obs_model`)
`m_pred`, `P_pred_sqr`	one-step predictions before each update (`None` for the adaptive save-at solver)
`G_back`, `d_back`, `P_back_sqr`	backward transitions consumed by `rts_smoother`
`sigma_sqr`	per-step calibrated diffusion `sigma_hat^2` (`None` for the adaptive solver)
`success`	adaptive only: whether sub-stepping reached every save time and the log-likelihood is finite (`None` for the fixed grid)
`m_bar`, `P_bar_sqr`, `T`	preconditioned-space internals (`None` for a plain prior)

gaussian_filter_adaptive is filtering-only by default (smoother=False): its backward-pass fields are then None and rts_smoother raises. Pass smoother=True to additionally compute a fixed-point smoothing backward pass (one composite conditional per save interval, O(#save points) memory); the result then carries G_back / d_back / P_back_sqr and rts_smoother(prior, result) works directly. smoother=True is not supported together with obs_model.

The low-level scan loops and step functions that gaussian_filter wraps live in the ode_filters.filters.ode_filter_loop / ode_filter_step / ode_filter_adaptive submodules; they are implementation detail, not part of the public API.

Filtering and smoothing routines for ODE models.

The public solver API is the consolidated :func:gaussian_filter / :func:gaussian_filter_adaptive / :func:rts_smoother (returning :class:FilterResult), plus the pluggable :class:Correction strategies and the adaptive step-size controllers. The historical sqr_loop* / *_step* matrix is now an implementation detail of those entry points; import it from the ode_filters.filters.ode_filter_loop / ode_filter_step / ode_filter_adaptive submodules if you need the low-level variants directly.

Modules:

Name	Description
`adaptive_controller`	Step-size controllers for the adaptive EKF loop.
`correction`	Pluggable linearization/correction strategies for the Gaussian ODE filter.
`gaussian_filter`	Consolidated public solver API: `gaussian_filter` + smoother.
`ode_filter_adaptive`	Adaptive-step EKF loop with online sigma calibration.
`ode_filter_loop`
`ode_filter_step`

Classes:

Name	Description
`Correction`	Base class for linearization/correction strategies.
`CorrectionResult`	Output of a single :meth:`Correction.correct` call.
`FilterResult`	Filtered solution of a Gaussian ODE filter.
`IteratedTaylorCorrection`	Iterated EKF (IEKF): per-step Gauss-Newton relinearization.
`PController`	Proportional (single-error) step-size controller.
`PIController`	Gustafsson-style proportional-integral controller.
`StepSizeController`	Duck-typed interface for step-size controllers.
`TaylorCorrection`	Taylor-series correction: EK1 (`order=1`) or EK0 (`order=0`).

Functions:

Name	Description
`gaussian_filter_adaptive`	Adaptive-step Gaussian filter, returning the solution at `save_at`.
`rts_smoother`	Rauch-Tung-Striebel smoothing of a Gaussian-filter result.

Classes¶

Correction ¶

Bases: Module

Base class for linearization/correction strategies.

A Correction is an :class:equinox.Module (static configuration), so it can be captured by jax.jit / jax.lax.scan without retracing. Subclasses own the full predicted -> posterior update and return a :class:CorrectionResult.

Methods:

Name	Description
`correct`	Update a predicted Gaussian with the measurement model at time `t`.

Functions¶

correct `abstractmethod` ¶

correct(
    measure: BaseODEInformation,
    m_pred: Array,
    P_pred_sqr: Array,
    *,
    t: ArrayLike = 0.0,
) -> CorrectionResult

Update a predicted Gaussian with the measurement model at time t.

Parameters:

Name	Type	Description	Default
`measure`	`BaseODEInformation`	ODE-information measurement model.	required
`m_pred`	`Array`	Predicted (prior) state mean.	required
`P_pred_sqr`	`Array`	Predicted state covariance, square-root form.	required
`t`	`ArrayLike`	Current time.	`0.0`

CorrectionResult ¶

Bases: NamedTuple

Output of a single :meth:Correction.correct call.

Attributes:

Name	Type	Description
`m`	`Array`	Updated posterior mean (shape `[state_dim]`).
`P_sqr`	`Array`	Updated posterior covariance in square-root form (`P = P_sqr.T @ P_sqr`).
`mz`	`Array`	Innovation (predicted-observation) mean (shape `[obs_dim]`).
`Pz_sqr`	`Array`	Innovation covariance in square-root form -- consumed by the Gaussian log-marginal-likelihood.

FilterResult ¶

Bases: NamedTuple

Filtered solution of a Gaussian ODE filter.

Attributes:

Name	Type	Description
`t`	`Array`	Time grid, shape `[K]`.
`m`	`Array`	Filtered state means at `t`, shape `[K, state_dim]`.
`P_sqr`	`Array`	Square-root covariances at `t`, shape `[K, state_dim, state_dim]` (`P = P_sqr.T @ P_sqr`).
`log_likelihood`	`Array`	Marginal log-likelihood of the ODE-information residuals, taken after diffusion calibration. Because each `calibration` mode defines a different generative model, this value is not comparable across calibration modes (for data-driven model comparison use `log_likelihood_obs`).
`m_pred`	`Array \| None`	Predicted (prior) means per step, `[K-1, state_dim]` (`None` for the adaptive save-at solver, which keeps no backward pass).
`P_pred_sqr`	`Array \| None`	Predicted square-root covariances per step.
`G_back`	`Array \| None`	Backward-pass gains per step (smoother input).
`d_back`	`Array \| None`	Backward-pass offsets per step.
`P_back_sqr`	`Array \| None`	Backward-pass square-root covariances per step.
`mz`	`Array \| None`	Predicted-observation (ODE-defect) innovation means per step; the input to post-hoc diffusion calibration. For the adaptive solver this is the innovation of the sub-step that lands on each save time (its `h` is clamped to hit the save time, so the raw magnitudes are not comparable across save points -- see :class:`AdaptiveSolveResult`; the whitened residual and NIS are unaffected).
`Pz_sqr`	`Array \| None`	Predicted-observation innovation square-root covariances per step.
`mz_obs`	`Array \| None`	External-observation innovation means per step, `h(m_ode_n) - y_n`, shape `[K-1, obs_dim]` (`None` when no `obs_model` was given). The innovation is taken at the ODE-updated predictive marginal: the sequential filter conditions on the ODE (and Conservation) information before the observation update, so `m_ode_n` is the post-ODE-update mean, not the raw prior prediction `m_pred`. Note the sign -- this is `h(m) - y`, not `y - h(m)` (the magnitude, hence NIS, is unaffected). The observation-channel analog of `mz`; together with `Pz_obs_sqr` it gives the innovation sequence used for filter-consistency tests (NIS/whitened residuals), outlier gating, and innovation-based noise tuning.
`Pz_obs_sqr`	`Array \| None`	External-observation innovation square-root covariances per step, `S_n = H P_ode_n H^T + R` in square-root form, shape `[K-1, obs_dim, obs_dim]` (`None` when no `obs_model`). `P_ode_n` is the post-ODE-update covariance (see `mz_obs`), not the prior-prediction `P_pred_sqr`.
`sigma_sqr`	`Array \| None`	Per-step calibrated diffusion `sigma_hat^2`.
`log_likelihood_obs`	`Array \| None`	Marginal log-likelihood of the external observations (`None` when no `obs_model` was given) -- the quantity to maximize for data-driven parameter inference. Fixed-grid path only: the adaptive solver always returns `None` here and folds the observation contribution into the combined `log_likelihood` instead.
`m_bar`	`Array \| None`	Preconditioned-space means (`None` unless the prior is preconditioned); internal, consumed by :func:`rts_smoother`.
`P_bar_sqr`	`Array \| None`	Preconditioned-space square-root covariances (`None` for plain).
`T`	`Array \| None`	Preconditioner matrix (`None` for plain); presence selects the preconditioned smoother.
`success`	`Array \| None`	Scalar boolean -- whether an adaptive solve reached every save time and produced a finite log-likelihood (see :class:`AdaptiveSolveResult`). `None` for the fixed-grid paths, which run a deterministic number of steps and always complete.

IteratedTaylorCorrection ¶

Bases: Correction

Iterated EKF (IEKF): per-step Gauss-Newton relinearization.

A single forward pass; within each update, the EK1 linearization is recomputed at the updated mean and the update redone, for a fixed max_iters passes (max_iters=1 reproduces EK1). A fixed iteration count -- rather than a convergence-based while_loop -- keeps the correction reverse-mode differentiable, so it can be used inside gradient-based parameter inference.

Attributes:

Name	Type	Description
`max_iters`	`int`	Number of relinearizations per step (>= 1).

PController `dataclass` ¶

PController(
    order: int,
    safety: float = 0.9,
    alpha: float | None = None,
    min_factor: float = 0.2,
    max_factor: float = 5.0,
)

Proportional (single-error) step-size controller.

The proposed step is

h_new = h * safety * err^(-alpha)

clipped to [min_factor, max_factor] * h. The previous error is ignored, so this controller has no memory.

Default alpha = 1 / order matches the standard "I-controller" of Hairer-Wanner-Norsett (1993) and the elementary controller in scipy's solve_ivp. For an EKF1 with IWP(q) prior, order = q.

Parameters:

Name	Type	Description	Default
`order`	`int`	Convergence order of the local error estimate.	required
`safety`	`float`	Safety factor on the proposed step.	`0.9`
`alpha`	`float \| None`	Gain on the current error; defaults to `1.0 / order`.	`None`
`min_factor`	`float`	Lower clip on `h_new / h`.	`0.2`
`max_factor`	`float`	Upper clip on `h_new / h`.	`5.0`

Methods:

Name	Description
`propose`	Return the next proposed step size.

Functions¶

propose ¶

propose(
    h: float, err: float, err_prev: float | None = None
) -> float

Return the next proposed step size.

Parameters:

Name	Type	Description	Default
`h`	`float`	Current step size.	required
`err`	`float`	Normalised error of the current step (1.0 = at tolerance).	required
`err_prev`	`float \| None`	Ignored; accepted for protocol compatibility with :class:`PIController`.	`None`

PIController `dataclass` ¶

PIController(
    order: int,
    safety: float = 0.9,
    alpha: float | None = None,
    beta: float | None = None,
    min_factor: float = 0.2,
    max_factor: float = 5.0,
)

Gustafsson-style proportional-integral controller.

The proposed step is

h_new = h * safety * err^(-alpha) * (err_prev / err)^(beta)

clipped to [min_factor, max_factor] * h. When err_prev is unknown (e.g. on the first step or right after a reject) the I-term is dropped and the update reduces to the :class:PController form h_new = h * safety * err^(-alpha).

Defaults follow Gustafsson (1991): alpha = 0.7 / order, beta = 0.4 / order, which are also the values used in Bosch et al. (2021). For an EKF1 with IWP(q) prior, order = q.

Parameters:

Name	Type	Description	Default
`order`	`int`	Convergence order of the local error estimate.	required
`safety`	`float`	Safety factor on the proposed step.	`0.9`
`alpha`	`float \| None`	Proportional gain; defaults to `0.7 / order`.	`None`
`beta`	`float \| None`	Integral gain; defaults to `0.4 / order`. Set `beta=0` for a pure proportional controller -- :class:`PController` is the standalone equivalent with the more natural default `alpha`.	`None`
`min_factor`	`float`	Lower clip on `h_new / h`.	`0.2`
`max_factor`	`float`	Upper clip on `h_new / h`.	`5.0`

Methods:

Name	Description
`propose`	Return the next proposed step size.

Functions¶

propose ¶

propose(
    h: float, err: float, err_prev: float | None
) -> float

Return the next proposed step size.

Parameters:

Name	Type	Description	Default
`h`	`float`	Current step size.	required
`err`	`float`	Normalised error of the current step (1.0 = at tolerance).	required
`err_prev`	`float \| None`	Normalised error of the previously accepted step, or `None` if unavailable.	required

StepSizeController ¶

Bases: Protocol

Duck-typed interface for step-size controllers.

err_prev may be None to indicate the absence of memory (first step, or right after a reject). Besides :meth:propose, controllers expose the static coefficients (safety, _alpha, min_factor, max_factor) that the adaptive loop reads to re-implement the proposal under jax.numpy tracing.

Methods:

Name	Description
`propose`	Return the proposed next step size.

Functions¶

propose ¶

propose(
    h: float, err: float, err_prev: float | None
) -> float

Return the proposed next step size.

TaylorCorrection ¶

Bases: Correction

Taylor-series correction: EK1 (order=1) or EK0 (order=0).

Attributes:

Name	Type	Description
`order`	`int`	`1` for the first-order (EK1) linearization with the full vector-field Jacobian (default; bit-identical to the historical filter step). `0` for the zeroth-order (EK0) linearization.

Functions¶

gaussian_filter_adaptive ¶

gaussian_filter_adaptive(
    mu_0: Array,
    P_0_sqr: Array,
    prior: BasePrior,
    measure: BaseODEInformation,
    save_at: Array,
    *,
    correction: Correction | None = None,
    obs_model: ObsModel | None = None,
    atol: float = 0.0001,
    rtol: float = 0.01,
    h_init: float | None = None,
    calibration: CalibrationMode = "dynamic",
    controller=None,
    min_sigma_sqr: float = 0.0,
    max_steps: int = 4096,
    smoother: bool = False,
) -> FilterResult

Adaptive-step Gaussian filter, returning the solution at save_at.

jit / vmap / reverse-grad-able (checkpointed adaptive loop). See :func:sqr_adaptive_solve for the full argument docs.

correction selects the linearization (EK0/EK1/IEKF), matching :func:gaussian_filter; result.success reports whether the adaptive sub-stepping reached every save time.

With an obs_model the observation likelihood is folded into the combined result.log_likelihood (summed over accepted steps); unlike the fixed-grid :func:gaussian_filter, result.log_likelihood_obs is always None on the adaptive path.

With smoother=True the result carries a fixed-point-smoothing backward pass (one composite conditional per save interval, O(#save points) memory), so :func:rts_smoother applies directly. Default False keeps the filtering-only path (no backward pass, lower cost); smoother=True is not supported together with obs_model.

rts_smoother ¶

rts_smoother(
    prior: BasePrior, result: FilterResult
) -> tuple[Array, Array]

Rauch-Tung-Striebel smoothing of a Gaussian-filter result.

Works on a :func:gaussian_filter result (backward pass over the fixed grid) or a :func:gaussian_filter_adaptive result run with smoother=True (the fixed-point backward pass over the save grid); both carry the required G_back / d_back / P_back_sqr. A filtering-only adaptive result (smoother=False, the default) does not, and raises.

Parameters:

Name	Type	Description	Default
`prior`	`BasePrior`	The prior used for the forward filter (selects the preconditioned smoother when preconditioned).	required
`result`	`FilterResult`	A :class:`FilterResult` carrying a backward pass.	required

Returns:

Type	Description
`Array`	Tuple `(m_smooth, P_smooth_sqr)` of smoothed means and square-root
`Array`	covariances, shapes `[K, state_dim]` and `[K, state_dim, state_dim]`.

Filters¶

The FilterResult¶

Classes¶

Correction ¶

Functions¶

correct abstractmethod ¶

CorrectionResult ¶

FilterResult ¶

IteratedTaylorCorrection ¶

PController dataclass ¶

Functions¶

propose ¶

PIController dataclass ¶

Functions¶

propose ¶

StepSizeController ¶

Functions¶

propose ¶

TaylorCorrection ¶

Functions¶

gaussian_filter_adaptive ¶

rts_smoother ¶

The `FilterResult`¶

correct `abstractmethod` ¶

PController `dataclass` ¶

PIController `dataclass` ¶