Coupling Mechanisms

Coupling Mechanisms#

Coupling mechanisms connect neural mass models across brain regions, enabling the simulation of large-scale brain networks. brainmass provides both class-based and functional coupling APIs.

Overview#

Coupling maps activity from source nodes to target nodes via a structural connectivity matrix. brainmass provides both linear and nonlinear coupling forms, matching the canonical TVB coupling library:

  1. Diffusive Coupling: Proportional to the difference between connected nodes (TVB Difference)

  2. Additive / Linear Coupling: Weighted sum of neighbour inputs, with optional bias (TVB Linear)

  3. Sigmoidal Coupling: Logistic nonlinearity applied after the network sum (TVB Sigmoidal)

  4. Hyperbolic-Tangent Coupling: Symmetric tanh saturation after the sum (TVB HyperbolicTangent)

  5. Sigmoidal Jansen-Rit Coupling: Jansen-Rit firing-rate sigmoid applied before the sum (TVB SigmoidalJansenRit)

Note

brainmass names the global coupling strength k; TVB / tvboptim call it G. The two are identical: G ≡ k. New trainable scalars (slope, midpoint, cmin/cmax, r, b) go through Param.init and are constrainable / fittable like k.

Coupling Types#

Coupling Type

Mathematical Form

Description

Diffusive

\(k \sum_j W_{ij} (x_j - x_i)\)

Drives nodes toward neighbors’ states

Additive / Linear

\(k \sum_j W_{ij} x_j + b\)

Weighted sum of neighbor inputs (+ bias)

Sigmoidal

\(k\,\sigma\!\big(s\,(a\sum_j W_{ij} x_j + b - m)\big)\)

Logistic saturation after the sum

Hyperbolic Tangent

\(k\,\tanh\!\big(s\sum_j W_{ij} x_j\big)\)

Symmetric saturation after the sum

Sigmoidal Jansen-Rit

\(k \sum_j W_{ij}\,\sigma_{\mathrm{JR}}(x_j)\)

Firing-rate sigmoid before the sum

Mathematical Details#

Diffusive Coupling

For a network with \(N\) nodes and connectivity matrix \(\mathbf{W}\), the diffusive coupling input to node \(i\) is:

\[C_i = k \sum_{j=1}^{N} W_{ij} (x_j - x_i)\]

This can be rewritten using the graph Laplacian \(\mathbf{L}\):

\[\mathbf{C} = -k \mathbf{L} \mathbf{x}\]

where \(L_{ij} = \delta_{ij} \sum_k W_{ik} - W_{ij}\).

Additive / Linear Coupling

The additive coupling input is:

\[C_i = k \sum_{j=1}^{N} W_{ij} x_j + b = k (\mathbf{W}^T \mathbf{x})_i + b\]

where \(b\) is an optional offset (default \(0\), which reproduces the bias-free coupling exactly). This is TVB’s Linear coupling.

Nonlinear Couplings

Three nonlinear forms close parity with the TVB coupling library. The post-nonlinearity forms apply a saturating function to the network sum:

\[\begin{split}\text{Sigmoidal:}\quad C_i &= k\,\sigma\!\Big(s\,\big(a\textstyle\sum_j W_{ij} x_j + b - m\big)\Big), \qquad \sigma(z) = \frac{1}{1 + e^{-z}} \\[4pt] \text{Hyperbolic Tangent:}\quad C_i &= k\,\tanh\!\Big(s\textstyle\sum_j W_{ij} x_j\Big)\end{split}\]

with slope \(s\), midpoint \(m\), input scaling \(a\) and offset \(b\). The pre-nonlinearity Jansen-Rit form applies a firing-rate sigmoid to each source before summation:

\[C_i = k \sum_j W_{ij}\,\sigma_{\mathrm{JR}}(x_j), \qquad \sigma_{\mathrm{JR}}(x) = c_{\min} + \frac{c_{\max} - c_{\min}}{1 + e^{\,r\,(m - x)}}\]

