Overview¶

Simulators run pure-JAX forward trajectories for a DynamicalModel on a chosen time grid, then optionally attach the realized outputs to a NumPyro trace when used through dsx.sample(...).

For pure-JAX forward generation, user-facing code should prefer dsx.simulate(...). The simulator classes documented here are the NumPyro-aware handler implementations used by dsx.sample(...) and by posterior rollout from Filter/Smoother.

When to use each time argument

predict_times: use this when you want rollout trajectories at specific times for simulation and/or post-filter rollout.
In posterior-rollout mode, predictions are generated at predict_times from inference-handler posteriors.
Typical use: forward simulation, forecasting, or dense trajectories for visualization.
obs_times / obs_values are consumed by observation-aware handlers such as LatentPathBuilder, Filter, and Smoother, not by the public simulator interface itself.
If predict_times is omitted: the simulator does not run and adds no deterministic sites.

Context and caveats

NumPyro context required for dsx.sample(...): simulator handlers draw randomness from the active NumPyro PRNG key, but the rollout itself is pure JAX and sites are registered only at the end. Use dsx.simulate(...) for the explicit pure-JAX API.
Generation-only public API: raw Simulator, DiscreteTimeSimulator, ODESimulator, and SDESimulator calls expect predict_times, not direct observation conditioning.
Inference lives elsewhere: use LatentPathBuilder for explicit latent paths, Filter for marginalized inference, and Smoother for smoothing. Simulators can then wrap those handlers for rollout with predict_times.

Deterministic sites

When simulator trajectories are produced, sites are recorded as "{name}_{key}" where name is the first argument to dsx.sample(name, dynamics, ...) (conventionally "f"):

"f_x_0": realized initial-state draw, shape (n_sim, state_dim),
"f_times": trajectory time grid, shape (n_sim, T),
"f_states": latent trajectory, shape (n_sim, T, state_dim),
"f_observations": sampled observations, shape (n_sim, T, obs_dim).

In filter-rollout mode (predict_times with filtered posteriors), additional keys "f_predicted_states", "f_predicted_times", and "f_predicted_observations" are recorded. Segment-level rollouts also register realized anchor-state sites such as "f_1_x_0" when applicable.

Under numpyro.infer.Predictive(model, num_samples=N), NumPyro prepends a leading num_samples axis, giving final shapes (num_samples, n_sim, T, dim). Use dynestyx.flatten_draws to collapse the (num_samples, n_sim) prefix into one axis for plotting or downstream analysis.

If predict_times is omitted, no public simulator rollout is performed and these sites are not added.

Simulators¶

Simulation backends and simulator handlers.

`DiscreteTimeSimulator` ¶

Bases: BaseSimulator

Generate trajectories from a discrete-time dynamical model.

For prediction times \(t_0,\ldots,t_{T-1}\), this simulator draws n_simulations independent paths according to

\[ x_0^{(m)} \sim p_0(x_0), \qquad x_{k+1}^{(m)} \sim p\!\left(x_{k+1}\mid x_k^{(m)},u_k,t_k,t_{k+1}\right), \qquad y_k^{(m)} \sim p(y_k\mid x_k^{(m)},u_k,t_k). \]

The first state in the returned path is the initial-condition draw at predict_times[0]; the simulator then makes one transition draw for each adjacent pair of prediction times and samples one observation conditional on every realized state. See DiscreteTimeStateEvolution for how a discrete transition model is represented in a DynamicalModel.

Use DiscreteTimeSimulator as a context manager around a model containing dsx.sample(name, dynamics, predict_times=...). The active NumPyro seed supplies randomness, while the realized paths are attached to the trace as deterministic sites. Use dsx.simulate for standalone pure-JAX generation without a NumPyro trace.

Examples:

