Coupling

Introduction

Coupling defines how network nodes interact through structural connectivity. While dynamics govern local temporal evolution at each node, coupling transmits information between nodes, enabling collective behaviors like synchronization, wave propagation, and functional integration.

The general form of coupling input is:

\[ c_i = \sum_{j} w_{ij} \cdot f(\text{state}_j, \text{state}_i, \theta) \]

where $w_{ij}$ represents connection weights from the graph, and $f$ is the coupling function with parameters $\theta$.

Coupling in Network Dynamics

Coupling integrates with dynamics through named input channels. The dynamics() method receives coupling as a Bunch object with named attributes:

def dynamics(self, t, state, params, coupling, external):
    # Access coupling inputs by name
    c_instant = coupling.instant[0]
    c_delayed = coupling.delayed[0]

    # Use in dynamics equations
    dX_dt = ... + c_instant + c_delayed
    return derivatives

Networks support multiple coupling types through a named dictionary:

import jax.numpy as jnp
import matplotlib.pyplot as plt
from tvboptim.experimental.network_dynamics import Network, solve
from tvboptim.experimental.network_dynamics.dynamics import ReducedWongWang
from tvboptim.experimental.network_dynamics.coupling import LinearCoupling
from tvboptim.experimental.network_dynamics.graph import DenseGraph
from tvboptim.experimental.network_dynamics.solvers import Euler

# Create simple network
dynamics = ReducedWongWang()
graph = DenseGraph(jnp.array([[0.0, 1.0], [1.0, 0.0]]))
coupling = LinearCoupling(source='S', G=0.5)

network = Network(
    dynamics=dynamics,
    coupling={'instant': coupling},  # Named coupling
    graph=graph
)

print(network)

Network(
  dynamics=ReducedWongWang
  nodes=2
  coupling=['instant']
)

Key point: If a dynamics model declares coupling inputs that are not provided in the network, those inputs are automatically filled with zeros. This allows flexible network configurations without requiring all possible couplings.

The Coupling API

Coupling objects follow a three-phase lifecycle that mirrors the simulation flow:

Phase 1: Preparation (prepare)

Called once before simulation starts:

coupling_data, coupling_state = coupling.prepare(network, dt, t0, t1)

Returns:

coupling_data (Bunch): Prepared metadata
- State indices that adapt the dynamics state into transmitted signals
- A prepared graph-topology view shared across coupling channels
- Edge-aligned parameter values and live delay read indices from precompute()
coupling_state (Bunch): Mutable internal state
- History buffers (for delayed coupling)
- Empty for instantaneous coupling

Phase 2: Computation (compute)

Called every time step during integration:

coupling_input = coupling.compute(t, state, coupling_data,
                                   coupling_state, params, graph)

Returns: - Coupling input array [n_coupling_dims, n_nodes]

Phase 3: State Update (update_state)

Called after each integration step:

coupling_state = coupling.update_state(coupling_data,
                                       coupling_state, new_state)

Returns: - Updated coupling_state (e.g., updated history buffer)

Flow pseudocode:

# Before simulation
coupling_data, coupling_state = coupling.prepare(network, dt, t0, t1)

# During simulation (at each time step)
for step in range(n_steps):
    # Compute coupling input
    coupling_input = coupling.compute(t, state, coupling_data,
                                      coupling_state, params, graph)

    # Integrate dynamics (uses coupling_input)
    new_state = integrate_step(state, coupling_input, dt)

    # Update coupling state
    coupling_state = coupling.update_state(coupling_data,
                                           coupling_state, new_state)
    state = new_state

TVB-Style Couplings: Pre and Post Pattern

Most brain network models follow a standard coupling pattern: transform states before aggregation (pre), perform weighted summation, then transform the result (post). This pattern is captured in the InstantaneousCoupling and DelayedCoupling base classes.

The computation flow is:

pre(states) → weighted sum → post(sum)

LinearCoupling Example

The simplest coupling applies a gain and offset to the weighted sum:

from tvboptim.experimental.network_dynamics.coupling import LinearCoupling

# Create linear coupling
coupling = LinearCoupling(source='x', G=2.0, b=0.1)

print(f"Coupling parameters: G={coupling.params.G}, b={coupling.params.b}")
print(f"Incoming states: {coupling.SOURCE_STATE_NAMES}")

Coupling parameters: G=2.0, b=0.1
Incoming states: x

The post() implementation is simple:

def post(self, summed_inputs, local_states, params):
    return params.G * summed_inputs + params.b

Since pre() is not overridden, it returns incoming states unchanged (identity function).

DifferenceCoupling: Overriding Pre

When coupling depends on differences between nodes, we override pre():

from tvboptim.experimental.network_dynamics.coupling import DifferenceCoupling

# Difference coupling: couples based on (x_j - x_i)
diff_coupling = DifferenceCoupling(
    source='x',
    local='x',  # Need local state for difference
    G=1.0
)

