Handlers¶

Contains the core dynestyx primitives and effectful handler utilities.

`plate` ¶

Bases: ObjectInterpretation

Hierarchical plate for batched trajectories.

dsx.plate wraps numpyro.plate for parameter sampling semantics and intercepts dsx.sample to pass plate sizes to simulator and filter handlers. Use it when a dynamical system has conditionally independent members, such as multiple trajectories, patients, groups, or treatment arms.

Shape semantics

Dynestyx treats plate axes as leading data-batch axes. Time axes come after plate axes in observation arrays, and state/observation event axes remain trailing axes.

For one plate of size N:

obs_values          # (N, T, obs_dim), or (N, T) for scalar observations
mu_i                # (N, state_dim), a vector parameter per trajectory
initial_mean        # (N, state_dim), or shared as (state_dim,)
initial_cov         # (N, state_dim, state_dim), or shared as (state_dim, state_dim)
prior["f_states"]   # (num_samples, N, n_sim, T, state_dim)

Distribution-valued model components use their NumPyro batch_shape/event_shape split: leading plate dimensions are batch dimensions, while state and observation sizes are inferred from event_shape. Thus a batched initial condition may have loc.shape == (N, state_dim) with either shared or batched covariance.

Built-in LTI vector fields, including transition/drift and observation biases, may be shared or plate-batched:

with dsx.plate("trajectories", N):
    mu_i = numpyro.sample(
        "mu_i",
        dist.Normal(mu_global, sigma).to_event(1),
    )  # (N, state_dim)

    dynamics = LTI_discrete(
        A=A,
        Q=Q,
        H=H,
        R=R,
        b=mu_i,              # plate-batched vector bias
        initial_mean=mu_i,   # plate-batched initial mean
    )

Ambiguous arrays are kept shared rather than sliced. In particular, a shared vector whose length happens to equal N is not treated as plate-batched unless it is a known vector-valued model field with an explicit event axis, such as (N, state_dim). For one-dimensional vector fields, prefer explicit singleton event axes like (N, 1). Nested plates follow the same rule with multiple leading plate axes.

Why event shapes drive sizing

Inside a plate, state_dim and observation_dim are inferred from a distribution's NumPyro event_shape, not from the full sample shape. The full sample shape includes leading plate batch axes, which are independent-member dimensions, not event dimensions; using it would misread (N, d) as state_dim == N. Sticking to event_shape keeps the per-member event size unambiguous.

The contract for a single plate of size N:

Sampled shape	event_shape	Interpretation
`(d,)`	`(d,)`	Shared vector event, broadcast.
`(N, d)`	`(d,)`	Per-member vector event of dim d.
`(N,)`	`()`	Per-member scalar event.

The third row is the subtle one: dist.Normal(mu, sigma) with mu.shape == (N,) produces event_shape == (), which we treat as state_dim == 1. If the intent is a 1-D vector state with one entry per member, wrap with .to_event(1) or use a vector-valued distribution (dist.MultivariateNormal) so the rank-1 axis is an event axis. This is the same ambiguity rule as the shared-vector case above, applied at the distribution level.

Nested plates extend this with multiple leading batch axes; the inner plate is the leftmost data batch axis, matching NumPyro's convention.

Output axis ordering

Predictive draws and filter outputs preserve a consistent axis order:

(num_samples, *plate_axes_inner_to_outer, n_sim, T, *event_shape)

For example, with one plate of size N and a vector state:

prior["f_states"]        # (num_samples, N, n_sim, T, state_dim)
prior["f_observations"]  # (num_samples, N, n_sim, T, obs_dim)

num_samples comes first (NumPyro Predictive), then plate axes from inner to outer (the inner plate is the leftmost data batch axis, so it appears immediately after num_samples), then n_sim from the simulator, then time, then the event axes. flatten_draws is the standard helper for collapsing (num_samples, n_sim) for plotting/credible intervals.

Examples:

>>> with dsx.plate("trajectories", M):
...     theta = numpyro.sample("theta", dist.Normal(0, 1))  # shape (M,)
...     dynamics = DynamicalModel(...)  # built from theta
...     dsx.sample("f", dynamics, obs_times=t, obs_values=y)

