Abmax is a general-purpose agent-based modeling(ABM) framework in Jax It provides:

• Two algorithms to apply unique updates to a dynamic number of agents selected during run-time. Both are JIT friendly and can be vectorized across different number of models
  • Rank-Match (RM)
  • Sort-Count-Iterate (SCI)
• JAX-transformation-friendly data structures and functions that can be used to define sets of agents and their manipulations.
  • Selecting agents based on a run-time determined key.
  • Searching and sorting agents based on their attributes.
  • Updating an arbitrary number of agents to a specific state.
  • Stepping agents in a vectorized way.
  • Running multiple such simulations in parallel.
• Implementation of common algorithms used in ABM implemented in vmap and jit friendly way.

pip install abmax

Dependencies:

• Python >= 3.10
• Jax >= 0.4.13
• Flax >= 0.7.4

A comparison of the performance of Abmax with other ABM frameworks: Agents.jl and Mesa based on the Wolf-Sheep (Grid space) and Bird-Flock (Continuous space) models. These simulations are run for 100 steps and the median time taken for 10 runs is logged in ms. The benchmark is run on a gcn GPU node(Intel Xeon Platinum 8360Y + Nvidia A100) of the Snelius cluster The number of initial agents for these simulations are as follows:

• Wolf-Sheep small: 600 sheep, 400 wolves on a 100 x 100 grid
• Wolf-Sheep large: 6000 sheep, 4000 wolves on a 1000 x 1000 grid

In Abmax, we can run multiple simulations in parallel because of automatic batching and vectorization. Here is a trend in running different numbers of wolf-sheep small models in parallel.

Note: All times that are reported, are excluding the model setup time.

A basic tutorial on how to use Abmax is available here[Outdated]

If you use Abmax in your work, please consider citing it as follows:

@misc{chaturvedi2025abmax,
    title={Abmax: A JAX-based Agent-based Modeling Framework},
    author={Siddharth Chaturvedi and Ahmed El-Gazzar and Marcel van Gerven},
    year={2025},
    eprint={2508.16508},
    archivePrefix={arXiv},
    primaryClass={cs.MA}
}