>>> def model(predict_times=None):
...     dynamics = DynamicalModel(
...         initial_condition=initial_dist,
...         state_evolution=transition,
...         observation_model=observation,
...     )
...     dsx.sample("f", dynamics, predict_times=predict_times)
>>> with DiscreteTimeSimulator(n_simulations=3):
...     predictive = Predictive(
...         model, num_samples=10, exclude_deterministic=False
...     )
...     draws = predictive(
...         jr.PRNGKey(0), predict_times=jnp.arange(20.0)
...     )
>>> draws["f_states"].shape
(10, 3, 20, state_dim)

For one direct, pure-JAX model execution:

>>> result = dsx.simulate(
...     dynamics,
...     rng_key=jr.PRNGKey(0),
...     predict_times=jnp.arange(20.0),
...     n_simulations=3,
... )

What this does¶

The transition distribution is evaluated with the current state, current control, and the two adjacent time values. Prediction times therefore need not be uniformly spaced, provided the model's transition accepts those intervals.

If controls are supplied, ctrl_times must contain every prediction time exactly. ctrl_values[k] is used for the transition beginning at \(t_k\) and for the observation at \(t_k\). The paired control arrays are validated before simulation.

This handler is generation-only and does not condition on obs_times or obs_values. Use LatentPathBuilder for explicit latent-path inference, or use Filter or Smoother for marginalized inference. Placing this simulator outside a compatible Filter or Smoother draws posterior rollouts at predict_times.

Configuration and defaults¶

Discrete simulation has no solver configuration: transitions are sampled directly from dynamics.state_evolution. n_simulations defaults to one and must be at least one. The simulation dimension is retained even when it has length one.

NumPyro trace¶

For a raw rollout from dsx.sample("f", ...), the following numpyro.deterministic sites are added:

"f_x_0": initial states, shape (*plate_shape, n_simulations, state_dim);
"f_times": prediction times, shape (*plate_shape, n_simulations, T);
"f_states": latent states, shape (*plate_shape, n_simulations, T, state_dim);
"f_observations": sampled observations, shape (*plate_shape, n_simulations, T, observation_dim).

Here "f" is replaced by the name passed to dsx.sample. Under Predictive(..., num_samples=N), NumPyro prepends an N axis to each shape. Because these sites are deterministic, pass exclude_deterministic=False to Predictive (or request the site names explicitly) to include them in its returned dictionary.

When this simulator wraps a Filter or Smoother, the inner handler records its own configured sites and the simulator's aggregate rollout sites are instead "f_predicted_times", "f_predicted_states", and "f_predicted_observations", with the corresponding time, state, and observation shapes above. Each nonempty prediction segment also records the state from which that segment starts, with shape (n_simulations, state_dim): "f_0_x_0" for a segment before the first posterior time, and "f_{j+1}_x_0" for a segment initialized from the posterior at inference-time index j. Only segments containing at least one requested prediction time are recorded. Inside dsx.plate, the segment name also identifies the plate member, for example "f_p0_1_x_0".

If predict_times is omitted, no simulator rollout or simulator trace sites are produced. Direct calls to DiscreteTimeSimulator().simulate return SimulatedResult without adding NumPyro sites.

Notes

Use Simulator instead when automatic selection among discrete, ODE, and SDE backends is desirable.
DiscreteTimeSimulator().simulate(...) consumes an already allocated simulation key. The public dsx.simulate function splits its root key before dispatch.

Attributes:

Name	Type	Description
`n_simulations`	`int`	Number of independent trajectories drawn per model execution. Defaults to one and must be greater than or equal to one.

`_simulate_forward_from_initial_state(dynamics: DynamicalModel, *, initial_state: Real[Array, 'n_simulations state_dim'] | Real[Array, ' n_simulations'], rng_key: PRNGKeyArray, times: Real[Array, ' time'], ctrl_values: Real[Array, 'time control_dim'] | Real[Array, ' time'] | None) -> SimulatedResult` ¶

Run pure forward simulation for a discrete-time model.

`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, **kwargs) -> SimulatedResult` ¶

Run pure-JAX forward simulation for a discrete-time model.

Unlike dsx.simulate, rng_key is consumed directly as an already-allocated simulation key and is not pre-split. Therefore, dsx.simulate(..., rng_key=root_key) is equivalent to DiscreteTimeSimulator().simulate(..., rng_key=split_key), where split_key = jax.random.split(root_key)[1].