>>> with dsx.plate("groups", G):
...     beta = numpyro.sample("beta", dist.Normal(0, 1))  # shape (G,)
...     with dsx.plate("trajectories", M):
...         alpha = numpyro.sample("alpha", dist.Normal(beta, 1))  # shape (M, G)
...         dynamics = DynamicalModel(...)  # built from alpha
...         dsx.sample("f", dynamics, obs_times=t, obs_values=y)

Sharp edges

Drift/diffusion must be sliceable pytrees. Plated parameters must be stored as array fields of an eqx.Module, not captured in a Python closure. A closure-captured variable is invisible to pytree munging, and can introduce shape errors. The built-in components (AffineDrift, LTI_continuous, FullDiffusion, etc.) follow this rule. See the hierarchical inference tutorial (08_hierarchical_inference.ipynb) for the full sharp-edges list including event-shape vs. sample-shape rules and the rank-1 shared/batched ambiguity.

Note

The dim argument is not currently supported for dynestyx plates.

`enter()` ¶

Enter both numpyro.plate context and dynestyx plate interpretation.

`exit(exc_type, exc, tb)` ¶

Exit both numpyro.plate context and dynestyx plate interpretation.

`init(name: str, size: int, dim: int | None = None)` ¶

Initialize the plate handler.

Parameters:

Name	Type	Description	Default
`name`	`str`	Name of the plate.	required
`size`	`int`	Size of the plate.	required
`dim`	`int \| None`	Dimension of the plate.	`None`

condition(name: str, dynamics: DynamicalModel, , obs_times: Real[Array, 'obs_time_plate obs_time'] | None = None, obs_values: Real[Array, 'obs_value_plate obs_time observation_dim'] | Real[Array, 'obs_value_plate obs_time'] | None = None, ctrl_times: Real[Array, 'ctrl_time_plate ctrl_time'] | None = None, ctrl_values: Real[Array, 'ctrl_value_plate ctrl_time control_dim'] | Real[Array, 'ctrl_value_plate ctrl_time'] | None = None, predict_times: Real[Array, 'predict_time_plate predict_time'] | None = None, **kwargs) ¶

Run inference on a dynamical model without registering numpyro sites.

This is the numpyro-free entry point. When a Filter or Smoother handler is active, returns a ConditionedResult dataclass with marginal_loglik, states, etc.

Parameters:

Name	Type	Description	Default
`name`	`str`	Name of the inference site.	required
`dynamics`	`DynamicalModel`	Dynamical model to infer.	required
`obs_times`	`Real[Array, '*obs_time_plate obs_time'] \| None`	Times at which observations are available.	`None`
`obs_values`	`Real[Array, 'obs_value_plate obs_time observation_dim'] \| Real[Array, 'obs_value_plate obs_time'] \| None`	Values of the observations at the given times.	`None`
`ctrl_times`	`Real[Array, '*ctrl_time_plate ctrl_time'] \| None`	Times at which controls are applied.	`None`
`ctrl_values`	`Real[Array, 'ctrl_value_plate ctrl_time control_dim'] \| Real[Array, 'ctrl_value_plate ctrl_time'] \| None`	Values of the controls at the given times.	`None`
`predict_times`	`Real[Array, '*predict_time_plate predict_time'] \| None`	Times at which to predict.	`None`
`**kwargs`		Additional keyword arguments.	`{}`

sample(name: str, dynamics: DynamicalModel, , obs_times: Real[Array, 'obs_time_plate obs_time'] | None = None, obs_values: Real[Array, 'obs_value_plate obs_time observation_dim'] | Real[Array, 'obs_value_plate obs_time'] | None = None, ctrl_times: Real[Array, 'ctrl_time_plate ctrl_time'] | None = None, ctrl_values: Real[Array, 'ctrl_value_plate ctrl_time control_dim'] | Real[Array, 'ctrl_value_plate ctrl_time'] | None = None, predict_times: Real[Array, 'predict_time_plate predict_time'] | None = None, **kwargs) ¶

Samples from a dynamical model. This is the main primitive of dynestyx.

The sample primitive is meant to mimic the numpyro.sample primitive in usage, but using a DynamicalModel instead of a Distribution.

