Skip to content

Simulate

jax_morph.simulate #

simulate(
    model: Model,
    state: BaseState,
    n_steps: int,
    dt: float | Array = 1.0,
    key: Array | None = None,
    *,
    history: bool = False,
    checkpoint: bool = False,
) -> BaseState

Roll out model for n_steps macro-steps of size dt; return the trajectory or final.

The forward pass is a pure sampler, pathwise-differentiable end to end. With history=True the result is the stacked complete state history s_0, ..., s_n (leading n_steps + 1 axis); otherwise it is the final state. key is required only when the model has a stochastic step - passing key=None then raises rather than silently freezing the seed. Trace fields on history frame zero belong to the input state, not a transition, so use transition_logp only when an intermediate state is intentionally an observed conditioning boundary.

Parameters:

  • model (Model) –

    The Model to advance.

  • state (BaseState) –

    Initial BaseState.

  • n_steps (int) –

    Number of macro-steps to scan.

  • dt (float | Array, default: 1.0 ) –

    Macro-step size; physical end time is n_steps * dt. Defaults to 1.0.

  • key (Array | None, default: None ) –

    PRNG key. Required for a model containing stochastic steps. Defaults to None.

  • history (bool, default: False ) –

    Whether to return the complete history s_0, ..., s_n instead of only the final state. Defaults to False.

  • checkpoint (bool, default: False ) –

    Whether to rematerialize macro-steps during backpropagation. Defaults to False.

Returns:

  • BaseState –

    The final BaseState, or a stacked complete history with leading axis n_steps + 1 when

  • BaseState –

    history is true.

Raises:

  • ValueError –

    If key is omitted for a model containing stochastic steps.

Scoring.

The result carries no aggregate log-probability. Request a complete history and pass it to trajectory_logp to score a sampled rollout.

Sample a trajectory 
trajectory = jxm.simulate(model, state, 20, dt=0.1, key=key, history=True)