ABMax: A JAX-based Agent-based Modeling Framework

TL;DR

ABMax is a JAX-based agent-based modeling framework that enables flexible agent updates and parallel simulations, overcoming shape immutability constraints for scalable complex system modeling.

Contribution

Introduces ABMax, a novel JAX-compatible ABM framework with algorithms for dynamic agent updates and parallel execution, enhancing scalability and flexibility.

Findings

Achieves runtime performance comparable to state-of-the-art ABM tools.

Enables vectorized execution of multiple models in parallel.

Demonstrates applicability with traffic-flow and financial market models.

Abstract

Agent-based modeling (ABM) is a principal approach for studying complex systems. By decomposing a system into simpler, interacting agents, agent-based modeling (ABM) allows researchers to observe the emergence of complex phenomena. High-performance array computing libraries like JAX can help scale such computational models to a large number of agents by using automatic vectorization and just-in-time (JIT) compilation. One of the caveats of using JAX to achieve such scaling is that the shapes of arrays used in the computational model should remain immutable throughout the simulation. In the context of agent-based modeling (ABM), this can pose constraints on certain agent manipulation operations that require flexible data structures. A subset of which is represented by the ability to update a dynamically selected number of agents by applying distinct changes to them during a simulation. To this effect, we introduce ABMax, an ABM framework based on JAX that implements multiple just-in-time (JIT) compilable algorithms to provide this functionality. On the canonical predation model benchmark, ABMax achieves runtime performance comparable to state-of-the-art implementations. Further, we show that this functionality can also be vectorized, making it possible to run many similar agent-based models in parallel. We also present two examples in the form of a traffic-flow model and a financial market model to show the use case of ABMax