API reference

Contents

API reference#

Core#

Core Benettin/QR Lyapunov-spectrum engine.

Scope: single-node (or otherwise uncoupled) systems given as a one-step map state -> new_state. Networked/coupled systems (lyapax.network) and delayed systems (lyapax.dde) reuse the same renormalization idea but need extra tangent bookkeeping (a delay ring buffer’s tangent, in the DDE case).

Method: propagate a (d, k) matrix of tangent vectors alongside the trajectory. Every renorm_every steps, QR-decompose the tangent matrix, accumulate log|diag(R)|, and replace the tangent matrix with the orthonormal factor Q (Benettin’s method).

Tangent propagation is jax.jvp-based, not jax.jacfwd-based: one jax.jvp call per tracked column, batched via jax.vmap - cost is O(k) forward-mode passes per raw step, not O(d) for a dense Jacobian, which matters whenever k < d (the partial-spectrum case jacfwd can’t exploit, since it always computes all d columns regardless of how many are actually tracked). See Matrix-free tangent propagation for the design rationale.

Differentiating an exponent w.r.t. a system parameter: lyapunov_spectrum is built entirely from jax.lax.scan (static trip count) and jnp.linalg.qr, both of which support reverse-mode AD, so jax.grad/jax.jacrev do not raise here (unlike lyapax.adaptive’s diffrax-backed integrator, whose internal dynamic-trip-count while_loop blocks reverse-mode outright). But not raising is not the same as useful: for a genuinely chaotic trajectory, step_fn closes over the parameter being differentiated, so jax.grad/jax.jacfwd differentiate through the entire unrolled state trajectory, and the result inherits that trajectory’s own exponential sensitivity to perturbation - the returned “gradient” grows roughly like exp(lambda_max * horizon) and is numerically meaningless well before it overflows (confirmed empirically: unusable already within a few hundred steps for the Lorenz system, worse for longer runs; see tests/test_differentiability.py). This is a known phenomenon in chaotic sensitivity analysis, not a lyapax-specific bug - naive trajectory-unrolling gradients of long-time-averaged chaotic quantities are fundamentally unreliable, which is why shadowing-based methods (e.g. least-squares shadowing) exist in that literature. Practical guidance: jax.grad/jax.jacfwd through this engine are reliable for non-chaotic or short-horizon systems (e.g. tuning a parameter that keeps the trajectory on a stable fixed point/limit cycle, or a genuinely short run) - do not trust a gradient computed through a long chaotic trajectory without independently checking it (e.g. against a finite-difference estimate) first.

class lyapax.core.ConvergenceDrift(absolute, relative, converged)[source]#
absolute: Array#

(k,) |history[-1] - history[-1 - window_rows]| per exponent – the raw change in the running estimate over the tail window.

converged: Array | None#

(k,) bool, relative <= tol per exponent – None if tol was not given to convergence_drift.

relative: Array#

(k,) absolute / |history[-1]| per exponent. Blows up for exponents near zero (e.g. a conservative system’s zero exponent, or any near-degenerate pair) – a near-zero reference makes “relative to it” not meaningful; use absolute for those columns instead.

class lyapax.core.LyapunovCheckpoint(state, Y, cum_log_growth, elapsed_time, dt)[source]#
Y: Array#

(d, k) QR-orthonormalized tangent basis at the end of the checkpointed run, in its raw (not display-sorted) column order – see cum_log_growth.

cum_log_growth: Array#

(k,) total accumulated log-growth (sum of log|diag(R)| over every renormalization block since accumulation started) in the same raw column order as Y – not LyapunovResult.history’s column order, which is re-sorted by the final row on every call and can reorder differently between two otherwise-identical resumed runs if exponents are near-degenerate. Keeping this in Y’s raw order is what makes resuming exact: it is the same “never reordered until the very end” bookkeeping lyapunov_spectrum already does internally within a single call, just carried across two calls instead of within one.

dt: float#

The dt of the run that produced this checkpoint (a plain Python float, not a traced array – kept out of the tangent computation). On resume=, lyapunov_spectrum checks this against the dt of the resuming call and raises if they differ, since silently resuming under a different dt would misinterpret elapsed_time/history without any structural (shape-based) sign of the mismatch. This does not validate step_fn/system parameters/integrator – those aren’t identifiable from an opaque JAX closure – so a same-dt checkpoint can still be resumed against a different system silently; dt is only the one stable field cheap enough to check automatically.

elapsed_time: Array#

total elapsed accumulation time since the end of the transient, across every resumed call so far – same units as LyapunovResult.times.

Type:

Scalar

state: Array#

(d,) trajectory state at the end of the checkpointed run.

class lyapax.core.LyapunovResult(exponents, history, times, checkpoint)[source]#
checkpoint: LyapunovCheckpoint | DDECheckpoint | None#

Enough state to continue this run with another lyapunov_spectrum(..., resume=result.checkpoint) / lyapunov_spectrum_dde(..., resume=result.checkpoint) call, picking up exactly where this one left off (no re-transient, no discontinuity in history/times) – the “run a fixed n_steps, eyeball history/convergence_drift, resume if not converged yet” workflow. Always set by both lyapunov_spectrum (the ODE engine, LyapunovCheckpoint) and lyapax.dde.lyapunov_spectrum_dde (the DDE engine, lyapax.dde.DDECheckpoint – also carries the delay ring buffer’s state, since a DDE’s Markovian state is (state, buf) together, not state alone).

exponents: jnp.ndarray#

(k,) final Lyapunov-exponent estimates, sorted descending.

history: jnp.ndarray#

(n_renorm, k) running estimate at each renormalization point, in the same column order as exponents - use to check convergence.

Columns are ordered once, by the final row (history[-1] == exponents), then that column order is applied to every row. Near- degenerate exponents can cross over during the run, so an early row’s per-column values are not guaranteed to be individually sorted or to track the same Oseledets direction throughout - only the last row is.

times: jnp.ndarray#

(n_renorm,) elapsed time (in units of dt) at each row of history, measured from the end of the transient.

class lyapax.core.ODEProblem(step_fn: Callable[[Array], Array], state0: Array, dt: float)[source]#

Owns the (step_fn, state0, dt) triple that lyapunov_spectrum otherwise asks the caller to keep in sync by hand – in particular, the dt baked into step_fn (e.g. via ode_step) and the dt passed separately to lyapunov_spectrum must agree, and nothing checks that when they’re two independent arguments. Build one with ode_problem (plain ODE) or network_problem (coupled network), or construct directly if you already have a step_fn.

Parameters:

  • step_fn – one fixed-time-step update, state (d,) -> new_state (d,).

  • state0 – (d,) initial state (pre-transient).

  • dt – time represented by one call to step_fn.

lyapax.core.convergence_drift(result: LyapunovResult, window: float = 0.1, tol: float | None = None) → ConvergenceDrift[source]#

Summarize how much each exponent’s running estimate has moved over the tail of the run, by comparing result.history[-1] (the final estimate) to the estimate window fraction of the run’s renormalization points earlier.

lyapunov_spectrum/lyapunov_spectrum_dde always run a fixed n_steps – there is no built-in stopping criterion, adaptive or otherwise. This is a diagnostic to help a caller judge, after the fact, whether that fixed-length run was long enough: a large drift means the estimate was probably still moving and n_steps should be increased; a small, stable drift is evidence (not proof) of convergence. Pairs with result.checkpoint/resume=: if not converged, continue the same run with lyapunov_spectrum(..., resume=result.checkpoint) instead of restarting from scratch.

Parameters:

  • result – output of lyapunov_spectrum/lyapunov_spectrum_dde.

  • window – fraction, in (0, 1], of result.history’s rows making up the tail comparison window. Default 0.1 compares the final estimate against the estimate from 10% of the run ago.

  • tol – if given, converged is relative <= tol per exponent; otherwise converged is None.

Returns:

ConvergenceDrift(absolute, relative, converged).

A single relative <= tol pass is evidence of settling, not a statistical stopping rule: pick tol as a mix of absolute and relative tolerance (relative alone blows up for exponents near zero – see ConvergenceDrift.relative’s docstring, and use absolute directly for those columns), and prefer requiring several consecutive resumed chunks to report converged before trusting it, rather than stopping at the first pass. CPU/GPU QR and reduction differences are not bitwise-identical, so the exact chunk where a tight tol first passes can differ by backend even for the same system and seed.

Usage#

>>> import jax
>>> jax.config.update("jax_enable_x64", True)
>>> import jax.numpy as jnp
>>> from lyapax import lyapunov_spectrum, ode_problem, systems
>>> from lyapax.core import convergence_drift
>>> rhs = systems.lorenz(sigma=10.0, rho=28.0, beta=8.0 / 3.0)
>>> problem = ode_problem(rhs, state0=jnp.array([1.0, 1.0, 1.0]), dt=1e-2)
>>> result = lyapunov_spectrum(
...     problem, n_steps=50_000, renorm_every=10, t_transient=100.0,
... )
>>> drift = convergence_drift(result, window=0.1, tol=1e-2)
>>> bool(drift.converged[0])
True
lyapax.core.kaplan_yorke_dimension(exponents: Array, d_total: int | None = None) → float[source]#

