LatentPathBuilder¶

LatentPathBuilder performs joint inference on model parameters and the state path. A state path is the sequence of states \(x_{0:T}\) over time. State values that are not given by the observations are called latent states.

Mental model¶

One way to perform Bayesian inference on a dynamical system is to unroll the system over time. Consider a model with state transitions

\[ x_{t+1} \sim p(x_{t+1} \mid x_t) \]

and observations

\[ y_t \sim p(y_t \mid x_t). \]

Unrolling means writing one state transition and one observation step for each time point. We can then infer the model parameters \(\theta\) and the complete state path together from the posterior distribution \(p(x_{0:T}, \theta \mid y_{1:T})\).

A direct NumPyro program would look like this:

def model(ys):
    theta = numpyro.sample("theta", p_theta)
    x_i = numpyro.sample("x_0", p_initial(theta))

    for i, y_i in enumerate(ys):
        x_i = numpyro.sample(
            f"x_{i + 1}",
            p_transition(x_i, theta),
        )
        numpyro.sample(
            f"y_{i + 1}",
            p_observation(x_i, theta),
            obs=y_i,
        )

LatentPathBuilder represents this same joint model, with support for arbitrary missingness. However, to support this arbitrary missingness, the actual implementation differs significantly, and the LatentPathBuilder does not build one NumPyro sample site for each time point. Instead, it stores the state values in one array so that JAX can evaluate many density terms together.

The state values that define the path are stored in state_path_params. The builder uses them to reconstruct the complete state path and then evaluates its joint log density with the observations.

Joint density¶

For a discrete state-space model, the log density is

\[ \log p(x_0) + \sum_{i=0}^{T-1} \log p(x_{i+1} \mid x_i, u_i) + \sum_{j=0}^{N-1} \log p(y_j \mid x(\tau_j), u(\tau_j)). \]

Here, \(x_i\) is a state, \(y_j\) is an observation at time \(\tau_j\), and \(u_i\) is a control input. The control terms are omitted when the model has no controls.

A deterministic ODE has no transition-density terms because the initial state determines the rest of the path. A DiracIdentityObservation has \(y_t = x_t\). Each observation requires the state to have the same value, so the builder does not calculate a separate observation-density term.

LatentPathBuilder is used when the state path should be represented directly in the NumPyro model. Filter and Smoother are used when the state path should instead be marginalized, which may be more efficient for large or complex state-space models. Simulator generates new state and observation paths without conditioning on observed data, and should not be used for inference.

The output names reflect this difference:

f_state_path and f_state_path_times contain the path reconstructed by LatentPathBuilder.
f_states and f_times contain a path generated by Simulator or dsx.simulate(...).

Supported models¶

LatentPathBuilder:

runs through dsx.sample(...) inside a NumPyro model;
requires both obs_times and obs_values;
supports discrete, discretized, and deterministic continuous-time models;
requires continuous-time SDE models to be discretized first; and
supports rollout only at or after the end of the state path: predict_times >= max(state_path_times).

The constructor accepts the following arguments:

