Architecture Overview

This page explains how brainmass is organised and why. It is the conceptual companion to the API reference: rather than listing every class, it describes the layers of the package, the contract every model obeys, and the design philosophy that keeps the library small. Understanding these layers makes the rest of the documentation — and the package’s source — easy to navigate.

The single sentence to remember: brainmass ships only models and delegates everything else to the rest of the BrainX ecosystem.

“Models-only, delegate the rest”

brainmass deliberately does not reinvent state management, units, integrators, optimisers, or the compile/transform layer. It builds on three companion libraries and contributes the neural-mass-specific pieces on top:

Concern	Delegated to	brainmass uses it for
State, modules, transforms, autodiff	brainstate	Dynamics / Module base classes, HiddenState / ParamState, Param with constraint Transforms, jit / grad / vmap / for_loop
Physical units	brainunit (u)	dimensional analysis on every quantity
Initializers, integrators, metrics, optimizers	braintools	braintools.init.*, braintools.quad.ode_*_step, braintools.metric.*, braintools.optim.*
Numerical backend	JAX	JIT compilation, GPU/TPU, automatic differentiation

What brainmass itself adds is the parts that are specific to neural mass modelling: the model equations, the coupling kernels, conduction delays, the forward/observation models, and a thin in-package orchestration layer (Simulator / Network / Fitter / objectives) so a user can run a whole study without hand-wiring the primitives.

The benefit of this discipline is a small, focused, auditable codebase whose improvements (e.g. a faster integrator in braintools, GPU support in JAX) are inherited for free.

The layers at a glance

brainmass is best understood as four cooperating layers. Activity flows down through them at simulation time; gradients flow back up at fitting time.

                ┌─────────────────────────────────────────────┐
   orchestration│  Simulator   Network   Fitter   objectives   │
                └─────────────────────────────────────────────┘
                                   │  drives / fits
                ┌─────────────────────────────────────────────┐
   models       │  *Step model contract  +  coupling kernels   │
                │  (Hopf, WilsonCowan, JansenRit, Montbrió…)    │
                └─────────────────────────────────────────────┘
                                   │  produces activity
                ┌─────────────────────────────────────────────┐
   observation  │  BOLD (Balloon / HRF)   EEG/MEG lead fields   │
                └─────────────────────────────────────────────┘
                                   │  built on
                ┌─────────────────────────────────────────────┐
   foundation   │  brainstate · brainunit · braintools · JAX   │
                └─────────────────────────────────────────────┘

The rest of this page walks each layer.

The model contract: the *Step convention

Every neural mass model in brainmass is a class named <Model>Step (e.g. HopfStep, WilsonCowanStep, MontbrioPazoRoxinStep) that subclasses brainstate.nn.Dynamics and follows one contract. The Step suffix is a reminder that the class implements one integration step of the dynamics; the Simulator repeats that step to produce a trajectory.

The contract has four parts:

1. __init__ wraps parameters through Param.init so they can be constants or trainable, and validates in_size (the number of regions / parallel units).

2. init_state(batch_size=None) allocates the model’s HiddenStates — the ODE variables — with their initial values.

3. Per-variable derivative methods (dx, dV, …) and an aggregator derivative(state, t, *inputs) define the right-hand side of the ODE.

4. update(*inputs) advances the state by one step (the default is an exponential-Euler step; alternatives like rk4 are available) and returns the observable(s).

Crucially, update returns the observable rather than the full state — for a Jansen–Rit column the EEG-like output is a derived quantity (a difference of two populations), and the Simulator records whatever update returns.

Let’s inspect a model to see the contract in the flesh.

import brainmass
import brainstate

node = brainmass.HopfStep(in_size=1, a=0.25, w=0.3)

# It is a brainstate Dynamics following the *Step contract.
print('class       :', type(node).__name__)
print('is Dynamics :', isinstance(node, brainstate.nn.Dynamics))
print('contract    :', [m for m in ('init_state', 'update', 'derivative')
                        if hasattr(node, m)])

# Hidden states are the ODE variables; they appear after init_state().
brainstate.nn.init_all_states(node)
states = node.states(brainstate.HiddenState)
print('state vars  :', sorted({k[0] for k in states.keys()}))

class       : HopfStep
is Dynamics : True
contract    : ['init_state', 'update', 'derivative']
state vars  : ['x', 'y']

Because every model shares this contract, the orchestration layer can drive any of them uniformly — and because the step is built from differentiable JAX ops, the whole rollout is differentiable.

Parameters: constants or trainables, in one wrapper

A model parameter is wrapped by brainstate’s Param. Passing a plain number gives a fixed constant; passing Param(value, t=Transform, fit=True) makes it a trainable ParamState whose value is stored in an unconstrained space (via the inverse transform) and read back constrained — so an optimiser can move freely while the parameter stays within physical bounds. The Fitter discovers exactly these trainable parameters via model.states(ParamState). brainmass does not reimplement any of this bounds handling; it consumes brainstate’s.

The orchestration layer