Kaplan-Yorke (Lyapunov) dimension: an estimate of an attractor’s fractal dimension from its Lyapunov spectrum (Kaplan & Yorke, 1979), j + sum(exponents[:j]) / |exponents[j]| where j is the largest prefix length with a non-negative partial sum. Pure post-processing of LyapunovResult.exponents – no tangent-propagation or QR involved.

Eager, host-side post-processing: the crossing index j is found by branching on Python float``s in a Python loop, not ``jnp.where/ lax.cond, so this is not jax.jit-, jax.vmap-, or jax.grad-compatible. Call it after (not inside) any transformed code, on a concrete exponents array.

Parameters:

  • exponents – (k,) exponents, sorted descending – exactly LyapunovResult.exponents’s own order, so kaplan_yorke_dimension(result.exponents) is the usual call.

  • d_total – the full system dimension, if exponents is only a partial spectrum (k < d_total). If given, and the partial sum never goes negative within the tracked k exponents, raises ValueError instead of silently returning k – the true crossing point lies beyond what was tracked, so k would understate the answer; track more exponents (increase k) instead. Omit (the default) when exponents is already the full spectrum – the “sum stays non-negative” case then correctly returns len(exponents) (the attractor fills the whole tracked phase space, e.g. a conservative system).

Returns:

the Kaplan-Yorke dimension, in [0, len(exponents)].

Usage#

>>> import jax.numpy as jnp
>>> from lyapax.core import kaplan_yorke_dimension
>>> float(kaplan_yorke_dimension(jnp.array([0.906, 0.0, -14.57])))  # Lorenz
2.0621...
lyapax.core.lyapunov_spectrum(step_fn_or_problem: Callable[[Array], Array] | ODEProblem, state0: Array | int | None = None, dt: float | None = None, n_steps: int | None = None, k: int | None = None, renorm_every: int = 1, t_transient: float = 0.0, seed: int = 0, check_finite: bool = False, resume: LyapunovCheckpoint | None = None) → LyapunovResult[source]#

Compute the (partial or full) Lyapunov spectrum of step_fn along the trajectory started from state0, via the Benettin/QR method.

Parameters:

  • step_fn_or_problem (either an ODEProblem (from ode_problem /) – network_problem) – in which case state0 and dt are read off it and the second positional argument is n_steps (lyapunov_spectrum(problem, n_steps) or lyapunov_spectrum(problem, n_steps=...)) – or a plain step_fn, state (d,) -> new_state (d,), in which case state0, dt, n_steps are all given explicitly (the original, lower-level call form). Must be a pure, differentiable JAX function of state alone - close over any parameters.

  • state0 ((d,) initial state (pre-transient). Ignored (and may be) – omitted) when an ODEProblem is passed.

  • dt (time represented by one call to step_fn. For discrete maps,) – pass dt=1.0 and interpret the exponents as per-iterate. When an ODEProblem is passed, this may be omitted; if provided, it must match the dt stored on the problem.

  • n_steps (number of steps to run after the transient. Must be a) – multiple of renorm_every.

  • k (number of leading exponents to track (k <= d). Defaults to the) – full spectrum (k = d). Cost scales with k, not d - this is the “only the first few largest exponents” case.

  • renorm_every (QR-renormalize every this many steps. Larger values) – reduce QR overhead but risk tangent-vector overflow/underflow for fast-growing/shrinking directions (see Choosing renorm_every) - keep small enough that exp(|lambda_max| * renorm_every * dt) stays well within float64 range.

  • t_transient (time to integrate (discarding tangent tracking) before) – starting the Lyapunov accumulation. Mutually exclusive with resume (raises ValueError if both are given and nonzero) – a resumed run is already transient-aligned.

  • seed (PRNG seed for the initial random tangent-vector basis. Unused) – when resume is given (the tangent basis comes from the checkpoint instead).

  • check_finite (if True, raise FloatingPointError when any running) – estimate in history is non-finite (a QR diagonal underflowed to 0 or overflowed to inf, e.g. from renorm_every too large). Off by default and only usable when this function is called eagerly (not wrapped in jax.jit).

  • resume (a previous call’s result.checkpoint, to continue that run) – instead of starting a fresh one – skips the random tangent-basis init and the transient, and history/times in the returned LyapunovResult continue the cumulative running estimate (not reset to zero), so concatenating two calls’ history reads as one continuous convergence curve. k must match the checkpoint’s tracked dimension if given explicitly (defaults to it otherwise). See 16_convergence_drift.py for the run-inspect-resume workflow this enables.

Return type:

LyapunovResult

Examples

>>> import jax
>>> jax.config.update("jax_enable_x64", True)
>>> import jax.numpy as jnp
>>> from lyapax import lyapunov_spectrum, ode_problem, systems
>>> rhs = systems.lorenz(sigma=10.0, rho=28.0, beta=8.0 / 3.0)
>>> problem = ode_problem(rhs, state0=jnp.array([1.0, 1.0, 1.0]), dt=1e-2)
>>> result = lyapunov_spectrum(
...     problem, n_steps=50_000, renorm_every=10, t_transient=100.0,
... )
>>> result.exponents
Array([ 0.906,  0.   , -14.57], dtype=float64)
lyapax.core.ode_problem(rhs, state0: Array, dt: float, integrator: str | Callable = 'rk4') → ODEProblem[source]#

Build an ODEProblem for a plain (uncoupled) ODE – bundles ode_step’s step_fn together with state0 and dt so lyapunov_spectrum(problem, n_steps=...) never needs dt (or state0) passed a second time.

Parameters:

  • rhs – right-hand side, state (d,) -> dstate (d,).

  • state0 – (d,) initial state (pre-transient).

  • dt – fixed step size.

  • integrator – "euler", "heun", "rk4", "rk6", or a callable (rhs, dt) -> step_fn – see lyapax.integrators.ode_step.

DDE#

Fixed-delay DDE Lyapunov engine, on top of the vendored ring buffer.

Reuses the vendored ring-buffer simulator (lyapax.simulator.step) as-is rather than a second, parallel history mechanism: the genuinely missing piece was never “how do we store delayed history” (the vendored _write_ring/_read_delayed_coupling already do that correctly), it was that lyapax.core.lyapunov_spectrum only differentiates a flat state, not the (state, buf) carry a DDE actually needs – a sensitivity to the delayed value (d f/d x(t-tau)) is just as real as the sensitivity to the current value, and dropping it silently gives wrong exponents, not merely imprecise ones.

Method, following Farmer’s (1982) approach to DDE Lyapunov spectra (the same method the established jitcdde_lyap package implements for adaptive-step DDEs via Hermite-interpolated history): augment the system with explicit tangent dynamics for both state and buf, propagate them alongside the primal trajectory, periodically re-orthonormalize via QR. Where this differs from jitcdde_lyap: our scope is fixed-step, with a finite-size ring buffer and a finite-dimensional tangent state (a plain jnp.linalg.qr, not jitcdde’s continuous-function inner product over a Hermite interpolant). By default a delayed value snaps to the nearest stored grid point (integer-step delay, no interpolation – resolve_tau_steps rounds tau to the nearest dt multiple); passing interpolate=True (see lyapax.simulator.make_step_fn, dde_problem, network_dde_problem) switches to cubic Hermite interpolation over the ring buffer’s stored (value, derivative) pairs, removing that rounding at the cost of storing/differentiating twice as much per buffer slot – see Grid-snapped vs. interpolated history reads for the design and tradeoffs.

Tangent propagation is jax.jvp-based (the same matrix-free pattern as lyapax.core, see Matrix-free tangent propagation): cost is O(k) forward passes per raw step (k = tracked exponents), not O(d_total) for the full augmented (state, buf) dimension – a dense-Jacobian approach does not scale to real coupled networks, since d_total grows with both network size and ring-buffer depth.

Scope note: lyapunov_spectrum_dde itself has no opinion on delay structure – it differentiates through whatever carry step_fn produces, so it already works correctly with a genuine per-edge (heterogeneous) delay matrix via the vendored step’s legacy, hardcoded-linear coupling path (coupling_fn=None, delay_steps=<(n_nodes,n_nodes) matrix>, verified directly against an asymmetric 2-node case). The real scope limit lives in lyapax.simulator.make_step_fn, not here: a custom coupling_fn (lyapax.coupling’s plain-callable style, e.g. for a delayed sigmoidal/Kuramoto network) is currently only wired up for zero-delay and uniform-delay (single global tau_steps) branches – combining a custom coupling_fn with per-edge delays needs an edge-aware coupling_fn signature, which does not exist yet.

