Installation

Installation#

This guide covers installing brainmass for different hardware configurations.

Quick Installation#

For most users, CPU installation is the simplest:

pip install -U brainmass[cpu]

Hardware-Specific Installation#

CPU#

For CPU-only systems (recommended for getting started):

pip install -U brainmass[cpu]

pip install -U brainmass[cpu]

GPU (CUDA)#

For NVIDIA GPUs with CUDA support:

pip install -U brainmass[cuda12]

pip install -U brainmass[cuda13]

Requirements:

  • NVIDIA GPU with compute capability ≥ 3.5

  • CUDA toolkit installed (12.x or 13.x)

  • cuDNN library

Troubleshooting CUDA:

If GPU is not detected:

import jax
print(jax.devices())  # Should show GPU devices

If only CPU is shown, reinstall JAX with CUDA support:

pip install --upgrade jax[cuda12_pip] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

TPU#

For Google Cloud TPU:

pip install -U brainmass[tpu]

Note: TPU support requires running on Google Cloud Platform with TPU VMs.

Development Installation#

To install from source for development:

# Clone the repository
git clone https://github.com/chaobrain/brainmass.git
cd brainmass

# Install in editable mode with development dependencies
pip install -e .[dev,doc]

Development dependencies include:

  • Testing: pytest, pytest-cov

  • Documentation: sphinx, sphinx-book-theme, etc.

  • Code quality: ruff, mypy

Verifying Installation#

After installation, verify that everything works:

import brainmass
import jax.numpy as jnp
import brainunit as u

print(f"brainmass version: {brainmass.__version__}")

# Test basic functionality
model = brainmass.HopfOscillator(in_size=10, omega=10 * u.Hz)
model.init_all_states()
output = model.update()

print(f"Model output shape: {output.shape}")
print("✓ Installation successful!")

Dependencies#

brainmass depends on:

Core:

  • jax ≥ 0.4.0: For numerical computations and automatic differentiation

  • brainstate ≥ 0.2.9: State management for dynamical systems

  • brainunit: Unit-aware array operations

  • braintools: Utility functions and transforms

  • numpy: NumPy array operations

Optional:

  • matplotlib: For visualization (examples)

  • nevergrad: Gradient-free optimization (parameter fitting)

  • scipy: SciPy optimizers (parameter fitting)

Install optional dependencies:

pip install brainmass[opt]  # Optimization libraries
pip install brainmass[viz]  # Visualization

Platform-Specific Notes#

macOS with Apple Silicon (M1/M2):

  • Use CPU installation

  • JAX has limited GPU support for Apple Silicon

  • Performance is still excellent on M1/M2 CPUs

Windows:

  • GPU support requires WSL2 for best compatibility

  • Native Windows GPU support is experimental

Linux:

  • Recommended platform for GPU/TPU

  • Most straightforward CUDA setup

Next Steps#

After installation:

  1. Follow the Quickstart tutorial for your first simulation

  2. Read the Choosing Models guide to select appropriate models

  3. Explore the Examples for practical applications

See Also#