`ODESimulator` ¶

Bases: BaseSimulator

Generate trajectories from deterministic continuous-time dynamics.

For an initial-condition distribution \(p_0\), drift function \(f\), and observation model \(p(y\mid x,u,t)\), ODESimulator draws n_simulations independent initial states and computes

\[ x_0^{(m)} \sim p_0(x_0), \qquad \frac{\mathrm{d}x^{(m)}(t)}{\mathrm{d}t} = f\!\left(x^{(m)}(t),u(t),t\right), \qquad y_k^{(m)} \sim p\!\left(y_k\mid x^{(m)}(t_k),u(t_k),t_k\right). \]

The ODE solution is evaluated at every value in predict_times. Conditional on the initial state and controls, the state path is deterministic; the initial-condition and observation distributions may still make the complete simulation stochastic. See ContinuousTimeStateEvolution for how an ODE is represented in a DynamicalModel by specifying its drift without a diffusion.

Use ODESimulator as a context manager around a model containing dsx.sample(name, dynamics, predict_times=...). The active NumPyro seed supplies randomness, and the computed arrays are then attached to the trace as deterministic sites. Pass an ODESimulatorConfig to choose the Diffrax solver, step-size controller, adjoint, step size, and step limit. Use dsx.simulate for standalone pure-JAX generation without a NumPyro trace.

Examples:

Prior-predictive ODE trajectories:

>>> def model(predict_times=None):
...     dynamics = DynamicalModel(
...         initial_condition=initial_dist,
...         state_evolution=ContinuousTimeStateEvolution(
...             drift=lambda x, u, t: -rate * x,
...         ),
...         observation_model=observation,
...     )
...     dsx.sample("f", dynamics, predict_times=predict_times)
>>> config = ODESimulatorConfig(dt0=1e-2)
>>> with ODESimulator(config, n_simulations=3):
...     predictive = Predictive(
...         model, num_samples=10, exclude_deterministic=False
...     )
...     draws = predictive(
...         jr.PRNGKey(0), predict_times=jnp.linspace(0.0, 5.0, 51)
...     )
>>> draws["f_states"].shape
(10, 3, 51, state_dim)

Standalone pure-JAX simulation uses the same ODE solver:

>>> result = dsx.simulate(
...     dynamics,
...     rng_key=jr.PRNGKey(0),
...     predict_times=times,
...     n_simulations=3,
...     simulator_config=ODESimulatorConfig(dt0=1e-2),
... )

What this does¶

Each initial-condition draw is integrated independently with Diffrax. The integration starts at dynamics.t0 when it is defined and otherwise at the first prediction time. The solved state is saved only at predict_times, after which the observation model is sampled independently at those states.

If controls are supplied, they form a right-continuous rectilinear path: the control at a knot ctrl_times[k] is ctrl_values[k], and that value is held until the next knot.

This handler is generation-only and does not condition on obs_times or obs_values. Use LatentPathBuilder for explicit latent-path inference, or use Filter or Smoother for marginalized inference. Placing this simulator outside a compatible continuous-time Filter or Smoother draws posterior rollouts at predict_times.

Configuration and defaults¶

ODEs are solved using Diffrax, and settings are controlled by ODESimulatorConfig. Its default settings are diffrax.Tsit5(), diffrax.ConstantStepSize(), diffrax.RecursiveCheckpointAdjoint(), dt0=1e-3, and max_steps=100_000. Pass different settings when the model requires them. If simulator_config=None, a default ODESimulatorConfig() is created.

n_simulations defaults to one and must be at least one. The simulation dimension is retained even when it has length one.

NumPyro trace¶

For a raw rollout from dsx.sample("f", ...), the following numpyro.deterministic sites are added:

"f_x_0": initial states, shape (*plate_shape, n_simulations, state_dim);
"f_times": prediction times, shape (*plate_shape, n_simulations, T);
"f_states": solved states, shape (*plate_shape, n_simulations, T, state_dim);
"f_observations": sampled observations, shape (*plate_shape, n_simulations, T, observation_dim).