lyapax.dde.CarryStepFn#

step(carry, _) -> (new_carry, new_state), carry = (state, buf, t, params) – exactly what lyapax.simulator.make_step_fn(…) returns.

alias of Callable

class lyapax.dde.DDECheckpoint(state: jnp.ndarray, buf: jnp.ndarray, t: jnp.ndarray, Y_state: jnp.ndarray, Y_buf: jnp.ndarray, cum_log_growth: jnp.ndarray, elapsed_time: jnp.ndarray, dt: float = None)[source]#

The DDE counterpart of lyapax.core.LyapunovCheckpoint: enough state to continue a lyapunov_spectrum_dde run with another lyapunov_spectrum_dde(..., resume=result.checkpoint) call. A DDE’s Markovian state is (state, buf) together (see the module docstring), not state alone, so unlike the ODE checkpoint this also carries buf, the ring-buffer tangent Y_buf, and the ring-buffer step counter t.

Y_buf: Array#

(horizon, n_cvar, n_nodes, k) (or the interpolate=True shape) tangent basis for the ring buffer, same raw column order as Y_state – together, (Y_state, Y_buf) is the full augmented tangent basis, mirroring how (state, buf) is the full augmented primal state.

Y_state: Array#

(n_sv, n_nodes, k) QR-orthonormalized tangent basis for state, in raw (not display-sorted) column order – see cum_log_growth.

buf: Array#

(horizon, n_cvar, n_nodes) (or the interpolate=True value/ derivative shape) ring buffer at the end of the checkpointed run.

cum_log_growth: Array#

(k,) total accumulated log-growth since accumulation started, in the same raw column order as Y_state/Y_buf – see lyapax.core.LyapunovCheckpoint.cum_log_growth.

dt: float#

The dt of the run that produced this checkpoint (a plain Python float, not a traced array). dt also determines the grid-snapped delay tau_steps/horizon for a DDE, so a dt mismatch on resume would be even more likely to silently corrupt a DDE run than an ODE one – see lyapax.core.LyapunovCheckpoint.dt, which this mirrors. lyapunov_spectrum_dde checks this on resume= and raises if it differs from the resuming call’s dt.

elapsed_time: Array#

total elapsed accumulation time since the end of the transient, across every resumed call so far – same units as LyapunovResult.times.

Type:

Scalar

state: Array#

(n_sv, n_nodes) trajectory state at the end of the checkpointed run.

t: Array#

Scalar int32 ring-buffer step counter at the end of the checkpointed run – must be threaded through unchanged (not reset to 0), since it drives the buffer’s modular read/write indices (see the module docstring’s note on why t lives in the scanned carry).

class lyapax.dde.DDEProblem(step_fn: Callable, state0: Array, buf0: Array, params: dict, dt: float, tau_steps: int | float)[source]#

Owns the carry-style plumbing (step_fn, ring buffer, delay length) that lyapunov_spectrum_dde otherwise asks the caller to assemble and pass by hand – buf0, tau_steps, and the carry layout become properties of one object instead of four separate positional arguments. Build one with dde_problem (single/scalar DDE) or network_dde_problem (coupled, uniformly-delayed network), or construct directly if you already have a carry step_fn.

Parameters:

  • step_fn – see lyapunov_spectrum_dde.

  • state0 – (n_sv, n_nodes) initial state (pre-transient).

  • buf0 – (horizon, n_cvar, n_nodes) initial ring buffer.

  • params – closed over for tangent propagation.

  • dt – fixed step size.

  • tau_steps – delay in units of dt (see resolve_tau_steps) – carried for reference/inspection (problem.tau_steps, tau_eff(problem.tau_steps, problem.dt)); lyapunov_spectrum_dde itself only needs buf0’s shape. An integer for the default (grid-snapped) delay; a non-integer float when built with interpolate=True.

lyapax.dde.DelayedRHS#

rhs_delayed(state_now, state_delayed) -> dstate, both shape (m,) – m=1 for a scalar DDE (e.g. Mackey-Glass), m>1 for a small non-networked vector DDE. See make_scalar_delayed_step_fn.

alias of Callable[[Array, Array], Array]

lyapax.dde.constant_history_buf0(cvar_state0: Array, horizon: int, interpolate: bool = False) → Array[source]#

Initial ring buffer under the constant-history DDE convention: the coupling-variable state is assumed equal to cvar_state0 for all t <= 0.

Parameters:

  • cvar_state0 – (n_cvar, n_nodes) coupling-variable state at t=0.

  • interpolate – build the (value, derivative) buffer shape lyapax.simulator.make_step_fn(..., interpolate=True) expects. The constant-history derivative is exactly 0 (a constant function has zero derivative), not an approximation.

Returns:

(horizon, n_cvar, n_nodes) buf0 (or (horizon, 2, n_cvar, n_nodes) when interpolate=True), ready for carry0 = (state0, buf0, jnp.int32(0), params).

lyapax.dde.dde_problem(rhs_delayed: Callable[[Array, Array], Array], state0: Array, tau: float, dt: float, history: Array | None = None, integrator: str | Callable = 'heun', interpolate: bool = False) → DDEProblem[source]#

Build a DDEProblem for a simple, non-networked scalar/vector fixed-delay DDE directly from rhs_delayed – the parallel-to-ODE front door for lyapunov_spectrum_dde, hiding the ring-buffer/carry construction that make_scalar_delayed_step_fn + scalar_delayed_history0 otherwise ask the caller to wire up by hand.

Parameters:

  • rhs_delayed – (state_now, state_delayed) -> dstate, both shape (m,) – see DelayedRHS.

  • state0 – (m,) initial state (pre-transient, pre-history).

  • tau – physical delay. Rounded to an integer number of dt steps (see resolve_tau_steps) unless interpolate=True, in which case tau is used exactly (via Hermite-interpolated history reads, see lyapax.simulator.make_step_fn) – no rounding, no warning.

  • dt – fixed step size.

  • history – optional initial ring buffer ((tau_steps + 1, m), or the interpolate=True value/derivative shape – see scalar_delayed_history0), for a non-constant history. Defaults to the constant-history convention (state0 held for all t <= 0).

  • integrator – "euler", "heun", "rk4", "rk6", or a callable – see lyapax.simulator.make_step_fn.

  • interpolate – see lyapax.simulator.make_step_fn.

Examples

>>> import jax
>>> jax.config.update("jax_enable_x64", True)
>>> import jax.numpy as jnp
>>> from lyapax import systems
>>> from lyapax.dde import dde_problem, lyapunov_spectrum_dde
>>> a, tau, dt = 0.5, 0.3, 1e-2
>>> problem = dde_problem(
...     systems.linear_scalar_dde(a=a), state0=jnp.array([1.0]),
...     tau=tau, dt=dt,
... )
>>> result = lyapunov_spectrum_dde(
...     problem, n_steps=20_000, k=1, renorm_every=5, t_transient=10.0,
... )
>>> result.exponents
Array([-0.5985], dtype=float64)  # dominant root of the Lambert-W characteristic equation
lyapax.dde.lyapunov_spectrum_dde(step_fn_or_problem: Callable | DDEProblem, state0: Array | int | None = None, buf0: Array | None = None, params: dict | None = None, dt: float | None = None, n_steps: int | None = None, k: int | None = None, renorm_every: int = 1, t_transient: float = 0.0, seed: int = 0, check_finite: bool = False, resume: DDECheckpoint | None = None) → LyapunovResult[source]#

Compute the (partial or full) Lyapunov spectrum of a fixed-delay DDE, via the Benettin/QR method generalized to a (state, buf) tangent pair (see module docstring).