Internally, sample calls dsx.condition(...) and then registers the results as numpyro sites (numpyro.factor, numpyro.deterministic).

Shape note

Inside dsx.plate, observation arrays use leading plate axes followed by time and event axes, e.g. (N, T, obs_dim). Model parameters follow the same leading-plate, trailing-event convention. See :class:plate for the full plated-shape contract.

Parameters:

Name	Type	Description	Default
`name`	`str`	Name of the sample site.	required
`dynamics`	`DynamicalModel`	Dynamical model to sample from.	required
`obs_times`	`Real[Array, '*obs_time_plate obs_time'] \| None`	Times at which to sample the observations.	`None`
`obs_values`	`Real[Array, 'obs_value_plate obs_time observation_dim'] \| Real[Array, 'obs_value_plate obs_time'] \| None`	Values of the observations at the given times.	`None`
`ctrl_times`	`Real[Array, '*ctrl_time_plate ctrl_time'] \| None`	Times at which to sample the controls.	`None`
`ctrl_values`	`Real[Array, 'ctrl_value_plate ctrl_time control_dim'] \| Real[Array, 'ctrl_value_plate ctrl_time'] \| None`	Values of the controls at the given times.	`None`
`predict_times`	`Real[Array, '*predict_time_plate predict_time'] \| None`	Times at which to predict the observations.	`None`
`**kwargs`		Additional keyword arguments.	`{}`

Top-level pure-JAX API for simulation and scoring. Consider using as an alternative to the NumPyro-based API if simulation and scoring are the only requirements.

log_prob(dynamics: DynamicalModel, , state_path_params: Real[Array, 'state_path_param_time state_dim'] | Real[Array, ' _'] | Real[Array, ''], state_path_param_times: Real[Array, ' state_path_param_time'], obs_times: Real[Array, ' obs_time'] | None = None, obs_values: Real[Array, 'obs_time observation_dim'] | Real[Array, ' obs_time'] | None = None, ctrl_times: Real[Array, ' ctrl_time'] | None = None, ctrl_values: Real[Array, 'ctrl_time control_dim'] | Real[Array, ' ctrl_time'] | None = None, missing_observation_strategy: MissingObservationStrategy = 'auto', missing_obs_values: Real[Array, ' n_missing_obs'] | Real[Array, ' obs_time'] | Real[Array, 'obs_time observation_dim'] | Real[Array, ''] | None = None, missing_obs_metadata: MissingObservationMetadata | None = None, chunk_size: int | None = 0, ode_diffeqsolve_settings: dict[str, Any] | None = None) -> Real[Array, 'log_prob_batch'] ¶

Evaluate the joint log density of a reconstructed state path.

The function reconstructs a complete state path from state_path_params, then evaluates its initial, transition, and observation terms. If observations are omitted, it evaluates only the state-path density.

Parameters:

Name	Type	Description	Default
`dynamics`	`DynamicalModel`	Dynamical model to score.	required
`state_path_params`	`Real[Array, 'state_path_param_time state_dim'] \| Real[Array, ' _'] \| Real[Array, '']`	Values used to reconstruct the latent state path. For a discrete or discretized model, provide the complete path. For a deterministic ODE, provide its initial state.	required
`state_path_param_times`	`Real[Array, ' state_path_param_time']`	Strictly increasing times associated with `state_path_params`. A deterministic ODE expects one time.	required
`obs_times`	`Real[Array, ' obs_time'] \| None`	Strictly increasing times associated with `obs_values`. Every observation time must occur in the reconstructed path.	`None`
`obs_values`	`Real[Array, 'obs_time observation_dim'] \| Real[Array, ' obs_time'] \| None`	Observation values, including any missing entries. Provide this argument together with `obs_times`.	`None`
`ctrl_times`	`Real[Array, ' ctrl_time'] \| None`	Strictly increasing times associated with `ctrl_values`. When controls are provided, these times must match the union of `obs_times` and `state_path_param_times`.	`None`
`ctrl_values`	`Real[Array, 'ctrl_time control_dim'] \| Real[Array, ' ctrl_time'] \| None`	Control values, or `None` for an uncontrolled model.	`None`
`missing_observation_strategy`	`MissingObservationStrategy`	Method used to handle missing observations. `"auto"` marginalizes supported observation distributions and otherwise uses augmentation for continuous distributions.	`'auto'`
`missing_obs_values`	`Real[Array, ' n_missing_obs'] \| Real[Array, ' obs_time'] \| Real[Array, 'obs_time observation_dim'] \| Real[Array, ''] \| None`	Values used to complete missing observations when augmentation is active. Supply either a flat vector ordered by `missing_obs_metadata`, a scalar for one missing entry, or a dense array shaped like `obs_values`; observed entries in a dense array are ignored. In this case, the result includes the density of these values instead of marginalizing them.	`None`
`missing_obs_metadata`	`MissingObservationMetadata \| None`	Positions, times, and component indices for `missing_obs_values`. Precompute this metadata before JIT-compiled augmentation when the missingness pattern cannot be inspected eagerly.	`None`
`chunk_size`	`int \| None`	Batch size passed to `jax.lax.map` while scoring transition and observation terms. The default, `0`, evaluates all terms with one `jax.vmap`. `None` maps one term at a time. A positive integer evaluates batches of that size with `jax.vmap`.	`0`
`ode_diffeqsolve_settings`	`dict[str, Any] \| None`	Diffrax settings used to reconstruct a deterministic ODE path.	`None`