Here "f" is replaced by the name passed to dsx.sample. Under Predictive(..., num_samples=N), NumPyro prepends an N axis to each shape. Because these sites are deterministic, pass exclude_deterministic=False to Predictive (or request the site names explicitly) to include them in its returned dictionary.

When this simulator wraps a Filter or Smoother, the inner handler records its own configured sites and the simulator's aggregate rollout sites are instead "f_predicted_times", "f_predicted_states", and "f_predicted_observations", with the corresponding time, state, and observation shapes above. Each nonempty prediction segment also records the state from which that segment starts, with shape (n_simulations, state_dim): "f_0_x_0" for a segment before the first posterior time, and "f_{j+1}_x_0" for a segment initialized from the posterior at inference-time index j. Only segments containing at least one requested prediction time are recorded. Inside dsx.plate, the segment name also identifies the plate member, for example "f_p0_1_x_0".

If predict_times is omitted, no simulator rollout or simulator trace sites are produced. Direct calls to ODESimulator().simulate return SimulatedResult without adding NumPyro sites.

Notes

Use Simulator instead when automatic selection among discrete, ODE, and SDE backends is desirable.
ODESimulator().simulate(...) consumes an already allocated simulation key. The public dsx.simulate function splits its root key before dispatch.

Attributes:

Name	Type	Description
`simulator_config`		ODE solver and integration settings. Defaults to `ODESimulatorConfig()`.
`n_simulations`	`int`	Number of independent initial states and trajectories drawn per model execution. Defaults to one and must be greater than or equal to one.
`diffeqsolve_settings`		Normalized settings passed to Diffrax.

`init(simulator_config: ODESimulatorConfig | None = None, *, n_simulations: int = 1) -> None` ¶

Configure ODE integration.

Parameters:

Name	Type	Description	Default
`simulator_config`	`ODESimulatorConfig \| None`	Structured simulator settings. Defaults to `ODESimulatorConfig()` when omitted.	`None`
`n_simulations`	`int`	Number of independent trajectories to simulate. State and observation paths have shape `(n_simulations, T, ...)`. Must be greater than or equal to one.	`1`

`_simulate_forward_from_initial_state(dynamics: DynamicalModel, *, initial_state: Real[Array, 'n_simulations state_dim'] | Real[Array, ' n_simulations'], rng_key: PRNGKeyArray, times: Real[Array, ' time'], ctrl_times: Real[Array, ' ctrl_time'] | None = None, ctrl_values: Real[Array, 'ctrl_time control_dim'] | Real[Array, ' ctrl_time'] | None = None) -> SimulatedResult` ¶

Run pure forward simulation for a deterministic continuous-time model.

`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, **kwargs) -> SimulatedResult` ¶

Run pure-JAX forward simulation for deterministic continuous-time models.

Unlike dsx.simulate, rng_key is consumed directly as an already-allocated simulation key and is not pre-split. Therefore, dsx.simulate(..., rng_key=root_key) is equivalent to ODESimulator().simulate(..., rng_key=jax.random.split(root_key)[1]).

`SDESimulator` ¶

Bases: BaseSimulator

Generate trajectories from stochastic continuous-time dynamics.

For an initial-condition distribution \(p(x_0)\), drift \(f\), diffusion coefficient \(g\), and observation model \(p(y\mid x,u,t)\), SDESimulator draws n_simulations independent paths satisfying

\[ \begin{aligned} x_0^{(m)} &\sim p(x_0), \\ dx_t^{(m)} &= f(x_t^{(m)},u_t,t)\,dt + g(x_t^{(m)},u_t,t)\,dW_t^{(m)}, \\ y_k^{(m)} &\sim p(y_k\mid x_k^{(m)},u_k,t_k). \end{aligned} \]

The numerical SDE solution is evaluated at every value in predict_times. Initial states, Brownian paths, and observations are drawn independently across the simulation dimension. See ContinuousTimeStateEvolution for how an SDE is represented in a DynamicalModel.