print(f"Incoming states: {diff_coupling.SOURCE_STATE_NAMES}")
print(f"Local states: {diff_coupling.LOCAL_STATE_NAMES}")

Incoming states: x
Local states: x

The pre() method computes differences:

PRE_USES_LOCAL = True

def pre(self, incoming_states, local_states, params):
    # Both operands are already aligned to the same message axis.
    return incoming_states - local_states

This enables synchronization dynamics where coupling strength depends on state mismatch.

Node and Edge Message Paths

The framework selects the scalable transport representation from the coupling contract:

Incoming-only instantaneous coupling without edge parameters transforms one value per source node. Dense graphs reduce with messages @ weights.T; sparse graphs gather source values at stored edges and reduce by target.
Delays, target-local operands (PRE_USES_LOCAL = True), or declared EDGE_PARAMS require an aligned message per connection. Sparse graphs execute these as [channels, E], while dense graphs use [channels, N_target, N_source].

The pre() body is identical in every case: it performs only elementwise math on operands the framework has already aligned. Do not add [:, :, None] or [None, :, :] reshapes. Incoming-only pre() receives local=None; post() continues to receive node-shaped local states.

LinearCoupling now uses the optimized incoming-only path directly. FastLinearCoupling is a deprecated compatibility wrapper that maps its historical local_states=... spelling to LinearCoupling(source=...); new code should use LinearCoupling.

Custom `pre()` migration and edge parameters

Custom couplings must declare the operands that force edge messages:

class EdgeScaledCoupling(InstantaneousCoupling):
    N_OUTPUT_STATES = 1
    DEFAULT_PARAMS = Bunch(edge_gain=1.0)
    EDGE_PARAMS = ("edge_gain",)

    def pre(self, incoming_states, local_states, params):
        return incoming_states * params.edge_gain

Each name in EDGE_PARAMS accepts either a graph-shaped [N_target, N_source] array or a prepared-edge [E] array. For a large sparse graph, construct the scalable form through the public graph API:

edge_gain_e = sparse_graph.gather_edges(edge_gain_matrix)
coupling = EdgeScaledCoupling(source="S", edge_gain=edge_gain_e)

sparse_graph.edge_indices documents how an existing E-vector is interpreted: rows are (target, source) in the same order used by gather_edges(). That topology remains fixed for one prepared solve; numerical weights, delays, and edge parameters may be swept or differentiated, but changing or reordering indices requires graph reconstruction and another prepare() call.

Preparation validates PRE_USES_LOCAL, EDGE_PARAMS, and the exact output shape of pre() before JIT compilation. If an older custom coupling fails with the elementwise-contract migration message, remove explicit message-axis reshapes and declare the operands above.

Dense graphs evaluate an edge-path pre() at every matrix cell, including zero-weight cells; sparse graphs omit cells that are not stored. These are equivalent for finite elementwise math, but a custom transform that deliberately emits NaN or raises at a zero-weight cell can observe the difference.

Creating Custom Couplings: Adaptive Gain

Let’s implement an adaptive coupling where coupling strength adjusts based on local activity - a homeostatic mechanism common in biological networks that allows nodes to maintain independent dynamics despite strong connectivity.

from tvboptim.experimental.network_dynamics.core.bunch import Bunch
from tvboptim.experimental.network_dynamics.coupling.base import InstantaneousCoupling

class AdaptiveGainCoupling(InstantaneousCoupling):
    """Coupling with activity-dependent gain modulation.

    Coupling strength decreases with local activity:
    G_effective = G * (1 - alpha * |local_state|)

    This implements homeostatic regulation where highly active
    nodes reduce their coupling sensitivity.
    """

    N_OUTPUT_STATES = 1
    DEFAULT_PARAMS = Bunch(
        G=1.0,      # Base coupling strength
        alpha=0.5   # Adaptation strength (0 = no adaptation)
    )

    def post(self, summed_inputs, local_states, params):
        """Apply activity-dependent gain modulation."""
        # Measure local activity (amplitude of oscillations)
        activity = jnp.abs(local_states[0])

        # Effective gain decreases with activity
        G_effective = params.G * (1.0 - params.alpha * activity)

        return G_effective * summed_inputs

print("Adaptive coupling class defined!")

Adaptive coupling class defined!

Demonstrating Adaptive vs Fixed Coupling

We’ll use two coupled oscillators with different initial conditions. Fixed coupling forces rapid synchronization, while adaptive coupling allows them to maintain more independent oscillations.

from tvboptim.experimental.network_dynamics.dynamics.tvb import Generic2dOscillator

# Bidirectional coupling between two nodes
graph_osc = DenseGraph(jnp.array([[0.0, 1.0], [1.0, 0.0]]))

