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:

@software{brainmass,
  title={BrainMass: Whole-brain modeling with differentiable neural mass models},
  author={BrainMass Developers},
  url={https://github.com/chaobrain/brainmass},
  version={0.1.0},
  year={2026}
}

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/