Parameters:

  • step_fn_or_problem (either a DDEProblem (from dde_problem /) – network_dde_problem) – in which case state0, buf0, params, and dt are read off it and the second positional argument is n_steps (lyapunov_spectrum_dde(problem, n_steps) or lyapunov_spectrum_dde(problem, n_steps=...)) – or a plain carry step_fn, step(carry, _) -> (new_carry, new_state), carry = (state, buf, t, params), in which case state0, buf0, params, dt, n_steps are all given explicitly (the original, lower-level call form).

  • state0 ((n_sv, n_nodes) initial state (pre-transient). Ignored (and) – may be omitted) when a DDEProblem is passed.

  • buf0 ((horizon, n_cvar, n_nodes) initial ring buffer (see) – constant_history_buf0). Ignored when a DDEProblem is passed.

  • params (closed over for tangent propagation (not differentiated --) – only state/buf are). Ignored when a DDEProblem is passed.

  • dt (see lyapax.core.lyapunov_spectrum) – (same meaning). When a DDEProblem is passed, dt may be omitted; if provided, it must match the dt stored on the problem.

  • n_steps (see lyapax.core.lyapunov_spectrum) – (same meaning). When a DDEProblem is passed, dt may be omitted; if provided, it must match the dt stored on the problem.

  • renorm_every (see lyapax.core.lyapunov_spectrum) – (same meaning). When a DDEProblem is passed, dt may be omitted; if provided, it must match the dt stored on the problem.

  • seed (see lyapax.core.lyapunov_spectrum) – (same meaning). When a DDEProblem is passed, dt may be omitted; if provided, it must match the dt stored on the problem.

  • k (number of leading exponents to track. Defaults to the full) – augmented spectrum, k = d_total = state0.size + buf0.size – note this is the ring-buffer-augmented dimension, not just the physical state dimension, and grows with horizon (delay length). Per-step cost and tangent-matrix memory are O(k), and QR is O(d_total * k^2), so for large horizon (long delays or many network nodes) pass an explicit small k rather than relying on the full-spectrum default.

  • t_transient (time to integrate (discarding tangent tracking) before) – starting the Lyapunov accumulation. Internally rounded up to cover at least one full ring cycle (horizon * dt): until every buffer slot has been written at least once from real dynamics, delayed- direction tangent information is incomplete, and (for large horizon) a random initial tangent basis is disproportionately concentrated in the buffer’s trivial “just shifted forward unchanged” directions – the same class of silent-bias risk M1’s transient fix addresses for the plain ODE case, amplified here.

  • check_finite (see lyapax.core.lyapunov_spectrum (same meaning);) – the augmented (state, buf) tangent dimension makes underflow more likely for large horizon, so this is worth enabling while tuning renorm_every for a new DDE system.

  • resume (a previous call’s result.checkpoint (a DDECheckpoint),) – to continue that run instead of starting a fresh one – skips the random tangent-basis init and the mandatory ring-cycle transient (a resumed run’s buf was already fully populated by the checkpointed run), and history/times in the returned LyapunovResult continue the cumulative running estimate, so concatenating two calls’ history reads as one continuous convergence curve – same contract as lyapax.core.lyapunov_spectrum’s resume. Mutually exclusive with a nonzero t_transient. k must match the checkpoint’s tracked dimension if given explicitly (defaults to it otherwise).

Return type:

LyapunovResult

Examples

See dde_problem for a full worked example (lyapunov_spectrum_dde(dde_problem(...), n_steps=...)); this function is the DDE analogue of lyapax.core.lyapunov_spectrum, generalized to the ring-buffer-augmented (state, buf) tangent pair.

lyapax.dde.make_scalar_delayed_step_fn(rhs_delayed: Callable[[Array, Array], Array], m: int, tau_steps: int | float, dt: float, integrator: str | Callable = 'heun', interpolate: bool = False) → Callable[source]#

Build a carry-based step function for a simple, non-networked fixed-delay DDE directly from a plain rhs_delayed – no ModelSpec/coupling_fn ceremony needed. Mirrors lyapax.integrators.rk4_step’s role on the ODE side: a lightweight front door for the “just one system, no network” case, sitting next to lyapax.simulator.make_step_fn (the general, ModelSpec-based front door used for real delayed networks, e.g. tests/test_dde.py’s benchmark test) – both build a step for the same underlying vendored ring-buffer simulator, so there is still only one DDE mechanism, just two ways to construct a step for it.

Internally: a trivial “coupled to your own delayed history” 1-node self-loop, with an identity coupling_fn (the delayed state passes through unchanged; any scaling/nonlinearity belongs in rhs_delayed itself, not a coupling formula – there’s no real network here).

Parameters:

  • rhs_delayed – (state_now, state_delayed) -> dstate, both shape (m,) – see DelayedRHS.

  • m – state dimension (e.g. 1 for a scalar DDE like Mackey-Glass).

  • tau_steps – delay in units of dt (see resolve_tau_steps). An integer when interpolate=False; may be a non-integer float when interpolate=True.

  • integrator – "euler", "heun", "rk4", "rk6", or a callable – see lyapax.simulator.make_step_fn.

  • interpolate – see lyapax.simulator.make_step_fn.

lyapax.dde.network_dde_problem(dfun: Callable, network: Network, coupling: Callable, params: dict, state0: Array, dt: float, tau: float | None = None, history: Array | None = None, integrator: str | Callable = 'heun', interpolate: bool = False) → DDEProblem[source]#

Build a DDEProblem for a coupled, uniformly-delayed network – the DDE counterpart of lyapax.network.network_step (see that function and Network for the topology/coupling/integrator split this keeps parallel between the ODE and DDE construction paths).

Only a single global delay is supported (a uniform tau_steps read via one ring-buffer lookup per node) – per-edge heterogeneous delays combined with a custom coupling callable are not yet supported, see lyapax.simulator.make_step_fn’s docstring.

Parameters:

  • dfun – (state, coupling, params) -> dstate.

  • network – topology; exactly one of tau or network.delay_steps must give the uniform delay.

  • coupling – (cvar_state, weights, params) -> coupling.

  • params – closed over for tangent propagation.

  • state0 – (n_sv, n_nodes) initial state.

  • dt – fixed step size.

  • tau – physical delay. Rounded via resolve_tau_steps unless interpolate=True (used exactly, via Hermite-interpolated history reads). Mutually exclusive with network.delay_steps.

  • history – optional initial ring buffer ((tau_steps + 1, n_cvar, n_nodes), or the interpolate=True value/derivative shape – see constant_history_buf0). Defaults to the constant-history convention.

  • integrator – "euler", "heun", "rk4", "rk6", or a callable – see lyapax.simulator.make_step_fn.

  • interpolate – see lyapax.simulator.make_step_fn. When True, tau (or network.delay_steps) need not be an integer number of dt steps.

Examples

Two linearly delay-coupled nodes, x_i' = gamma*x_i + G*x_j(t-tau), whose symmetric/antisymmetric modes have closed-form Lambert-W exponents:

>>> import jax
>>> jax.config.update("jax_enable_x64", True)
>>> import jax.numpy as jnp
>>> import lyapax
>>> from lyapax.coupling import linear_coupling
>>> from lyapax.simulator import ModelSpec, StateVar, Parameter, build_jax_dfun
>>> gamma, G, tau, dt = -1.0, 0.3, 0.2, 1e-3
>>> 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)
>>> network = lyapax.Network(
...     weights=jnp.array([[0., 1.], [1., 0.]]), cvar_indices=model.cvar_indices,
... )
>>> problem = lyapax.network_dde_problem(
...     dfun, network, linear_coupling(a=1.0, b=0.0),
...     params={"gamma": gamma, "G": G}, state0=jnp.array([[0.3, -0.2]]),
...     dt=dt, tau=tau,
... )
>>> result = lyapax.lyapunov_spectrum_dde(
...     problem, n_steps=20_000, k=2, renorm_every=10, t_transient=5.0,
... )
>>> result.exponents
Array([-0.6578, -1.3967], dtype=float64)  # symmetric, antisymmetric mode
lyapax.dde.resolve_tau_steps(tau: float, dt: float, warn_tol: float = 1e-06) → int[source]#

Round a physical delay tau to an integer number of dt steps.

Integer-step only, no sub-step interpolation – see Grid-snapped vs. interpolated history reads for the accuracy tradeoff this implies. Every downstream Lyapunov spectrum is computed for the rounded delay tau_eff(tau_steps, dt), not the exact tau passed in – use tau_eff to see what delay was actually used, and decrease dt to shrink the gap.

warn_tolrelative tolerance (of tau) above which a mismatch

between tau and the rounded tau_eff triggers a warning.

lyapax.dde.scalar_delayed_history0(state0_now: Array, tau_steps: int | float, interpolate: bool = False)[source]#

(state0, buf0) for make_scalar_delayed_step_fn, from a flat (m,) initial condition, under the constant-history DDE convention – the scalar-case counterpart to constant_history_buf0, hiding the (n_sv, n_nodes)/(horizon, n_cvar, n_nodes) reshaping the general (networked) path needs.

Parameters:

interpolate – see lyapax.simulator.make_step_fn.

lyapax.dde.tau_eff(tau_steps: int, dt: float) → float[source]#

The effective (rounded-to-grid) delay actually used by a DDE run built with integer tau_steps, in the same units as dt. See resolve_tau_steps.

Simulator#

Trimmed model-specification dataclasses.

Adapted from vbi/simulator/spec/model.py - see NOTICE.md in this directory for provenance and what was dropped.

class lyapax.simulator.model_spec.ModelSpec(name: str, state_variables: tuple[StateVar, ...], parameters: tuple[Parameter, ...], cvar: tuple[str, ...], dfun_str: dict[str, str])[source]#

Backend-agnostic description of a dynamical system’s right-hand side.

dfun_str maps each state variable name to a bare math expression string. Only these symbols may appear: state variable names, parameter names, ‘c’ (coupling input, scalar per node - zero for uncoupled/single-node systems), and the math functions: exp, log, sin, cos, tanh, sqrt, abs, pi. No ‘np.’ / ‘jnp.’ prefix - the code generator injects the namespace.