# Create oscillatory dynamics with different initial conditions per node
dynamics_fixed = Generic2dOscillator(
    a=2.0,    # Parameter controlling oscillation frequency
    b=-10.0,   # Damping
    c=0.0,
    d=0.02,
    tau=1.0,
    I=0.0,
    INITIAL_STATE=(1.0, 0.0)  # Node 0 starts here
)

dynamics_adaptive = Generic2dOscillator(
    a=2.0,
    b=-10.0,
    c=0.0,
    d=0.02,
    tau=1.0,
    I=0.0,
    INITIAL_STATE=(1.0, 0.0)
)

# Fixed coupling: strong constant gain
fixed_coupling = LinearCoupling(source='V', G=0.8)
network_fixed = Network(
    dynamics=dynamics_fixed,
    coupling={'instant': fixed_coupling},
    graph=graph_osc
)

# Adaptive coupling: same base gain but adapts
adaptive_coupling = AdaptiveGainCoupling(
    source='V',
    local='V',  # Need local state for adaptation
    G=0.8,
    alpha=0.9
)
network_adaptive = Network(
    dynamics=dynamics_adaptive,
    coupling={'instant': adaptive_coupling},
    graph=graph_osc
)

# Manually set different initial conditions for second node
initial_state_fixed = network_fixed.initial_state.at[:, 1].set(jnp.array([-1.0, 0.0]))
initial_state_adaptive = network_adaptive.initial_state.at[:, 1].set(jnp.array([-1.0, 0.0]))

# Update network history to use custom initial states
from tvboptim.experimental.network_dynamics.result import NativeSolution
network_fixed.update_history(NativeSolution(
    ts=jnp.array([0.0]),
    ys=initial_state_fixed[None, :, :],
    variable_names=tuple(network_fixed.dynamics.STATE_NAMES),
))
network_adaptive.update_history(NativeSolution(
    ts=jnp.array([0.0]),
    ys=initial_state_adaptive[None, :, :],
    variable_names=tuple(network_adaptive.dynamics.STATE_NAMES),
))

# Simulate both
result_fixed = solve(
    network_fixed, Euler(),
    t0=0.0, t1=500.0, dt=1.0
)

result_adaptive = solve(
    network_adaptive, Euler(),
    t0=0.0, t1=500.0, dt=1.0
)

print(f"Simulation complete: {len(result_fixed.ts)} time steps")

Simulation complete: 500 time steps

Visualization

fig, axes = plt.subplots(2, 2, figsize=(8.1, 4.63))

# Time series - Fixed coupling
axes[0, 0].plot(result_fixed.ts, result_fixed.ys[:, 0, 0], 'b-', label='Node 0', linewidth=2)
axes[0, 0].plot(result_fixed.ts, result_fixed.ys[:, 0, 1], 'r-', label='Node 1', linewidth=2, alpha=0.7)
axes[0, 0].set_ylabel('V (fast variable)')
axes[0, 0].set_title('Fixed Coupling (G=0.8): Rapid Synchronization')
axes[0, 0].legend(loc='upper right')
axes[0, 0].grid(True, alpha=0.3)
# axes[0, 0].set_xlim(0, 100)

# Time series - Adaptive coupling
axes[1, 0].plot(result_adaptive.ts, result_adaptive.ys[:, 0, 0], 'b-', label='Node 0', linewidth=2)
axes[1, 0].plot(result_adaptive.ts, result_adaptive.ys[:, 0, 1], 'r-', label='Node 1', linewidth=2, alpha=0.7)
axes[1, 0].set_xlabel('Time')
axes[1, 0].set_ylabel('V (fast variable)')
axes[1, 0].set_title('Adaptive Coupling (G=0.8, α=0.5): Maintained Independence')
axes[1, 0].legend(loc='upper right')
axes[1, 0].grid(True, alpha=0.3)
# axes[1, 0].set_xlim(0, 100)

# Phase portraits - Fixed coupling
axes[0, 1].plot(result_fixed.ys[:, 0, 0], result_fixed.ys[:, 1, 0],
                'b-', linewidth=1, alpha=0.6, label='Node 0')
axes[0, 1].plot(result_fixed.ys[:, 0, 1], result_fixed.ys[:, 1, 1],
                'r-', linewidth=1, alpha=0.6, label='Node 1')
axes[0, 1].set_xlabel('V (fast variable)')
axes[0, 1].set_ylabel('W (slow variable)')
axes[0, 1].set_title('Phase Portrait: Fixed Coupling')
axes[0, 1].legend()
axes[0, 1].grid(True, alpha=0.3)

# Phase portraits - Adaptive coupling
axes[1, 1].plot(result_adaptive.ys[:, 0, 0], result_adaptive.ys[:, 1, 0],
                'b-', linewidth=1, alpha=0.6, label='Node 0')
axes[1, 1].plot(result_adaptive.ys[:, 0, 1], result_adaptive.ys[:, 1, 1],
                'r-', linewidth=1, alpha=0.6, label='Node 1')
