Skip to content

Specialized Models

Observation models

Observation model implementations.

DiracIdentityObservation

Bases: ObservationModel

Noise-free identity observation model.

Observations are modeled as

\[ y_t \sim \delta(x_t), \]

i.e., the observation equals the latent state almost surely.

GaussianObservation

Bases: ObservationModel

Nonlinear Gaussian observation model.

Observations are modeled as

\[ y_t \sim \mathcal{N}(h(x_t, u_t, t), R), \]

where \(h\) is a user-provided measurement function and \(R\) is the observation noise covariance.

__init__(h: Callable[[Real[Array, ' state_dim'] | Real[Array, ''], Real[Array, ' control_dim'] | Real[Array, ''] | None, Real[Array, '']], Real[Array, ' observation_dim'] | Real[Array, '']], R: Float[Array, '*plate observation_dim observation_dim'])

Parameters:

Name Type Description Default
h Callable[[State, Control, Time], Array]

Measurement function mapping \((x, u, t)\) to the mean observation.

 required
R Array

Observation noise covariance with shape \((d_y, d_y)\).

 required

LinearGaussianObservation

Bases: ObservationModel

Linear-Gaussian observation model.

Observations are modeled as

\[ y_t \sim \mathcal{N}(H x_t + D u_t + b, R). \]

Here, \(H\) is the observation matrix, \(D\) is an optional control-input matrix, \(b\) is an optional observation bias, and \(R\) is the observation noise covariance.

Each parameter may be a constant array (time-invariant) or a callable (t,) -> value evaluated at each observation time (time-varying); constant and callable parameters may be mixed freely. The single time argument mirrors the observation model contract \(p(y_t | x_t, u_t, t)\) (transitions span an interval; observations happen at one time).

Note
  • Callable parameters receive only the observation time t; they must not depend on state or controls (use GaussianObservation for nonlinear measurement functions).
  • Callables must be pure, JAX-traceable functions returning a fixed shape.
  • Backend support: time-varying parameters work with the simulators and the filter_source="cuthbert" filters/smoothers; the cd_dynamax backend requires constant arrays and raises TypeError otherwise.

is_time_invariant: bool property

True iff every parameter is a constant array (no callables).

__init__(H: Float[Array, '*h_plate observation_dim state_dim'] | Callable[[float | int | Real[Array, '']], Float[Array, '*h_plate observation_dim state_dim']], R: Float[Array, '*r_plate observation_dim observation_dim'] | Callable[[float | int | Real[Array, '']], Float[Array, '*r_plate observation_dim observation_dim']], D: Float[Array, '*d_matrix_plate observation_dim control_dim'] | Callable[[float | int | Real[Array, '']], Float[Array, '*d_matrix_plate observation_dim control_dim']] | None = None, bias: Float[Array, '*bias_plate observation_dim'] | Callable[[float | int | Real[Array, '']], Float[Array, '*bias_plate observation_dim']] | None = None)

Parameters:

Name Type Description Default
H Array | Callable

Observation matrix with shape \((d_y, d_x)\), or a callable (t,) returning it.

 required
R Array | Callable

Observation noise covariance with shape \((d_y, d_y)\), or a callable (t,) returning it.

 required
D Array | Callable | None

Optional control matrix with shape \((d_y, d_u)\), or a callable (t,) returning it. If None, no control contribution is used.

 None
bias Array | Callable | None

Optional additive bias with shape \((d_y,)\), or a callable (t,) returning it.

 None

params_at(t: float | int | Real[Array, '']) -> LinearGaussianObservationParams

Resolve (H, D, bias, R) at one observation time.

Constant parameters are returned unchanged; callable parameters are evaluated at t.

LinearGaussianObservationParams

Bases: NamedTuple

Linear-Gaussian observation parameters resolved at one time.

Returned by LinearGaussianObservation.params_at: any callable (time-varying) parameter has been evaluated at the requested time, so every entry is a plain array (or None for an absent optional term).

Expected shapes match the LinearGaussianObservation fields; they are deliberately not enforced here because plate slicing can legally hand a member-sliced (reduced-rank) parameter to __call__.

State evolution models

State evolution implementations.

Specialty implementations for discrete-time systems. Structure allows future extension to LTI factories, Neural SDEs, etc.

AffineDrift

Bases: Module

Affine drift function for continuous-time models.

This implements an affine map of the form

\[f(x, u, t) = A x + B u + b,\]

where \(A \in \mathbb{R}^{d_x \times d_x}\), \(B \in \mathbb{R}^{d_x \times d_u}\) (optional), and \(b \in \mathbb{R}^{d_x}\) (optional). The time argument \(t\) is accepted for compatibility with the Drift protocol but is not used.

This is commonly used as the drift term \(\mu(x_t, u_t, t)\) inside ContinuousTimeStateEvolution, and is a building block for LTI models such as LTI_continuous.

Attributes:

Name Type Description
A Array

Drift matrix with shape \((d_x, d_x)\).
B Array | None