Warning

dfun_str is compiled and run via Python’s exec() (build_jax_dfun) with no sanitization - it can execute arbitrary code. Only build a ModelSpec from strings you (or your own code) wrote; never from untrusted input such as user uploads, network responses, or config files from unknown sources.

class lyapax.simulator.model_spec.Parameter(name: 'str', default: 'float')[source]#
class lyapax.simulator.model_spec.StateVar(name: 'str', default_init: 'float' = 0.0)[source]#

dfun codegen + the ring-buffer step(carry, _) factory.

Adapted from vbi/simulator/backend/jax_/codegen.py and vbi/simulator/backend/jax_/simulator.py - see NOTICE.md in this directory for provenance and what was dropped (stochastic noise, state clipping, sigmoidal/kuramoto coupling, monitors/stimuli).

This is the piece that unifies ODE and DDE: has_delays=False reads coupling from the instantaneous state; has_delays=True reads it from a ring buffer of past coupling-variable states via a flat-index gather. Both paths produce the same step(carry, _) -> (new_carry, new_state) shape consumed by jax.lax.scan, and - critically for the Lyapunov engines - both are plain, differentiable JAX functions of carry.

lyapax.simulator.step.build_jax_dfun(spec: ModelSpec) → Callable[source]#

Compile dfun_str expressions into a JAX-traceable function.

Generated signature:

fn(state, coupling, params) -> jnp.ndarray shape (n_sv, n_nodes)

where:

state : (n_sv, n_nodes) coupling : (n_cvar, n_nodes) - one row per coupling variable params : dict[str, scalar | jnp.ndarray]

Compiled once via exec() on the spec’s own strings, same as the vbi original - nothing about the generated function blocks jax.jacfwd / jax.jvp.

Warning

The dfun_str expressions are code-generated and executed via exec() with no sanitization. This is fine for specs a caller writes themselves, but never build a ModelSpec from untrusted input (user uploads, network data, config files from unknown sources) - the strings can execute arbitrary code.

lyapax.simulator.step.make_step_fn(dfun: Callable, weights: Array, has_delays: bool, horizon: int, n_nodes: int, cvar_indices: tuple[int, ...], dt: float, delay_steps: Array | None = None, G_default: float = 1.0, coup_a: float = 1.0, coup_b: float = 0.0, integrator: str | Callable = 'heun', coupling_fn: Callable | None = None, tau_steps: int | float | None = None, interpolate: bool = False) → Callable[source]#

Returns step(carry, _) -> (new_carry, new_state).

carry = (state, buf, step_int32, params)

params travels inside the carry (not closed over) so that, later, jax.vmap over swept parameters and jax.jacfwd/jvp w.r.t. state both see it as data rather than a baked-in constant.

buf is a no-op array when has_delays=False - the ODE and DDE cases share this exact function; only the coupling branch differs.

integrator"euler", "heun", "rk4", "rk6", or a callable

(state, dfun, coupling_at, dt, params) -> new_state, where coupling_at(y_stage, c) -> coupling is called fresh at each internal stage (y_stage: that stage’s intra-step state estimate; c: its Butcher node, 0 at the first stage, 1 at the last) rather than being frozen once per step – see the “Integrators” section’s module comment, and Delays, coupling, and Runge–Kutta stage order for why freezing it instead silently caps the whole scheme at O(dt) whenever coupling is nonzero. The zero-delay case reads coupling from y_stage (the evolving state) fresh at each stage; the interpolated-delay case (interpolate=True) reads the delayed history at each stage’s own intra-step time step + c via the Hermite interpolant, so both realize the integrator’s nominal order (the delayed case up to the cubic Hermite history’s own ~4th-order ceiling). Only the grid-snapped delayed default (interpolate=False) still ignores c and reads once per step – without an interpolant there is no sub-step history to read – and remains O(dt); pass interpolate=True for higher-order DDE stepping.

coupling_fnoptional lyapax.coupling-style plain callable,

(cvar_state, weights, params) -> coupling (see lyapax.coupling.CouplingFn). When given, replaces the hardcoded G_default/coup_a/coup_b linear formula for both the zero-delay and (uniform-delay-only, see tau_steps) delayed branches, unifying this step with lyapax.coupling’s plain-callable design. G_default/coup_a/coup_b are ignored when coupling_fn is given. Default None preserves the exact original vendored behavior.

tau_stepsrequired alongside coupling_fn when has_delays=True

– a single global delay (in steps), read via _read_uniform_delayed_cvar (O(1) ring-buffer read) or, when interpolate=True, _read_uniform_delayed_cvar_interp (in which case tau_steps need not be an integer – see interpolate). Per-edge heterogeneous delays (the general delay_steps matrix) are not supported together with a custom coupling_fn yet – that combination needs an edge-aware coupling_fn signature, which does not exist yet; use the legacy coupling_fn=None path (which does support per-edge delay_steps, just with the hardcoded linear formula) until then.

interpolatereconstruct the delayed coupling-variable value via cubic

Hermite interpolation (using the stored value and derivative at the two grid points bracketing the delayed time) instead of snapping to the nearest stored grid sample – see Grid-snapped vs. interpolated history reads. Requires has_delays=True, coupling_fn given, and tau_steps given (the uniform-delay path only; not supported with the legacy per-edge delay_steps path). tau_steps may then be a non-integer float: the physical delay no longer needs rounding to a whole number of dt steps. Requires tau_steps >= 1 (tau >= dt): the per-stage delayed reads at intra-step times up to step + 1 (and the derivative stored at write time) must only ever land on ring-buffer slots already finalized by earlier steps, which fails once the delay is shorter than one step. Costs one extra dfun evaluation per step (to compute the derivative to store) and doubles buf’s per-slot size (value and derivative, instead of just value).

Trimmed connectivity dataclass.

Adapted from vbi/simulator/spec/connectivity.py - see NOTICE.md in this directory for provenance and what was dropped. vbi’s CouplingSpec (a fixed-kind enum) is not vendored here - lyapax’s coupling is a plain callable instead (see lyapax.coupling).

class lyapax.simulator.coupling.Connectivity(weights: ndarray, tract_lengths: ndarray | None = None, speed: float = 1.0)[source]#

Structural connectivity: weights + optional per-edge delay (tract length / speed). tract_lengths is None (or all-zero) means a pure ODE (zero-delay) network; any non-zero entry makes it a DDE network.

delay_steps(dt: float) → ndarray[source]#

(n, n) int32 array of delay expressed in integration steps.

Integer-step only (no sub-step interpolation) - see Grid-snapped vs. interpolated history reads for the accuracy tradeoff this implies.

horizon(dt: float) → int[source]#

Ring-buffer depth = max delay (in steps) + 1.

Integrators#

General-purpose fixed-step ODE integrators.

Independent of the vendored (Euler/Heun) integrators in lyapax.simulator.step - those are tied to the coupling/ring-buffer plumbing for network models (M3+). These are plain state -> new_state maps for standalone benchmark systems (Lorenz, Rössler, …), where RK4’s better accuracy at a given dt matters for getting a clean Lyapunov-exponent estimate (see Tier 2 of the validation guide).

lyapax.integrators.RK6_STAGE_C = (0.0, 0.06, 0.09593333333333333, 0.1439, 0.4973, 0.9725, 0.9995, 1.0)#

Butcher c_i (fractional position within the step) for each of rk6_combine’s 8 stages – verified via each stage’s row-sum consistency condition (sum_j a_ij == c_i) against the source tableau, see rk6_combine’s docstring. Exposed so a coupled/delayed caller can recompute a state- or time-dependent input (e.g. network coupling, or a DDE’s delayed history lookup) at the correct intra-step point for each stage, instead of freezing it at the step’s start – see Delays, coupling, and Runge–Kutta stage order.

lyapax.integrators.get_integrator(name: str) → Callable[[Callable[[Array], Array], float], Callable[[Array], Array]][source]#

Look up one of the built-in fixed-step ODE integrator builders by name ("euler", "heun", "rk4", or "rk6"), each mapping (rhs, dt) -> step_fn.

lyapax.integrators.heun_step(rhs: Callable[[Array], Array], dt: float) → Callable[[Array], Array][source]#

Fixed-step 2nd-order Heun (explicit trapezoidal).

lyapax.integrators.ode_step(rhs: Callable[[Array], Array], dt: float, integrator: str | Callable[[Callable[[Array], Array], float], Callable[[Array], Array]] = 'rk4') → Callable[[Array], Array][source]#

Build a single-step state -> new_state map for a plain (uncoupled) ODE, with the integration method as an explicit, swappable argument.

Parameters:

  • rhs – right-hand side, state (d,) -> dstate (d,).

  • dt – fixed step size.

  • integrator – one of "euler", "heun", "rk4", "rk6", or a callable (rhs, dt) -> step_fn (e.g. rk4_step itself, or a user-supplied builder with the same signature).