Use SDESimulator as a context manager around a model containing dsx.sample(name, dynamics, predict_times=...). The active NumPyro seed supplies randomness, and the realized paths are then attached to the trace as deterministic sites. Pass an SDESimulatorConfig to choose the SDE backend, solver, step-size controller, adjoint, step size, and Brownian-tree tolerance. Use dsx.simulate for standalone pure-JAX generation without a NumPyro trace.

Examples:

Fast fixed-step Euler--Maruyama simulation:

>>> def model(predict_times=None):
...     dynamics = DynamicalModel(
...         initial_condition=initial_dist,
...         state_evolution=ContinuousTimeStateEvolution(
...             drift=drift,
...             diffusion=diffusion,
...         ),
...         observation_model=observation,
...     )
...     dsx.sample("f", dynamics, predict_times=predict_times)
>>> config = SDESimulatorConfig(source="em_scan", dt0=1e-3)
>>> with SDESimulator(config, n_simulations=4):
...     predictive = Predictive(
...         model, num_samples=10, exclude_deterministic=False
...     )
...     draws = predictive(
...         jr.PRNGKey(0), predict_times=jnp.linspace(0.0, 5.0, 51)
...     )
>>> draws["f_states"].shape
(10, 4, 51, state_dim)

Use the Diffrax backend when a different SDE solver or step-size controller is required:

>>> config = SDESimulatorConfig(
...     source="diffrax",
...     solver=diffrax.EulerHeun(),
...     dt0=1e-3,
... )

What this does¶

Each initial-condition draw is integrated with an independent Brownian path. Integration starts at dynamics.t0 when it is defined and otherwise at the first prediction time. After the state paths are solved, the observation model is sampled independently at the requested times.

If controls are supplied, they form a right-continuous rectilinear path: the control at a knot ctrl_times[k] is ctrl_values[k], and that value is held until the next knot.

This handler is generation-only and does not condition on obs_times or obs_values. Native SDE latent-path inference should go through a Discretizer and LatentPathBuilder; use Filter or Smoother for marginalized inference. Placing this simulator outside a compatible continuous-time Filter or Smoother draws posterior rollouts at predict_times.

Configurations and defaults¶

SDESimulatorConfig selects one of two backends:

source="em_scan" uses a fixed-step Euler--Maruyama jax.lax.scan. It is the default and uses dt0=1e-4.
source="diffrax" uses the configured Diffrax solver, step-size controller, adjoint, and virtual Brownian tree. Defaults include diffrax.Heun(), diffrax.ConstantStepSize(), diffrax.RecursiveCheckpointAdjoint(), and dt0=1e-4. tol_vbt=None resolves to dt0 / 2.

Solver choice determines the stochastic integral represented by the numerical solution. In particular, the default Diffrax Heun solver converges to a Stratonovich solution, while Euler--Maruyama converges to an Itô solution. This distinction matters for state-dependent diffusion.

If simulator_config=None, a default SDESimulatorConfig() is created. n_simulations defaults to one and must be at least one. The simulation dimension is retained even when it has length one.

NumPyro trace¶

For a raw rollout from dsx.sample("f", ...), the following numpyro.deterministic sites are added:

"f_x_0": initial states, shape (*plate_shape, n_simulations, state_dim);
"f_times": prediction times, shape (*plate_shape, n_simulations, T);
"f_states": solved states, shape (*plate_shape, n_simulations, T, state_dim);
"f_observations": sampled observations, shape (*plate_shape, n_simulations, T, observation_dim).

Here "f" is replaced by the name passed to dsx.sample. Under Predictive(..., num_samples=N), NumPyro prepends an N axis to each shape. Because these sites are deterministic, pass exclude_deterministic=False to Predictive (or request the site names explicitly) to include them in its returned dictionary.

