Skip to content

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

Examples

Predictive with SDESimulator 
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,
    FullDiffusion,
    SDESimulator,
)
from numpyro.infer import Predictive

state_dim = 1
observation_dim = 1
bm_dim = 1

def model(predict_times=None):
    theta = numpyro.sample("theta", dist.LogNormal(-0.5, 0.2))
    sigma_x = numpyro.sample("sigma_x", dist.LogNormal(-1.0, 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,
            diffusion=FullDiffusion(sigma_x * jnp.eye(state_dim, bm_dim)),
        ),
        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 SDESimulator():
    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_x', '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

SDESimulator is generation-only. Native SDE explicit latent-path inference should currently go through a discretization step first, then LatentPathBuilder; otherwise prefer Filter for marginalized inference.