axes[1, 1].set_xlabel('V (fast variable)')
axes[1, 1].set_ylabel('W (slow variable)')
axes[1, 1].set_title('Phase Portrait: Adaptive Coupling')
axes[1, 1].legend()
axes[1, 1].grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

The key difference:

Fixed coupling: Nodes quickly synchronize and oscillate in phase (overlapping trajectories)
Adaptive coupling: Nodes maintain more independent oscillations with persistent phase differences

When oscillation amplitude increases, adaptive coupling reduces its strength, preventing complete synchronization. This demonstrates how local_states in post() enables state-dependent coupling mechanisms that can maintain network diversity despite strong structural connectivity.

Delays

Brain networks exhibit transmission delays due to finite axonal conduction speeds. The DelayedCoupling base class handles delayed interactions automatically.

Delay Computation

Delays are computed from tract lengths and conduction speed:

\[ \tau_{ij} = \frac{\text{length}_{ij}}{\text{speed}} \]

The default conduction speed is 3 m/s (typical for unmyelinated axons). Tract lengths are stored in the graph’s delay matrix.

Implementation Details

Delayed coupling maintains a history buffer - a circular buffer storing past states. During compute(), the appropriate delayed states are extracted based on connection-specific delays.

Computational cost: Delayed coupling requires more memory access for indexing into the history buffer at different time offsets per connection. This is more expensive than instantaneous coupling but essential for biological realism.

Example

from tvboptim.experimental.network_dynamics.coupling import DelayedLinearCoupling
from tvboptim.experimental.network_dynamics.graph import DenseDelayGraph
from tvboptim.experimental.network_dynamics.dynamics import Lorenz

# Create graph with delays (convert tract lengths to delays at 3 m/s)
weights_small = jnp.array([[0.0, 1.0], [1.0, 0.0]])
lengths_small = jnp.array([[0.0, 30.0], [30.0, 0.0]])  # 30 mm tract length
delays_small = lengths_small / 3.0

graph_delayed = DenseDelayGraph(weights_small, delays_small)

# Delayed coupling
delayed_coupling = DelayedLinearCoupling(source='x', G=0.5)

# Create network
network_delayed = Network(
    dynamics=Lorenz(),
    coupling=delayed_coupling,
    graph=graph_delayed
)

print(f"Tract length: {lengths_small[0, 1]} mm")
print(f"Conduction speed: 3 m/s (default)")
print(f"Delay: {lengths_small[0, 1] / 3:.1f} ms")

Tract length: 30.0 mm
Conduction speed: 3 m/s (default)
Delay: 10.0 ms

Delayed coupling requires DelayGraph (or SparseDelayGraph) which stores both weights and delays. See the Graph section for details on delay graph construction.

Multiple Couplings

Networks can have multiple coupling mechanisms simultaneously. Dynamics models declare which coupling inputs they expect via COUPLING_INPUTS:

class MyDynamics(AbstractDynamics):
    COUPLING_INPUTS = {
        'instant': 1,   # Expects 1D instantaneous coupling
        'delayed': 1    # Expects 1D delayed coupling
    }

    def dynamics(self, t, state, params, coupling, external):
        c_instant = coupling.instant[0]
        c_delayed = coupling.delayed[0]
        # Use both in dynamics
        ...

When creating a network, provide couplings as a named dictionary:

network = Network(
    dynamics=dynamics,
    coupling={
        'instant': LinearCoupling(source='S', G=1.0),
        'delayed': DelayedLinearCoupling(source='S', G=0.5)
    },
    graph=graph  # Must be DelayGraph for delayed coupling
)

Each input name accepts exactly one coupling. Ordinary Network does not add multiple couplings into one input; define separate COUPLING_INPUTS when the dynamics must receive separate mechanisms. Duplicate keys in a Python dictionary are discarded before Network can validate them.

Key behavior: If a dynamics model declares coupling inputs that are not provided, those inputs are automatically filled with zeros. This allows flexible configurations:

# Model declares both instant and delayed, but only provide instant
network = Network(
    dynamics=dynamics,
    coupling={'instant': coupling_instant},  # 'delayed' will be zeros
    graph=graph
)

This zero-filling enables gradual model building and testing without requiring all coupling types upfront.

Summary

See the API Reference for complete parameter descriptions and available coupling types.

Coupling enables inter-region communication in brain networks:

Lifecycle: Three-phase pattern (prepare → compute → update_state) mirrors simulation flow
TVB pattern: Most couplings use pre() → weighted sum → post() pattern; typically only post() needs implementation
Performance: Incoming-only transforms use node messages; delays, target-local math, and edge parameters use aligned edge messages (O(E) on sparse graphs)
Delays: Computed from tract lengths and conduction speed; managed automatically via history buffers with increased computational cost
Multiple couplings: Named dictionary with automatic zero-filling for missing inputs
Custom couplings: Subclass InstantaneousCoupling or DelayedCoupling and override pre()/post()

