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

Examples¶

Predictive with ODESimulator

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

state_dim = 1
observation_dim = 1

def model(predict_times=None):
    theta = numpyro.sample("theta", dist.LogNormal(-0.5, 0.2))
    sigma_y = numpyro.sample("sigma_y", dist.LogNormal(-1.5, 0.2))
    dynamics = DynamicalModel(
        control_dim=0,
        initial_condition=dist.MultivariateNormal(
            loc=jnp.zeros(state_dim),
            covariance_matrix=jnp.eye(state_dim),
        ),
        state_evolution=ContinuousTimeStateEvolution(
            drift=lambda x, u, t: -theta * x,
        ),
        observation_model=lambda x, u, t: dist.MultivariateNormal(
            x,
            sigma_y**2 * jnp.eye(observation_dim),
        ),
    )
    return dsx.sample("f", dynamics, predict_times=predict_times)

predict_times = jnp.linspace(0.0, 5.0, 51)
with ODESimulator():
    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', 'sigma_y', 'theta', ...]
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

ODESimulator is generation-only. For explicit ODE latent-state inference, use LatentPathBuilder. For marginalized continuous-time inference, use Filter. Posterior rollouts can still be generated with ODESimulator outside those handlers.