Argument	Purpose
`ode_simulator_config`	An `ODESimulatorConfig` containing the solver and integration settings used when a deterministic continuous-time state path is reconstructed.
`missing_observation_strategy`	Chooses how missing observations are handled. The options are `"auto"`, `"marginalize"`, `"augment"`, and `"error"`.
`chunk_size`	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`.

State-path parameters¶

state_path_params contains the state values needed to construct the complete path. If the user does not provide these values, NumPyro infers them. Their meaning depends on the model:

Model	Values in `state_path_params`	Complete state path
Discrete or discretized	One state at each observation time	The parameters are already the complete path.
Deterministic ODE	The initial state at `dynamics.t0`	The builder solves the ODE at the observation times.
`DiracIdentityObservation`	State components for missing observations	The builder combines these components with the observed state components.

A component is one entry of a vector-valued state. For exact identity observations, several missing components can belong to the same time. state_path_param_coordinate_indices records the component for each parameter, and state_path_param_times records its time. The time array can therefore contain repeated values.

Execution¶

For one trajectory, _sample_single(...) performs these steps:

Validate the model and read the filled observations and observation mask prepared by dsx.sample(...).
Use the mask to determine which observations are missing.
Create the "{name}_state_path_params" sample site.
Reconstruct the complete state path.
If NumPyro must infer missing observation values, create the "{name}_missing_obs_values" sample site and fill those values into the observation array.
Call compute_state_path_log_prob(...) and add the result to the model as "{name}_joint_log_prob_factor".
Record the reconstructed path and its metadata. If future times were requested, pass the final state to the simulator for rollout.

For dsx.plate(...), _sample_ds(...) processes each plate member separately and stacks the results into the requested plate shape. A call named "f" uses member names such as "f_p0_1".

NumPyro sites¶

A sample site contains a value that NumPyro can infer. A deterministic site records a calculated value. A factor adds a log-density value to the model.

For an unplated call named "f", the builder can create the following sites:

Site	Type	When it is created
`f_state_path_params`	Sample	Always.
`f_missing_obs_values`	Sample	When NumPyro infers missing observations and the observation model is not `DiracIdentityObservation`.
`f_joint_log_prob_factor`	Factor	Always. This is the builder's contribution to the model density.
`f_state_path_param_times`	Deterministic	Always.
`f_state_path_param_coordinate_indices`	Deterministic	For `DiracIdentityObservation`.
`f_state_path`	Deterministic	Always.
`f_state_path_times`	Deterministic	Always.
`f_missing_obs_times`	Deterministic	When NumPyro infers missing observations.
`f_missing_obs_coordinate_indices`	Deterministic	When NumPyro infers missing observations.
`f_completed_obs_values`	Deterministic	For exact observations or when NumPyro infers missing observations.
`f_joint_log_prob`	Deterministic	Always. It records the value added by `f_joint_log_prob_factor`.

Sampling and density¶

The sample sites use _ForwardSimulationImproperUniform. This distribution has a log density of zero, so it does not add a second density term to the model. Its sample(...) method still returns useful initial values:

discrete paths are sampled by forward simulation;
ODE paths start with a draw from the initial-condition distribution;
for exact identity observations, the complete path is simulated and only the components that correspond to missing observations are kept; and
missing observation values are sampled from the observation model given the reconstructed state path.

The builder adds the actual joint density once through f_joint_log_prob_factor. This prevents the state and observation densities from being counted twice.

A NumPyro Predictive call without posterior samples therefore draws paths from the model prior. When posterior samples are provided, Predictive uses their stored state-path values.

Missing observations¶

The missing_observation_strategy options have the following behavior:

"marginalize" evaluates the density of the observed components without creating values for the missing components;
"augment" asks NumPyro to infer the missing values, fills them into the observation array, and evaluates the density of the completed observation;
"auto" selects between these two methods based on the observation model; and
"error" raises an error when observations are missing.

DiracIdentityObservation follows a different path. An observed value fixes the matching state component, while a missing value leaves that state component to be inferred through state_path_params. The builder does not create a separate missing_obs_values site in this case. Exact missing observations support "auto" and "augment".

The missingness pattern can determine NumPyro site shapes, so it must remain static while JAX traces the model. With concrete observations, the builder caches the layout by sample site and observation shape.

There are three supported paths: run concretely, reuse a cache from a concrete run, or pass missing_obs_metadata directly. The metadata can come from dsx.prepare_missing_observation_metadata(...) or a manually constructed dsx.MissingObservationMetadata object.

The recommended pattern defines the model first and applies LatentPathBuilder when the model runs. Cache reuse requires the same builder object for the concrete and traced calls. NumPyro MCMC routines typically make the concrete call automatically.

The test suite covers all three paths and ragged MCMC. The builder derives missing-observation times from the current obs_times, so values and times can change while the layout stays fixed.

For plate members with different layouts, use one builder for all traced calls. Ragged aggregate fields are flat lists of per-member arrays. Member sites keep their shapes, and the rightmost plate index varies fastest.

Implementation modules¶

dynestyx.inference.latent.builder creates the sites, handles plates, records results, and starts future rollout.
dynestyx.inference.state_paths.reconstruct checks parameter shapes and reconstructs the state path.
dynestyx.inference.state_paths.score evaluates the initial, transition, and observation density terms.
dynestyx.observation_missingness prepares observation masks, records missing components, selects a missing-data strategy, and completes observations.
dynestyx.inference.utils.distribution_utils defines _ForwardSimulationImproperUniform.
dynestyx.inference.utils.plate_utils slices and stacks plate values.
dynestyx.simulation.discrete and dynestyx.simulation.utils provide the state and observation samplers used for initialization.

Use with dsx.LatentPathBuilder(): dsx.sample(...) for joint inference with an explicit state path. Use dsx.log_prob(...) to evaluate a state path with pure JAX.

NumPyro-facing explicit latent-path inference.

`LatentPathBuilder` ¶

Bases: ObjectInterpretation, HandlesSelf

Construct and score explicit latent state paths in a NumPyro model.

Use this handler as a context manager around dsx.sample(...). The builder creates array-valued sample sites, reconstructs the complete state path, and adds the joint state-observation log density to the NumPyro model.

Several different strategies are provided for missing data, specified by the missing_observation_strategy parameter:

"marginalize" evaluates the observation density using only the observed components. The observation distribution must support this calculation (namely, a MultivariateNormal or IndependentDistribution).
"augment" creates a "{name}_missing_obs_values" sample site, fills the missing components, and evaluates the completed observation. This method requires a continuous observation distribution.
"auto" uses marginalization when supported and otherwise uses augmentation for a continuous observation distribution.
"error" rejects partially observed vectors.

DiracIdentityObservation uses a separate path. Its missing components are inferred through state_path_params, and missing data support only "auto" or "augment".

A missingness layout can determine NumPyro site shapes. The builder infers and caches this layout whenever observations are concrete. Before an outer jax.jit passes obs_values dynamically, either reuse a builder that has seen concrete observations or pass eagerly prepared missing_obs_metadata to dsx.sample(...). JAX itself does not run the function eagerly before tracing.

Attributes:

Name	Type	Description
`ode_simulator_config`		ODE solver and integration settings used during deterministic continuous-time path reconstruction.
`missing_observation_strategy`		Method used to handle missing entries in `obs_values`.
`chunk_size`		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`.

