Custom coupling functions: no registry needed

lyapax coupling is a plain callable coupling_fn(cvar_state, weights, params) -> coupling – not a fixed set of named coupling “kinds” dispatched through hardcoded if/elif branches internally. That means adding a new coupling rule never requires touching the library’s source: any function with the right signature is a first-class coupling, exactly like the built-in linear_coupling/ sigmoidal_coupling/kuramoto_coupling in lyapax.coupling. This demo writes one from scratch – no import from lyapax.coupling at all – and plugs it directly into network_problem, reproducing the exact linear-network result from 04_linear_network.py.

The system. Same idea as 04_linear_network.py – scalar linear node dynamics x_i' = gamma * x_i + c_i coupled linearly, so the whole network reduces to x' = (gamma*I + G*W) x – but scaled down to 2 nodes joined by a single edge (weights = [[0,1],[1,0]]). That adjacency matrix has eigenvalues {1, -1}, so A’s eigenvalues are exactly gamma +/- G = {-0.8, -2.2}. The only thing this example changes versus 04_linear_network.py is the coupling term itself: my_linear_coupling below reimplements lyapax.coupling.linear_coupling by hand (same G * W @ x formula) purely to demonstrate that network_problem accepts any callable of the right signature – the result should (and does) match 04_linear_network.py’s exact-eigenvalue check bit-for-bit in spirit, just for this smaller network.

import os

os.environ["JAX_PLATFORMS"] = "cpu"

import jax
import jax.numpy as jnp
import numpy as np

jax.config.update("jax_enable_x64", True)

from lyapax.core import lyapunov_spectrum
from lyapax.network import Network, network_problem
from lyapax.simulator import ModelSpec, Parameter, StateVar, build_jax_dfun

def my_linear_coupling(cvar_state, weights, params):
    """A user-written coupling -- deliberately not imported from
    lyapax.coupling, to prove the extension point needs no library
    changes and no registration step."""
    G = params["G"]
    return G * jnp.einsum("ts,cs->ct", weights, cvar_state)

weights = np.array([[0., 1.], [1., 0.]])
gamma, G = -1.5, 0.7
A = gamma * np.eye(2) + G * weights
expected = np.sort(np.linalg.eigvalsh(A))[::-1]

model = ModelSpec(
    name="linear_node",
    state_variables=(StateVar("x", default_init=0.0),),
    parameters=(Parameter("gamma", gamma),),
    cvar=("x",),
    dfun_str={"x": "gamma * x + c"},
)
dfun = build_jax_dfun(model)
params = {"gamma": gamma, "G": G}
dt = 1e-3
network = Network(weights=jnp.array(weights), cvar_indices=model.cvar_indices)
problem = network_problem(
    dfun, network, my_linear_coupling,  # <- user function, not a lyapax builder
    params=params, state0=jnp.array([0.2, -0.3]), dt=dt,
)

result = lyapunov_spectrum(
    problem, n_steps=20_000, renorm_every=10, t_transient=5.0,
)

estimate = np.array(result.exponents)
print(f"exact eigenvalues:  {expected}")
print(f"lyapax estimate:    {estimate}")
print(f"max abs error:      {np.max(np.abs(estimate - expected)):.2e}")

exact eigenvalues:  [-0.8 -2.2]
lyapax estimate:    [-0.79999995 -2.19999819]
max abs error:      1.81e-06