Modeling tools for brain simulation

October 9, 2025 · View on GitHub

Header image of braintools.

braintools is a lightweight, JAX-friendly toolbox with practical utilities for brain modeling.

Highlights

Composable connectivity: declarative builders for point, multi-compartment, and population networks with spatial kernels, degree constraints, and unit-aware metadata
Visualization suite: publication plots, interactive dashboards, 3D viewers, and animation helpers in braintools.visualize
Metrics and solvers: losses, evaluation metrics, and PyTree-aware ODE/SDE/DDE integrators ready for jit/vmap
Signal and optimization helpers: reusable generators and lightweight optimizers to prototype models quickly

braintools integrates smoothly with the broader ecosystem (e.g., brainstate, brainunit) while keeping a simple, functional style.

pip install -U braintools

Optional extras are published for hardware-specific builds:

pip install -U braintools[cpu]
# CUDA 12.x wheels
pip install -U braintools[cuda12]
# TPU runtime
pip install -U braintools[tpu]

Alternatively, install the curated BrainX bundle that ships with braintools and related projects:

pip install -U BrainX

The full documentation is available at https://braintools.readthedocs.io

braintools is one part of our brain simulation ecosystem: https://brainmodeling.readthedocs.io/

Contributions and issue reports are welcome! See CONTRIBUTING.md for guidelines.

Apache 2.0. See LICENSE for details.

If you use braintools in your work, please cite the Zenodo DOI: 10.5281/zenodo.17110064