Coupling integrates with Dynamics through named inputs and with Graph through connectivity structure. Together they define collective network behavior.

--- title: "Coupling" format: html: code-fold: false toc: true toc-depth: 3 fig-width: 8 out-width: "100%" jupyter: python3 execute: cache: true --- # Introduction Coupling defines how network nodes interact through structural connectivity. While dynamics govern local temporal evolution at each node, coupling transmits information between nodes, enabling collective behaviors like synchronization, wave propagation, and functional integration. The general form of coupling input is: $$ c_i = \sum_{j} w_{ij} \cdot f(\text{state}_j, \text{state}_i, \theta) $$ where $w_{ij}$ represents connection weights from the graph, and $f$ is the coupling function with parameters $\theta$. # Coupling in Network Dynamics Coupling integrates with dynamics through named input channels. The `dynamics()` method receives coupling as a `Bunch` object with named attributes: ```python def dynamics(self, t, state, params, coupling, external): # Access coupling inputs by name c_instant = coupling.instant[0] c_delayed = coupling.delayed[0] # Use in dynamics equations dX_dt = ... + c_instant + c_delayed return derivatives ``` Networks support multiple coupling types through a named dictionary: ```{python} import jax.numpy as jnp import matplotlib.pyplot as plt from tvboptim.experimental.network_dynamics import Network, solve from tvboptim.experimental.network_dynamics.dynamics import ReducedWongWang from tvboptim.experimental.network_dynamics.coupling import LinearCoupling from tvboptim.experimental.network_dynamics.graph import DenseGraph from tvboptim.experimental.network_dynamics.solvers import Euler # Create simple network dynamics = ReducedWongWang() graph = DenseGraph(jnp.array([[0.0, 1.0], [1.0, 0.0]])) coupling = LinearCoupling(source='S', G=0.5) network = Network( dynamics=dynamics, coupling={'instant': coupling}, # Named coupling graph=graph ) print(network) ``` **Key point**: If a dynamics model declares coupling inputs that are not provided in the network, those inputs are automatically filled with zeros. This allows flexible network configurations without requiring all possible couplings. # The Coupling API Coupling objects follow a three-phase lifecycle that mirrors the simulation flow: ## Phase 1: Preparation (*prepare*) Called once before simulation starts: ```python coupling_data, coupling_state = coupling.prepare(network, dt, t0, t1) ``` **Returns**: - `coupling_data` (Bunch): Prepared metadata - State indices that adapt the dynamics state into transmitted signals - A prepared graph-topology view shared across coupling channels - Edge-aligned parameter values and live delay read indices from `precompute()` - `coupling_state` (Bunch): Mutable internal state - History buffers (for delayed coupling) - Empty for instantaneous coupling ## Phase 2: Computation (*compute*) Called every time step during integration: ```python coupling_input = coupling.compute(t, state, coupling_data, coupling_state, params, graph) ``` **Returns**: - Coupling input array `[n_coupling_dims, n_nodes]` ## Phase 3: State Update (*update_state*) Called after each integration step: ```python coupling_state = coupling.update_state(coupling_data, coupling_state, new_state) ``` **Returns**: - Updated `coupling_state` (e.g., updated history buffer) **Flow pseudocode**: ```python # Before simulation coupling_data, coupling_state = coupling.prepare(network, dt, t0, t1) # During simulation (at each time step) for step in range(n_steps): # Compute coupling input coupling_input = coupling.compute(t, state, coupling_data, coupling_state, params, graph) # Integrate dynamics (uses coupling_input) new_state = integrate_step(state, coupling_input, dt) # Update coupling state coupling_state = coupling.update_state(coupling_data, coupling_state, new_state) state = new_state ``` # TVB-Style Couplings: Pre and Post Pattern Most brain network models follow a standard coupling pattern: transform states before aggregation (`pre`), perform weighted summation, then transform the result (`post`). This pattern is captured in the `InstantaneousCoupling` and `DelayedCoupling` base classes. The computation flow is: ``` pre(states) → weighted sum → post(sum) ``` ## LinearCoupling Example The simplest coupling applies a gain and offset to the weighted sum: ```{python} from tvboptim.experimental.network_dynamics.coupling import LinearCoupling # Create linear coupling coupling = LinearCoupling(source='x', G=2.0, b=0.1) print(f"Coupling parameters: G={coupling.params.G}, b={coupling.params.b}") print(f"Incoming states: {coupling.SOURCE_STATE_NAMES}") ``` The `post()` implementation is simple: ```python def post(self, summed_inputs, local_states, params): return params.G * summed_inputs + params.b ``` Since `pre()` is not overridden, it returns incoming states unchanged (identity function). ## DifferenceCoupling: Overriding Pre When coupling depends on differences between nodes, we override `pre()`: ```{python} from tvboptim.experimental.network_dynamics.coupling import DifferenceCoupling # Difference coupling: couples based on (x_j - x_i) diff_coupling = DifferenceCoupling( source='x', local='x', # Need local state for difference G=1.0 ) print(f"Incoming states: {diff_coupling.SOURCE_STATE_NAMES}") print(f"Local states: {diff_coupling.LOCAL_STATE_NAMES}") ``` The `pre()` method computes differences: ```python PRE_USES_LOCAL = True def pre(self, incoming_states, local_states, params): # Both operands are already aligned to the same message axis. return incoming_states - local_states ``` This enables synchronization dynamics where coupling strength depends on state mismatch. # Node and Edge Message Paths The framework selects the scalable transport representation from the coupling contract: - Incoming-only instantaneous coupling without edge parameters transforms one value per source node. Dense graphs reduce with `messages @ weights.T`; sparse graphs gather source values at stored edges and reduce by target. - Delays, target-local operands (`PRE_USES_LOCAL = True`), or declared `EDGE_PARAMS` require an aligned message per connection. Sparse graphs execute these as `[channels, E]`, while dense graphs use `[channels, N_target, N_source]`. The `pre()` body is identical in every case: it performs only elementwise math on operands the framework has already aligned. Do not add `[:, :, None]` or `[None, :, :]` reshapes. Incoming-only `pre()` receives `local=None`; `post()` continues to receive node-shaped local states. `LinearCoupling` now uses the optimized incoming-only path directly. `FastLinearCoupling` is a deprecated compatibility wrapper that maps its historical `local_states=...` spelling to `LinearCoupling(source=...)`; new code should use `LinearCoupling`. ## Custom `pre()` migration and edge parameters Custom couplings must declare the operands that force edge messages: ```python class EdgeScaledCoupling(InstantaneousCoupling): N_OUTPUT_STATES = 1 DEFAULT_PARAMS = Bunch(edge_gain=1.0) EDGE_PARAMS = ("edge_gain",) def pre(self, incoming_states, local_states, params): return incoming_states * params.edge_gain ``` Each name in `EDGE_PARAMS` accepts either a graph-shaped `[N_target, N_source]` array or a prepared-edge `[E]` array. For a large sparse graph, construct the scalable form through the public graph API: ```python edge_gain_e = sparse_graph.gather_edges(edge_gain_matrix) coupling = EdgeScaledCoupling(source="S", edge_gain=edge_gain_e) ``` `sparse_graph.edge_indices` documents how an existing E-vector is interpreted: rows are `(target, source)` in the same order used by `gather_edges()`. That topology remains fixed for one prepared solve; numerical weights, delays, and edge parameters may be swept or differentiated, but changing or reordering indices requires graph reconstruction and another `prepare()` call. Preparation validates `PRE_USES_LOCAL`, `EDGE_PARAMS`, and the exact output shape of `pre()` before JIT compilation. If an older custom coupling fails with the elementwise-contract migration message, remove explicit message-axis reshapes and declare the operands above. Dense graphs evaluate an edge-path `pre()` at every matrix cell, including zero-weight cells; sparse graphs omit cells that are not stored. These are equivalent for finite elementwise math, but a custom transform that deliberately emits NaN or raises at a zero-weight cell can observe the difference. # Creating Custom Couplings: Adaptive Gain Let's implement an adaptive coupling where coupling strength adjusts based on local activity - a homeostatic mechanism common in biological networks that allows nodes to maintain independent dynamics despite strong connectivity. ```{python} from tvboptim.experimental.network_dynamics.core.bunch import Bunch from tvboptim.experimental.network_dynamics.coupling.base import InstantaneousCoupling class AdaptiveGainCoupling(InstantaneousCoupling): """Coupling with activity-dependent gain modulation. Coupling strength decreases with local activity: G_effective = G * (1 - alpha * |local_state|) This implements homeostatic regulation where highly active nodes reduce their coupling sensitivity. """ N_OUTPUT_STATES = 1 DEFAULT_PARAMS = Bunch( G=1.0, # Base coupling strength alpha=0.5 # Adaptation strength (0 = no adaptation) ) def post(self, summed_inputs, local_states, params): """Apply activity-dependent gain modulation.""" # Measure local activity (amplitude of oscillations) activity = jnp.abs(local_states[0]) # Effective gain decreases with activity G_effective = params.G * (1.0 - params.alpha * activity) return G_effective * summed_inputs print("Adaptive coupling class defined!") ``` ## Demonstrating Adaptive vs Fixed Coupling We'll use two coupled oscillators with different initial conditions. Fixed coupling forces rapid synchronization, while adaptive coupling allows them to maintain more independent oscillations. ```{python} from tvboptim.experimental.network_dynamics.dynamics.tvb import Generic2dOscillator # Bidirectional coupling between two nodes graph_osc = DenseGraph(jnp.array([[0.0, 1.0], [1.0, 0.0]])) # Create oscillatory dynamics with different initial conditions per node dynamics_fixed = Generic2dOscillator( a=2.0, # Parameter controlling oscillation frequency b=-10.0, # Damping c=0.0, d=0.02, tau=1.0, I=0.0, INITIAL_STATE=(1.0, 0.0) # Node 0 starts here ) dynamics_adaptive = Generic2dOscillator( a=2.0, b=-10.0, c=0.0, d=0.02, tau=1.0, I=0.0, INITIAL_STATE=(1.0, 0.0) ) # Fixed coupling: strong constant gain fixed_coupling = LinearCoupling(source='V', G=0.8) network_fixed = Network( dynamics=dynamics_fixed, coupling={'instant': fixed_coupling}, graph=graph_osc ) # Adaptive coupling: same base gain but adapts adaptive_coupling = AdaptiveGainCoupling( source='V', local='V', # Need local state for adaptation G=0.8, alpha=0.9 ) network_adaptive = Network( dynamics=dynamics_adaptive, coupling={'instant': adaptive_coupling}, graph=graph_osc ) # Manually set different initial conditions for second node initial_state_fixed = network_fixed.initial_state.at[:, 1].set(jnp.array([-1.0, 0.0])) initial_state_adaptive = network_adaptive.initial_state.at[:, 1].set(jnp.array([-1.0, 0.0])) # Update network history to use custom initial states from tvboptim.experimental.network_dynamics.result import NativeSolution network_fixed.update_history(NativeSolution( ts=jnp.array([0.0]), ys=initial_state_fixed[None, :, :], variable_names=tuple(network_fixed.dynamics.STATE_NAMES), )) network_adaptive.update_history(NativeSolution( ts=jnp.array([0.0]), ys=initial_state_adaptive[None, :, :], variable_names=tuple(network_adaptive.dynamics.STATE_NAMES), )) # Simulate both result_fixed = solve( network_fixed, Euler(), t0=0.0, t1=500.0, dt=1.0 ) result_adaptive = solve( network_adaptive, Euler(), t0=0.0, t1=500.0, dt=1.0 ) print(f"Simulation complete: {len(result_fixed.ts)} time steps") ``` ```{python} #| code-fold: true #| code-summary: "Visualization" fig, axes = plt.subplots(2, 2, figsize=(8.1, 4.63)) # Time series - Fixed coupling axes[0, 0].plot(result_fixed.ts, result_fixed.ys[:, 0, 0], 'b-', label='Node 0', linewidth=2) axes[0, 0].plot(result_fixed.ts, result_fixed.ys[:, 0, 1], 'r-', label='Node 1', linewidth=2, alpha=0.7) axes[0, 0].set_ylabel('V (fast variable)') axes[0, 0].set_title('Fixed Coupling (G=0.8): Rapid Synchronization') axes[0, 0].legend(loc='upper right') axes[0, 0].grid(True, alpha=0.3) # axes[0, 0].set_xlim(0, 100) # Time series - Adaptive coupling axes[1, 0].plot(result_adaptive.ts, result_adaptive.ys[:, 0, 0], 'b-', label='Node 0', linewidth=2) axes[1, 0].plot(result_adaptive.ts, result_adaptive.ys[:, 0, 1], 'r-', label='Node 1', linewidth=2, alpha=0.7) axes[1, 0].set_xlabel('Time') axes[1, 0].set_ylabel('V (fast variable)') axes[1, 0].set_title('Adaptive Coupling (G=0.8, α=0.5): Maintained Independence') axes[1, 0].legend(loc='upper right') axes[1, 0].grid(True, alpha=0.3) # axes[1, 0].set_xlim(0, 100) # Phase portraits - Fixed coupling axes[0, 1].plot(result_fixed.ys[:, 0, 0], result_fixed.ys[:, 1, 0], 'b-', linewidth=1, alpha=0.6, label='Node 0') axes[0, 1].plot(result_fixed.ys[:, 0, 1], result_fixed.ys[:, 1, 1], 'r-', linewidth=1, alpha=0.6, label='Node 1') axes[0, 1].set_xlabel('V (fast variable)') axes[0, 1].set_ylabel('W (slow variable)') axes[0, 1].set_title('Phase Portrait: Fixed Coupling') axes[0, 1].legend() axes[0, 1].grid(True, alpha=0.3) # Phase portraits - Adaptive coupling axes[1, 1].plot(result_adaptive.ys[:, 0, 0], result_adaptive.ys[:, 1, 0], 'b-', linewidth=1, alpha=0.6, label='Node 0') axes[1, 1].plot(result_adaptive.ys[:, 0, 1], result_adaptive.ys[:, 1, 1], 'r-', linewidth=1, alpha=0.6, label='Node 1') axes[1, 1].set_xlabel('V (fast variable)') axes[1, 1].set_ylabel('W (slow variable)') axes[1, 1].set_title('Phase Portrait: Adaptive Coupling') axes[1, 1].legend() axes[1, 1].grid(True, alpha=0.3) plt.tight_layout() plt.show() ``` The key difference: - **Fixed coupling**: Nodes quickly synchronize and oscillate in phase (overlapping trajectories) - **Adaptive coupling**: Nodes maintain more independent oscillations with persistent phase differences When oscillation amplitude increases, adaptive coupling reduces its strength, preventing complete synchronization. This demonstrates how `local_states` in `post()` enables state-dependent coupling mechanisms that can maintain network diversity despite strong structural connectivity. # Delays Brain networks exhibit transmission delays due to finite axonal conduction speeds. The `DelayedCoupling` base class handles delayed interactions automatically. ## Delay Computation Delays are computed from tract lengths and conduction speed: $$ \tau_{ij} = \frac{\text{length}_{ij}}{\text{speed}} $$ The default conduction speed is **3 m/s** (typical for unmyelinated axons). Tract lengths are stored in the graph's delay matrix. ## Implementation Details Delayed coupling maintains a **history buffer** - a circular buffer storing past states. During `compute()`, the appropriate delayed states are extracted based on connection-specific delays. **Computational cost**: Delayed coupling requires more memory access for indexing into the history buffer at different time offsets per connection. This is more expensive than instantaneous coupling but essential for biological realism. ## Example ```{python} from tvboptim.experimental.network_dynamics.coupling import DelayedLinearCoupling from tvboptim.experimental.network_dynamics.graph import DenseDelayGraph from tvboptim.experimental.network_dynamics.dynamics import Lorenz # Create graph with delays (convert tract lengths to delays at 3 m/s) weights_small = jnp.array([[0.0, 1.0], [1.0, 0.0]]) lengths_small = jnp.array([[0.0, 30.0], [30.0, 0.0]]) # 30 mm tract length delays_small = lengths_small / 3.0 graph_delayed = DenseDelayGraph(weights_small, delays_small) # Delayed coupling delayed_coupling = DelayedLinearCoupling(source='x', G=0.5) # Create network network_delayed = Network( dynamics=Lorenz(), coupling=delayed_coupling, graph=graph_delayed ) print(f"Tract length: {lengths_small[0, 1]} mm") print(f"Conduction speed: 3 m/s (default)") print(f"Delay: {lengths_small[0, 1] / 3:.1f} ms") ``` Delayed coupling requires `DelayGraph` (or `SparseDelayGraph`) which stores both weights and delays. See the [Graph](graph.qmd) section for details on delay graph construction. # Multiple Couplings Networks can have multiple coupling mechanisms simultaneously. Dynamics models declare which coupling inputs they expect via `COUPLING_INPUTS`: ```python class MyDynamics(AbstractDynamics): COUPLING_INPUTS = { 'instant': 1, # Expects 1D instantaneous coupling 'delayed': 1 # Expects 1D delayed coupling } def dynamics(self, t, state, params, coupling, external): c_instant = coupling.instant[0] c_delayed = coupling.delayed[0] # Use both in dynamics ... ``` When creating a network, provide couplings as a named dictionary: ```python network = Network( dynamics=dynamics, coupling={ 'instant': LinearCoupling(source='S', G=1.0), 'delayed': DelayedLinearCoupling(source='S', G=0.5) }, graph=graph # Must be DelayGraph for delayed coupling ) ``` Each input name accepts exactly one coupling. Ordinary `Network` does not add multiple couplings into one input; define separate `COUPLING_INPUTS` when the dynamics must receive separate mechanisms. Duplicate keys in a Python dictionary are discarded before `Network` can validate them. **Key behavior**: If a dynamics model declares coupling inputs that are not provided, those inputs are **automatically filled with zeros**. This allows flexible configurations: ```python # Model declares both instant and delayed, but only provide instant network = Network( dynamics=dynamics, coupling={'instant': coupling_instant}, # 'delayed' will be zeros graph=graph ) ``` This zero-filling enables gradual model building and testing without requiring all coupling types upfront. # Summary See the [API Reference](../reference/index.qmd) for complete parameter descriptions and available coupling types. Coupling enables inter-region communication in brain networks: - **Lifecycle**: Three-phase pattern (`prepare` → `compute` → `update_state`) mirrors simulation flow - **TVB pattern**: Most couplings use `pre()` → weighted sum → `post()` pattern; typically only `post()` needs implementation - **Performance**: Incoming-only transforms use node messages; delays, target-local math, and edge parameters use aligned edge messages (O(E) on sparse graphs) - **Delays**: Computed from tract lengths and conduction speed; managed automatically via history buffers with increased computational cost - **Multiple couplings**: Named dictionary with automatic zero-filling for missing inputs - **Custom couplings**: Subclass `InstantaneousCoupling` or `DelayedCoupling` and override `pre()`/`post()` Coupling integrates with [Dynamics](dynamics.qmd) through named inputs and with [Graph](graph.qmd) through connectivity structure. Together they define collective network behavior.