Examples:

>>> builder = dsx.LatentPathBuilder()
>>> with builder:
...     result = dsx.sample(
...         "f",
...         dynamics,
...         obs_times=obs_times,
...         obs_values=obs_values,
...     )

`init(ode_simulator_config: ODESimulatorConfig | None = None, missing_observation_strategy: MissingObservationStrategy = 'auto', chunk_size: int | None = 0) -> None` ¶

Initialize explicit latent-path inference.

Parameters:

Name	Type	Description	Default
`ode_simulator_config`	`ODESimulatorConfig \| None`	ODE solver and integration settings used during deterministic continuous-time path reconstruction. Defaults to `ODESimulatorConfig()` when omitted.	`None`
`missing_observation_strategy`	`MissingObservationStrategy`	Method used to handle missing entries in `obs_values`, as described in the class documentation.	`'auto'`
`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`

_sample_ds(name: str, dynamics: DynamicalModel, , plate_shapes: tuple[int, ...] = (), 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, _obs_values_filled: Real[Array, 'obs_value_plate obs_time observation_dim'] | Real[Array, 'obs_value_plate obs_time'] | None = None, _obs_mask: Bool[Array, 'obs_value_plate obs_time observation_dim'] | Bool[Array, 'obs_value_plate obs_time'] | None = None, _obs_has_missing: bool | None = None, missing_obs_metadata: MissingObservationMetadata | 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, state_path_params: Real[Array, 'state_path_param_shape'] | None = None, missing_obs_values: Real[Array, 'missing_obs_shape'] | None = None, _dsx_sample_mode: bool = False, **kwargs) -> LatentStateResult ¶

Construct latent paths for one trajectory or a plate of trajectories.

Parameters:

