Differentiable simulation (JAX)¶
JAX-based differentiable simulation engine unsim.unsim_diff.
UNsim-diff: JAX-based differentiable Network Link Transmission Model simulator.
All simulation internals are pure JAX functions, enabling jax.grad, jax.jit, and jax.lax.scan for automatic differentiation and compilation.
Usage¶
>>> from unsim import World
>>> from unsim.unsim_diff import world_to_jax, simulate, total_travel_time
>>> W = World(...); W.addNode(...); W.addLink(...); W.adddemand(...)
>>> params, config, lengths = world_to_jax(W)
>>> state = simulate(params, config, lengths)
>>> ttt = total_travel_time(state, config)
>>> grad_fn = jax.grad(lambda p: total_travel_time(simulate(p, config, lengths), config))
>>> grads = grad_fn(params)
- unsim.unsim_diff.simulate(params, config, differentiable=True, checkpoint_every=None)¶
Run full LTM simulation.
- Parameters:
params (Params)
config (NetworkConfig)
differentiable (bool, optional) – If True (default), use windowed carry suitable for jax.grad. If False, use full-array carry for faster forward-only evaluation (not compatible with reverse-mode AD).
checkpoint_every (int or None, optional) – If None (default), no gradient checkpointing (current behavior). If a positive integer, split the time loop into segments of this many steps and apply jax.checkpoint to each, reducing peak GPU memory of reverse-mode AD at the cost of extra recomputation in the backward pass. Smaller values save more memory. Only effective when
differentiable=True.
- Returns:
Final simulation state.
- Return type:
SimState
- unsim.unsim_diff.simulate_duo(params, config, differentiable=True, checkpoint_every=None)¶
Run DUO simulation.
- Parameters:
params (Params)
config (NetworkConfig)
differentiable (bool, optional) – If True (default), use windowed carry suitable for jax.grad. If False, use full-array carry for faster forward-only evaluation (not compatible with reverse-mode AD).
checkpoint_every (int or None, optional) – If None (default), no gradient checkpointing (current behavior). If a positive integer, split the time loop into segments of this many steps and apply jax.checkpoint to each, reducing peak GPU memory of reverse-mode AD at the cost of extra recomputation in the backward pass. Smaller values save more memory. Only effective when
differentiable=True.
- Return type:
SimState
- unsim.unsim_diff.total_travel_time(state, config)¶
Total travel time (s). Differentiable.
- TTT = sum of (N_U(t) - N_D(t)) * dt over all links and timesteps
sum of demand_queue * dt over all origins and timesteps.
- Parameters:
state (SimState)
config (NetworkConfig)
- Return type:
jnp scalar