lyapax.integrators.rk4_step(rhs: Callable[[Array], Array], dt: float) → Callable[[Array], Array][source]#

Classic fixed-step 4th-order Runge-Kutta.

lyapax.integrators.rk6_combine(f: Callable[[Array, float], Array], state: Array, dt: float) → Array[source]#

One fixed-step 6th-order Runge-Kutta update, state -> new_state, against a derivative evaluator f(y, c) -> dy, where c is that stage’s Butcher node (RK6_STAGE_C) – the fractional position within the step, 0 at the first stage, 1 at the last. Purely autonomous callers (no coupling, no explicit time-dependence) can ignore c; lyapax.simulator.step’s coupled-network RK6 uses it to recompute coupling (or a delayed history lookup) fresh at each stage’s own intra-step state/time, rather than freezing it once per step (see Delays, coupling, and Runge–Kutta stage order) – this and rk6_step share the same tableau instead of duplicating ~30 lines of coefficients twice.

The 8-stage tableau is the a_ij coefficients (stages 1-8) of Verner’s “efficient order 6(5)” embedded pair – the coefficients underlying Vern6 in Julia’s OrdinaryDiffEq.jl – with the 6th-order solution taken as the FSAL point these stages build towards (c=1, i.e. what would be stage 9’s argument), rather than a separately weighted sum over 9 evaluated stages: an earlier draft of this function used the b1..b9 weights given as comments in the Julia source directly, which turned out to be a red herring – symbolically expanding that combination’s stability function against e^z (exact rational arithmetic, state' = z*state) showed it matches only through z^5, i.e. order 5, not 6. The 8-stage FSAL combination below was verified the same way and matches e^z through z^6 (mismatching first at z^7, the expected order-6 signature), and separately confirmed with an empirical convergence-order check against a genuinely nonlinear system with a known exact solution – see tests/test_integrators.py.

lyapax.integrators.rk6_step(rhs: Callable[[Array], Array], dt: float) → Callable[[Array], Array][source]#

Fixed-step 6th-order Runge-Kutta, for when RK4’s O(dt^4) local error is the accuracy bottleneck (e.g. wanting a much larger dt at the same error, or a much smaller error at the same dt, than RK4 gives). See rk6_combine for the tableau, provenance, and correctness checks.

Adaptive#

Optional adaptive-step ODE integration via the adaptive extra (pip install lyapax[adaptive]); not imported by lyapax’s top-level __init__.

Adaptive-step ODE integration via diffrax.

ODE-only: diffrax has no DDE support (patrick-kidger/diffrax#406, open since April 2024, re-checked 2026-07-07) – lyapax.dde stays fixed-step, see that module.

Not imported by lyapax’s top-level __init__, since diffrax is an optional dependency (the adaptive extra) – import this module directly: from lyapax.adaptive import diffrax_adaptive_step.

This is not a speed optimization – measured slower, not faster, than fixed-step integration. At matched accuracy (not matched dt/rtol, which isn’t a fair comparison), this integrator was 2-4x slower than rk4_step/rk6_step on Lorenz, and up to ~4.5x slower on GPU for that same small system, across every system tested including large (d=1500-3000) ones – the gap narrows at scale but a fixed-step win was never overturned. Root cause: the internal accept/reject jax.lax.while_loop pays real per-step overhead (PID controller update, solver.step bookkeeping) that a fixed-step raw step doesn’t. Checked and still true even on a relaxation oscillator (Van der Pol, mu=30, sharply time-varying stiffness unlike Lorenz’s roughly uniform one) where fixed-step “wastes” steps during slow phases: RK4 at matched accuracy was still ~4x faster (0.31s vs 1.28s), because avoiding those extra steps only pays off if each one was expensive – for a small system they’re nearly free regardless of stiffness, so adaptive’s fixed control-flow overhead per step is never amortized. A win would need both sharply time-varying stiffness and an expensive per-step cost (large state dimension) together; not tested, and not assumed here. Reach for this integrator for tolerance-driven accuracy control (no more manual dt-convergence sweeps) and for its being the only forward-mode-differentiable adaptive option here, not for throughput. See Adaptive integration is not a speed optimization for a summary of these findings.

Design: diffrax_adaptive_step(...) returns a builder with the same (rhs, dt) -> step_fn signature as the fixed-step builders in lyapax.integrators (rk4_step, etc.), so it drops straight into lyapax.ode_problem(rhs, state0, dt, integrator=diffrax_adaptive_step(...)) with no changes to lyapax.core – core.py’s _advance already treats step_fn as an opaque state -> new_state map called once per raw step and differentiated with jax.jvp; whether that one call does a single fixed-size update or an internal accept/reject/step-size-adaptation loop is invisible to it. dt keeps its usual lyapax meaning here too – the fixed sampling/renormalization interval lyapunov_spectrum advances by per raw step – and is not the adaptive integrator’s own internal step size, which varies within [t, t + dt] under rtol/atol control.

The internal accept/reject loop is a jax.lax.while_loop over diffrax’s low-level solver.step primitive (the same primitive diffrax.diffeqsolve is built on), mirroring diffrax._integrate.loop’s own step/clip-to-boundary structure closely enough to inherit its correctness, but scoped to a single dt advance with no SaveAt/dense-output machinery. See tests/test_adaptive_diffrax.py for the isolated primitive check this is built on, and two findings that shaped this implementation:

  • ``scan_kind=”lax”`` is required. A diffrax Runge-Kutta solver’s default construction (scan_kind=None, internally "checkpointed") iterates its stages via an equinox-checkpointed while_loop wrapped in a custom_vjp that rejects forward-mode jax.jvp outright (TypeError). scan_kind="lax" uses a plain jax.lax.scan instead, which is jvp-compatible. Enforced here, not left to the caller.

  • Reverse-mode differentiation (``jax.grad``/``jax.jacrev``) does not work end-to-end through this integrator, even though the Lyapunov engine’s own tangent propagation is forward-mode (jax.jvp, see lyapax.core’s module docstring): the outer while_loop here has a data-dependent trip count (however many accept/reject iterations a given rtol/atol/system needs), and JAX’s reverse-mode AD cannot replay a lax.while_loop with a dynamic trip count backward at all – a fundamental JAX limitation, not a diffrax quirk or a precision issue. Differentiating a Lyapunov exponent w.r.t. a handful of system parameters still works, but requires jax.jacfwd/jax.jvp (forward-mode over the parameter), not jax.grad. Confirmed empirically: jax.grad raises ValueError: Reverse-mode differentiation does not work for lax.while_loop ... with dynamic start/stop values; jax.jacfwd of the same function matches a central finite-difference reference to ~1e-8.

diffrax’s own step-size controller (PIDController) already applies jax.lax.stop_gradient internally to its step-size-adjustment factor and initial-step guess (see its source) – the “stop_gradient on the step-size decision” concern from the design note is therefore already handled by using PIDController, not something this module adds on top.

lyapax.adaptive.diffrax_adaptive_step(solver=None, rtol: float = 1e-06, atol: float = 1e-09, dt0: float | None = None, max_steps: int = 4096) → Callable[[Callable[[Array], Array], float], Callable[[Array], Array]][source]#

Build a lyapax.ode_step-compatible adaptive integrator: (rhs, dt) -> step_fn where step_fn(state) -> new_state advances by exactly dt using an adaptive-step diffrax solver internally.

Parameters:

  • solver – a diffrax solver instance, e.g. diffrax.Dopri5(), constructed with scan_kind="lax". Defaults to diffrax.Dopri5(scan_kind="lax"). Raises ValueError if given a solver without scan_kind="lax" – see the module docstring.

  • atol (rtol,) – diffrax.PIDController step-size tolerances.

  • dt0 – initial internal step-size guess. Defaults to diffrax’s own heuristic initial-step selection (PIDController.init).

  • max_steps – safety cap on internal accept+reject iterations per dt advance, so a misbehaving rhs/tolerance combination loops a bounded number of times inside jax.lax.while_loop rather than running away. Unlike diffrax.diffeqsolve (which returns whatever was reached when max_steps is hit unless throw=True), exhausting max_steps here always raises via equinox.error_if – returning a state from an earlier time than the requested t1 while lyapunov_spectrum advances its elapsed-time denominator by the full dt regardless would silently mis-associate exponent with elapsed time, which check_finite=True cannot detect (a truncated step is still finite). equinox.error_if raises immediately in eager mode and under jax.jit/jax.vmap; see its docstring for the EQX_ON_ERROR environment variable if you need to disable this check or substitute NaN instead. Raise max_steps or loosen rtol/atol if this triggers on a well-posed system.

Usage#

>>> import lyapax
>>> from lyapax.adaptive import diffrax_adaptive_step
>>> problem = lyapax.ode_problem(
...     rhs, state0, dt=0.1, integrator=diffrax_adaptive_step(rtol=1e-8, atol=1e-10),
... )
>>> result = lyapax.lyapunov_spectrum(problem, n_steps=...)

Coupling#

Network coupling as plain callables.

Design note: coupling here is deliberately not a closed enum dispatched through hardcoded if/elif branches, unlike vbi’s CouplingSpec.kind + build_coupling (duplicated across its numpy and JAX backends, with jr_sigmoidal even special-cased to skip the G multiply). Each builder below just returns a plain function:

coupling_fn(cvar_state, weights, params) -> coupling

that lyapax.network.make_network_step_fn calls directly. A user’s own function with that exact signature is a first-class coupling - no registry, no “kind” string to add to a dispatch table, no library changes.

param cvar_state:

(n_cvar, n_nodes) coupling-variable state – the instantaneous state for zero-delay networks, a delayed array of the same shape for DDE networks.

param weights:

(n_nodes, n_nodes) weights[tgt, src].

param params:

dict - the global coupling strength G is read from here (params.get("G", G_default)), the same convention the vendored step function uses, so G can be swept or differentiated without being baked into a closure constant.

returns:

(n_cvar, n_nodes)

lyapax.coupling.kuramoto_coupling(alpha: float = 0.0, G_default: float = 1.0) → Callable[[Array, Array, dict], Array][source]#

c[0, tgt] = (G/N) * sum_src w[tgt, src] * sin(theta_src - theta_tgt + alpha).

Requires exactly one coupling variable (the phase); cvar_state[0] is taken as theta.

lyapax.coupling.linear_coupling(a: float = 1.0, b: float = 0.0, G_default: float = 1.0) → Callable[[Array, Array, dict], Array][source]#

c[cvar, tgt] = G * a * sum_src(w[tgt, src] * x[cvar, src]) + b.

lyapax.coupling.sigmoidal_coupling(a: float = 1.0, b: float = 0.0, midpoint: float = 0.0, sigma: float = 1.0, G_default: float = 1.0) → Callable[[Array, Array, dict], Array][source]#

c[cvar, tgt] = G * a * sum_src(w[tgt, src] * sigm(x[cvar, src])) + b, sigm(x) = 1 / (1 + exp(-(x - midpoint) / sigma)).

Network#

Coupled, zero-delay networks - wires a dfun + a coupling callable into the flat state -> new_state shape lyapax.core.lyapunov_spectrum expects.

Built on the vendored ring-buffer step (lyapax.simulator.make_step_fn, has_delays=False) via a thin carry-to-flat adapter, not a second, independent Euler/Heun implementation – the adapter is the single place that bridges the carry-based step to the flat map lyapax.core.lyapunov_spectrum wants, so there is only one copy of the integrators to maintain.

Why this still doesn’t need lyapax.dde’s carry-aware tangent engine: for has_delays=False the vendored step’s buf never changes (it’s a structural no-op) and t is never read, so closing over a fixed placeholder buf/t and only threading state through lyapunov_spectrum’s jacfwd is exactly equivalent to differentiating the whole carry – there is no delay-buffer tangent information to lose.

class lyapax.network.Network(weights: Array, cvar_indices: tuple[int, ...], delay_steps: int | float | Array | None = None)[source]#

Topology for a coupled network: what M3+ call weights/ cvar_indices (and, for delayed networks, delay_steps), gathered into one named object instead of separate positional arguments repeated across network_step/network_dde_problem call sites.

Parameters:

  • weights – (n_nodes, n_nodes), weights[tgt, src].

  • cvar_indices – indices into the state-variable axis selecting which rows of state feed the coupling.

  • delay_steps – None for a zero-delay network, a number (int, or float when interpolating – see lyapax.simulator.make_step_fn) for a single global delay (in units of dt), or an (n_nodes, n_nodes) integer matrix for per-edge delays. Only relevant to the DDE path (network_dde_problem); ignored by network_step.

lyapax.network.make_network_step_fn(dfun: Callable, weights: Array, cvar_indices: tuple[int, ...], params: dict, dt: float, coupling_fn: Callable[[Array, Array, dict], Array], integrator: str | Callable = 'heun') → Callable[[Array], Array][source]#

Build a flat state_flat (n_sv * n_nodes,) -> new_state_flat step function for use with lyapax.core.lyapunov_spectrum.

Parameters:

  • dfun – (state, coupling, params) -> dstate, e.g. from lyapax.simulator.build_jax_dfun(model_spec). state and coupling are (n_sv, n_nodes) / (n_cvar, n_nodes).

  • weights – (n_nodes, n_nodes), weights[tgt, src].

  • cvar_indices – indices into the state-variable axis selecting which rows of state feed the coupling (ModelSpec.cvar_indices).

  • params – closed over for this step function. To sweep over params via jax.vmap instead, use make_parametrized_network_step_fn + lyapax.sweep.sweep_lyapunov_spectrum (M6).

  • coupling_fn – see lyapax.coupling - any callable with signature (cvar_state, weights, params) -> coupling works, including a user-defined one.

  • integrator – "euler", "heun", "rk4", "rk6", or a callable (state, dfun, coupling, dt, params) -> new_state – see lyapax.simulator.make_step_fn.

lyapax.network.make_parametrized_network_step_fn(dfun: Callable, weights: Array, cvar_indices: tuple[int, ...], dt: float, coupling_fn: Callable[[Array, Array, dict], Array], integrator: str | Callable = 'heun') → Callable[[Array, dict], Array][source]#

Same wiring as make_network_step_fn, except params is a plain call-time argument, (state_flat, params) -> new_state_flat, instead of being closed over at construction time – for M6’s jax.vmap parameter sweeps (lyapax.sweep.sweep_lyapunov_spectrum), where params needs to be data the vmapped function can batch over, not a baked-in Python constant. Works because lyapax.simulator.make_step_fn already threads params through the scanned carry rather than closing over it (see that function’s docstring) – this adapter was the only piece still closing params in early, so it’s the only piece that needed to change.

dfun, weights, cvar_indices, coupling_fn, integratorsee

make_network_step_fn.

lyapax.network.network_problem(dfun: Callable, network: Network, coupling: Callable[[Array, Array, dict], Array], params: dict, state0: Array, dt: float, integrator: str | Callable = 'heun') → ODEProblem[source]#

Build an ODEProblem for a zero-delay coupled network – bundles network_step’s step_fn together with state0 and dt, the same lyapunov_spectrum(problem, n_steps=...) recipe ode_problem gives the plain-ODE case (see ODEProblem).

Parameters:

  • dfun – (state, coupling, params) -> dstate.

  • network – topology, see Network.

  • coupling – any lyapax.coupling-style callable, (cvar_state, weights, params) -> coupling.

  • params – closed over for this step function.

  • state0 – initial state, either already flat (n_sv * n_nodes,) or in the network’s natural (n_sv, n_nodes) shape – flattened here either way, since network_step’s step function operates on flat state.

  • dt – fixed step size.

  • integrator – "euler", "heun", "rk4", "rk6", or a callable – see lyapax.simulator.make_step_fn.

Examples

>>> import jax
>>> jax.config.update("jax_enable_x64", True)
>>> import jax.numpy as jnp
>>> from lyapax.core import lyapunov_spectrum
>>> from lyapax.coupling import linear_coupling
>>> from lyapax.network import Network, network_problem
>>> from lyapax.simulator import ModelSpec, StateVar, Parameter, build_jax_dfun
>>> weights = jnp.array([[0., 1., 0., 1.],
...                      [1., 0., 1., 0.],
...                      [0., 1., 0., 1.],
...                      [1., 0., 1., 0.]])  # 4-cycle graph
>>> model = ModelSpec(
...     name="linear_node",
...     state_variables=(StateVar("x", default_init=0.0),),
...     parameters=(Parameter("gamma", -2.0),),
...     cvar=("x",),
...     dfun_str={"x": "gamma * x + c"},
... )
>>> dfun = build_jax_dfun(model)
>>> network = Network(weights=weights, cvar_indices=model.cvar_indices)
>>> problem = network_problem(
...     dfun, network, linear_coupling(a=1.0, b=0.0),
...     params={"gamma": -2.0, "G": 0.5},
...     state0=jnp.array([0.3, -0.1, 0.2, -0.4]), dt=1e-3,
... )
>>> result = lyapunov_spectrum(problem, n_steps=20_000, renorm_every=10)
>>> result.exponents
Array([-1., -2., -2., -3.], dtype=float64)
lyapax.network.network_step(dfun: Callable, network: Network, coupling: Callable[[Array, Array, dict], Array], params: dict, dt: float, integrator: str | Callable = 'heun') → Callable[[Array], Array][source]#

Build a flat state_flat -> new_state_flat step for a zero-delay coupled network, naming the four concepts make_network_step_fn otherwise mixes into one long argument list separately: model dynamics (dfun), topology (network), the coupling rule (coupling), and the numerical method (integrator).

A thin wrapper – it calls straight through to make_network_step_fn, reading weights/cvar_indices off network. Prefer this over make_network_step_fn in new code; the latter remains for direct, lower-level use.

Parameters:

  • dfun – (state, coupling, params) -> dstate, e.g. from lyapax.simulator.build_jax_dfun(model_spec).

  • network – topology, see Network.

  • coupling – any lyapax.coupling-style callable, (cvar_state, weights, params) -> coupling.

  • params – closed over for this step function.

  • dt – fixed step size.

  • integrator – "euler", "heun", "rk4", "rk6", or a callable – see lyapax.simulator.make_step_fn.

lyapax.network.network_step_parametrized(dfun: Callable, network: Network, coupling: Callable[[Array, Array, dict], Array], dt: float, integrator: str | Callable = 'heun') → Callable[[Array, dict], Array][source]#

network_step’s counterpart for jax.vmap parameter sweeps – see make_parametrized_network_step_fn, which this wraps.

Systems#

Benchmark systems for validating the Lyapunov engine.

Plain flat-vector JAX functions, deliberately independent of the vendored ModelSpec/coupling machinery in lyapax.simulator - these are single-node, uncoupled textbook systems, so routing them through the string-expression codegen built for networked models would only add indirection. Each entry here corresponds to a tier of the validation suite; see Validation: how lyapax’s results are checked for the reference Lyapunov values and citations.

Tier 4 (DDE) entries have a different shape than the rest: rhs_delayed (state_now, state_delayed) -> dstate instead of rhs(state) -> dstate, since they need both the current and the delayed state (see lyapax.dde.DelayedRHS). Pair with lyapax.dde.make_scalar_delayed_step_fn, not the plain integrators in lyapax.integrators – mirrors this module’s relationship to lyapax.integrators.rk4_step for the ODE case: plain math here, a step-function builder elsewhere.

lyapax.systems.henon_map(a: float = 1.4, b: float = 0.3) → Callable[[Array], Array][source]#

Constant Jacobian determinant -b => sum(LE) = ln|b| exactly (Tier 0.3).

lyapax.systems.linear_scalar_dde(a: float = 1.0)[source]#

x’(t) = -a*x(t-tau). Dominant Lyapunov exponent solves the transcendental characteristic equation lambda = -a*exp(-lambda*tau) (Lambert W: lambda = W(-a*tau)/tau) – a closer-to-analytic reference for isolating DDE tangent-propagation bugs from Mackey-Glass’s nonlinear chaos (Tier 4.2).

lyapax.systems.linear_system(A: Array) → Callable[[Array], Array][source]#

rhs(x) = A @ x. Exact Lyapunov spectrum = Re(eigvals(A)) (Tier 0.1).

lyapax.systems.logistic_map(r: float) → Callable[[Array], Array][source]#

x_{n+1} = r x (1-x). At r=4, exact LE = ln(2) (Tier 0.2).

lyapax.systems.lorenz(sigma: float = 10.0, rho: float = 28.0, beta: float = 2.6666666666666665) → Callable[[Array], Array][source]#

Constant trace(J) = -(sigma+1+beta) => sum(LE) exactly known (Tier 1.1); lambda1 has a well-documented published value (Tier 2).

lyapax.systems.mackey_glass(beta: float = 0.2, gamma: float = 0.1, n: float = 10.0)[source]#

x’(t) = beta*x(t-tau)/(1+x(t-tau)^n) - gamma*x(t). The classic beta=0.2, gamma=0.1, n=10, tau=17 parameter set is chaotic – the system Farmer (1982) used to introduce the discretized-map method for DDE Lyapunov spectra (Tier 4.1).

lyapax.systems.rossler(a: float = 0.2, b: float = 0.2, c: float = 5.7) → Callable[[Array], Array][source]#

trace(J) = a + x - c (state-dependent) => sum(LE) = a - c + <x>_t (Tier 1.2); lambda1 has a published order-of-magnitude value (Tier 2).

lyapax.systems.tent_map() → Callable[[Array], Array][source]#

Exact LE = ln(2) (Tier 0.2).

Sweep#

Batched parameter/initial-condition sweeps via jax.vmap.

Mirrors vbi’s JaxSweeper pattern (batch a simulation over a grid of parameter values) without reusing its code, per the vendoring decision. The surface needed here is much smaller than a general sweeper: lyapax.simulator.make_step_fn’s carry already threads params through as data rather than closing over it (see that function’s docstring), so a parameter sweep is just jax.vmap applied to a function that calls lyapax.core.lyapunov_spectrum once, with params (and optionally state0) as the vmapped argument. No new tangent-propagation or QR machinery – this is a batching wrapper around the existing engine, not a third one.

The one piece of the network path that still closed over params at construction time was lyapax.network.make_network_step_fn (a thin adapter, not the tangent engine); make_parametrized_network_step_fn (in that module) is the call-time-params sibling this needs.

lyapax.sweep.sweep_lyapunov_spectrum(step_fn: Callable[[Array, dict], Array], state0: Array, params_batch: dict, dt: float, n_steps: int, k: int | None = None, renorm_every: int = 1, t_transient: float = 0.0, seed: int = 0, state0_batch: Array | None = None) → LyapunovResult[source]#

Batched lyapunov_spectrum over a grid of parameter values (and, optionally, a matching grid of initial conditions) via one jax.vmap call, rather than a Python loop making one call per grid point (as in e.g. examples/05_kuramoto_sync.py/09_kuramoto_delayed_network.py).

step_fn(state, params) -> new_state – params must be a

call-time argument, not closed over. See lyapax.network.make_parametrized_network_step_fn for the network-coupling case; a hand-written function with this signature works too (same “plain callable, no registry” extension point as lyapax.coupling).

state0(d,) initial state, shared across the whole sweep, unless

state0_batch is given.

params_batchpytree matching one params dict’s structure, but

every leaf carries an extra leading batch axis of size n_sweep – e.g. {"G": jnp.linspace(0, 4, 13), "omega": jnp.broadcast_to(omega, (13, n_nodes))} to sweep G at fixed omega.

dt, n_steps, k, renorm_every, t_transient, seedsee

lyapax.core.lyapunov_spectrum - shared (not swept) across the whole grid.

state0_batchoptional (n_sweep, d) initial conditions, one row

per grid point, if the initial condition should vary alongside params_batch rather than staying fixed at state0.

Returns:

  • LyapunovResult, every field with an extra leading batch axis of size

  • n_sweep – e.g. result.exponents.shape == (n_sweep, k).

Examples

Sweep the coupling strength G of a Kuramoto network in one jax.vmap call instead of a Python loop over lyapunov_spectrum:

>>> import jax
>>> jax.config.update("jax_enable_x64", True)
>>> import jax.numpy as jnp
>>> from lyapax.coupling import kuramoto_coupling
>>> from lyapax.network import Network, network_step_parametrized
>>> from lyapax.simulator import ModelSpec, StateVar, Parameter, build_jax_dfun
>>> from lyapax.sweep import sweep_lyapunov_spectrum
>>> n_nodes = 6
>>> omega = jnp.linspace(-1.0, 1.0, n_nodes)
>>> weights = jnp.ones((n_nodes, n_nodes)) - jnp.eye(n_nodes)
>>> model = ModelSpec(
...     name="kuramoto",
...     state_variables=(StateVar("theta", default_init=0.0),),
...     parameters=(Parameter("omega", 0.0),),
...     cvar=("theta",),
...     dfun_str={"theta": "omega + c"},
... )
>>> dfun = build_jax_dfun(model)
>>> network = Network(weights=weights, cvar_indices=model.cvar_indices)
>>> dt = 1e-2
>>> state0 = jnp.linspace(0.0, 2 * jnp.pi, n_nodes, endpoint=False)
>>> G_values = jnp.linspace(0.0, 4.0, 13)
>>> step_p = network_step_parametrized(dfun, network, kuramoto_coupling(alpha=0.0), dt)
>>> params_batch = {"omega": jnp.broadcast_to(omega, (13, n_nodes)), "G": G_values}
>>> result = sweep_lyapunov_spectrum(
...     step_p, state0, params_batch, dt=dt, n_steps=5_000,
...     renorm_every=10, k=2, t_transient=50.0,
... )
>>> result.exponents.shape
(13, 2)

Utils#

Small helpers for demos/examples – not part of the core LE engine.

lyapax.utils.simulate_trajectory(step_fn: Callable[[Array], Array], state0: Array, n_steps: int, dt: float | None = None) → Array | tuple[Array, Array][source]#

Run step_fn for n_steps iterations via lax.scan (fast, compiled – unlike a Python for-loop calling a JAX function repeatedly).

Returns the trajectory including the initial state: shape (n_steps + 1, d). If dt is given, also returns the matching time axis as (t, traj), t = arange(n_steps + 1) * dt – dt is otherwise opaque to step_fn (baked into its closure), so this is the only place that can hand it back to the caller.

Contents