SubspaceCoupling

experimental.network_dynamics.coupling.SubspaceCoupling(
    inner_coupling,
    region_mapping,
    regional_graph,
    use_sparse=True,
    **kwargs,
)

Coupling on regional subspace: aggregate nodes → couple regions → distribute.

Performs three-stage computation: 1. Aggregate node states to regional states (default: mean) 2. Apply inner coupling on regional graph 3. Distribute regional results to nodes (default: broadcast)

Parameters

Name	Type	Description	Default
inner_coupling	AbstractCoupling	Coupling applied at regional level	required
region_mapping	(ndarray, shape(n_nodes))	Maps each node to region ID (0 to n_regions-1)	required
regional_graph	AbstractGraph	Regional connectivity (n_nodes must equal n_regions)	required
use_sparse	bool	Use sparse BCOO format for aggregation (memory efficient)	True
**kwargs		Additional parameters for custom subclasses	{}

Attributes

Name	Type	Description
n_regions	int	Number of regions from region_mapping
N_OUTPUT_STATES	int	Output coupling dimensions from inner_coupling

Notes

Customization: Override aggregate() or distribute() for custom aggregation/ distribution strategies beyond mean/broadcast.

Examples

Surface-level network with regional delayed coupling:

>>> region_mapping = jnp.array([...])  # [n_vertices] -> region IDs
>>> regional_graph = DenseDelayGraph.from_weights_delays(SC, delays)
>>>
>>> coupling = SubspaceCoupling(
...     inner_coupling=DelayedLinearCoupling(source='S', G=0.5),
...     region_mapping=region_mapping,
...     regional_graph=regional_graph
... )

Methods

Name	Description
aggregate	Aggregate node states to regional states (default: mean).
compute	Compute coupling using cached aggregated state.
describe	Generate description for network printer.
distribute	Distribute regional coupling to nodes (default: broadcast).
precompute	Forward precompute() to the inner coupling on the regional graph.
prepare	Prepare aggregation matrices, regional context, and inner coupling.
update_state	Update inner coupling state and cache aggregated state.

aggregate

experimental.network_dynamics.coupling.SubspaceCoupling.aggregate(
    node_state,
    coupling_data,
)

Aggregate node states to regional states (default: mean).

Override for custom aggregation strategies (sum, weighted, etc).

Parameters

Name	Type	Description	Default
node_state	(ndarray, shape(n_states, n_nodes))	Node-level states	required
coupling_data	Bunch	Contains region_one_hot_normalized for aggregation	required

Returns

Name	Type	Description
	(ndarray, shape(n_states, n_regions))	Regional states

Notes

Default uses normalized one-hot matrix: node_state @ region_one_hot_normalized Supports both dense and sparse (BCOO) matrices.

compute

experimental.network_dynamics.coupling.SubspaceCoupling.compute(
    t,
    state,
    coupling_data,
    coupling_state,
    params,
    graph,
)

Compute coupling using cached aggregated state.

Parameters

Name	Type	Description	Default
t	float	Current time	required
state	(ndarray, shape(n_states, n_nodes))	Not used - cached state from previous update_state() avoids aggregation	required
coupling_data	Bunch	Static regional structures	required
coupling_state	Bunch	Contains cached_regional_state, inner_state	required
params	Bunch	Not used - inner coupling has own params	required
graph	AbstractGraph	Not used - regional_graph used instead	required

Returns

Name	Type	Description
	(ndarray, shape(n_coupling_inputs, n_nodes))	Node-level coupling input

describe

experimental.network_dynamics.coupling.SubspaceCoupling.describe()

Generate description for network printer.

Returns

Name	Type	Description
	dict	Coupling metadata with regional structure info

distribute

experimental.network_dynamics.coupling.SubspaceCoupling.distribute(
    regional_coupling,
    coupling_data,
)

Distribute regional coupling to nodes (default: broadcast).

Override for custom distribution strategies (scaled, weighted, etc).

Parameters

Name	Type	Description	Default
regional_coupling	(ndarray, shape(n_coupling_inputs, n_regions))	Regional coupling values	required
coupling_data	Bunch	Contains region_mapping for distribution	required

Returns

Name	Type	Description
	(ndarray, shape(n_coupling_inputs, n_nodes))	Node-level coupling

Notes

Default broadcasts: each node receives its region’s value via indexing.

precompute

experimental.network_dynamics.coupling.SubspaceCoupling.precompute(
    coupling_data,
    params,
    graph,
)

Forward precompute() to the inner coupling on the regional graph.

compute() hands coupling_data.inner_data straight to inner_coupling.compute(), so it must be the enriched data the inner coupling’s own precompute() produces, not the raw prepare() output. A DelayedCoupling inner coupling resolves its delay read indices there, against self.regional_graph rather than the node graph passed in here.

The solver’s stage-time fields are written into the outer coupling_data by Network.prepare(), which never sees inner_data (it comes from a nested prepare() call on the regional network). Forward them so the inner coupling’s delayed read gets the same stage-time shift the node-level couplings get.

prepare

experimental.network_dynamics.coupling.SubspaceCoupling.prepare(
    network,
    dt,
    t0,
    t1,
)

Prepare aggregation matrices, regional context, and inner coupling.

Parameters

Name	Type	Description	Default
network	Network	Node-level network instance	required
dt	float	Integration timestep	required
t0	float	Simulation time window	required
t1	float	Simulation time window	required

Returns

Name	Type	Description
coupling_data	Bunch	Static data with region_one_hot_normalized, inner_data
coupling_state	Bunch	Mutable state with inner_state, cached_regional_state

Raises

Name	Type	Description
	ValueError	If len(region_mapping) != network.graph.n_nodes

update_state

experimental.network_dynamics.coupling.SubspaceCoupling.update_state(
    coupling_data,
    coupling_state,
    new_state,
)

Update inner coupling state and cache aggregated state.

Parameters

Name	Type	Description	Default
coupling_data	Bunch	Static regional structures	required
coupling_state	Bunch	Current state with inner_state, cached_regional_state	required
new_state	(ndarray, shape(n_states, n_nodes))	Network state after integration step	required

Returns

Name	Type	Description
	Bunch	Updated inner_state and cached_regional_state for next compute()