LeadFieldModel

LeadFieldModel#

class brainmass.LeadFieldModel(in_size, out_size, L, sensor_unit=Unit('mV'), dipole_unit=Unit('m * nA'), scale=None, noise_cov=None)#

Lead-field model.

The lead-field matrix is a mathematical representation that describes the relationship between the electrical activity generated by neural sources in the brain and the resulting electrical measurements observed on the scalp’s surface during EEG recordings.

A differentiable, unit-safe forward (lead-field) model that maps Neural Mass Model (NMM) region-level outputs to EEG/MEG sensor space:

\[\mathbf{y}(t) \;=\; \mathbf{s}(t) \, \mathbf{L} \;+\; \boldsymbol{\varepsilon}(t),\]

where

  • \(\mathbf{s}(t)\in\mathbb{R}^{R}\) is the equivalent current dipole moment (ECD) per region (typically in nA·m), obtained from NMM observables via a scale/biophysical mapping.

  • \(\mathbf{L}\in\mathbb{R}^{R\times M}\) is the (region-aggregated) lead field matrix (EEG unit: V/(nA·m), MEG unit: T/(nA·m)).

  • \(\mathbf{y}(t)\in\mathbb{R}^{M}\) are the sensor measurements (EEG in V, MEG in T).

  • \(\boldsymbol{\varepsilon}(t)\sim\mathcal{N}(\mathbf{0},\boldsymbol{\Sigma})\) is sensor noise.

This implementation is unit-aware by using brainunit:

  • If you pass unitless arrays, you can attach/declare units via constructor arguments.

  • If you pass Quantities (with units), the module will validate and convert to consistent units.

Notes on Units#

  • Consistency: The class internally converts to a consistent set of units so that \(\mathbf{L} [\text{sensor}/(\text{nA·m})] \times \mathbf{s} [\text{nA·m}] \Rightarrow \mathbf{y} [\text{sensor}]\).

  • Two common workflows:

    1. Unitless workflow: provide unitless L, scale, and inputs. Specify leadfield_unit, sensor_unit, dipole_unit in the constructor to attach units. The model ensures consistency.

    2. Quantity workflow: provide L (with proper units), and pass NMM observables/scale as brainunit Quantities. The model converts as needed and outputs a Quantity in sensor_unit.

  • Sensor unit: "V" for EEG, "T" for MEG

  • Dipole unit: "nA*m", ECD unit

Shape Conventions#

  • Let M = #sensors, R = #regions, T = #time steps.

  • L: (R, M).

  • nmm_obs_or_dipoles: (T, R).

  • Output y: (T, M).

Biophysical Guidance#

  • For EEG, \(\mathbf{L}\) is usually computed with a BEM/forward solver in units of V/(nA·m) under a specific reference (e.g., average reference).

  • For MEG, \(\mathbf{L}\) is typically in T/(nA·m), accounting for sensor type (magnetometer/gradiometer).

  • Mapping NMM observables to ECD:

    If your NMM provides (population) membrane potentials in mV, use a scale such as scale = alpha * u("nA*m")/u("mV") so that s = scale * V_m produces dipole moments. alpha encodes column geometry, synchrony, etc.

Examples

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

>>> # constants
>>> M, R, T = 64, 68, 2000

>>> # EEG example: lead field in V/(nA*m)
>>> L = jnp.ones((R, M)) * u.volt / (u.nA * u.meter)
>>> # Suppose NMM observable is in mV; choose scale to map mV -> nA*m
>>> scale = 0.1 * u.nA * u.meter / u.volt     # toy number
>>> model = LeadFieldModel(L=L, scale=scale, sensor_unit=u.volt, dipole_unit=u.nA * u.meter)
>>> nmm_obs = 50.0 * jnp.ones((T, R)) * u.mV # (T, R) in mV
>>> y = model(nmm_obs)  # Quantity with unit "V", shape (T, M)

>>> # MEG example: lead field in T/(nA*m)
>>> L_u = jnp.ones((R, M)) * u.tesla * (u.nA * u.meter)
>>> model_u = LeadFieldModel(L=L_u,  sensor_unit=u.volt, dipole_unit=u.nA*u.meter, scale=u.nA*u.meter/u.tesla)
>>> y2 = model_u(50.0 * jnp.ones((T, R)) * u.tesla)
param L:

