BrainMass: whole-brain modeling with differentiable neural mass models

BrainMass is a Python library for whole-brain computational modeling using differentiable neural mass models. Built on JAX for high-performance computing, it provides tools for simulating brain dynamics, fitting neural signal data, and training cognitive tasks.

Installation

From PyPI (recommended)

pip install brainmass

From Source

git clone https://github.com/chaobrain/brainmass.git
cd brainmass
pip install -e .

GPU Support

For CUDA support:

pip install brainmass[cuda12]
pip install brainmass[cuda13]

For TPU support:

pip install brainmass[tpu]

Ecosystem

For whole brain modeling ecosystem:

pip install BrainX 

# GPU support
pip install BrainX[cuda12]
pip install BrainX[cuda13]

# TPU support
pip install BrainX[tpu]

Quick Start

Simulate a single brain region with the Hopf oscillator model. The high-level Simulator drives the compiled run loop and collects the monitored trajectories into a unit-aware result dict -- no hand-written stepping required. The integration time step dt is supplied to the Simulator:

import brainmass
import brainunit as u

# Create a single-region Hopf oscillator in the limit-cycle regime (a > 0)
node = brainmass.HopfStep(in_size=1, a=0.25, w=0.3)

# Run for 200 ms, dropping the first 20 ms transient, recording x and y
sim = brainmass.Simulator(node, dt=0.1 * u.ms)
res = sim.run(200 * u.ms, monitors=['x', 'y'], transient=20 * u.ms)

print(res['x'].shape)  # (1800, 1)

# Plot the limit cycle (matplotlib via the optional [viz] extra)
brainmass.viz.plot_timeseries(res['x'], ts=res['ts'])

Wiring a whole-brain network and fitting a parameter are just as direct:

import brainmass
import brainunit as u
from brainstate.nn import Param

# A delay-coupled network from a bundled example connectome
conn = brainmass.datasets.load_dataset('example_connectome')
net = brainmass.Network(
    brainmass.HopfStep(in_size=conn.weights.shape[0], a=0.1, w=0.3),
    conn=conn.weights, distance=conn.distances, speed=10 * u.mm / u.ms,
    coupling='diffusive', coupled_var='x', k=0.5,
)

# Fit a model parameter to data with gradients (or swap to Nevergrad / SciPy)
node = brainmass.HopfStep(in_size=1, a=Param(0.1, fit=True), w=0.3)
fitter = brainmass.Fitter(node, loss_fn=my_loss)   # backend='grad' | 'nevergrad' | 'scipy'
result = fitter.fit(n_steps=50)

Documentation

The full documentation is organized for different goals:

• Getting Started — installation, a five-minute quickstart, key concepts, and persona learning paths.
• Tutorials — a sequential path from a first simulation to building networks, forward modeling, fitting, and training.
• How-To Guides — task-focused recipes (choose a model, work with units, accelerate, custom coupling/objective, sweeps, analysis).
• Concepts — the why: neural mass models, differentiable programming, architecture, coupling/delays, forward models.
• Data-Driven Modeling — the flagship guided path through the differentiable, data-driven workflow.
• Gallery — a runnable model zoo (one demo per model family) and end-to-end case studies.
• API Reference — every public symbol, organized by category.
• Developer Guide — contributing and the extension playbooks (custom models, couplings, objectives, workflows).

Citation

If you use BrainMass in your research, please cite the archived release on Zenodo (DOI: 10.5281/zenodo.20762192):

@software{chaoming_wang_2026_20762193,
  author       = {Chaoming Wang and
                  Catttyfyx and
                  xinzhu-L},
  title        = {chaobrain/brainmass: Version 0.1.0},
  month        = jun,
  year         = 2026,
  publisher    = {Zenodo},
  version      = {v0.1.0},
  doi          = {10.5281/zenodo.20762193},
  url          = {https://doi.org/10.5281/zenodo.20762193},
  swhid        = {swh:1:dir:38b86b9216559a6cfd2a4f6b7325fc9c444fbaa9
                   ;origin=https://doi.org/10.5281/zenodo.20762192;vi
                   sit=swh:1:snp:b47f4abfd6c7e4fb4ac20f2367b46100c865
                   01f8;anchor=swh:1:rel:ae0f8da161588d4af993a702181a
                   1e2307916d2e;path=chaobrain-brainmass-e09e9f6
                  },
}

License

BrainMass is licensed under the Apache License 2.0. See LICENSE for details.

See also the ecosystem

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