Skip to content

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(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].

Examples

Predictive with DiscreteTimeSimulator 
import dynestyx as dsx
import jax.numpy as jnp
import jax.random as jr
import numpyro
import numpyro.distributions as dist
from dynestyx import DynamicalModel, DiscreteTimeSimulator
from numpyro.infer import Predictive

state_dim = 1
observation_dim = 1

def model(phi=None, predict_times=None):
    phi = numpyro.sample("phi", dist.Uniform(0.0, 1.0), obs=phi)
    dynamics = DynamicalModel(
        control_dim=0,
        initial_condition=dist.MultivariateNormal(
            loc=jnp.zeros(state_dim),
            covariance_matrix=jnp.eye(state_dim),
        ),
        state_evolution=lambda x, u, t_now, t_next: dist.MultivariateNormal(
            loc=phi * x + 0.1 * jnp.sin(x),
            covariance_matrix=0.2**2 * jnp.eye(state_dim),
        ),
        observation_model=lambda x, u, t: dist.MultivariateNormal(
            x,
            0.3**2 * jnp.eye(observation_dim),
        ),
    )
    return dsx.sample("f", dynamics, predict_times=predict_times)

predict_times = jnp.arange(20.0)
with DiscreteTimeSimulator():
    prior_pred = Predictive(model, num_samples=5)(
        jr.PRNGKey(0),
        predict_times=predict_times,
    )
print("Predictive keys:", sorted(prior_pred.keys()))  # e.g. ['f_observations', 'f_states', 'f_times', 'phi', ...]
print("Predictive shapes:", {k: v.shape for k, v in prior_pred.items()})  # trajectory arrays: (num_samples, n_sim, T, dim); here num_samples=5, n_sim=1

Note

DiscreteTimeSimulator is generation-only. For explicit latent-state inference use LatentPathBuilder; for marginalized inference use Filter. You can then wrap the fitted model in DiscreteTimeSimulator again to generate rollouts at predict_times.