When this simulator wraps a Filter or Smoother, the inner handler records its own configured sites and the simulator's aggregate rollout sites are instead "f_predicted_times", "f_predicted_states", and "f_predicted_observations", with the corresponding time, state, and observation shapes above. Each nonempty prediction segment also records the state from which that segment starts, with shape (n_simulations, state_dim): "f_0_x_0" for a segment before the first posterior time, and "f_{j+1}_x_0" for a segment initialized from the posterior at inference-time index j. Only segments containing at least one requested prediction time are recorded. Inside dsx.plate, the segment name also identifies the plate member, for example "f_p0_1_x_0".

If predict_times is omitted, no simulator rollout or simulator trace sites are produced. Direct calls to SDESimulator().simulate return SimulatedResult without adding NumPyro sites.

Notes

Use Simulator instead when automatic selection among discrete, ODE, and SDE backends is desirable.
SDESimulator().simulate(...) consumes an already allocated simulation key. The public dsx.simulate function splits its root key before dispatch.

Attributes:

Name	Type	Description
`simulator_config`		SDE backend and integration settings. Defaults to `SDESimulatorConfig()`.
`n_simulations`	`int`	Number of independent initial states, Brownian paths, and trajectories drawn per model execution. Defaults to one and must be greater than or equal to one.
`source`		Active SDE backend, either `"em_scan"` or `"diffrax"`.
`diffeqsolve_settings`		Normalized Diffrax-style solver settings.
`tol_vbt`		Resolved virtual-Brownian-tree tolerance for the Diffrax backend, or `None` for `"em_scan"`.

`init(simulator_config: SDESimulatorConfig | None = None, *, n_simulations: int = 1) -> None` ¶

Configure SDE integration settings.

Parameters:

Name	Type	Description	Default
`simulator_config`	`SDESimulatorConfig \| None`	Structured simulator settings. Defaults to `SDESimulatorConfig()` when omitted.	`None`
`n_simulations`	`int`	Number of independent trajectories to simulate. State and observation paths have shape `(n_simulations, T, ...)`. Must be greater than or equal to one.	`1`

`_simulate_forward_from_initial_state(dynamics: DynamicalModel, *, initial_state: Real[Array, 'n_simulations state_dim'] | Real[Array, ' n_simulations'], rng_key: PRNGKeyArray, times: Real[Array, ' time'], ctrl_times: Real[Array, ' ctrl_time'] | None = None, ctrl_values: Real[Array, 'ctrl_time control_dim'] | Real[Array, ' ctrl_time'] | None = None) -> SimulatedResult` ¶

Run pure forward SDE simulation from provided initial states.

`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, **kwargs) -> SimulatedResult` ¶

Run pure-JAX forward simulation for stochastic continuous-time models.

Unlike dsx.simulate, rng_key is consumed directly as an already-allocated simulation key and is not pre-split. Therefore, dsx.simulate(..., rng_key=root_key) is equivalent to SDESimulator().simulate(..., rng_key=jax.random.split(root_key)[1]).

`Simulator` ¶

Bases: BaseSimulator

Generate trajectories using the simulator appropriate for the model.

Simulator is the auto-routing simulator handler. Given prediction times \(t_{0:T-1}\), it draws n_simulations independent trajectories from the dynamical model:

\[ x_0^{(m)} \sim p_0(x_0), \qquad x_{1:T-1}^{(m)} \sim p(x_{1:T-1}\mid x_0^{(m)},u_{1:T-1},t_{1:T-1}), \qquad y_k^{(m)} \sim p(y_k\mid x_k^{(m)}, u_k, t_k). \]

For continuous-time models, the middle term denotes a numerical ODE or SDE solve evaluated at predict_times. For discrete-time models, it denotes successive draws from the transition distribution between adjacent prediction times.

Use Simulator as a context manager around model execution when the model contains dsx.sample(name, dynamics, predict_times=...). The active NumPyro seed supplies randomness; the rollout itself is computed with pure JAX and its realized arrays are then attached to the NumPyro trace. For continuous-time models, pass an ODESimulatorConfig or SDESimulatorConfig as simulator_config to control how the differential equation is solved. Use dsx.simulate instead when no NumPyro trace is needed and an explicit rng_key is more convenient.

Examples:

