kellax: continuation in JAX

Quickstart

The figure above is this program. The residual is one lambda; the rest follows from it by automatic differentiation.

import jax; jax.config.update("jax_enable_x64", True)
import jax.numpy as np
from kellax import arclength_continuation, refine_fold, fold_sensitivity

R = lambda x, p: np.array([x[0]**3 - x[0] + p])          # canonical S-curve
br = arclength_continuation(R, x0 = np.array([-1.2]), p0 = 0.7, ds = 0.05,
                            ds_max = 0.05, direction = -1.0)   # through both folds
for i in br.turning_points:
    x_f, p_f, v_f, res = refine_fold(R, br.x[i], float(br.p[i]))
    print(p_f)                                            # -+2/(3*sqrt(3)), to 1e-10

R3 = lambda x, p, th: np.array([x[0]**3 - th[0]*x[0] + p])
x, p, v, dp, res = fold_sensitivity(R3, np.array([0.6]), 0.4, np.array([1.0]))
dp                                                        # d(fold)/d(theta) = sqrt(1/3), exactly

The landscape

Numerical continuation is the standard tool for mapping out multiplicity and hysteresis: phase transitions, ignition thresholds, pattern selection. Its mature implementations live outside the modern Python stack: AUTO-07p in Fortran; MatCont, COCO and pde2path in MATLAB; LOCA/Trilinos in C++. The one modern autodiff-native package is BifurcationKit.jl in Julia. It remains considerably more complete than kellax on periodic orbits and kellax does not attempt to displace it. In the Python/JAX ecosystem a comparable tool is elusive. The options are young finite-difference codes (pycont-lite), stale wrappers around AUTO (PyDSTool/PyCont) and the matrix-free but autodiff-less pacopy. kellax fills this gap.

kellax began as the isotherm engine of a classical density-functional-theory code. It traced capillary condensation loops and located spinodals as the external conditions varied. Nothing in the library is specific to that origin: any smooth R(x, p) will do, from steady states of discretised PDEs to reaction networks, phase equilibria and fixed points of learned dynamics.

Why JAX

Continuation is built out of derivatives, and JAX supplies every order of them from the residual alone. The first derivatives drive the two engines: jax.jacfwd assembles the dense bordered solves, and jax.linearize provides the JVPs for the matrix-free GMRES path. The exact fold and Hopf systems require second derivatives. These are equally automatic: the Moore–Spence block d(R_xv)/dx and the Jacobian of the Hopf system are autodiff of code written once for the residual. In the classical packages the same objects are hand-assembled operator calculus.

The Jacobian of the deflation operator is one more jax.jacfwd call; that product-rule term is the reason deflated-Newton implementations grow long. Finally, implicit differentiation of the converged Moore–Spence system turns the fold location into a differentiable function of the model parameters (fold_sensitivity). It costs one extra linear solve against the Jacobian Newton has already used. In return the fold acquires an exact gradient: a model can be optimised or learned against its own bifurcation diagram. The whole predictor–corrector jit-compiles and runs unchanged on CPU or GPU, in float64 throughout.

The toolbox v0.4.0

arclength_continuation: The dense Keller trace with adaptive steps and fold detection. Returns a Branch (states, parameters, tangents, turning points).
mf_arclength_continuation: The matrix-free counterpart: preconditioned GMRES over jax.linearize JVPs with a precond hook and p_stop landing. For fields of 10⁴–10⁶ dof the Jacobian is never formed.
refine_fold / track_fold: Moore–Spence refinement of a detected fold to Newton precision, and continuation of the fold itself in a second parameter. The augmented system is continued by arclength, so cusps are passed and reported.
analyze_branch / branch_eigenvalues: The spectrum along the branch: stability, and the classification of every axis crossing into fold, branch point or Hopf candidate.
refine_hopf: A Hopf point pinned by the standard (3N+2) augmented system, with second derivatives supplied by autodiff.
deflated_newton / deflated_search: Farrell-style deflation: Newton converges away from the solutions already found. This is the route to disconnected branches.
branch_off / bifurcation_diagram: Switching onto the bifurcating branch at a simple branch point, and a bounded-depth recursive driver: trace, classify, switch, recurse (equilibria only).
fold_sensitivity: The exact gradient of a fold location with respect to the model parameters, by implicit differentiation of the converged Moore–Spence system.
bordered_newton / newton: The generic (N+k) bordered primitive and plain Newton for seeding.

Validated, number by number

Every claim in the toolbox is validated against an exact result or the literature.

Cubic normal-form folds	recovered to 1e‑10
Two-parameter fold law (the cusp)	to 1e‑8
Fold gradient vs the closed form dp*/dθ = √(θ/3)	to 1e‑9
Hopf, normal form and Brusselator	b* = 1 + a², ω = a, exactly
Pitchfork: both arms on x² = p	to 1e‑8
Bratu ignition, 1-D and 2-D	λ* = 3.5138 · 6.808
CSTR ignition/extinction pair	Uppal–Ray–Poore values
Predator–prey folds vs MatCont	five digits
Swift–Hohenberg snaking, one continuation	38 folds

Homoclinic snaking branch of the Swift-Hohenberg equation: the solution norm oscillates upward through 38 folds as the localised pattern grows, with solution profiles shown alongside — The homoclinic snaking branch of the Swift–Hohenberg equation: 38 folds passed in one continuation. examples/swift_hohenberg.py regenerates the figure.

The book

kellax by solved problems: nine chapters, each a single worked problem. Each chapter's script lives in examples/ and regenerates its figure; the printed numbers are real output. The TUTORIAL is the quick tour of the API.

The foldthe canonical S-curve
The cuspthe two-parameter fold law
Bratu–Gelfandignition in 1-D
Matrix-free scalingGMRES over JVPs
Homoclinic snaking38 folds in one continuation
CSTR hysteresisignition and extinction
A predator–prey fold pairfolds matched to MatCont
Bratu in 2-Dmatrix-free in two dimensions
Differentiable continuationgradients of the fold location

Install & test

uv venv .venv && uv pip install --python .venv/bin/python -e ".[test,examples]"
.venv/bin/python tests/test_kellax.py && .venv/bin/python tests/test_bifurcations.py
python examples/cubic_fold.py          # -> figures/cubic_fold.png

kellax requires float64: jax.config.update("jax_enable_x64", True) once at the top of every script.

Citing

Cite via CITATION.cff. Releases are archived on Zenodo; 10.5281/zenodo.21433166 resolves to the latest version.

Roadmap

Periodic-orbit continuation from Hopf points (shooting first).
Hopf tracking in a second parameter.
Matrix-free eigenvalues (Arnoldi over JVPs) so the classification layer reaches 10⁴+ dof.
Deflated continuation.
A jax.custom_jvp rule for the fold location, so that jax.grad composes through the diagram directly.

The long-term goal is unchanged: models optimised or learned against their bifurcation diagrams.