where \(x_j\) is the presynaptic source (e.g. the Jansen-Rit \(y_1 - y_2\)), \(r\) the steepness and \(m\) the half-activation potential. Because the transcendental functions require dimensionless arguments, their input is reduced to its magnitude; the post-nonlinearity output then carries the units of \(k\), and the Jansen-Rit output the units of \(k\,W\).

Connectivity Matrix Conventions#

Shape and Indexing: - Connectivity matrix W has shape (N_out, N_in) or (N_regions, N_regions) - W[i, j] represents the connection strength from source node j to target node i - For symmetric networks (undirected graphs), W is symmetric - Row sums W[i, :].sum() give the total input weights to node i

Common Structures: - Structural connectivity: From DTI tractography (weighted by fiber density) - Functional connectivity: From fMRI correlations (correlation matrices) - Random networks: Erdős-Rényi, scale-free, small-world topologies

API Reference#

Class-Based Coupling#

DiffusiveCoupling

Diffusive coupling.

AdditiveCoupling

Additive (linear) coupling.

SigmoidalCoupling

Sigmoidal coupling (TVB Sigmoidal, post-nonlinearity).

HyperbolicTangentCoupling

Hyperbolic-tangent coupling (TVB HyperbolicTangent, post-nonlinearity).

SigmoidalJansenRitCoupling

Sigmoidal Jansen-Rit coupling (TVB SigmoidalJansenRit, pre-nonlinearity).

Example: Diffusive Coupling

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

N = 10  # number of regions

# Create structural connectivity
W = jnp.ones((N, N)) * 0.1  # uniform weak coupling
W = W.at[jnp.diag_indices(N)].set(0.)  # no self-connections

# Create coupling module
coupling = brainmass.DiffusiveCoupling(
    conn=W,
    k=0.5,  # global coupling strength
)
coupling.init_all_states()

# Create node dynamics
nodes = brainmass.HopfStep(in_size=N, w=0.3)
nodes.init_all_states()

# Simulation loop
def step(i):
    local_activity = nodes.update()
    coupled_input = coupling(local_activity, local_activity)
    # Add coupling to node state
    nodes.x.value += coupled_input
    return local_activity

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

Functional Coupling#

diffusive_coupling

Diffusive coupling kernel (function form).

additive_coupling

Additive (linear) coupling kernel (function form).

sigmoidal_coupling

Sigmoidal coupling kernel (function form, post-nonlinearity).

hyperbolic_tangent_coupling

Hyperbolic-tangent coupling kernel (function form, post-nonlinearity).

sigmoidal_jansen_rit_coupling

Sigmoidal Jansen-Rit coupling kernel (function form, pre-nonlinearity).

Functional APIs provide stateless coupling for imperative use:

# Functional diffusive coupling
coupled_input = brainmass.diffusive_coupling(
    source=source_activity,  # shape (..., N)
    target=target_activity,  # shape (..., N)
    conn=W,                  # shape (N, N)
    k=0.5,
)

# Functional additive coupling
coupled_input = brainmass.additive_coupling(
    source=source_activity,
    conn=W,
    k=0.5,
)

Connectivity Utilities#

laplacian_connectivity

Build graph Laplacian matrix from adjacency/connectivity matrix.

LaplacianConnParam

Graph Laplacian connectivity module.

Laplacian Matrix:

The graph Laplacian is commonly used in diffusive coupling:

W = ...  # connectivity matrix

# Compute unnormalized Laplacian
L = brainmass.laplacian_connectivity(W, normalize=False)
# L[i,i] = sum_j W[i,j], L[i,j] = -W[i,j] for i≠j

# Compute normalized Laplacian (symmetric normalization)
L_norm = brainmass.laplacian_connectivity(W, normalize=True)

Usage Patterns#

Basic Network Simulation#

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

# Parameters
N_regions = 90
coupling_strength = 0.1

# Load structural connectivity (example: random)
rng = jax.random.PRNGKey(0)
W = jax.random.uniform(rng, (N_regions, N_regions)) * 0.1
W = W.at[jnp.diag_indices(N_regions)].set(0.)

# Create network components
nodes = brainmass.WilsonCowanStep(in_size=N_regions)
coupling = brainmass.DiffusiveCoupling(conn=W, k=coupling_strength)

