AbstractDynamics

experimental.network_dynamics.dynamics.AbstractDynamics(**kwargs)

Abstract base class for neural dynamics models with multi-coupling support.

This base class extends the original AbstractDynamics with the ability to declare and receive multiple named coupling inputs via Bunch objects, enabling flexible network architectures with different coupling mechanisms.

Attributes

Name	Type	Description
STATE_NAMES	tuple of str	Names of integrated state variables
INITIAL_STATE	tuple of float	Default initial values for integrated states
AUXILIARY_NAMES	tuple of str	Names of auxiliary (non-integrated) variables
DEFAULT_PARAMS	Bunch	Default parameter values as Bunch object
COUPLING_INPUTS	dict	Dictionary of expected coupling inputs {name: n_dims}
EXTERNAL_INPUTS	dict	Dictionary of external input specifications {name: n_dims}
VARIABLES_OF_INTEREST	tuple	Variables to record (empty tuple = record all state variables)

Methods

Name	Description
dynamics	Compute state derivatives and auxiliary variables.
get_default_initial_state	Get default initial state array.
get_variables_of_interest_indices	Get indices of variables to record.
name_to_index	Convert state names to indices.
plot	Plot dynamics timeseries for state variables.
simulate	Simulate dynamics using Diffrax.
verify	Verify that dynamics implementation is correct.

dynamics

experimental.network_dynamics.dynamics.AbstractDynamics.dynamics(
    t,
    state,
    params,
    coupling,
    external,
)

Compute state derivatives and auxiliary variables.

Parameters

Name	Type	Description	Default
t	float	Current time	required
state	jnp.ndarray	Current state with shape [N_STATES, n_nodes]	required
params	Bunch	Model parameters as Bunch object (supports broadcasting)	required
coupling	Bunch	Named coupling inputs accessed as attributes (e.g., coupling.structural, coupling.modulatory). Missing couplings are automatically filled with zeros.	required
external	Bunch	Named external inputs accessed as attributes (e.g., external.stimulus, external.perturbation). Missing inputs are automatically filled with zeros.	required

Returns

Name	Type	Description
derivatives	jnp.ndarray or tuple	If N_AUXILIARIES == 0: Returns derivatives array with shape [N_STATES, n_nodes] If N_AUXILIARIES > 0: Returns tuple (derivatives, auxiliaries) where: - derivatives has shape [N_STATES, n_nodes] - auxiliaries has shape [N_AUXILIARIES, n_nodes]

Notes

• If COUPLING_INPUTS is empty, the model has no coupling
• Coupling arrays have shape [n_dims, n_nodes] where n_dims is declared in COUPLING_INPUTS
• If EXTERNAL_INPUTS is empty, the model has no external inputs
• External input arrays have shape [n_dims, n_nodes] where n_dims is declared in EXTERNAL_INPUTS

get_default_initial_state

experimental.network_dynamics.dynamics.AbstractDynamics.get_default_initial_state(
    n_nodes=1,
)

Get default initial state array.

Args: n_nodes: Number of network nodes

Returns: Initial state array of shape [N_STATES, n_nodes]

get_variables_of_interest_indices

experimental.network_dynamics.dynamics.AbstractDynamics.get_variables_of_interest_indices(
)

Get indices of variables to record.

Returns: Tuple of indices into all_variable_names. If VARIABLES_OF_INTEREST is empty, returns only state variable indices (default behavior).

name_to_index

experimental.network_dynamics.dynamics.AbstractDynamics.name_to_index(names)

Convert state names to indices.

Args: names: Single state name, tuple/list of state names, or empty list

Returns: Array of indices corresponding to the state names

Examples: For Lorenz with STATE_NAMES = (“x”, “y”, “z”): name_to_index(“x”) -> jnp.array([0]) name_to_index([“y”, “x”]) -> jnp.array([1, 0]) name_to_index((“z”, “y”)) -> jnp.array([2, 1]) name_to_index([]) -> jnp.array([], dtype=int)

plot

experimental.network_dynamics.dynamics.AbstractDynamics.plot(
    t0=0.0,
    t1=10.0,
    dt=0.01,
    n_nodes=1,
    figsize=(10, 6),
    **simulate_kwargs,
)

Plot dynamics timeseries for state variables.

Note: Currently only plots state variables (not auxiliaries) as auxiliaries are not computed during simulation.

Args: t0: Start time t1: End time dt: Time step n_nodes: Number of nodes to simulate figsize: Figure size (width, height) **simulate_kwargs: Additional arguments for simulate()

simulate

experimental.network_dynamics.dynamics.AbstractDynamics.simulate(
    t0=0.0,
    t1=10.0,
    dt=0.01,
    n_nodes=1,
    params=None,
    initial_state=None,
    solver=None,
    **solver_kwargs,
)

Simulate dynamics using Diffrax.

Args: t0: Start time t1: End time dt: Time step for output n_nodes: Number of nodes to simulate params: Parameters (default: self.params) initial_state: Initial state (default: from get_default_initial_state) solver: Diffrax solver (default: Dopri5) **solver_kwargs: Additional solver arguments

Returns: times: [n_timesteps] - Time points trajectory: [n_timesteps, n_states, n_nodes] - State trajectory

verify

experimental.network_dynamics.dynamics.AbstractDynamics.verify(
    n_nodes=1,
    verbose=True,
)

Verify that dynamics implementation is correct.

Tests the dynamics function with default initial state and parameters at t=0 with zero coupling to ensure it runs without errors and returns the expected output format.

Args: n_nodes: Number of nodes to test with verbose: Whether to print verification details

Returns: True if verification passes, False otherwise