Modeling tools for brain simulation

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
• Cognitive task framework: phase-based, composable trial construction (braintools.cogtask) for training and analyzing neural networks on classic systems-neuroscience paradigms — decision making, working memory, reasoning, motor — plus a registry of pre-built tasks
• 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.

Installation

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

Cognitive task framework (braintools.cogtask)

braintools.cogtask lets you build cognitive tasks from composable, named phases (fixation, stimulus, delay, response, ...) connected by simple operators:

• a >> b runs b after a (sequential)
• a * n repeats a n times
• a | b runs a and b in parallel (multi-modal stimuli)
• If, Switch, While add trial-by-trial branching

Trials are generated lazily via task.sample(i) / task.batch_sample(B), both JIT- and vmap-friendly, with per-trial keys derived from seed so batches are fully reproducible.

import brainunit as u
from braintools.cogtask import PerceptualDecisionMaking

task = PerceptualDecisionMaking(t_stimulus=1500 * u.ms, num_choices=2, seed=0)
X, Y = task.batch_sample(32)   # X: (T, B, n_in), Y: (T, B) labels

The package also ships a registry of canonical tasks: PerceptualDecisionMaking, ContextDecisionMaking, DelayMatchSample, DelayComparison, GoNoGo, ReadySetGo, PostDecisionWager, Reaching1D, EvidenceAccumulation, HierarchicalReasoning, and more — see the API reference for the full list.

Documentation

The full documentation is available at https://brainx.chaobrain.com/braintools/

Ecosystem

braintools is one part of our brain simulation ecosystem: https://brainx.chaobrain.com/

Contributing

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

License

Apache 2.0. See LICENSE for details.

Citation

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