JaxABM Documentation
JaxABM is a high-performance agent-based modeling framework built on JAX, designed for fast, scalable, and differentiable simulations. It provides powerful tools for parameter calibration, sensitivity analysis, and model optimization using modern machine learning techniques.
Key Features
High Performance: Built on JAX for GPU acceleration and JIT compilation
Advanced Calibration: Multiple optimization methods including reinforcement learning
Sensitivity Analysis: Comprehensive tools for parameter importance analysis
Differentiable: Full compatibility with JAX’s automatic differentiation
Scalable: Handle large-scale agent populations efficiently
Flexible: Support for custom agent behaviors and model architectures
Quick Start
Install JaxABM:
pip install jaxabm
Basic usage:
import jaxabm as jx
# Create a simple model
model = jx.Model()
# Add agents
agents = jx.AgentCollection("traders", 1000)
model.add_agent_collection(agents)
# Run simulation
results = model.run(steps=100)
Documentation Contents
User Guide
Core Features
API Reference
Development