Returns:

Name	Type	Description
`Array`	`Real[Array, '*log_prob_batch']`	Joint log density, retaining any distribution batch axes.

Raises:

Type	Description
`ValueError`	If time, path, control, observation, or missing-observation inputs are inconsistent; if the model is not supported for scoring; or if a native SDE has not been discretized.
`EquinoxRuntimeError`	If a time array is not strictly increasing, `dynamics.t0` does not match the earliest supplied time, or a required observation or control time is absent.
`NotImplementedError`	If the selected missing-observation strategy is unsupported by the observation distribution.

`simulate(dynamics: DynamicalModel, *, rng_key: PRNGKeyArray, ctrl_times: Real[Array, ' ctrl_time'] | None = None, ctrl_values: Real[Array, 'ctrl_time control_dim'] | Real[Array, ' ctrl_time'] | None = None, predict_times: Real[Array, ' predict_time'] | None = None, n_simulations: int = 1, simulator_config: SimulatorConfig | None = None) -> SimulatedResult` ¶

Simulate states and observations without registering NumPyro sites.

The simulation runs on the grid specified by predict_times.

Parameters:

Name	Type	Description	Default
`dynamics`	`DynamicalModel`	Dynamical model to simulate.	required
`rng_key`	`PRNGKeyArray`	JAX pseudorandom number generator key.	required
`ctrl_times`	`Real[Array, ' ctrl_time'] \| None`	Times associated with `ctrl_values`. If controls are provided, these times must match `predict_times`.	`None`
`ctrl_values`	`Real[Array, 'ctrl_time control_dim'] \| Real[Array, ' ctrl_time'] \| None`	Control values, or `None` for an uncontrolled model.	`None`
`predict_times`	`Real[Array, ' predict_time'] \| None`	Times at which to simulate states and observations.	`None`
`n_simulations`	`int`	Number of independent trajectories to simulate.	`1`
`simulator_config`	`SimulatorConfig \| None`	ODE or SDE solver configuration. Its type must match the model's state evolution. Discrete-time models do not accept a simulator configuration.	`None`

Returns:

Name	Type	Description
`SimulatedResult`	`SimulatedResult`	Simulated times, initial states, state paths, and observations.

Raises:

Type	Description
`ValueError`	If `predict_times` is not provided, controls are incomplete or incompatible with the model, or the simulator configuration does not match the model.
`EquinoxRuntimeError`	If a time array is not strictly increasing, `ctrl_times` does not match the required time grid, or `dynamics.t0` does not match the first prediction time.

Handlers¶

plate ¶

__enter__() ¶

__exit__(exc_type, exc, tb) ¶

__init__(name: str, size: int, dim: int | None = None) ¶

`plate` ¶

`enter()` ¶

`exit(exc_type, exc, tb)` ¶

`init(name: str, size: int, dim: int | None = None)` ¶