Name	Type	Description	Default
`name`	`str`	Prefix used for NumPyro site names.	required
`dynamics`	`DynamicalModel`	Dynamical model to condition on observations.	required
`plate_shapes`	`tuple[int, ...]`	Leading plate dimensions for independent trajectories.	`()`
`obs_times`	`Real[Array, '*obs_time_plate obs_time'] \| None`	Observation times.	`None`
`obs_values`	`Real[Array, 'obs_value_plate obs_time observation_dim'] \| Real[Array, 'obs_value_plate obs_time'] \| None`	Observation values.	`None`
`_obs_values_filled`	`Real[Array, 'obs_value_plate obs_time observation_dim'] \| Real[Array, 'obs_value_plate obs_time'] \| None`	Internal observation values with missing entries replaced by shape-preserving filler values.	`None`
`_obs_mask`	`Bool[Array, 'obs_value_plate obs_time observation_dim'] \| Bool[Array, 'obs_value_plate obs_time'] \| None`	Internal boolean array that marks observed entries.	`None`
`_obs_has_missing`	`bool \| None`	Internal flag indicating whether any observations are missing.	`None`
`missing_obs_metadata`	`MissingObservationMetadata \| None`	Optional concrete missingness layout. One layout is shared by all plate members.	`None`
`ctrl_times`	`Real[Array, '*ctrl_time_plate ctrl_time'] \| None`	Times associated with `ctrl_values`.	`None`
`ctrl_values`	`Real[Array, 'ctrl_value_plate ctrl_time control_dim'] \| Real[Array, 'ctrl_value_plate ctrl_time'] \| None`	Control values, or `None` for an uncontrolled model.	`None`
`predict_times`	`Real[Array, '*predict_time_plate predict_time'] \| None`	Future times used for posterior rollout. Each time must be at or after the end of the reconstructed path.	`None`
`state_path_params`	`Real[Array, '*state_path_param_shape'] \| None`	Optional state-path values used to condition the corresponding NumPyro sample site.	`None`
`missing_obs_values`	`Real[Array, '*missing_obs_shape'] \| None`	Optional values used to condition the missing-observation sample site when augmentation is active.	`None`
`_dsx_sample_mode`	`bool`	Internal flag set by `dsx.sample(...)`.	`False`
`**kwargs`		Additional arguments forwarded to the next handler.	`{}`

Returns:

Name	Type	Description
`LatentStateResult`	`LatentStateResult`	Path values and metadata. Plate dimensions are
	`LatentStateResult`	leading dimensions in array-valued fields.

Raises:

Type	Description
`ValueError`	If called outside `dsx.sample(...)`, if observations or latent values are invalid, or if `predict_times` contains an unsupported in-window time.

_sample_single(name: str, dynamics: DynamicalModel, , obs_times: Real[Array, ' obs_time'] | None, obs_values: Real[Array, 'obs_time observation_dim'] | Real[Array, ' obs_time'] | None, obs_values_filled: Real[Array, 'obs_time observation_dim'] | Real[Array, ' obs_time'] | None, obs_mask: Bool[Array, 'obs_time observation_dim'] | Bool[Array, ' obs_time'] | None, missing_obs_metadata: MissingObservationMetadata | None, ctrl_times: Real[Array, ' ctrl_time'] | None, ctrl_values: Real[Array, 'ctrl_time control_dim'] | Real[Array, ' ctrl_time'] | None, state_path_params: Real[Array, 'state_path_param_shape'] | None, missing_obs_values: Real[Array, ' n_missing_obs'] | Real[Array, ''] | None) -> LatentStateResult ¶

Construct and score one latent state path.

Parameters:

Name	Type	Description	Default
`name`	`str`	Prefix used for NumPyro site names.	required
`dynamics`	`DynamicalModel`	Dynamical model to condition on observations.	required
`obs_times`	`Real[Array, ' obs_time'] \| None`	Observation times. This argument is required.	required
`obs_values`	`Real[Array, 'obs_time observation_dim'] \| Real[Array, ' obs_time'] \| None`	Observation values. This argument is required.	required
`obs_values_filled`	`Real[Array, 'obs_time observation_dim'] \| Real[Array, ' obs_time'] \| None`	Observation values with missing entries replaced by shape-preserving filler values.	required
`obs_mask`	`Bool[Array, 'obs_time observation_dim'] \| Bool[Array, ' obs_time'] \| None`	Boolean array that marks observed entries.	required
`missing_obs_metadata`	`MissingObservationMetadata \| None`	Optional concrete missingness layout.	required
`ctrl_times`	`Real[Array, ' ctrl_time'] \| None`	Times associated with `ctrl_values`.	required
`ctrl_values`	`Real[Array, 'ctrl_time control_dim'] \| Real[Array, ' ctrl_time'] \| None`	Control values, or `None` for an uncontrolled model.	required
`state_path_params`	`Real[Array, '*state_path_param_shape'] \| None`	State values used to construct the path. `None` creates an unconditioned NumPyro sample site.	required
`missing_obs_values`	`Real[Array, ' n_missing_obs'] \| Real[Array, ''] \| None`	Values used to fill missing observations when augmentation is active. `None` creates an unconditioned sample site when these values are required.	required

Returns:

