# BrainMass
**Whole-brain modeling with differentiable neural mass models**
[](https://opensource.org/licenses/Apache-2.0)
[](https://www.python.org/downloads/)
[](https://brainmass.readthedocs.io/)
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, analyzing neural networks, and modeling hemodynamic responses.
## Features
- **Neural Mass Models**: Wilson-Cowan model for excitatory-inhibitory population dynamics
- **Hemodynamic Modeling**: BOLD signal simulation using the Balloon-Windkessel model
- **Network Coupling**: Diffusive and additive coupling mechanisms for brain connectivity
- **Noise Modeling**: Ornstein-Uhlenbeck processes for realistic neural noise
- **JAX-Powered**: GPU acceleration and automatic differentiation
- **Real Brain Data**: Example datasets from HCP and other neuroimaging studies
## Installation
### From PyPI (recommended)
```bash
pip install brainmass
```
### From Source
```bash
git clone https://github.com/chaobrain/brainmass.git
cd brainmass
pip install -e .
```
### GPU Support
For CUDA 12 support:
```bash
pip install brainmass[cuda12]
```
For TPU support:
```bash
pip install brainmass[tpu]
```
For whole brain modeling ecosystem:
```bash
pip install BrainX
# GPU support
pip install BrainX[cuda12]
# TPU support
pip install BrainX[tpu]
```
## Quick Start
### Single Node Simulation
```python
import brainstate
import brainmass
import numpy as np
import matplotlib.pyplot as plt
# Set simulation parameters
brainstate.environ.set(dt=0.1)
# Create a Wilson-Cowan model with noise
node = brainmass.WilsonCowanModel(
1,
noise_E=brainmass.OUProcess(1, sigma=0.01),
noise_I=brainmass.OUProcess(1, sigma=0.01),
)
# Initialize model states
brainstate.nn.init_all_states(node)
# Define simulation function
def step_run(i):
with brainstate.environ.context(i=i, t=i * brainstate.environ.get_dt()):
return node.update()
# Run simulation
indices = np.arange(10000)
activity = brainstate.transform.for_loop(step_run, indices)
# Plot results
plt.plot(indices * brainstate.environ.get_dt(), activity)
plt.xlabel("Time [ms]")
plt.ylabel("Neural Activity")
plt.show()
```
### Brain Network Simulation
```python
import brainstate
import brainmass
from datasets import Dataset
import numpy as np
# Load brain connectivity data
hcp = Dataset('hcp')
class BrainNetwork(brainstate.nn.Module):
def __init__(self, signal_speed=2., k=1.):
super().__init__()
# Get connectivity and distance matrices
conn_weight = hcp.Cmat
np.fill_diagonal(conn_weight, 0)
delay_time = hcp.Dmat / signal_speed
np.fill_diagonal(delay_time, 0)
indices_ = np.tile(np.arange(conn_weight.shape[1]), conn_weight.shape[0])
# Create network nodes
self.node = brainmass.WilsonCowanModel(
80, # 80 brain regions
noise_E=brainmass.OUProcess(80, sigma=0.01),
noise_I=brainmass.OUProcess(80, sigma=0.01),
)
# Define coupling between regions
self.coupling = brainmass.DiffusiveCoupling(
self.node.prefetch_delay('rE', (delay_time.flatten(), indices_),
init=brainstate.init.Uniform(0, 0.05)),
self.node.prefetch('rE'),
conn_weight,
k=k
)
def update(self):
current = self.coupling()
rE = self.node(current)
return rE
def step_run(self, i):
with brainstate.environ.context(i=i, t=i * brainstate.environ.get_dt()):
return self.update()
# Create and run network simulation
net = BrainNetwork()
brainstate.nn.init_all_states(net)
indices = np.arange(0, 6000 // brainstate.environ.get_dt())
brain_activity = brainstate.transform.for_loop(net.step_run, indices)
```
## Examples and Tutorials
The `examples/` directory contains Jupyter notebooks demonstrating:
- **Single Node Dynamics**: Basic Wilson-Cowan model simulation
- **Bifurcation Analysis**: Parameter space exploration
- **Brain Network Modeling**: Whole-brain connectivity simulations
- **Parameter Optimization**: Using Nevergrad for parameter tuning
- **BOLD Signal Analysis**: Hemodynamic response modeling
## Dependencies
Core dependencies:
- `jax`: High-performance computing and automatic differentiation
- `numpy`: Numerical computations
- `brainstate`: State management and neural dynamics
- `brainunit`: Unit system for neuroscience
- `brainscale`: Scaling utilities
Optional dependencies:
- `matplotlib`: Plotting and visualization
- `braintools`: Additional analysis tools
- `nevergrad`: Parameter optimization
## Documentation
Full documentation is available at [brainmass.readthedocs.io](https://brainmass.readthedocs.io/).
## Contributing
We welcome contributions! Please see [CONTRIBUTING.md](CONTRIBUTING.md) for guidelines.
## Citation
If you use BrainMass in your research, please cite:
```bibtex
@software{brainmass,
title={BrainMass: Whole-brain modeling with differentiable neural mass models},
author={BrainMass Developers},
url={https://github.com/chaobrain/brainmass},
version={0.0.1},
year={2024}
}
```
## License
BrainMass is licensed under the Apache License 2.0. See [LICENSE](LICENSE) for details.
## Support
- **Issues**: [GitHub Issues](https://github.com/chaobrain/brainmass/issues)
- **Documentation**: [ReadTheDocs](https://brainmass.readthedocs.io/)
- **Contact**: chao.brain@qq.com
---
**Keywords**: neural mass model, brain modeling, computational neuroscience, JAX, differentiable programming
Raw data
{
"_id": null,
"home_page": "https://github.com/chaobrain/brainmass",
"name": "brainmass",
"maintainer": null,
"docs_url": null,
"requires_python": ">=3.10",
"maintainer_email": null,
"keywords": "neural mass model, brain modeling, rate model, mean field model, neuroscience, computational neuroscience, brain simulation",
"author": "BranMass Developers",
"author_email": "brainmass Developers <chao.brain@qq.com>",
"download_url": null,
"platform": null,
"description": "# BrainMass\r\n\r\n**Whole-brain modeling with differentiable neural mass models**\r\n\r\n[](https://opensource.org/licenses/Apache-2.0)\r\n[](https://www.python.org/downloads/)\r\n[](https://brainmass.readthedocs.io/)\r\n\r\nBrainMass 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, analyzing neural networks, and modeling hemodynamic responses.\r\n\r\n## Features\r\n\r\n- **Neural Mass Models**: Wilson-Cowan model for excitatory-inhibitory population dynamics\r\n- **Hemodynamic Modeling**: BOLD signal simulation using the Balloon-Windkessel model\r\n- **Network Coupling**: Diffusive and additive coupling mechanisms for brain connectivity\r\n- **Noise Modeling**: Ornstein-Uhlenbeck processes for realistic neural noise\r\n- **JAX-Powered**: GPU acceleration and automatic differentiation\r\n- **Real Brain Data**: Example datasets from HCP and other neuroimaging studies\r\n\r\n## Installation\r\n\r\n### From PyPI (recommended)\r\n```bash\r\npip install brainmass\r\n```\r\n\r\n### From Source\r\n```bash\r\ngit clone https://github.com/chaobrain/brainmass.git\r\ncd brainmass\r\npip install -e .\r\n```\r\n\r\n### GPU Support\r\nFor CUDA 12 support:\r\n```bash\r\npip install brainmass[cuda12]\r\n```\r\n\r\nFor TPU support:\r\n```bash\r\npip install brainmass[tpu]\r\n```\r\n\r\nFor whole brain modeling ecosystem:\r\n```bash\r\npip install BrainX \r\n\r\n# GPU support\r\npip install BrainX[cuda12]\r\n\r\n# TPU support\r\npip install BrainX[tpu]\r\n```\r\n\r\n\r\n\r\n\r\n## Quick Start\r\n\r\n### Single Node Simulation\r\n\r\n```python\r\nimport brainstate\r\nimport brainmass\r\nimport numpy as np\r\nimport matplotlib.pyplot as plt\r\n\r\n# Set simulation parameters\r\nbrainstate.environ.set(dt=0.1)\r\n\r\n# Create a Wilson-Cowan model with noise\r\nnode = brainmass.WilsonCowanModel(\r\n 1,\r\n noise_E=brainmass.OUProcess(1, sigma=0.01),\r\n noise_I=brainmass.OUProcess(1, sigma=0.01),\r\n)\r\n\r\n# Initialize model states\r\nbrainstate.nn.init_all_states(node)\r\n\r\n# Define simulation function\r\ndef step_run(i):\r\n with brainstate.environ.context(i=i, t=i * brainstate.environ.get_dt()):\r\n return node.update()\r\n\r\n# Run simulation\r\nindices = np.arange(10000)\r\nactivity = brainstate.transform.for_loop(step_run, indices)\r\n\r\n# Plot results\r\nplt.plot(indices * brainstate.environ.get_dt(), activity)\r\nplt.xlabel(\"Time [ms]\")\r\nplt.ylabel(\"Neural Activity\")\r\nplt.show()\r\n```\r\n\r\n### Brain Network Simulation\r\n\r\n```python\r\nimport brainstate\r\nimport brainmass\r\nfrom datasets import Dataset\r\nimport numpy as np\r\n\r\n# Load brain connectivity data\r\nhcp = Dataset('hcp')\r\n\r\nclass BrainNetwork(brainstate.nn.Module):\r\n def __init__(self, signal_speed=2., k=1.):\r\n super().__init__()\r\n \r\n # Get connectivity and distance matrices\r\n conn_weight = hcp.Cmat\r\n np.fill_diagonal(conn_weight, 0)\r\n delay_time = hcp.Dmat / signal_speed\r\n np.fill_diagonal(delay_time, 0)\r\n indices_ = np.tile(np.arange(conn_weight.shape[1]), conn_weight.shape[0])\r\n \r\n # Create network nodes\r\n self.node = brainmass.WilsonCowanModel(\r\n 80, # 80 brain regions\r\n noise_E=brainmass.OUProcess(80, sigma=0.01),\r\n noise_I=brainmass.OUProcess(80, sigma=0.01),\r\n )\r\n \r\n # Define coupling between regions\r\n self.coupling = brainmass.DiffusiveCoupling(\r\n self.node.prefetch_delay('rE', (delay_time.flatten(), indices_), \r\n init=brainstate.init.Uniform(0, 0.05)),\r\n self.node.prefetch('rE'),\r\n conn_weight,\r\n k=k\r\n )\r\n \r\n def update(self):\r\n current = self.coupling()\r\n rE = self.node(current)\r\n return rE\r\n \r\n def step_run(self, i):\r\n with brainstate.environ.context(i=i, t=i * brainstate.environ.get_dt()):\r\n return self.update()\r\n\r\n# Create and run network simulation\r\nnet = BrainNetwork()\r\nbrainstate.nn.init_all_states(net)\r\nindices = np.arange(0, 6000 // brainstate.environ.get_dt())\r\nbrain_activity = brainstate.transform.for_loop(net.step_run, indices)\r\n```\r\n\r\n\r\n## Examples and Tutorials\r\n\r\nThe `examples/` directory contains Jupyter notebooks demonstrating:\r\n\r\n- **Single Node Dynamics**: Basic Wilson-Cowan model simulation\r\n- **Bifurcation Analysis**: Parameter space exploration\r\n- **Brain Network Modeling**: Whole-brain connectivity simulations\r\n- **Parameter Optimization**: Using Nevergrad for parameter tuning\r\n- **BOLD Signal Analysis**: Hemodynamic response modeling\r\n\r\n\r\n## Dependencies\r\n\r\nCore dependencies:\r\n- `jax`: High-performance computing and automatic differentiation\r\n- `numpy`: Numerical computations\r\n- `brainstate`: State management and neural dynamics\r\n- `brainunit`: Unit system for neuroscience\r\n- `brainscale`: Scaling utilities\r\n\r\nOptional dependencies:\r\n- `matplotlib`: Plotting and visualization\r\n- `braintools`: Additional analysis tools\r\n- `nevergrad`: Parameter optimization\r\n\r\n## Documentation\r\n\r\nFull documentation is available at [brainmass.readthedocs.io](https://brainmass.readthedocs.io/).\r\n\r\n## Contributing\r\n\r\nWe welcome contributions! Please see [CONTRIBUTING.md](CONTRIBUTING.md) for guidelines.\r\n\r\n## Citation\r\n\r\nIf you use BrainMass in your research, please cite:\r\n\r\n```bibtex\r\n@software{brainmass,\r\n title={BrainMass: Whole-brain modeling with differentiable neural mass models},\r\n author={BrainMass Developers},\r\n url={https://github.com/chaobrain/brainmass},\r\n version={0.0.1},\r\n year={2024}\r\n}\r\n```\r\n\r\n## License\r\n\r\nBrainMass is licensed under the Apache License 2.0. See [LICENSE](LICENSE) for details.\r\n\r\n## Support\r\n\r\n- **Issues**: [GitHub Issues](https://github.com/chaobrain/brainmass/issues)\r\n- **Documentation**: [ReadTheDocs](https://brainmass.readthedocs.io/)\r\n- **Contact**: chao.brain@qq.com\r\n\r\n---\r\n\r\n**Keywords**: neural mass model, brain modeling, computational neuroscience, JAX, differentiable programming\r\n",
"bugtrack_url": null,
"license": "Apache-2.0 license",
"summary": null,
"version": "0.0.2",
"project_urls": {
"Homepage": "https://github.com/chaobrain/brainmass",
"homepage": "http://github.com/chaobrain/brainmass",
"repository": "http://github.com/chaobrain/brainmass"
},
"split_keywords": [
"neural mass model",
" brain modeling",
" rate model",
" mean field model",
" neuroscience",
" computational neuroscience",
" brain simulation"
],
"urls": [
{
"comment_text": "",
"digests": {
"blake2b_256": "80ab2fc0b7cf99732a0bfe179690d5f3ee95d4b35780052eba814fa032558618",
"md5": "7df1e33e8b8ebbdbc5f4735f0aafa6eb",
"sha256": "24655bcf623af1f0893ea53679f82ddec3f76f389d1490f420a08bdc36025fbc"
},
"downloads": -1,
"filename": "brainmass-0.0.2-py2.py3-none-any.whl",
"has_sig": false,
"md5_digest": "7df1e33e8b8ebbdbc5f4735f0aafa6eb",
"packagetype": "bdist_wheel",
"python_version": "py2.py3",
"requires_python": ">=3.10",
"size": 19224,
"upload_time": "2025-09-04T13:49:24",
"upload_time_iso_8601": "2025-09-04T13:49:24.898626Z",
"url": "https://files.pythonhosted.org/packages/80/ab/2fc0b7cf99732a0bfe179690d5f3ee95d4b35780052eba814fa032558618/brainmass-0.0.2-py2.py3-none-any.whl",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2025-09-04 13:49:24",
"github": true,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"github_user": "chaobrain",
"github_project": "brainmass",
"travis_ci": false,
"coveralls": false,
"github_actions": true,
"requirements": [
{
"name": "numpy",
"specs": []
},
{
"name": "jax",
"specs": []
},
{
"name": "brainunit",
"specs": []
},
{
"name": "brainscale",
"specs": []
},
{
"name": "brainstate",
"specs": [
[
">=",
"0.1.8"
]
]
}
],
"lcname": "brainmass"
}
JSON
