Important
Early development stage. There may be bugs and inconsistencies. The code and API will change significantly in the future, especially for AD part.
- Simple, lightweight, and easy-to-use Python implementation of modern standard models of dynamic network traffic flow
- An end-to-end differentiable simulation using JAX
- Lightning-fast JAX mode mode with a GPU: 0.3 sec for forward simulation on the Chicago-Sketch dataset (2500 links, 1 million vehicles, 3 hours), and 0.5 sec for backward differentiation
- The basic features and syntax are almost identical to those of the UXsim traffic flow simulator
60000 vehicles travel through a 10 km grid network over 3 hours. Dark colors indicate congestion (slow speeds). The simulation wall-clock time was 5 seconds on a 2.0 GHz CPU in pure Python mode.
One million vehicles travel through the Chicago network (approximately 2500 links) over a 3-hour period. The simulation time for a single run was 0.2 seconds using a GPU.
Additionally, we solve a dynamic congestion pricing optimization problem on this network. The number of decision variables is 15000, corresponding to the number of links and tolling periods. This problem is very difficult to solve using conventional approaches (e.g., SPSA in the figure above), but our simulator quickly produced a high-quality solution (AD in the figure).
Simple scenario in a Y-shaped merge network:
from unsim import World
import matplotlib.pyplot as plt
# Define the main simulation
# Units are standardized to seconds (s) and meters (m)
W = World(name="merge", deltat=5, tmax=1200,
print_mode=1, save_mode=1, show_mode=1)
# Define the network
W.addNode("orig1", x=0, y=0)
W.addNode("orig2", x=0, y=2)
W.addNode("merge", x=1, y=1)
W.addNode("dest", x=2, y=1)
W.addLink("link1", "orig1", "merge", length=1000, free_flow_speed=20, capacity=0.8, merge_priority=1)
W.addLink("link2", "orig2", "merge", length=1000, free_flow_speed=20, capacity=0.8, merge_priority=1)
W.addLink("link3", "merge", "dest", length=1000, free_flow_speed=20)
# Define the vehicle demand
W.adddemand("orig1", "dest", t_start=0, t_end=1000, flow=0.45)
W.adddemand("orig2", "dest", t_start=400, t_end=1000, flow=0.6)
# Run the simulation
W.exec_simulation()
# Analysis
W.analyzer.print_simple_stats()
W.analyzer.network(t=200)
W.analyzer.network(t=800)
plt.show()Results:
Simulation completed. merge
Simulation Results:
Total trips: 810.0
Completed trips: 740.0
Total travel time: 136825.0 s
Avg travel time: 184.9 s
Avg delay: 84.9 s
For further usage including Autodiff, please see the Jupyter notebook tutorial and example scripts.
pip install unsimIf you want to use JAX acceleration, install your preferred JAX build such as jax[cpu] and jax[cuda13].
The optimal installation depends on your hardware and software configuration, so please check the JAX official document.
This simulator implements the following transportation scientific models:
- Link Transmission Model: very efficient and accurate macroscopic traffic flow model
- Incremental Node Model: general intersection or junction model
- Dynamic User Optimum: route choice model in which each traveler chooses the real-time minimum travel time path
For the details, please see the arXiv preprint.
UNsim is released under the MIT License. You are free to use it as long as the source is acknowledged.
If you use the code, please cite the arXiv article:
- Toru Seo. End-to-end differentiable network traffic simulation with dynamic route choice. arXiv preprint arXiv:2604.11380, 2026.
@Article{seo2026unsim_arxiv,
author = {Toru Seo},
journal = {arXiv preprint arXiv: 2604.11380},
title = {End-to-end differentiable network traffic simulation with dynamic route choice},
year = {2026},
doi = {10.48550/arXiv.2604.11380},
}