braintools is a lightweight, JAX-friendly toolbox with practical utilities for brain modeling.
- 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 braintoolsOptional 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 BrainXThe 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