UNsim: Python traffic simulation with Autodiff

Main Features

• 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 1.5 sec for backward differentiation, with 2GB VRAM (on RTX 5060 Ti).
• The basic features and syntax are almost identical to those of the UXsim traffic flow simulator

Simulation Examples

Simple Simulation

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.

Autodiff-based large-scale optimization

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

Usage

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.

Install

pip install unsim

If 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.

Technical Note

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.

Terms of Use & License

UNsim is released under the Apache License 2.0. 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},
}

Related Links

• Toru Seo (Author)
• Seo Laboratory, Institute of Science Tokyo