Architecture

Architecture#

Overview of brainmass package structure and design.

Package Structure#

brainmass/
├── __init__.py          # Public API exports
├── _fhn.py              # FitzHugh-Nagumo model
├── _hopf.py             # Hopf oscillator
├── _jansen_rit.py       # Jansen-Rit model
├── _wilson_cowan.py     # Wilson-Cowan model
├── _wong_wang.py        # Wong-Wang model
├── _vdp.py              # Van der Pol oscillator
├── _sl.py               # Stuart-Landau oscillator
├── _qif.py              # QIF model
├── _linear.py           # Threshold-linear model
├── _kuramoto.py         # Kuramoto network
├── _xy_model.py         # XY oscillator base class
├── _noise.py            # Noise processes
├── _coupling.py         # Coupling mechanisms
├── _forward_model.py    # BOLD, EEG, MEG models
├── _horn.py             # HORN models
├── _utils.py            # Utility functions
└── _typing.py           # Type aliases

Design Principles#

  1. Consistency: All models follow same API pattern

  2. Modularity: Components can be mixed and matched

  3. Unit Safety: Use brainunit for dimensional analysis

  4. JAX First: Designed for JIT compilation and autodiff

  5. Simplicity: Clear, readable code over cleverness

Dependencies#

Core:

  • brainstate: State management for dynamical systems

  • brainunit: Unit-aware computations

  • braintools: Transforms and utilities

  • jax: Numerical computing backend

Rationale:

  • Separates concerns (state, units, computation)

  • Reusable across BrainX ecosystem

  • Leverages JAX for performance

Model Base Classes#

All models inherit from brainstate.nn.Dynamics:

class MyModel(brainstate.nn.Dynamics):
    def __init__(self, in_size, **params):
        super().__init__()
        # Initialize parameters
        # Create state variables with HiddenState

    def init_state(self, batch_size=None):
        # Initialize state variables
        pass

    def update(self, **inputs):
        # Update dynamics by one time step
        # Return observable(s)
        pass

State Management#

Uses brainstate.HiddenState for internal variables:

self.x = brainstate.HiddenState(
    brainstate.init.Constant(0.),
    sharding=brainstate.ShardingMode.REPLICATED
)

Coupling Architecture#

Coupling as separate modules enables:

  • Reusable across models

  • Swappable coupling mechanisms

  • Clear separation of node vs edge dynamics

Forward Models#

Designed as transformations:

Neural Activity → Forward Model → Observable Signal
(hidden)         (biophysics)    (measured)

Optimization Support#

  • ArrayParam wraps parameters with transforms

  • Compatible with JAX autodiff

  • Works with gradient-free optimizers (Nevergrad)

Extension Points#

Extend brainmass by:

  1. New models: Inherit from Dynamics

  2. New noise: Inherit from noise base classes

  3. New coupling: Implement coupling interface

  4. New forward models: Follow BOLD/EEG pattern

Code Organization#

  • Private modules: _module.py (implementation)

  • Public API: __init__.py (exports)

  • Tests: tests/test_module.py

  • Examples: examples/###-name.ipynb

See Also#