Specialized Models

Observation models

For exact identity measurements (y = x), use DiracIdentityObservation(). With LatentPathBuilder, observed values fix those state coordinates, and only missing state coordinates are inferred. Relaxing the observation turns these into noisy Gaussian measurements: observed values contribute a likelihood, and the latent state coordinates are inferred.

Observation model implementations.

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'] | jax_sparse.JAXSparse | Callable[[float | int | Real[Array, '']], Float[Array, '*h_plate observation_dim state_dim']], R: Covariance | 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 | JAXSparse | Callable	Observation matrix with shape \((d_y, d_x)\), or a callable (t,) returning it. May be a sparse (e.g. BCOO) array for EnKFConfig/PFConfig/EKFConfig (EKF works but likely gives no efficiency gain, warns) and for KFConfig(filter_source="cd_dynamax"); raises for KFConfig(filter_source="cuthbert"), which cannot support a sparse H.	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.

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: Covariance | 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

DeterministicObservation

Bases: ObservationModel

Noise-free observations y = h(x, u, t) as Delta distributions.

Scalar outputs are scalar events; vector outputs use the trailing axis as the event axis, preserving leading batch/plate axes.

__call__(x: Real[Array, '*batch state_dim'] | Real[Array, ''], u: Real[Array, '*control_batch control_dim'] | Real[Array, ''] | None, t: float | int | Real[Array, '']) -> dist.Delta

Evaluate h and preserve its scalar or trailing vector event shape.

DiracIdentityObservation

Bases: DeterministicObservation

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.

__init__() -> None

Construct the argument-free deterministic identity observation.

State evolution models

State evolution implementations.

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

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: Covariance | 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).

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: Covariance | 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

DeterministicStateEvolution

Bases: DiscreteTimeStateEvolution

Noise-free discrete transition x_next = F(x, u, t_now, t_next).

Returns a Delta distribution with scalar events or a trailing vector event axis. Leading axes remain batch/plate dimensions.

__call__(x: Real[Array, '*batch state_dim'] | Real[Array, ''], u: Real[Array, '*control_batch control_dim'] | Real[Array, ''] | None, t_now: float | int | Real[Array, ''], t_next: float | int | Real[Array, '']) -> dist.Delta

Evaluate F and preserve its scalar or trailing vector event shape.

Covariance specifications

Scalar, diagonal, and full covariance objects are reusable in Gaussian models, LTI_discrete, and Gaussian relaxation:

• ScalarCovariance(sd=... | variance=...) represents isotropic covariance.
• DiagonalCovariance(sd=... | variance=...) holds one variance per event coordinate.
• FullCovariance(matrix) holds a dense covariance matrix.

Scalar and diagonal objects require exactly one of sd or variance. All three preserve leading batch/plate axes as Equinox pytrees; their trailing covariance event axes have rank zero, one, or two respectively.

Relaxation settings also accept shorthand: a scalar is a variance, a vector contains diagonal variances, and a matrix is a full covariance. Use explicit covariance objects for batched settings to distinguish plate axes from event axes.

Covariance objects check shape, finite values, and symmetry. Supply positive-semidefinite covariances and a positive-definite effective covariance for Gaussian densities. Zero covariance can be added to an already positive-definite Gaussian. Relaxation uses the specified covariance exactly; the EnKF setting recorded_filtered_states_cov_jitter only adds jitter to returned state distributions, leaving the filter recursion and marginal likelihood unchanged.

Structured, sliceable covariance specifications for Gaussian models.

Covariance

Bases: Module

Base interface for scalar, diagonal, and full covariance specifications.

Numeric fields are JAX pytree leaves. Leading axes are batch/plate axes; event_rank identifies the trailing axes belonging to one covariance.

value: Float[Array, ...] abstractmethod property

Return the numeric leaf with event_rank trailing covariance axes.

as_matrix(event_dim: int) -> Float[Array, '*batch {event_dim} {event_dim}'] abstractmethod

Return a dense covariance with trailing shape (event_dim, event_dim).

ScalarCovariance

Bases: Covariance

Isotropic covariance variance * I; specify sd or variance.

An array of scalars represents leading batch/plate dimensions.

value: Float[Array, '*batch'] property

Return one variance per batch member, with no trailing event axes.

__init__(*, sd: Real[ArrayLike, '*batch'] | None = None, variance: Real[ArrayLike, '*batch'] | None = None) -> None

Treat every input axis as a batch axis of isotropic variances.

as_matrix(event_dim: int) -> Float[Array, '*batch {event_dim} {event_dim}']

Expand each scalar variance to an isotropic event_dim covariance.

DiagonalCovariance

Bases: Covariance

Diagonal covariance; the trailing axis holds variances or standard deviations.

value: Float[Array, '*batch event_dim'] property

Return the diagonal variances, retaining their trailing event axis.

__init__(*, sd: Real[ArrayLike, ...] | None = None, variance: Real[ArrayLike, ...] | None = None) -> None

Use the trailing input axis for event variances and leading axes for batches.

as_matrix(event_dim: int) -> Float[Array, '*batch {event_dim} {event_dim}']

Build diagonal matrices after checking the trailing event dimension.

FullCovariance

Bases: Covariance

A full covariance matrix with two trailing event axes.

Leading axes represent batches/plates. Construction checks shape, finite values, and symmetry. Supply a positive-semidefinite matrix, with a positive-definite effective covariance when a Gaussian density is needed.

value: Float[Array, '*batch event_dim event_dim'] property

Return the dense matrix with two trailing covariance event axes.

__init__(matrix: Real[ArrayLike, ...]) -> None

Validate square trailing event axes, finite values, and symmetry.

as_matrix(event_dim: int) -> Float[Array, '*batch {event_dim} {event_dim}']

Return the matrix after checking both trailing event dimensions.

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 drift control term is used.	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

Notes

control_dim is inferred from B, or from D when B is None, and defaults to 0 when both are 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: Covariance | Float[Array, '*q_plate state_dim state_dim'], H: Float[Array, '*h_plate observation_dim state_dim'], R: Covariance | 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: Covariance | Float[Array, '*init_cov_plate state_dim state_dim'] | None = None, observation_control_alignment: ObservationControlAlignment | str | 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 transition control term is used.	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
observation_control_alignment	ObservationControlAlignment | str | None	Forwarded to DynamicalModel; see its docstring. Defaults to None (unspecified).	None

Notes

control_dim is inferred from B, or from D when B is None, and defaults to 0 when both are None.

Returns:

Name	Type	Description
DynamicalModel	DynamicalModel	A discrete-time LTI state-space model.