{ "cells": [ { "cell_type": "code", "execution_count": 1, "id": "16818731", "metadata": { "execution": { "iopub.execute_input": "2026-06-19T07:56:37.218243Z", "iopub.status.busy": "2026-06-19T07:56:37.218017Z", "iopub.status.idle": "2026-06-19T07:56:42.463838Z", "shell.execute_reply": "2026-06-19T07:56:42.462972Z" }, "tags": [ "remove-cell" ] }, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.\n" ] } ], "source": [ "%matplotlib inline\n", "import brainmass\n", "import brainstate\n", "import brainunit as u\n", "import jax.numpy as jnp\n", "import numpy as np\n", "import matplotlib.pyplot as plt\n", "brainstate.random.seed(0)\n", "brainstate.environ.set(dt=0.1 * u.ms)" ] }, { "cell_type": "markdown", "id": "5edc4a6f", "metadata": {}, "source": [ "# Architecture Overview\n", "\n", "This page explains how ``brainmass`` is *organised* and why. It is the conceptual\n", "companion to the API reference: rather than listing every class, it describes the\n", "**layers** of the package, the contract every model obeys, and the design\n", "philosophy that keeps the library small. Understanding these layers makes the rest\n", "of the documentation — and the package's source — easy to navigate.\n", "\n", "The single sentence to remember: **brainmass ships only models and delegates\n", "everything else** to the rest of the BrainX ecosystem." ] }, { "cell_type": "markdown", "id": "3692becf", "metadata": {}, "source": [ "## \"Models-only, delegate the rest\"\n", "\n", "``brainmass`` deliberately does *not* reinvent state management, units,\n", "integrators, optimisers, or the compile/transform layer. It builds on three\n", "companion libraries and contributes the neural-mass-specific pieces on top:\n", "\n", "| Concern | Delegated to | brainmass uses it for |\n", "|---|---|---|\n", "| State, modules, transforms, autodiff | **brainstate** | ``Dynamics`` / ``Module`` base classes, ``HiddenState`` / ``ParamState``, ``Param`` with constraint ``Transform``s, ``jit`` / ``grad`` / ``vmap`` / ``for_loop`` |\n", "| Physical units | **brainunit** (``u``) | dimensional analysis on every quantity |\n", "| Initializers, integrators, metrics, optimizers | **braintools** | ``braintools.init.*``, ``braintools.quad.ode_*_step``, ``braintools.metric.*``, ``braintools.optim.*`` |\n", "| Numerical backend | **JAX** | JIT compilation, GPU/TPU, automatic differentiation |\n", "\n", "What ``brainmass`` itself adds is the parts that are *specific to neural mass\n", "modelling*: the model equations, the coupling kernels, conduction delays, the\n", "forward/observation models, and a thin in-package orchestration layer\n", "(``Simulator`` / ``Network`` / ``Fitter`` / ``objectives``) so a user can run a\n", "whole study without hand-wiring the primitives.\n", "\n", "The benefit of this discipline is a small, focused, auditable codebase whose\n", "improvements (e.g. a faster integrator in braintools, GPU support in JAX) are\n", "inherited for free." ] }, { "cell_type": "markdown", "id": "5ad71a05", "metadata": {}, "source": [ "## The layers at a glance\n", "\n", "``brainmass`` is best understood as four cooperating layers. Activity flows *down*\n", "through them at simulation time; gradients flow *back up* at fitting time.\n", "\n", "```\n", " ┌─────────────────────────────────────────────┐\n", " orchestration│ Simulator Network Fitter objectives │\n", " └─────────────────────────────────────────────┘\n", " │ drives / fits\n", " ┌─────────────────────────────────────────────┐\n", " models │ *Step model contract + coupling kernels │\n", " │ (Hopf, WilsonCowan, JansenRit, Montbrió…) │\n", " └─────────────────────────────────────────────┘\n", " │ produces activity\n", " ┌─────────────────────────────────────────────┐\n", " observation │ BOLD (Balloon / HRF) EEG/MEG lead fields │\n", " └─────────────────────────────────────────────┘\n", " │ built on\n", " ┌─────────────────────────────────────────────┐\n", " foundation │ brainstate · brainunit · braintools · JAX │\n", " └─────────────────────────────────────────────┘\n", "```\n", "\n", "The rest of this page walks each layer." ] }, { "cell_type": "markdown", "id": "1fbec217", "metadata": {}, "source": [ "## The model contract: the ``*Step`` convention\n", "\n", "Every neural mass model in ``brainmass`` is a class named ``Step`` (e.g.\n", "``HopfStep``, ``WilsonCowanStep``, ``MontbrioPazoRoxinStep``) that subclasses\n", "``brainstate.nn.Dynamics`` and follows one contract. The ``Step`` suffix is a\n", "reminder that the class implements **one integration step** of the dynamics; the\n", "``Simulator`` repeats that step to produce a trajectory.\n", "\n", "The contract has four parts:\n", "\n", "1. **``__init__``** wraps parameters through ``Param.init`` so they can be\n", " constants *or* trainable, and validates ``in_size`` (the number of regions /\n", " parallel units).\n", "2. **``init_state(batch_size=None)``** allocates the model's ``HiddenState``s — the\n", " ODE variables — with their initial values.\n", "3. **Per-variable derivative methods** (``dx``, ``dV``, …) and an aggregator\n", " ``derivative(state, t, *inputs)`` define the right-hand side of the ODE.\n", "4. **``update(*inputs)``** advances the state by one step (the default is an\n", " exponential-Euler step; alternatives like ``rk4`` are available) and returns\n", " the observable(s).\n", "\n", "Crucially, ``update`` *returns* the observable rather than the full state — for a\n", "Jansen–Rit column the EEG-like output is a *derived* quantity (a difference of two\n", "populations), and the ``Simulator`` records whatever ``update`` returns.\n", "\n", "Let's inspect a model to see the contract in the flesh." ] }, { "cell_type": "code", "execution_count": 2, "id": "0af9c9bc", "metadata": { "execution": { "iopub.execute_input": "2026-06-19T07:56:42.466660Z", "iopub.status.busy": "2026-06-19T07:56:42.466142Z", "iopub.status.idle": "2026-06-19T07:56:42.511735Z", "shell.execute_reply": "2026-06-19T07:56:42.510828Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "class : HopfStep\n", "is Dynamics : True\n", "contract : ['init_state', 'update', 'derivative']\n", "state vars : ['x', 'y']\n" ] } ], "source": [ "import brainmass\n", "import brainstate\n", "\n", "node = brainmass.HopfStep(in_size=1, a=0.25, w=0.3)\n", "\n", "# It is a brainstate Dynamics following the *Step contract.\n", "print('class :', type(node).__name__)\n", "print('is Dynamics :', isinstance(node, brainstate.nn.Dynamics))\n", "print('contract :', [m for m in ('init_state', 'update', 'derivative')\n", " if hasattr(node, m)])\n", "\n", "# Hidden states are the ODE variables; they appear after init_state().\n", "brainstate.nn.init_all_states(node)\n", "states = node.states(brainstate.HiddenState)\n", "print('state vars :', sorted({k[0] for k in states.keys()}))" ] }, { "cell_type": "markdown", "id": "6baa8ff7", "metadata": {}, "source": [ "Because every model shares this contract, the orchestration layer can drive *any*\n", "of them uniformly — and because the step is built from differentiable JAX ops, the\n", "whole rollout is differentiable.\n", "\n", "### Parameters: constants or trainables, in one wrapper\n", "\n", "A model parameter is wrapped by ``brainstate``'s ``Param``. Passing a plain number\n", "gives a fixed constant; passing ``Param(value, t=Transform, fit=True)`` makes it a\n", "**trainable** ``ParamState`` whose value is stored in an *unconstrained* space (via\n", "the inverse transform) and read back constrained — so an optimiser can move freely\n", "while the parameter stays within physical bounds. The ``Fitter`` discovers exactly\n", "these trainable parameters via ``model.states(ParamState)``. ``brainmass`` does not\n", "reimplement any of this bounds handling; it consumes brainstate's." ] }, { "cell_type": "markdown", "id": "d86fc551", "metadata": {}, "source": [ "## The orchestration layer\n", "\n", "These four pieces turn the model contract into a usable workflow. They *compose*\n", "the existing brainstate idiom — they do not reimplement integration.\n", "\n", "- **``Simulator(model, dt)``** runs the loop. ``run(duration, monitors=…,\n", " transient=…, sample_every=…, batch_size=…)`` sets ``dt`` in the environment,\n", " initialises states, repeats the step with ``brainstate.transform.for_loop`` (one\n", " compiled XLA program, not a Python loop), and returns a plain ``dict`` of stacked\n", " trajectories plus a unit-aware ``ts`` time axis. Returning a dict (a clean\n", " pytree) keeps the result differentiable through ``jit`` / ``grad`` / ``vmap``.\n", "- **``Network(node, conn=…, distance=…, speed=…, coupling=…, coupled_var=…)``**\n", " wires many regions into a delay-coupled whole-brain model: it zeroes the SC\n", " diagonal, turns ``distance / speed`` into conduction delays, prefetches the\n", " delayed source state, and feeds a coupling kernel. The coupling current becomes\n", " the node's first input. See {doc}`/concepts/coupling_and_delays`.\n", "- **``Fitter(model, optimizer, …)``** fits parameters. Write the objective once and\n", " pick a ``backend``: ``'grad'`` (backprop), ``'nevergrad'``, or ``'scipy'``. It\n", " finds trainable parameters automatically and returns a ``FitResult``. See\n", " {doc}`/concepts/why_differentiable`.\n", "- **``objectives``** is a small library of *loss builders* — ``timeseries_rmse``,\n", " ``fc_corr``, ``fcd_wasserstein``, ``combine`` — each returning a unit-aware,\n", " jit/grad/vmap-safe ``callable(prediction, target)``. They wrap\n", " ``braintools.metric`` without re-implementing the maths.\n", "\n", "Here is the entire run loop in three lines — the idiom the whole package is built\n", "to make trivial." ] }, { "cell_type": "code", "execution_count": 3, "id": "b970a0e6", "metadata": { "execution": { "iopub.execute_input": "2026-06-19T07:56:42.513793Z", "iopub.status.busy": "2026-06-19T07:56:42.513600Z", "iopub.status.idle": "2026-06-19T07:56:42.632736Z", "shell.execute_reply": "2026-06-19T07:56:42.632024Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "result is a dict : True\n", "keys : ['ts', 'x']\n", "x trajectory : (2000, 1)\n", "ts is unit-aware : s\n" ] } ], "source": [ "import brainmass\n", "import brainunit as u\n", "\n", "node = brainmass.HopfStep(in_size=1, a=0.25, w=0.3)\n", "res = brainmass.Simulator(node, dt=0.1 * u.ms).run(200 * u.ms, monitors=['x'])\n", "\n", "print('result is a dict :', isinstance(res, dict))\n", "print('keys :', sorted(res.keys()))\n", "print('x trajectory :', res['x'].shape)\n", "print('ts is unit-aware :', res['ts'].dim) # carries time units" ] }, { "cell_type": "markdown", "id": "0f2f3412", "metadata": {}, "source": [ "## The observation / forward layer\n", "\n", "Neural activity is not what an experiment measures. The **observation layer** maps\n", "activity to neuroimaging signals, and these maps are differentiable too:\n", "\n", "- **BOLD (fMRI).** Two routes: ``BOLDSignal``, the four-state Balloon–Windkessel\n", " ODE (biophysical realism), and ``HRFBold``, a fast linear convolution with a\n", " closed-form hemodynamic response function (ideal for fitting). A family of HRF\n", " kernels and a ``TemporalAverage`` monitor round this out.\n", "- **EEG / MEG.** ``LeadFieldModel`` (and the ``EEGLeadFieldModel`` /\n", " ``MEGLeadFieldModel`` specialisations) project region-level dipole currents to\n", " sensor space through a lead-field matrix, unit-correctly.\n", "\n", "The theory of these maps — hemodynamics and lead fields — is the subject of\n", "{doc}`/concepts/from_activity_to_signals`. Architecturally, the point is that the\n", "forward model is just another differentiable module in the chain, so gradients flow\n", "from a *signal-space* loss all the way back to neural parameters." ] }, { "cell_type": "markdown", "id": "92be22d9", "metadata": {}, "source": [ "## Units everywhere\n", "\n", "Every quantity in ``brainmass`` can carry a physical unit via ``brainunit``. A\n", "``dt`` is ``0.1 * u.ms``; a conduction speed is ``10 * u.mm / u.ms``; a Jansen–Rit\n", "output is in ``u.mV``; a lead field carries ``V / (nA·m)``. Dimensional analysis is\n", "enforced at runtime, so an accidental mismatch — adding millivolts to a firing rate\n", "— raises immediately instead of silently producing nonsense.\n", "\n", "The orchestration layer respects this: ``Simulator`` returns a unit-aware ``ts``,\n", "and the ``viz`` helpers strip units internally for plotting (you can pass\n", "``res['x']`` straight to ``brainmass.viz.plot_timeseries``). When you need a raw\n", "array, ``u.get_magnitude(q)`` strips the unit explicitly — note that\n", "``np.asarray`` on a ``Quantity`` *raises*, by design, rather than silently dropping\n", "the unit. The hands-on guide is {doc}`/howto/work_with_units`." ] }, { "cell_type": "code", "execution_count": 4, "id": "b9f262f1", "metadata": { "execution": { "iopub.execute_input": "2026-06-19T07:56:42.635807Z", "iopub.status.busy": "2026-06-19T07:56:42.635567Z", "iopub.status.idle": "2026-06-19T07:56:42.985490Z", "shell.execute_reply": "2026-06-19T07:56:42.983461Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "r unit : Hz\n", "ts unit: ms\n" ] } ], "source": [ "import brainmass\n", "import brainunit as u\n", "\n", "node = brainmass.MontbrioPazoRoxinStep(in_size=1)\n", "res = brainmass.Simulator(node, dt=0.01 * u.ms).run(50 * u.ms, monitors=['r'])\n", "\n", "# The firing-rate observable comes back carrying its physical unit.\n", "print('r unit :', u.get_unit(res['r']))\n", "print('ts unit:', u.get_unit(res['ts']))" ] }, { "cell_type": "markdown", "id": "b6153955", "metadata": {}, "source": [ "## How the layers fit together\n", "\n", "A complete whole-brain study threads the four layers in order:\n", "\n", "1. **Foundation** provides states, units, integrators, autodiff.\n", "2. **Models** define per-region dynamics (the ``*Step`` contract) and how regions\n", " couple.\n", "3. **Orchestration** runs them (``Simulator`` / ``Network``) and fits them\n", " (``Fitter`` / ``objectives``).\n", "4. **Observation** turns activity into BOLD / EEG / MEG to compare against data.\n", "\n", "Because every layer is differentiable, the *whole stack* is differentiable: a loss\n", "in signal space backpropagates through observation, orchestration, and models down\n", "to any parameter. That end-to-end differentiability — the topic of\n", "{doc}`/concepts/why_differentiable` — is what the architecture exists to deliver.\n", "\n", "## See also\n", "\n", "- {doc}`/concepts/what_is_a_neural_mass_model` — the models the contract describes.\n", "- {doc}`/concepts/why_differentiable` — why the stack is differentiable end-to-end.\n", "- {doc}`/concepts/coupling_and_delays` — the ``Network`` wiring layer in depth.\n", "- {doc}`/concepts/from_activity_to_signals` — the observation layer in depth.\n", "- {doc}`/reference/index` — the full API reference.\n", "- {doc}`/developer/index` — building and extending models.\n", "\n", "## References\n", "\n", "- The BrainX ecosystem: https://brainx.chaobrain.com/\n", "- brainstate, brainunit, braintools documentation (linked from the ecosystem hub)." ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.13.11" } }, "nbformat": 4, "nbformat_minor": 5 }