braintools documentation

braintools implements a modern toolbox for brain simulation and analysis.

Overview

braintools is a lightweight, JAX-friendly collection of utilities used across computational neuroscience workflows:

• Composable synaptic connectivity builders (braintools.conn) with point, compartment, and population abstractions

• A composable cognitive task framework (braintools.cogtask) for constructing trial-structured tasks — decision making, working memory, reasoning, motor — with JIT/vmap-friendly batch sampling and a registry of pre-built paradigms

• Rich visualization helpers (braintools.visualize) covering static plots, interactive dashboards, and animation utilities

• Metric functions for model training and evaluation (classification, regression, ranking, correlation, LFP helpers)

• Numerical integration utilities for ODE/SDE/DDE one-step steppers (braintools.quad) designed for PyTrees and functional APIs

• Input generators, optimization helpers, and reusable data-structure utilities

• And more …

The project favors a simple, well-typed functional style that works seamlessly with brainstate, brainunit, and just-in-time compilation (jit/vmap).

Installation

CPU

pip install -U braintools[cpu]

GPU

pip install -U braintools[cuda12]
pip install -U braintools[cuda13]

TPU

pip install -U braintools[tpu]

---

See also the ecosystem

BrainTools is one part of our brain simulation ecosystem.

---