These four pieces turn the model contract into a usable workflow. They compose the existing brainstate idiom — they do not reimplement integration.

• Simulator(model, dt) runs the loop. run(duration, monitors=…, transient=…, sample_every=…, batch_size=…) sets dt in the environment, initialises states, repeats the step with brainstate.transform.for_loop (one compiled XLA program, not a Python loop), and returns a plain dict of stacked trajectories plus a unit-aware ts time axis. Returning a dict (a clean pytree) keeps the result differentiable through jit / grad / vmap.

• Network(node, conn=…, distance=…, speed=…, coupling=…, coupled_var=…) wires many regions into a delay-coupled whole-brain model: it zeroes the SC diagonal, turns distance / speed into conduction delays, prefetches the delayed source state, and feeds a coupling kernel. The coupling current becomes the node’s first input. See Coupling and Delays.

• Fitter(model, optimizer, …) fits parameters. Write the objective once and pick a backend: 'grad' (backprop), 'nevergrad', or 'scipy'. It finds trainable parameters automatically and returns a FitResult. See Why Differentiable?.

• objectives is a small library of loss builders — timeseries_rmse, fc_corr, fcd_wasserstein, combine — each returning a unit-aware, jit/grad/vmap-safe callable(prediction, target). They wrap braintools.metric without re-implementing the maths.

Here is the entire run loop in three lines — the idiom the whole package is built to make trivial.

import brainmass
import brainunit as u

node = brainmass.HopfStep(in_size=1, a=0.25, w=0.3)
res = brainmass.Simulator(node, dt=0.1 * u.ms).run(200 * u.ms, monitors=['x'])

print('result is a dict :', isinstance(res, dict))
print('keys             :', sorted(res.keys()))
print('x trajectory     :', res['x'].shape)
print('ts is unit-aware :', res['ts'].dim)   # carries time units

result is a dict : True
keys             : ['ts', 'x']
x trajectory     : (2000, 1)
ts is unit-aware : s

The observation / forward layer

Neural activity is not what an experiment measures. The observation layer maps activity to neuroimaging signals, and these maps are differentiable too:

• BOLD (fMRI). Two routes: BOLDSignal, the four-state Balloon–Windkessel ODE (biophysical realism), and HRFBold, a fast linear convolution with a closed-form hemodynamic response function (ideal for fitting). A family of HRF kernels and a TemporalAverage monitor round this out.

• EEG / MEG. LeadFieldModel (and the EEGLeadFieldModel / MEGLeadFieldModel specialisations) project region-level dipole currents to sensor space through a lead-field matrix, unit-correctly.

The theory of these maps — hemodynamics and lead fields — is the subject of From Activity to Signals. Architecturally, the point is that the forward model is just another differentiable module in the chain, so gradients flow from a signal-space loss all the way back to neural parameters.

Units everywhere

Every quantity in brainmass can carry a physical unit via brainunit. A dt is 0.1 * u.ms; a conduction speed is 10 * u.mm / u.ms; a Jansen–Rit output is in u.mV; a lead field carries V / (nA·m). Dimensional analysis is enforced at runtime, so an accidental mismatch — adding millivolts to a firing rate — raises immediately instead of silently producing nonsense.

The orchestration layer respects this: Simulator returns a unit-aware ts, and the viz helpers strip units internally for plotting (you can pass res['x'] straight to brainmass.viz.plot_timeseries). When you need a raw array, u.get_magnitude(q) strips the unit explicitly — note that np.asarray on a Quantity raises, by design, rather than silently dropping the unit. The hands-on guide is Working with Units.

import brainmass
import brainunit as u

node = brainmass.MontbrioPazoRoxinStep(in_size=1)
res = brainmass.Simulator(node, dt=0.01 * u.ms).run(50 * u.ms, monitors=['r'])

# The firing-rate observable comes back carrying its physical unit.
print('r unit :', u.get_unit(res['r']))
print('ts unit:', u.get_unit(res['ts']))

r unit : Hz
ts unit: ms

How the layers fit together

A complete whole-brain study threads the four layers in order:

1. Foundation provides states, units, integrators, autodiff.

2. Models define per-region dynamics (the *Step contract) and how regions couple.

3. Orchestration runs them (Simulator / Network) and fits them (Fitter / objectives).

4. Observation turns activity into BOLD / EEG / MEG to compare against data.

Because every layer is differentiable, the whole stack is differentiable: a loss in signal space backpropagates through observation, orchestration, and models down to any parameter. That end-to-end differentiability — the topic of Why Differentiable? — is what the architecture exists to deliver.

See also

• What Is a Neural Mass Model? — the models the contract describes.

• Why Differentiable? — why the stack is differentiable end-to-end.

• Coupling and Delays — the Network wiring layer in depth.

• From Activity to Signals — the observation layer in depth.

• API Reference — the full API reference.

• Developer Guide — building and extending models.

References

• The BrainX ecosystem: https://brainx.chaobrain.com/

• brainstate, brainunit, braintools documentation (linked from the ecosystem hub).