Optional control matrix with shape \((d_x, d_u)\).
b Array | None

Optional additive bias with shape \((d_x,)\).

GaussianStateEvolution

Bases: DiscreteTimeStateEvolution

Nonlinear Gaussian discrete-time state transition.

The next state is modeled as

\[ x_{t_{k+1}} \sim \mathcal{N}(F(x_{t_k}, u_{t_k}, t_k, t_{k+1}), Q), \]

where \(F\) is a user-provided transition function and \(Q\) is the process-noise covariance (either constant or state/time dependent).

__init__(F: Callable[[Real[Array, ' state_dim'] | Real[Array, ''], Real[Array, ' control_dim'] | Real[Array, ''] | None, float | int | Real[Array, ''], float | int | Real[Array, '']], Real[Array, ' state_dim'] | Real[Array, '']], cov: Float[Array, '*plate state_dim state_dim'] | Callable[[Real[Array, ' state_dim'] | Real[Array, ''], Real[Array, ' control_dim'] | Real[Array, ''] | None, float | int | Real[Array, ''], float | int | Real[Array, '']], Float[Array, '*plate state_dim state_dim']])

Parameters:

Name Type Description Default
F Callable[[State, Control, Time, Time], State]

Transition function mapping \((x, u, t_k, t_{k+1})\) to the conditional mean.

 required
cov Array | Callable[[State, Control, Time, Time], Array]

Process-noise covariance with shape \((d_x, d_x)\), or a callable mapping \((x, u, t_k, t_{k+1})\) to that covariance.

 required

LinearGaussianParams

Bases: NamedTuple

Linear-Gaussian transition parameters resolved at one time interval.

Returned by LinearGaussianStateEvolution.params_at: any callable (time-varying) parameter has been evaluated at the requested interval, so every entry is a plain array (or None for an absent optional term).

Expected shapes match the LinearGaussianStateEvolution fields; they are deliberately not enforced here because plate slicing can legally hand a member-sliced (reduced-rank) parameter to __call__.

LinearGaussianStateEvolution

Bases: DiscreteTimeStateEvolution

Linear-Gaussian discrete-time state transition.

The next state is modeled as

\[ x_{t_{k+1}} \sim \mathcal{N}(A x_{t_k} + B u_{t_k} + b, Q), \]

where \(A\) is the state transition matrix, \(B\) is an optional control-input matrix, \(b\) is an optional transition bias, and \(Q\) is the process-noise covariance.

Each parameter may be a constant array (time-invariant) or a callable (t_now, t_next) -> value evaluated per transition interval (time-varying); constant and callable parameters may be mixed freely.

Note
  • Callable parameters receive only the interval endpoints (t_now, t_next); they must not depend on state or controls (use GaussianStateEvolution for nonlinear transitions).
  • Callables must be pure, JAX-traceable functions returning a fixed shape.
  • Backend support: time-varying parameters work with the simulators and the filter_source="cuthbert" filters/smoothers; the cd_dynamax backend requires constant arrays and raises TypeError otherwise.

is_time_invariant: bool property

True iff every parameter is a constant array (no callables).

__init__(A: Float[Array, '*a_plate state_dim state_dim'] | Callable[[float | int | Real[Array, ''], float | int | Real[Array, '']], Float[Array, '*a_plate state_dim state_dim']], cov: Float[Array, '*cov_plate state_dim state_dim'] | Callable[[float | int | Real[Array, ''], float | int | Real[Array, '']], Float[Array, '*cov_plate state_dim state_dim']], B: Float[Array, '*b_matrix_plate state_dim control_dim'] | Callable[[float | int | Real[Array, ''], float | int | Real[Array, '']], Float[Array, '*b_matrix_plate state_dim control_dim']] | None = None, bias: Float[Array, '*bias_plate state_dim'] | Callable[[float | int | Real[Array, ''], float | int | Real[Array, '']], Float[Array, '*bias_plate state_dim']] | None = None)

Parameters:

Name Type Description Default
A Array | Callable

State transition matrix with shape \((d_x, d_x)\), or a callable (t_now, t_next) returning it.

 required
cov Array | Callable

Process-noise covariance with shape \((d_x, d_x)\), or a callable (t_now, t_next) returning it.

 required
B Array | Callable | None

Optional control matrix with shape \((d_x, d_u)\), or a callable (t_now, t_next) returning it.

 None
bias Array | Callable | None

Optional additive bias with shape \((d_x,)\), or a callable (t_now, t_next) returning it.

 None

params_at(t_now: float | int | Real[Array, ''], t_next: float | int | Real[Array, '']) -> LinearGaussianParams

Resolve (A, B, bias, cov) at one transition interval.

Constant parameters are returned unchanged; callable parameters are evaluated at (t_now, t_next).

LTI model factories

