Quickstart

Quickstart#

This 5-minute tutorial demonstrates the basic brainmass workflow: create a neural mass model, add noise, simulate dynamics, and visualize results.

Your First Simulation#

Let’s simulate a single brain region using the Hopf oscillator model:

import brainmass
import brainunit as u
import brainstate
import jax.numpy as jnp
import matplotlib.pyplot as plt

# 1. Create a Hopf oscillator (10 Hz oscillation)
model = brainmass.HopfOscillator(
    in_size=1,              # single region
    omega=2 * jnp.pi * 10 * u.Hz,  # 10 Hz frequency
    a=0.1,                  # bifurcation parameter (>0 for oscillations)
)

# 2. Initialize model state
model.init_all_states()

# 3. Run simulation for 1000 time steps
def step(i):
    return model.update()

time_series = brainstate.transform.for_loop(step, jnp.arange(1000))

# 4. Visualize
plt.figure(figsize=(10, 4))
plt.plot(time_series[:, 0])  # plot x coordinate
plt.xlabel('Time step')
plt.ylabel('Activity')
plt.title('Hopf Oscillator Dynamics')
plt.show()

What’s happening:

We create a 10 Hz oscillator (omega = 2π × 10)
a > 0 puts the oscillator in the limit cycle regime
for_loop efficiently runs 1000 simulation steps
The output is a (1000, 1) array of the oscillator’s x-coordinate

Adding Noise#

Real neural activity is noisy. Let’s add stochastic fluctuations:

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

# Create oscillator
model = brainmass.HopfOscillator(in_size=1, omega=2 * jnp.pi * 10 * u.Hz, a=0.1)

# Add Ornstein-Uhlenbeck noise
model.noise = brainmass.OUProcess(
    in_size=1,
    sigma=0.5 * u.Hz,  # noise amplitude
    tau=20. * u.ms,    # correlation time
)

# Initialize
model.init_all_states()

# Simulate with noise
time_series = brainstate.transform.for_loop(
    lambda i: model.update(),
    jnp.arange(1000)
)

The noise is automatically added during update(), creating more realistic fluctuations.

Multi-Region Networks#

Simulate multiple brain regions simultaneously:

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

N_regions = 10

# Create network of Wilson-Cowan models
network = brainmass.WilsonCowanModel(
    in_size=N_regions,
    tau_E=10. * u.ms,
    tau_I=20. * u.ms,
)

# Add noise to each population
network.noise_E = brainmass.OUProcess(
    in_size=N_regions,
    sigma=0.3 * u.Hz,
    tau=20. * u.ms
)

network.noise_I = brainmass.OUProcess(
    in_size=N_regions,
    sigma=0.2 * u.Hz,
    tau=30. * u.ms
)

# Initialize
network.init_all_states()

# Simulate
def step(i):
    return network.update(rE_inp=0.5, rI_inp=0.2)

exc_rates = brainstate.transform.for_loop(step, jnp.arange(2000))

# exc_rates has shape (2000, 10) - excitatory rates for 10 regions

Note: in_size=N_regions creates N independent copies of the model (no coupling yet).

Adding Coupling#

Connect regions with diffusive coupling:

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

N = 10

# Create connectivity matrix (random for demo)
W = jax.random.uniform(jax.random.PRNGKey(0), (N, N)) * 0.1
W = W.at[jnp.diag_indices(N)].set(0.)  # no self-connections

# Create models
nodes = brainmass.HopfOscillator(in_size=N, omega=10 * u.Hz)
coupling = brainmass.DiffusiveCoupling(conn=W, k=0.2)

# Initialize
nodes.init_all_states()
coupling.init_all_states()

# Network simulation with coupling
def network_step(i):
    x = nodes.x.value
    coupled_input = coupling(x, x)

    output = nodes.update()
    nodes.x.value += coupled_input  # add coupling to state

    return output

network_activity = brainstate.transform.for_loop(
    network_step,
    jnp.arange(1000)
)

Forward Modeling (BOLD Signal)#

Map neural activity to fMRI BOLD signal:

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

N_regions = 5

# Neural mass model
nmm = brainmass.WilsonCowanModel(in_size=N_regions)
nmm.init_all_states()

# BOLD hemodynamic model
bold_model = brainmass.BOLDSignal(in_size=N_regions)
bold_model.init_all_states()

# Simulate neural activity
def sim_neural(i):
    return nmm.update(rE_inp=0.5, rI_inp=0.2)

neural_activity = brainstate.transform.for_loop(
    sim_neural,
    jnp.arange(5000)  # 5000 time steps
)

# Generate BOLD signal from neural activity
def sim_bold(z):
    bold_model.update(z=z)
    return bold_model.bold()

bold_signal = brainstate.transform.for_loop(
    sim_bold,
    neural_activity
)

# bold_signal shape: (5000, 5)
# Downsample to TR ~ 2s if needed: bold_signal[::2000]

Common Patterns#

Pattern 1: Batched Simulations

Run multiple simulations in parallel:

batch_size = 32

model = brainmass.HopfOscillator(in_size=10, omega=10 * u.Hz)
model.init_all_states(batch_size=batch_size)

output = model.update()  # shape: (32, 10)

Pattern 2: Accessing Internal States

model = brainmass.WilsonCowanModel(in_size=5)
model.init_all_states()

model.update(rE_inp=0.5, rI_inp=0.2)

# Access internal state variables
exc_rate = model.rE.value  # excitatory firing rate
inh_rate = model.rI.value  # inhibitory firing rate

Pattern 3: Custom Time Steps

# Default dt is typically 1 ms
# To change, set during model creation or manually integrate

model = brainmass.HopfOscillator(in_size=1, omega=10 * u.Hz, dt=0.5 * u.ms)

Quick Reference#

Task	Code
Create model	`model = brainmass.ModelName(in_size=N, **params)`
Initialize state	`model.init_all_states(batch_size=None)`
Single step	`output = model.update(**inputs)`
Simulate loop	`brainstate.transform.for_loop(lambda i: model.update(), jnp.arange(T))`
Add noise	`model.noise = brainmass.OUProcess(in_size=N, sigma=..., tau=...)`
Access state	`value = model.state_var.value`
Reset state	`model.reset_state()`

Next Steps#

Now that you’ve run your first simulations:

Choosing Models - Learn which model to use for your application
Building Multi-Region Networks - Create realistic multi-region networks
Parameter Fitting - Optimize parameters to match empirical data
Examples - Explore complete examples

Tips for Beginners#

Start simple: Single region, no noise, no coupling
Visualize often: Plot time series to understand dynamics
Check units: Use brainunit quantities to avoid unit errors
Read docstrings: Use help(brainmass.ModelName) for parameter details
Use examples: The Examples are your best resource