Inference¶

The square-root Gaussian algebra (sqr_marginalization, sqr_inversion) is the QR-based numerical core every filter step is built on: it propagates covariances as upper-triangular factors A_sqr (with A = A_sqr.T @ A_sqr), never forming a dense covariance and keeping it symmetric and positive-semi-definite (no indefiniteness from round-off; see notation).

Built on top of these primitives, this module also exposes a parameter-inference layer — ODEFilter, fit, marginal_loglik, and InferenceProblem, plus the Real / PositiveReal parameter wrappers — for fitting ODE parameters by maximizing the marginal log-likelihood of external observations under the ODE-constrained Gauss-Markov model. (fit minimizes the negative log-likelihood via a caller-supplied Optax optimizer.)

Inference routines for ODE filtering.

Modules:

Name	Description
`model`	Ergonomic object API for ODE-parameter inference (Layer 2).
`parameter_inference`	Differentiable marginal likelihood for ODE-parameter inference (Layer 1).
`parameters`	Unconstrained-parameter wrappers for gradient-based inference.
`sqr_gaussian_inference`

Classes:

Name	Description
`AbstractParameter`	A parameter stored in unconstrained space.
`InferenceProblem`	Static specification mapping a parameter pytree to solver inputs.
`ODEFilter`	A fittable first-order probabilistic ODE solver.
`PositiveReal`	A strictly-positive parameter, stored via the softplus bijection.
`Real`	An unconstrained real parameter (identity transform).

Functions:

Name	Description
`fit`	Fit an :class:`ODEFilter` by maximizing the data marginal log-likelihood.
`marginal_loglik`	Data marginal log-likelihood for gradient-based ODE-parameter inference.
`sqr_inversion`	Numerically stable Bayesian inference using square-root representations.
`sqr_marginalization`	Marginalize out the linear transformation in a Gaussian model using square-root form.
`unwrap`	Replace every :class:`AbstractParameter` in a pytree with its value.

Classes¶

AbstractParameter ¶

Bases: Module

A parameter stored in unconstrained space.

Subclasses hold the raw (unconstrained) value as their array leaf and map it to the constrained value in :meth:unwrap. Being an equinox.Module they are registered pytrees, so the unconstrained leaf is what jax.grad and Optax see.

Methods:

Name	Description
`unwrap`	Return the constrained value.

Functions¶

unwrap `abstractmethod` ¶

unwrap() -> Array

Return the constrained value.

InferenceProblem ¶

Bases: NamedTuple

Static specification mapping a parameter pytree to solver inputs.

Attributes:

Name	Type	Description
`build`	`Callable[[Any], tuple[Array, Array, BaseODEInformation]]`	Callable `theta -> (mu_0, Sigma_0_sqr, measure)` constructing the (differentiable) initial Gaussian and measurement model from the parameters. This is the single problem-specific bridge; everything else is static.
`prior`	`Any`	Gauss-Markov prior (e.g. :class:`IWP`).
`tspan`	`tuple[float, float]`	Time interval `(t0, t1)` (a tuple, for jit-static hashing).
`N`	`int`	Number of fixed-grid steps.
`calibration`	`str`	Diffusion calibration mode. Defaults to `"none"` (fixed diffusion) -- recommended for parameter inference.
`min_sigma_sqr`	`float`	Lower bound passed through to the loop.
`correction`	`Any`	Linearization strategy (a :class:`~ode_filters.Correction`); `None` defaults to EK1. Use e.g. `TaylorCorrection(order=0)` or `IteratedTaylorCorrection()` to fit with EK0 / IEKF.

ODEFilter ¶

Bases: Module

A fittable first-order probabilistic ODE solver.

Attributes:

Name	Type	Description
`vf`	`Callable`	Vector field `vf(x, params, *, t) -> dx/dt` (`params` is the differentiable parameter pytree, `ode_params`).
`init_fn`	`Callable`	Callable `params -> (mu_0, Sigma_0_sqr)` producing the initial Gaussian (square-root covariance). May ignore `params` for a fixed initial condition.
`prior`	`Any`	Gauss-Markov prior (e.g. :class:`IWP`).
`tspan`	`tuple[float, float]`	Time interval `(t0, t1)`.
`N`	`int`	Number of fixed-grid steps.
`ode_params`	`Array`	The differentiable parameters (the only trainable leaf).
`calibration`	`str`	Diffusion calibration mode (default `"none"`).

Methods:

Name	Description
`loglik`	Data marginal log-likelihood at the current parameters.

Functions¶

loglik ¶

loglik(data: ObsModel) -> Array

Data marginal log-likelihood at the current parameters.

PositiveReal ¶

PositiveReal(value: Any)

Bases: AbstractParameter

A strictly-positive parameter, stored via the softplus bijection.

Construct it with the constrained (positive) value; it stores the corresponding unconstrained value and returns the positive value from :meth:unwrap.

Attributes:

Name	Type	Description
`unconstrained`	`Array`	The stored unconstrained value (`softplus_inv(value)`).

Real ¶

Real(value: Any)

Bases: AbstractParameter

An unconstrained real parameter (identity transform).

Attributes:

Name	Type	Description
`value`	`Array`	The (unconstrained) value, used as-is.

Functions¶

fit ¶

fit(
    model: ODEFilter,
    data: ObsModel,
    optim: Any,
    *,
    steps: int = 200,
) -> tuple[ODEFilter, Array]

