Metadata-Version: 2.4
Name: chaotax
Version: 0.0.1
Summary: Simple JAX-based tool to compute Lyapunov spectra and exponents of dynamical systems.
Project-URL: repository, https://github.com/Ceyron/chaotax
Author: Felix Koehler
License-Expression: MIT
License-File: LICENSE
Keywords: chaos,dynamical-systems,jax,lyapunov,sciml
Requires-Python: >=3.10
Requires-Dist: jax>=0.4.13
Requires-Dist: jaxtyping>=0.2.20
Provides-Extra: docs
Requires-Dist: exponax; extra == 'docs'
Requires-Dist: furo; extra == 'docs'
Requires-Dist: ipykernel; extra == 'docs'
Requires-Dist: matplotlib; extra == 'docs'
Requires-Dist: nbsphinx; extra == 'docs'
Requires-Dist: sphinx; extra == 'docs'
Requires-Dist: sphinx-autodoc-typehints; extra == 'docs'
Provides-Extra: gpu
Requires-Dist: jax[cuda12]>=0.8.0; extra == 'gpu'
Provides-Extra: test
Requires-Dist: equinox; extra == 'test'
Requires-Dist: exponax; extra == 'test'
Requires-Dist: pytest; extra == 'test'
Requires-Dist: pytest-cov; extra == 'test'
Description-Content-Type: text/markdown

<p align="center">
  <img src="docs/_static/logo.svg" alt="chaotax logo" width="200">
</p>

**🚨 This is a pre-release for me to cite in my dissertation. It is based on
software I once wrote for [APEBench](https://github.com/tum-pbs/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.**

<h1 align="center">
    <span style="color: #f2614f;">chao</span><span style="color: #3b59b4;">tax</span>
</h1>
<h4 align="center">Simple <a href="https://github.com/google/jax" target="_blank">JAX</a>-based tool to compute Lyapunov spectra and exponents of dynamical systems using autodiff.</h4>

<p align="center">
<a href="https://pypi.org/project/chaotax/">
  <img src="https://img.shields.io/pypi/v/chaotax.svg" alt="PyPI">
</a>
<a href="https://github.com/ceyron/chaotax/actions/workflows/tests.yml">
  <img src="https://github.com/ceyron/chaotax/actions/workflows/tests.yml/badge.svg" alt="Tests">
</a>
<a href="https://fkoehler.site/chaotax/">
  <img src="https://img.shields.io/badge/docs-latest-green" alt="docs-latest">
</a>
<a href="https://github.com/ceyron/chaotax/blob/main/LICENSE">
  <img src="https://img.shields.io/badge/license-MIT-blue" alt="License">
</a>
</p>

<p align="center">
  <a href="#installation">Installation</a> •
  <a href="#quickstart">Quickstart</a> •
  <a href="#documentation">Documentation</a> •
  <a href="#related">Related</a>
</p>

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

## Installation

```bash
pip install chaotax
```

Requires Python 3.10+ and JAX 0.4.13+. See the [JAX install guide](https://jax.readthedocs.io/en/latest/installation.html).

## Quickstart

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

```python
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:

```python
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](https://fkoehler.site/chaotax/).

## Related

- [Exponax](https://github.com/Ceyron/exponax) — Exponential and explicit time steppers for PDEs in JAX
- [Picardax](https://github.com/Ceyron/picardax) — Time-implicit differentiable PDE solvers in JAX

## License

MIT, see [here](https://github.com/Ceyron/chaotax/blob/main/LICENSE).

---

> [fkoehler.site](https://fkoehler.site/) &nbsp;&middot;&nbsp;
> GitHub [@ceyron](https://github.com/ceyron) &nbsp;&middot;&nbsp;
> X [@felix_m_koehler](https://twitter.com/felix_m_koehler) &nbsp;&middot;&nbsp;
> LinkedIn [Felix Koehler](https://www.linkedin.com/in/felix-koehler)