Prior-predictive trajectories with automatic backend selection:

>>> def model(predict_times=None):
...     dynamics = DynamicalModel(...)
...     dsx.sample("f", dynamics, predict_times=predict_times)
>>> with Simulator(n_simulations=4):
...     predictive = Predictive(
...         model, num_samples=20, exclude_deterministic=False
...     )
...     draws = predictive(jr.PRNGKey(0), predict_times=times)
>>> draws["f_states"].shape
(20, 4, T, state_dim)

Pure-JAX forward simulation returns the same trajectory fields without adding NumPyro sites:

>>> result = dsx.simulate(
...     dynamics,
...     rng_key=jr.PRNGKey(0),
...     predict_times=times,
...     n_simulations=4,
... )

For posterior state rollouts, place the simulator outside the inference handler and retain deterministic sites from Predictive:

>>> with Simulator(n_simulations=100):
...     with Filter():
...         predictive = Predictive(
...             model,
...             posterior_samples=posterior_samples,
...             exclude_deterministic=False,
...         )
...         forecast = predictive(
...             jr.PRNGKey(1),
...             obs_times=obs_times,
...             obs_values=obs_values,
...             predict_times=forecast_times,
...         )

What this does¶

The concrete backend is selected from dynamics.state_evolution:

StochasticContinuousTimeStateEvolution uses SDESimulator.
DeterministicContinuousTimeStateEvolution uses ODESimulator.
Discrete-time state evolution uses DiscreteTimeSimulator.

A simulator is generation-only: raw simulator calls accept predict_times, not obs_times or obs_values. For observation-conditioned inference, use LatentPathBuilder, Filter, or Smoother. A simulator may wrap a Filter or Smoother to draw posterior rollouts at predict_times; in that composition, the inference handler consumes the observations and passes posterior state distributions outward to the simulator.

Configurations and defaults¶

simulator_config accepts a SimulatorConfig and forwards it to the selected continuous-time backend:

stochastic continuous-time models accept SDESimulatorConfig;
deterministic continuous-time models accept ODESimulatorConfig;
discrete-time models do not accept a simulator config.

If simulator_config=None, discrete transitions are sampled directly, deterministic continuous-time models use the ODESimulatorConfig defaults (diffrax.Tsit5() with fixed dt0=1e-3), and stochastic continuous-time models use the SDESimulatorConfig defaults (the fixed-step "em_scan" Euler--Maruyama backend with dt0=1e-4). n_simulations defaults to one and must be at least one. The simulation dimension is retained even when it has length one.

NumPyro trace¶

For a raw rollout from dsx.sample("f", ...), the following numpyro.deterministic sites are added:

"f_x_0": initial states, shape (*plate_shape, n_simulations, state_dim);
"f_times": prediction times, shape (*plate_shape, n_simulations, T);
"f_states": latent states, shape (*plate_shape, n_simulations, T, state_dim);
"f_observations": sampled observations, shape (*plate_shape, n_simulations, T, observation_dim).

Here "f" is replaced by the name passed to dsx.sample, T is the length of the prediction grid, and plate_shape is absent outside a dsx.plate. Under Predictive(..., num_samples=N), NumPyro prepends an N axis to every shape above. Because these sites are deterministic, pass exclude_deterministic=False to Predictive (or request the site names explicitly) to include them in its returned dictionary.

When the simulator wraps a Filter or Smoother, the inner handler records its own configured sites and the simulator's aggregate rollout sites are instead "f_predicted_times", "f_predicted_states", and "f_predicted_observations", with the corresponding time, state, and observation shapes above. Each nonempty prediction segment also records the state from which that segment starts, with shape (n_simulations, state_dim): "f_0_x_0" for a segment before the first posterior time, and "f_{j+1}_x_0" for a segment initialized from the posterior at inference-time index j. Only segments containing at least one requested prediction time are recorded. Inside dsx.plate, the segment name also identifies the plate member, for example "f_p0_1_x_0".

