Skip to content

Overview

Simulators (also called unrollers) generate forward trajectories from a DynamicalModel on a provided time grid, and can also sit outside inference handlers to produce posterior rollouts.

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 the realized sites are registered only at the end. Use dsx.simulate(...) when you want a pure-JAX API with an explicit rng_key.
  • 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.

User code will usually choose between:

  • with Simulator(): ... or a concrete simulator handler inside a NumPyro model
  • dsx.simulate(...) for pure-JAX forward simulation without NumPyro sites

BaseSimulator

Bases: ObjectInterpretation, HandlesSelf

Base class for generation-only simulator handlers.

Interprets dsx.sample(name, dynamics, predict_times=..., ...) by running a pure-JAX forward simulation on the requested prediction grid, then registering the realized simulator outputs as deferred NumPyro sites only when the NumPyro-style API is used.

When the simulator runs, it records the solved trajectories as deterministic sites (conventionally "x_0", "times", "states", and "observations").

Notes
  • Raw simulator handlers are generation-only and therefore require predict_times.
  • Observation-conditioned latent-state inference now belongs to LatentPathBuilder (explicit latent paths) or Filter / Smoother (marginalized inference).
  • Posterior rollout remains supported because Filter / Smoother consume obs_times / obs_values before forwarding rollout metadata to the simulator.

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 without registering NumPyro sites.