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 >> brunsbaftera(sequential)a * nrepeatsantimesa | brunsaandbin parallel (multi-modal stimuli)If,Switch,Whileadd 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
1 maintainer
Metadata
1 maintainer
Release files for braintools 0.3.0
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| braintools-0.3.0.tar.gz | 770.0 kB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| braintools-0.3.0-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 1.7 MB
Release files / braintools-0.3.0.tar.gz
| Download URL | braintools-0.3.0.tar.gz |
|---|---|
| Size | 770.0 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
8cb9f0ed3501b0ba5cc08f9643b0cf5455875aef8e4167b61824ed585342d873
|
|
BLAKE2b-256 checksum How to use checksums |
63d6c242a5d29d9f15551383cd7c29b2362d4162575beb79804d99c45f70603a
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/6.1.0 CPython/3.13.12
|
Release files / braintools-0.3.0-py3-none-any.whl
| Download URL | braintools-0.3.0-py3-none-any.whl |
|---|---|
| Size | 931.7 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
7846be0329b9d1927d8d8ead899ca36aa45e9a6ea64ffd094125f159271d10cf
|
|
BLAKE2b-256 checksum How to use checksums |
74677e2ed918a71beb89a13b4471e87aeb9d41d573758c7bfd943534766615d3
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/6.1.0 CPython/3.13.12
|