Lead field matrix. Accepts either:

  • A brainunit.Quantity with unit sensor_unit / dipole_unit and shape (R, M); or

  • A unitless (R, M) array, in which case leadfield_unit must be provided.

type L:

Array | Quantity | Callable

param sensor_unit:

Unit of the sensor signals, e.g. "V" for EEG or "T" for MEG.

type sensor_unit:

Unit

param dipole_unit:

Unit of dipole moment (ECD), typically "nA*m".

type dipole_unit:

Unit

param scale:

Optional mapping from the NMM observable to dipole moment. This can be:

  • A unit-aware scalar Quantity, e.g. 10.0 * "nA*m" / u.mV, meaning: s = scale * nmm_obs will convert mV to nA·m;

  • A dimensionless float if your NMM output is already in dipole units.

If scale is not provided, the model expects inputs to update to already be in dipole_unit.

type scale:

float | Quantity | None

param noise_cov:

Optional sensor noise covariance ((M, M)). If provided, it can be a Quantity with unit sensor_unit^2, or a unitless array in which case its unit is assumed compatible with sensor_unit**2. Noise is sampled i.i.d. across time (same covariance at each step).

type noise_cov:

Quantity | Array | None

static build_fixed_orientation_weights(normals_V3, parcel_VxR, area_weights_V=None)[source]#

Build a fixed-orientation vertex->region weight matrix with per-vertex surface normals.

Parameters:

  • normals_V3 (Array) – Array of shape (V, 3) with unit surface normals (x, y, z) at each cortical vertex.

  • parcel_VxR (Array) – Array of shape (V, R) giving vertex-to-region assignment (one-hot or soft).

  • area_weights_V (Array | None) – Optional array of shape (V,) with vertex area (or voxel volume) weights. If None, uses ones.

Returns:

Array of shape (3V, R). For each vertex v and region r, the (x,y,z) blocks receive: area[v] * parcel[v,r] * normals[v,k].

Return type:

Array

Notes

This constructs a direction-weighted, area-aware aggregation matrix suitable for compressing a vertex-level lead field (M,3V) into region-level (M,R) with:

W = build_fixed_orientation_weights(normals_V3, parcel_VxR, area_weights_V)
L_region = L_vertex_3V @ W
static compress_vertex_leadfield(L_vertex_3V, W_vertex_to_region_3VxR)[source]#

Compress a vertex-level lead field to region-level:

\[\mathbf{L}_{\text{region}} \;=\; \mathbf{L}_{\text{vertex}} \, \mathbf{W},\]

where L_vertex_3V has shape (M, 3V) (x/y/z dipole components per vertex) and W_vertex_to_region_3VxR has shape (3V, R) (direction/area/parcel weights).

Unit behavior#

  • If either input is a Quantity, unit algebra is applied automatically.

  • Typically, L_vertex_3V is in sensor_unit/(nA·m) and W is dimensionless, yielding region-level sensor_unit/(nA·m).

returns:

Region-level lead field of shape (R, M) (Quantity if any input had units).

rtype:

Array | Quantity

Parameters:

  • L_vertex_3V (Array | Quantity)

  • W_vertex_to_region_3VxR (Array | Quantity)

Return type:

Array | Quantity

update(nmm_obs_or_dipoles)[source]#

Forward mapping from region signals to sensors.

Parameters:

nmm_obs_or_dipoles (Array | Quantity) – (R, T) region time series. If already_in_dipole_unit=False, this is the NMM observable (e.g., membrane potential in mV) which will be mapped to dipole moment via scale. If already_in_dipole_unit=True, it is interpreted directly as dipole moment (ECD) in dipole_unit. Accepts unitless arrays (units are attached via scale/constructor) or Quantity.

Returns:

Sensor-space time series Quantity of shape (M, T) with unit sensor_unit.

Return type:

Quantity

Notes

The mapping uses the instantaneous linear model:

\[\mathbf{Y} = \mathbf{S}\,\mathbf{L} + \mathbf{E}.\]

where \(\mathbf{S}\) stacks dipoles over time and \(\mathbf{E}\) is sampled per time step from \(\mathcal{N}(\mathbf{0}, \boldsymbol{\Sigma})\) if enabled.

Return type:

Any