Fit an :class:ODEFilter by maximizing the data marginal log-likelihood.

Parameters:

Name	Type	Description	Default
`model`	`ODEFilter`	Initial :class:`ODEFilter`.	required
`data`	`ObsModel`	Observations as an :class:`ObsModel`.	required
`optim`	`Any`	An Optax `GradientTransformation` (anything exposing `init` / `update`); the library does not import Optax itself.	required
`steps`	`int`	Number of optimization steps.	`200`

Returns:

Type	Description
`ODEFilter`	Tuple `(fitted_model, losses)` where `losses` is the per-step negative
`Array`	log-likelihood (shape `[steps]`).

marginal_loglik ¶

marginal_loglik(
    theta: Any, data: ObsModel, *, model: InferenceProblem
) -> Array

Data marginal log-likelihood for gradient-based ODE-parameter inference.

Parameters:

Name	Type	Description	Default
`theta`	`Any`	Differentiable parameter pytree, consumed by `model.build`. May contain :class:`~ode_filters.PositiveReal` / :class:`~ode_filters.Real` wrappers, which are unwrapped automatically.	required
`data`	`ObsModel`	Observations as an :class:`ObsModel` (e.g. from :func:`prepare_observations`); the measured values live in `data.c_seq`.	required
`model`	`InferenceProblem`	Static :class:`InferenceProblem` (closed over, not traced).	required

Returns:

Type	Description
`Array`	Scalar observation marginal log-likelihood. Maximize over `theta` (or
`Array`	minimize its negative) to fit parameters. `jax.jit` / `jax.grad` /
`Array`	`jax.vmap` over `theta` are all supported.

Raises:

Type	Description
`ValueError`	If `data` is `None` (observations are required).

sqr_inversion ¶

sqr_inversion(
    A: Array,
    mu: Array,
    Sigma_sqr: Array,
    mu_z: Array,
    Sigma_z_sqr: Array,
    Q_sqr: Array | None = None,
) -> tuple[Array, Array, Array]

Numerically stable Bayesian inference using square-root representations.

Performs Bayesian update of a Gaussian given a linear observation model, using Cholesky factors and QR decomposition to maintain numerical stability.

Given p(x) ~ N(mu, Sigma) and p(z|x) ~ N(Ax + b, Q), computes the posterior p(x|z) in square-root form.

Parameters:

Name	Type	Description	Default
`A`	`Array`	Observation matrix (shape [n_obs, n_state]).	required
`mu`	`Array`	Prior mean (shape [n_state]).	required
`Sigma_sqr`	`Array`	Square root of prior covariance (shape [n_state, n_state]).	required
`mu_z`	`Array`	Marginal observation mean (shape [n_obs]).	required
`Sigma_z_sqr`	`Array`	Square root of marginal observation covariance (shape [n_obs, n_obs]).	required
`Q_sqr`	`Array \| None`	Square root of measurement noise covariance (optional).	`None`

Returns:

Type	Description
`Array`	Tuple of (G, d, Lambda_sqr) where:
`Array`	G is the Kalman gain matrix (shape [n_state, n_obs])
`Array`	d is the posterior offset/mean correction (shape [n_state])
`tuple[Array, Array, Array]`	Lambda_sqr is the posterior covariance square root (shape [n_state, n_state])

sqr_marginalization ¶

sqr_marginalization(
    A: Array,
    b: Array,
    Q_sqr: Array,
    mu: Array,
    Sigma_sqr: Array,
) -> tuple[Array, Array]

Marginalize out the linear transformation in a Gaussian model using square-root form.

Computes the marginal distribution of z = Ax + b given p(x) ~ N(mu, Sigma) and p(z|x) ~ N(Ax + b, Q). The result is p(z) = N(mu_z, Sigma_z) where: - mu_z = A @ mu + b - Sigma_z = A @ Sigma @ A.T + Q

The square-root form is preserved to maintain numerical stability.

Parameters:

Name	Type	Description	Default
`A`	`Array`	Linear transformation matrix (shape [n_obs, n_state]).	required
`b`	`Array`	Observation offset (shape [n_obs]).	required
`Q_sqr`	`Array`	Square root of observation noise covariance. Shape [n_obs, n_obs] or [n_obs]. If 1D array, will be converted to 2D.	required
`mu`	`Array`	Prior mean (shape [n_state]).	required
`Sigma_sqr`	`Array`	Square root of prior covariance (shape [n_state, n_state]).	required

Returns:

Type	Description
`Array`	Tuple of (mu_z, Sigma_z_sqr) where:
`Array`	mu_z is the marginal mean of z (shape [n_obs])
`tuple[Array, Array]`	Sigma_z_sqr is the square root of marginal covariance of z (shape [n_obs, n_obs])

Raises:

Type	Description
`ValueError`	If input shapes are incompatible or invalid.

unwrap ¶

unwrap(tree: Any) -> Any

Replace every :class:AbstractParameter in a pytree with its value.

Identity on leaves that are not :class:AbstractParameter, so it is safe to call on any theta (plain arrays pass through unchanged).

Inference¶

Classes¶

AbstractParameter ¶

Functions¶

unwrap abstractmethod ¶

InferenceProblem ¶

ODEFilter ¶

Functions¶

loglik ¶

PositiveReal ¶

Real ¶

Functions¶

fit ¶

marginal_loglik ¶

sqr_inversion ¶

sqr_marginalization ¶

unwrap ¶

unwrap `abstractmethod` ¶