Name	Type	Description
`LatentStateResult`	`LatentStateResult`	Reconstructed path, joint log density, and
	`LatentStateResult`	metadata used for missing observations and posterior rollout.

Raises:

Type	Description
`ValueError`	If the model is unsupported, required observations are absent, the missing-observation strategy is invalid, or the supplied latent values have incompatible shapes or semantics.

`_build_state_path_distributions(dynamics: DynamicalModel, state_path: Real[Array, 'state_path_plate state_path_time state_dim'] | Real[Array, 'state_path_plate state_path_time']) -> list[dist.Distribution]` ¶

Create one Delta distribution for each state in a reconstructed path.

Parameters:

Name	Type	Description	Default
`dynamics`	`DynamicalModel`	Dynamical model that defines the state event shape.	required
`state_path`	`Real[Array, 'state_path_plate state_path_time state_dim'] \| Real[Array, 'state_path_plate state_path_time']`	Reconstructed state values.	required

Returns:

Type	Description
`list[Distribution]`	list[dist.Distribution]: Delta distributions used for posterior rollout.

`_resolve_missing_observation_metadata(*, name: str, dynamics: DynamicalModel, obs_times: Real[Array, ' obs_time'], obs_values: Real[Array, 'obs_time observation_dim'] | Real[Array, ' obs_time'], missing_obs_metadata: MissingObservationMetadata | None, cache: _MissingObservationMetadataCache) -> MissingObservationMetadata` ¶

Prepare static missingness layout while retaining the current times.

`_sample_missing_observation_prior(dynamics: DynamicalModel, state_path: Real[Array, 'state_path_time state_dim'] | Real[Array, ' state_path_time'], state_path_times: Real[Array, ' state_path_time'], obs_times: Real[Array, ' obs_time'], ctrl_times: Real[Array, ' ctrl_time'] | None, ctrl_values: Real[Array, 'ctrl_time control_dim'] | Real[Array, ' ctrl_time'] | None, missing_flat_indices: Int[Array, ' n_missing_obs'], key: PRNGKeyArray) -> Real[Array, ' n_missing_obs']` ¶

Sample missing observations conditional on the reconstructed state path.

Parameters:

Name	Type	Description	Default
`dynamics`	`DynamicalModel`	Dynamical model that defines the observation distribution.	required
`state_path`	`Real[Array, 'state_path_time state_dim'] \| Real[Array, ' state_path_time']`	Reconstructed state values.	required
`state_path_times`	`Real[Array, ' state_path_time']`	Times associated with `state_path`.	required
`obs_times`	`Real[Array, ' obs_time']`	Times at which observations are required.	required
`ctrl_times`	`Real[Array, ' ctrl_time'] \| None`	Times associated with `ctrl_values`.	required
`ctrl_values`	`Real[Array, 'ctrl_time control_dim'] \| Real[Array, ' ctrl_time'] \| None`	Control values, or `None` for an uncontrolled model.	required
`missing_flat_indices`	`Int[Array, ' n_missing_obs']`	Positions of missing values in the flattened observation array.	required
`key`	`PRNGKeyArray`	JAX pseudorandom key.	required

Returns:

Name	Type	Description
`Array`	`Real[Array, ' n_missing_obs']`	Sampled missing values ordered by `missing_flat_indices`.

LatentPathBuilder¶

Mental model¶

Joint density¶

Related handlers¶

Supported models¶

State-path parameters¶

Execution¶

NumPyro sites¶

Sampling and density¶

Missing observations¶

Implementation modules¶

LatentPathBuilder ¶

__init__(ode_simulator_config: ODESimulatorConfig | None = None, missing_observation_strategy: MissingObservationStrategy = 'auto', chunk_size: int | None = 0) -> None ¶

_build_state_path_distributions(dynamics: DynamicalModel, state_path: Real[Array, '*state_path_plate state_path_time state_dim'] | Real[Array, '*state_path_plate state_path_time']) -> list[dist.Distribution] ¶

`LatentPathBuilder` ¶

`init(ode_simulator_config: ODESimulatorConfig | None = None, missing_observation_strategy: MissingObservationStrategy = 'auto', chunk_size: int | None = 0) -> None` ¶

`_build_state_path_distributions(dynamics: DynamicalModel, state_path: Real[Array, 'state_path_plate state_path_time state_dim'] | Real[Array, 'state_path_plate state_path_time']) -> list[dist.Distribution]` ¶