Installation

Requirements

JaxABM requires Python 3.8 or later and has the following dependencies:

Core Dependencies

JAX: For high-performance computing and automatic differentiation
NumPy: For numerical operations
SciPy: For scientific computing utilities

Optional Dependencies

Matplotlib: For plotting and visualization
Seaborn: For advanced statistical plots
Pandas: For data manipulation and analysis
NetworkX: For network-based models

Installation Methods

From PyPI (Recommended)

pip install jaxabm

For GPU support with CUDA:

# Install JAX with CUDA support first
pip install "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

# Then install JaxABM
pip install jaxabm

From Source

For development or the latest features:

git clone https://github.com/username/jaxabm.git
cd jaxabm
pip install -e .

Development Installation

For contributing to JaxABM:

git clone https://github.com/username/jaxabm.git
cd jaxabm
pip install -e ".[dev]"

This installs additional development dependencies for testing and documentation.

Verification

Test your installation:

import jaxabm as jx
import jax.numpy as jnp

# Check JAX is working
print(f"JAX version: {jx.__version__}")
print(f"JAX devices: {jx.devices()}")

# Simple test
model = jx.Model()
print("✅ JaxABM installed successfully!")

GPU Support

To use GPU acceleration, ensure you have:

NVIDIA GPU with CUDA support
CUDA toolkit installed
JAX with CUDA installed

Check GPU availability:

import jax

print(f"JAX devices: {jax.devices()}")
print(f"GPU available: {len(jax.devices('gpu')) > 0}")

Troubleshooting

Common Issues

JAX Installation Issues

If you encounter issues with JAX installation, refer to the JAX installation guide.

Import Errors

If you get import errors, ensure all dependencies are installed:

pip install --upgrade jaxabm

GPU Not Detected

For GPU issues, check your CUDA installation and JAX GPU setup:

python -c "import jax; print(jax.devices())"

Performance Tips

For optimal performance:

Use JAX transformations: jit, vmap, pmap
Batch operations: Process multiple agents simultaneously
GPU acceleration: Use CUDA for large-scale simulations
Memory management: Use JAX’s memory-efficient operations

Next Steps

Read the Quick Start Guide guide
Explore Tutorials
Check out Examples