lyapax

JAX-native Lyapunov exponent computation for ODEs and DDEs, via the Benettin/QR method with jax.jvp/jax.vmap tangent propagation.

📚 Documentation

Looking for the documentation?

➡️ https://lyapax.readthedocs.io/en/latest/

It includes:

• Installation guide
• Quick start tutorial
• API reference
• Examples and gallery
• Validation, benchmarks, and implementation details

Install

pip install lyapax

Development installation

If you plan to contribute, run the test suite, or build the documentation, clone the repository and install the appropriate extras:

pip install -e ".[dev]"       # testing
pip install -e ".[examples]"  # example scripts
pip install -e ".[docs]"      # documentation
pip install -e ".[adaptive]"  # Diffrax backend

Requires jax>=0.10, Python >=3.11.

Enable float64 before anything else. Lyapunov exponents are averages of log-growth rates accumulated over many steps; JAX's default float32 silently degrades long-horizon estimates. Do this first, before creating any jax.numpy arrays:

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

lyapunov_spectrum/lyapunov_spectrum_dde warn if called with a float32 state0 and x64 is off.

Minimal example

import jax
jax.config.update("jax_enable_x64", True)
import jax.numpy as jnp

from lyapax import lyapunov_spectrum, ode_problem

sigma, rho, beta = 10.0, 28.0, 8.0 / 3.0

def rhs(state):
    x, y, z = state
    return jnp.array([sigma * (y - x), x * (rho - z) - y, x * y - beta * z])

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,
)
print(result.exponents)  # ~ [0.906, 0.0, -14.57]

result.history gives the running per-column estimate at each renormalization point, for checking convergence.

rhs is just a plain state -> jnp.ndarray function - write your own for any system. lyapax.systems has this Lorenz system and other standard test systems (Rössler, Kuramoto, ...) prebuilt, if you'd rather not retype them.

Method and scope

lyapunov_spectrum propagates a (d, k) matrix of tangent vectors alongside the trajectory via jax.jvp (one forward-mode pass per tracked column, batched with jax.vmap), QR-decomposes it every renorm_every steps, accumulates log|diag(R)|, and divides by elapsed time (Benettin's method). k defaults to the full spectrum (k = d); cost scales with k, not d, so partial spectra of high-dimensional systems are cheap.

lyapunov_spectrum_dde generalizes this to fixed-delay DDEs by differentiating through an augmented (state, ring_buffer) carry.

What this does not do:

• No stiff (implicit) ODE solvers. step_fn is whatever integrator you hand it - a fixed-step map (ode_problem(..., integrator="rk4") by default, direct rk4_step/euler_step step functions, or Euler/Heun/RK6 via lyapax.simulator), or an adaptive-step explicit one via the optional adaptive extra (lyapax.adaptive.diffrax_adaptive_step, backed by diffrax; ode_problem only - not network_problem or DDEs). Exponents are for the numerical time-dt map (or accepted-step sequence, for the adaptive integrator), not an exact flow - check dt/tolerance convergence for anything you report. The adaptive integrator is for tolerance-driven accuracy control, not speed - measured 2-4x slower than rk4/rk6 at matched accuracy, not faster; see docs/background/capabilities.md.
• DDE delays must be known and fixed, resolved one of two ways. By default (grid-snapped), a physical delay tau is rounded to the nearest multiple of dt (lyapax.dde.resolve_tau_steps); use lyapax.dde.tau_eff to see the delay actually used, and shrink dt to converge it toward tau. Passing interpolate=True instead reconstructs the delayed history with a cubic-Hermite interpolant, removing that rounding on the uniform-delay path at some extra cost - see 13_dde_history_interpolation.py and docs/background/lyapax_implementation.md.
• history columns are ordered once, by the final row. Near- degenerate exponents can swap order over the run - see LyapunovResult.history's docstring.

Examples

Runnable, sphinx-gallery-formatted demos in examples/ (pip install -e ".[examples]", then python examples/01_linear_ode.py):

Building the docs

pip install -e ".[docs,adaptive]"
sphinx-build -b html docs docs/_build/html

This re-runs every numbered examples/*.py file to render it into a sphinx-gallery gallery (code, output, and plots), including 15_adaptive_ode.py (hence the adaptive extra above), alongside the API reference; open docs/_build/html/index.html when it's done.

Further reading

• docs/background/validation.md - the correctness tests this package is held to (exact values, structural invariants, literature figures).
• docs/background/benchmarks.md - comparison against jitcode/jitcdde and ChaosTools.jl.
• docs/background/capabilities.md - a candid summary of what lyapax does and does not do, including measured performance characteristics.