# Add noise
nodes.noise_E = brainmass.OUProcess(
    in_size=N_regions,
    sigma=0.5 * u.Hz,
    tau=20. * u.ms
)

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

# Simulate
def network_step(i):
    rE = nodes.rE.value
    rI = nodes.rI.value

    # Update nodes with coupling
    coupling_E = coupling(rE, rE)
    output = nodes.update(rE_inp=coupling_E, rI_inp=0.)

    return output

time_series = brainstate.transform.for_loop(
    network_step,
    jnp.arange(5000)
)

Heterogeneous Coupling#

Different regions can have different coupling parameters:

# Region-specific coupling strengths
k_per_region = jnp.array([0.1, 0.2, 0.3, ...])  # shape (N,)

# Apply in coupling
coupled = brainmass.diffusive_coupling(source, target, conn=W, k=1.0)
coupled = coupled * k_per_region[:, None]  # scale per-region

Multi-Modal Coupling#

Combine multiple coupling mechanisms:

# Diffusive coupling on excitatory population
coupling_E = brainmass.DiffusiveCoupling(conn=W_excitatory, k=0.2)

# Additive coupling on inhibitory population
coupling_I = brainmass.AdditiveCoupling(conn=W_inhibitory, k=0.1)

def step(i):
    rE = nodes.rE.value
    rI = nodes.rI.value

    coupled_E = coupling_E(rE, rE)
    coupled_I = coupling_I(rI)

    nodes.update(rE_inp=coupled_E, rI_inp=coupled_I)

Time-Delayed Coupling#

For long-range connections with transmission delays:

# Simple delay implementation with circular buffer
delay_steps = 5  # time steps
history_buffer = jnp.zeros((delay_steps, N_regions))

def step_with_delay(i, buffer):
    current_activity = nodes.update()

    # Get delayed activity
    delayed_activity = buffer[0]  # oldest in buffer

    # Apply coupling with delay
    coupled = coupling(delayed_activity, current_activity)
    nodes.x.value += coupled

    # Update buffer (shift and append)
    buffer = jnp.roll(buffer, shift=-1, axis=0)
    buffer = buffer.at[-1].set(current_activity)

    return current_activity, buffer

Unit Handling#

Coupling respects brainunit units:

# Activity in Hz
activity = jnp.array([10., 20., 30.]) * u.Hz

# Coupling returns Hz
coupled = brainmass.diffusive_coupling(
    source=activity,
    target=activity,
    conn=W,
    k=0.5,  # unitless coupling strength
)
# coupled has units Hz

Performance Tips#

  1. Pre-compute Laplacian: For diffusive coupling, pre-compute the Laplacian if conn is static:

    L = brainmass.laplacian_connectivity(W, normalize=False)
coupled = -k * (L @ source_activity)

  2. Sparse Connectivity: For large sparse networks, consider sparse matrix operations (JAX BCOO)

  3. Batched Simulations: All coupling functions support batched inputs:

    # Batch of 32 simulations
source = jnp.zeros((32, N_regions)) * u.Hz
coupled = coupling(source, source)  # shape (32, N_regions)

Common Issues#

Sign Convention: - Diffusive coupling typically has k > 0, producing C_i = k * sum(x_j - x_i) - Negative k inverts the coupling (anti-phase synchronization)

Normalization: - Connectivity matrices can be normalized by row sums, column sums, or total sum - Normalization affects the effective coupling strength

Self-Connections: - Diagonal elements W[i, i] represent self-connections - Often set to zero to avoid redundant self-coupling

References#

  • Jansen, B. H., & Rit, V. G. (1995). Electroencephalogram and visual evoked potential generation in a mathematical model of coupled cortical columns. Biological Cybernetics, 73(4), 357-366. (Sigmoidal Jansen-Rit coupling.)

  • Sanz-Leon, P., Knock, S. A., Spiegler, A., & Jirsa, V. K. (2015). Mathematical framework for large-scale brain network modeling in The Virtual Brain. NeuroImage, 111, 385-430. (TVB Sigmoidal / HyperbolicTangent / Linear / Difference couplings.)

See Also#