🚨 This is a pre-release for me to cite in my dissertation. It is based on software I once wrote for APEBench (but did not yet have the time to release) slightly rewritten and enhanced. Please check back end of May 2026 for a full release with extended features, documentation, and examples.

chaotax

Simple JAX-based tool to compute Lyapunov spectra and exponents of dynamical systems using autodiff.

Installation • Quickstart • Documentation • Related

Chaotax estimates Lyapunov exponents and spectra of dynamical systems using JAX's automatic differentiation (JVP-based linearisation). It works with any one-step state transition function, is JIT-compilable, and GPU/TPU-ready.

Installation

pip install chaotax

Requires Python 3.10+ and JAX 0.4.13+. See the JAX install guide.

Quickstart

Estimate the largest Lyapunov exponent of a system given a one-step transition function:

import jax
import jax.numpy as jnp
import chaotax

# Define a one-step stepper (e.g., Lorenz via RK4)
def lorenz_stepper(u):
    ...
    return u_next

u_0 = jnp.array([1.0, 0.0, 0.0])
eps_0 = jax.random.normal(jax.random.PRNGKey(0), u_0.shape)

estimate_le = chaotax.lyapunov(lorenz_stepper, rollout_steps=20_000, dt=0.01)
le = estimate_le(u_0, eps_0)

Compute the full Lyapunov spectrum:

estimate_spectrum = chaotax.lyapunov_spectrum(lorenz_stepper, rollout_steps=20_000, dt=0.01)
spectrum = estimate_spectrum(u_0, eps_0=jnp.eye(3))
# Returns exponents in descending order, e.g. [+, 0, -] for Lorenz

Documentation

Documentation is available at fkoehler.site/chaotax.

Related

• Exponax — Exponential and explicit time steppers for PDEs in JAX
• Picardax — Time-implicit differentiable PDE solvers in JAX

License

MIT, see here.

fkoehler.site  ·  GitHub @ceyron  ·  X @felix_m_koehler  ·  LinkedIn Felix Koehler