LTI_continuous(A: Float[Array, '*a_plate state_dim state_dim'], L: Float[Array, '*diffusion_plate state_dim bm_dim'], H: Float[Array, '*h_plate observation_dim state_dim'], R: Float[Array, '*r_plate observation_dim observation_dim'], B: Float[Array, '*b_matrix_plate state_dim control_dim'] | None = None, b: Float[Array, '*state_bias_plate state_dim'] | None = None, D: Float[Array, '*d_matrix_plate observation_dim control_dim'] | None = None, d: Float[Array, '*obs_bias_plate observation_dim'] | None = None, initial_mean: Float[Array, '*init_mean_plate state_dim'] | None = None, initial_cov: Float[Array, '*init_cov_plate state_dim state_dim'] | None = None) -> DynamicalModel

Build a continuous-time linear time-invariant (LTI) DynamicalModel.

The state evolves according to the SDE and observation model

\[ \begin{aligned} x_0 &\sim \mathcal{N}(m_0, C_0), \\ dx_t &= (A x_t + B u_t + b) \, dt + L \, dW_t, \\ y_t &\sim \mathcal{N}(H x_t + D u_t + d, R). \end{aligned} \]

Here, \(L\) is a diffusion coefficient (not a covariance) with shape \((d_x, d_w)\). It multiplies a \(d_w\)-dimensional Brownian motion \(W_t\) whose increments have identity covariance: \(dW_t \sim \mathcal{N}(0, I_{d_w} \, dt)\). The Brownian motion dimension \(d_w\) is determined by the second dimension of \(L\). Under this convention, the infinitesimal state covariance contributed by the noise term is \(L L^\top \, dt\).

Parameters:

Name Type Description Default
A Array

Drift matrix with shape \((d_x, d_x)\).

 required
L Array

Diffusion coefficient with shape \((d_x, d_w)\).

 required
H Array

Observation matrix with shape \((d_y, d_x)\).

 required
R Array

Observation-noise covariance with shape \((d_y, d_y)\).

 required
B Array | None

Optional control matrix in the drift with shape \((d_x, d_u)\). If None, no control term is used and control_dim is set to 0.

 None
b Array | None

Optional additive drift bias with shape \((d_x,)\).

 None
D Array | None

Optional control matrix in the observation model with shape \((d_y, d_u)\).

 None
d Array | None

Optional additive observation bias with shape \((d_y,)\).

 None
initial_mean Array | None

Optional initial-state mean \(m_0\) with shape \((d_x,)\). Defaults to zeros.

 None
initial_cov Array | None

Optional initial-state covariance \(C_0\) with shape \((d_x, d_x)\). Defaults to identity.

 None

Returns:

Name Type Description
DynamicalModel DynamicalModel

A continuous-time LTI state-space model.

LTI_discrete(A: Float[Array, '*a_plate state_dim state_dim'], Q: Float[Array, '*q_plate state_dim state_dim'], H: Float[Array, '*h_plate observation_dim state_dim'], R: Float[Array, '*r_plate observation_dim observation_dim'], B: Float[Array, '*b_matrix_plate state_dim control_dim'] | None = None, b: Float[Array, '*state_bias_plate state_dim'] | None = None, D: Float[Array, '*d_matrix_plate observation_dim control_dim'] | None = None, d: Float[Array, '*obs_bias_plate observation_dim'] | None = None, initial_mean: Float[Array, '*init_mean_plate state_dim'] | None = None, initial_cov: Float[Array, '*init_cov_plate state_dim state_dim'] | None = None) -> DynamicalModel

Build a discrete-time linear time-invariant (LTI) DynamicalModel.

The model has transition and observation distributions

\[ \begin{aligned} x_0 &\sim \mathcal{N}(m_0, C_0), \\ x_{t_{k+1}} &\sim \mathcal{N}(A x_{t_k} + B u_{t_k} + b, Q), \\ y_{t_k} &\sim \mathcal{N}(H x_{t_k} + D u_{t_k} + d, R). \end{aligned} \]

This factory composes LinearGaussianStateEvolution and LinearGaussianObservation into a core DynamicalModel.

Parameters:

Name Type Description Default
A Array

State transition matrix with shape \((d_x, d_x)\).

 required
Q Array

Process-noise covariance with shape \((d_x, d_x)\).

 required
H Array

Observation matrix with shape \((d_y, d_x)\).

 required
R Array

Observation-noise covariance with shape \((d_y, d_y)\).

 required
B Array | None

Optional control matrix in the transition model with shape \((d_x, d_u)\). If None, no control term is used and control_dim is set to 0.

 None
b Array | None

Optional additive transition bias with shape \((d_x,)\).

 None
D Array | None

Optional control matrix in the observation model with shape \((d_y, d_u)\).

 None
d Array | None

Optional additive observation bias with shape \((d_y,)\).

 None
initial_mean Array | None

Optional initial-state mean \(m_0\) with shape \((d_x,)\). Defaults to zeros.

 None
initial_cov Array | None

Optional initial-state covariance \(C_0\) with shape \((d_x, d_x)\). Defaults to identity.

 None

Returns:

Name Type Description
DynamicalModel DynamicalModel

A discrete-time LTI state-space model.