If predict_times is omitted, the simulator performs no rollout and adds no simulator sites. Direct calls to Simulator().simulate and dsx.simulate return a SimulatedResult and do not add NumPyro sites.

Notes

Controls are passed through to the selected backend. See the concrete simulator for its control interpolation or alignment rules.
A Simulator instance selects and caches its concrete backend on the first dynamics object it handles. Use separate instances for models requiring different backend types.
Simulator().simulate(...) consumes an already allocated simulation key. The public dsx.simulate function splits its root key before dispatch.

Attributes:

Name	Type	Description
`simulator_config`		Optional ODE or SDE solver configuration. Its type must match the model selected at first use.
`n_simulations`	`int`	Number of independent trajectories drawn per model execution. Defaults to one and must be greater than or equal to one.
`simulator`	`BaseSimulator \| None`	Concrete auto-selected simulator cached on first use.

`_ensure_simulator(dynamics: DynamicalModel) -> BaseSimulator` ¶

Instantiate and cache the concrete simulator for dynamics.

`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, **kwargs) -> SimulatedResult` ¶

Auto-route to the appropriate pure-JAX simulator backend.

Unlike dsx.simulate, rng_key is consumed directly as an already-allocated simulation key and is not pre-split. Therefore, dsx.simulate(..., rng_key=root_key) is equivalent to Simulator().simulate(..., rng_key=jax.random.split(root_key)[1]).

`_slice_tree_for_plate_member(tree, plate_shapes: tuple[int, ...], plate_idx: PlateIndex)` ¶

Select one plate member from every matching leaf in a pytree.

Member-specific array leaves are indexed by plate_idx. Classification uses both shape and the leaf's location in the tree so shared vectors are not sliced only because their length matches a plate size. For a constant-coefficient Diffusion, only the coefficient is sliced. Shared leaves are returned unchanged.

Parameters:

Name	Type	Description	Default
`tree`		Pytree whose leaves may contain leading plate dimensions.	required
`plate_shapes`	`tuple[int, ...]`	Sizes of the leading plate dimensions.	required
`plate_idx`	`PlateIndex`	One index for each plate dimension.	required

Returns:

Name	Type	Description
`PyTree`		Copy of `tree` containing values for the selected plate member.

Overview¶

Simulators¶

DiscreteTimeSimulator ¶

What this does¶

Configuration and defaults¶

NumPyro trace¶

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, **kwargs) -> SimulatedResult ¶

ODESimulator ¶

What this does¶

Configuration and defaults¶

NumPyro trace¶

__init__(simulator_config: ODESimulatorConfig | None = None, *, n_simulations: int = 1) -> None ¶

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, **kwargs) -> SimulatedResult ¶

SDESimulator ¶

What this does¶

Configurations and defaults¶

NumPyro trace¶

__init__(simulator_config: SDESimulatorConfig | None = None, *, n_simulations: int = 1) -> None ¶

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, **kwargs) -> SimulatedResult ¶

Simulator ¶

What this does¶

Configurations and defaults¶

NumPyro trace¶

_ensure_simulator(dynamics: DynamicalModel) -> BaseSimulator ¶

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, **kwargs) -> SimulatedResult ¶

_slice_tree_for_plate_member(tree, plate_shapes: tuple[int, ...], plate_idx: PlateIndex) ¶

`DiscreteTimeSimulator` ¶

`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, **kwargs) -> SimulatedResult` ¶

`ODESimulator` ¶

`init(simulator_config: ODESimulatorConfig | None = None, *, n_simulations: int = 1) -> None` ¶

`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, **kwargs) -> SimulatedResult` ¶

`SDESimulator` ¶

`init(simulator_config: SDESimulatorConfig | None = None, *, n_simulations: int = 1) -> None` ¶

`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, **kwargs) -> SimulatedResult` ¶

`Simulator` ¶

`_ensure_simulator(dynamics: DynamicalModel) -> BaseSimulator` ¶

`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, **kwargs) -> SimulatedResult` ¶

`_slice_tree_for_plate_member(tree, plate_shapes: tuple[int, ...], plate_idx: PlateIndex)` ¶