Efficient Differentiable n-d PDE Solvers in JAX.
See the codeREADME
Efficient Differentiable n-d PDE solvers built on top of JAX & Equinox.
Installation • Quickstart • Equations • Features • Documentation • Background • Citation
Exponax solves partial differential equations in 1D, 2D, and 3D on periodic
domains highly efficiently using Fourier spectral methods and exponential time
differencing. It ships more than 46 PDE solvers covering linear, nonlinear, and
reaction-diffusion dynamics. Built entirely on
JAX and
Equinox, every solver is
automatically differentiable, JIT-compilable, and GPU/TPU-ready — making it
ideal for physics-based deep learning workflows.
Installation
pip install exponax
Requires Python 3.10+ and JAX 0.4.13+. 👉 JAX install guide.
Quickstart
Simulate the chaotic Kuramoto-Sivashinsky equation in 1D — a single stepper object, one line to roll out 500 time steps:
import jax
import exponax as ex
import matplotlib.pyplot as plt
ks_stepper = ex.stepper.KuramotoSivashinskyConservative(
num_spatial_dims=1, domain_extent=100.0,
num_points=200, dt=0.1,
)
u_0 = ex.ic.RandomTruncatedFourierSeries(
num_spatial_dims=1, cutoff=5
)(num_points=200, key=jax.random.PRNGKey(0))
trajectory = ex.rollout(ks_stepper, 500, include_init=True)(u_0)
plt.imshow(trajectory[:, 0, :].T, aspect='auto', cmap='RdBu', vmin=-2, vmax=2, origin="lower")
plt.xlabel("Time"); plt.ylabel("Space"); plt.show()
Because every stepper is a differentiable JAX function, you can freely compose
it with jax.grad, jax.vmap, and jax.jit:
# Jacobian of the stepper function
jacobian = jax.jacfwd(ks_stepper)(u_0)
For a next step, check out this tutorial on 1D
Advection
that explains the basics of Exponax.
Built-in Equations
Linear
| Equation | Stepper | Dimensions |
|---|---|---|
| Advection: $u_t + c \cdot \nabla u = 0$ | Advection | 1D, 2D, 3D |
| Diffusion: $u_t = \nu \Delta u$ | Diffusion | 1D, 2D, 3D |
| Advection-Diffusion: $u_t + c \cdot \nabla u = \nu \Delta u$ | AdvectionDiffusion | 1D, 2D, 3D |
| Dispersion: $u_t = \xi \nabla^3 u$ | Dispersion | 1D, 2D, 3D |
| Hyper-Diffusion: $u_t = -\zeta \Delta^2 u$ | HyperDiffusion | 1D, 2D, 3D |
| Wave: $u_{tt} = c^2 \Delta u$ | Wave | 1D, 2D, 3D |
Nonlinear
| Equation | Stepper | Dimensions |
|---|---|---|
| Burgers: $u_t + \frac{1}{2} \nabla \cdot (u \otimes u) = \nu \Delta u$ | Burgers | 1D, 2D, 3D |
| Korteweg-de Vries: $u_t + \frac{1}{2} \nabla \cdot (u \otimes u) - \nabla^3 u = \mu \Delta u$ | KortewegDeVries | 1D, 2D, 3D |
| Kuramoto-Sivashinsky: $u_t + \frac{1}{2} |\nabla u|^2 + \Delta u + \Delta^2 u = 0$ | KuramotoSivashinsky | 1D, 2D, 3D |
| KS (conservative): $u_t + \frac{1}{2} \nabla \cdot (u \otimes u) + \Delta u + \Delta^2 u = 0$ | KuramotoSivashinskyConservative | 1D, 2D, 3D |
| Navier-Stokes (vorticity): $\omega_t + (u \cdot \nabla)\omega = \nu \Delta \omega$ | NavierStokesVorticity | 2D |
| Kolmogorov Flow (vorticity): $\omega_t + (u \cdot \nabla)\omega = \nu \Delta \omega + f$ | KolmogorovFlowVorticity | 2D |
| Navier-Stokes (velocity): $u_t = \nu \Delta u + \mathcal{P}(u \times \omega)$ | NavierStokesVelocity | 3D |
| Kolmogorov Flow (velocity): $u_t = \nu \Delta u + \mathcal{P}(u \times \omega) + f$ | KolmogorovFlowVelocity | 3D |
Reaction-Diffusion
| Equation | Stepper | Dimensions |
|---|---|---|
| Fisher-KPP: $u_t = \nu \Delta u + r, u(1 - u)$ | reaction.FisherKPP | 1D, 2D, 3D |
| Allen-Cahn: $u_t = \nu \Delta u + c_1 u + c_3 u^3$ | reaction.AllenCahn | 1D, 2D, 3D |
| Cahn-Hilliard: $u_t = \nu \Delta(u^3 + c_1 u - \gamma \Delta u)$ | reaction.CahnHilliard | 1D, 2D, 3D |
| Gray-Scott: $u_t = \nu_1 \Delta u + f(1-u) - uv^2, \quad v_t = \nu_2 \Delta v - (f+k)v + uv^2$ | reaction.GrayScott | 1D, 2D, 3D |
| Swift-Hohenberg: $u_t = ru - (k + \Delta)^2 u + g(u)$ | reaction.SwiftHohenberg | 1D, 2D, 3D |
Generic stepper families (for advanced / custom dynamics)
These parametric families generalize the concrete steppers above. Each comes in three flavors: physical coefficients, normalized, and difficulty-based.
| Family | Nonlinearity | Generalizes |
|---|---|---|
GeneralLinearStepper | None | Advection, Diffusion, Dispersion, etc. |
GeneralConvectionStepper | Quadratic convection | Burgers, KdV, KS Conservative |
GeneralGradientNormStepper | Gradient norm | Kuramoto-Sivashinsky |
GeneralVorticityConvectionStepper | Vorticity convection (2D only) | Navier-Stokes, Kolmogorov Flow |
GeneralPolynomialStepper | Arbitrary polynomial | Fisher-KPP, Allen-Cahn, etc. |
GeneralNonlinearStepper | Convection + gradient norm + polynomial | Most of the above |
See the normalized & difficulty interface docs for details.
Features
- Hardware-agnostic — run on CPU, GPU, or TPU in single or double precision.
- Fully differentiable — compute gradients of solutions w.r.t. initial conditions, PDE parameters, or neural network weights when composed with PDE solvers via
jax.grad. - Vectorized batching — advance multiple states or sweep over parameter grids in parallel using
jax.vmap(andeqx.filter_vmap). - Deep-learning native — every stepper is an Equinox Module, composable with neural networks out of the box.
- Lightweight design — no custom grid or state objects; everything is plain
jax.numpyarrays and callable PyTrees. - Initial conditions — library of random IC distributions (truncated Fourier series, Gaussian random fields, etc.).
- Utilities — spectral derivatives, grid creation, autoregressive rollout, interpolation, and more.
- Extensible — add new PDEs by subclassing
BaseStepper.
Documentation
Documentation is available at fkoehler.site/exponax. Key pages:
- 1D Advection Tutorial — learn the basics
- Solver Showcase 1D / 2D / 3D — visual gallery of all dynamics
- Creating Your Own Solvers — extend Exponax with custom PDEs
- Training a Neural Operator — use
Exponaxfor synthetic data generation and training of a neural emulator - Stepper Overview — API reference for all steppers
- Performance Hints — tips for fast simulations
Background
Exponax solves semi-linear PDEs of the form
$$ \partial u / \partial t = Lu + N(u), $$
where $L$ is a linear differential operator and $N$ is a nonlinear differential operator. The linear part is solved exactly via a matrix exponential in Fourier space, while the nonlinear part is integrated using exponential time differencing Runge-Kutta (ETDRK) schemes of order 1 through 4. The complex contour integral method of Kassam & Trefethen is used for numerical stability.
By restricting to periodic domains on scaled hypercubes with uniform Cartesian grids, all transforms reduce to FFTs — yielding blazing-fast simulations. For example, 50 trajectories of the 2D Kuramoto-Sivashinsky equation (200 time steps, 128x128 grid) are generated in under a second on a modern GPU.
References
- Cox, S.M. and Matthews, P.C. "Exponential time differencing for stiff systems." Journal of Computational Physics 176.2 (2002): 430-455. doi:10.1006/jcph.2002.6995
- Kassam, A.K. and Trefethen, L.N. "Fourth-order time-stepping for stiff PDEs." SIAM Journal on Scientific Computing 26.4 (2005): 1214-1233. doi:10.1137/S1064827502410633
- Montanelli, H. and Bootland, N. "Solving periodic semilinear stiff PDEs in 1D, 2D and 3D with exponential integrators." Mathematics and Computers in Simulation 178 (2020): 307-327. doi:10.1016/j.matcom.2020.06.008
Related & Motivation
This package is greatly inspired by the
spinX module of the
ChebFun package in MATLAB. spinX served as a
reliable data generator for early works in physics-based deep learning, e.g.,
DeepHiddenPhysics
and Fourier Neural
Operators.
However, due to the two-language barrier, dynamically calling MATLAB solvers
from Python-based deep learning workflows is hard to impossible. This also
excludes the option to differentiate through them — ruling out
differentiable-physics approaches like solver-in-the-loop correction or
diverted-chain training.
We view Exponax as a spiritual successor of spinX. JAX, as the
computational backend, elevates the power of this solver type with automatic
vectorization (jax.vmap), backend-agnostic execution (CPU/GPU/TPU), and tight
integration for deep learning via its versatile automatic differentiation
engine. With reproducible randomness in JAX, datasets can be re-created in
seconds — no need to ever write them to disk.
Beyond ChebFun, other popular pseudo-spectral implementations include Dedalus in the Python world and FourierFlows.jl in the Julia ecosystem (the latter was especially helpful for verifying our implementation of the contour integral method and dealiasing).
Citation
Exponax was developed as part of the
APEBench benchmark suite for
autoregressive neural emulators of PDEs. The accompanying paper was accepted at
NeurIPS 2024. If you find this package useful for your research, please
consider citing it:
@article{koehler2024apebench,
title={Apebench: A benchmark for autoregressive neural emulators of pdes},
author={Koehler, Felix and Niedermayr, Simon and Westermann, R{\"u}diger and Thuerey, Nils},
journal={Advances in Neural Information Processing Systems},
volume={37},
pages={120252--120310},
year={2024}
}
If you enjoy the project, feel free to give it a star on GitHub!
Funding
The main author (Felix Koehler) is a PhD student in the group of Prof. Thuerey at TUM and his research is funded by the Munich Center for Machine Learning.
License
MIT, see here
fkoehler.site · GitHub @ceyron · X @felix_m_koehler · LinkedIn Felix Köhler
Contributors
Similar repos
Languages
Jupyter Notebook
79.6%
Python
20.4%
Efficient Differentiable n-d PDE Solvers in JAX.
See the codeREADME
Efficient Differentiable n-d PDE solvers built on top of JAX & Equinox.
Installation • Quickstart • Equations • Features • Documentation • Background • Citation
Exponax solves partial differential equations in 1D, 2D, and 3D on periodic
domains highly efficiently using Fourier spectral methods and exponential time
differencing. It ships more than 46 PDE solvers covering linear, nonlinear, and
reaction-diffusion dynamics. Built entirely on
JAX and
Equinox, every solver is
automatically differentiable, JIT-compilable, and GPU/TPU-ready — making it
ideal for physics-based deep learning workflows.
Installation
pip install exponax
Requires Python 3.10+ and JAX 0.4.13+. 👉 JAX install guide.
Quickstart
Simulate the chaotic Kuramoto-Sivashinsky equation in 1D — a single stepper object, one line to roll out 500 time steps:
import jax
import exponax as ex
import matplotlib.pyplot as plt
ks_stepper = ex.stepper.KuramotoSivashinskyConservative(
num_spatial_dims=1, domain_extent=100.0,
num_points=200, dt=0.1,
)
u_0 = ex.ic.RandomTruncatedFourierSeries(
num_spatial_dims=1, cutoff=5
)(num_points=200, key=jax.random.PRNGKey(0))
trajectory = ex.rollout(ks_stepper, 500, include_init=True)(u_0)
plt.imshow(trajectory[:, 0, :].T, aspect='auto', cmap='RdBu', vmin=-2, vmax=2, origin="lower")
plt.xlabel("Time"); plt.ylabel("Space"); plt.show()
Because every stepper is a differentiable JAX function, you can freely compose
it with jax.grad, jax.vmap, and jax.jit:
# Jacobian of the stepper function
jacobian = jax.jacfwd(ks_stepper)(u_0)
For a next step, check out this tutorial on 1D
Advection
that explains the basics of Exponax.
Built-in Equations
Linear
| Equation | Stepper | Dimensions |
|---|---|---|
| Advection: $u_t + c \cdot \nabla u = 0$ | Advection | 1D, 2D, 3D |
| Diffusion: $u_t = \nu \Delta u$ | Diffusion | 1D, 2D, 3D |
| Advection-Diffusion: $u_t + c \cdot \nabla u = \nu \Delta u$ | AdvectionDiffusion | 1D, 2D, 3D |
| Dispersion: $u_t = \xi \nabla^3 u$ | Dispersion | 1D, 2D, 3D |
| Hyper-Diffusion: $u_t = -\zeta \Delta^2 u$ | HyperDiffusion | 1D, 2D, 3D |
| Wave: $u_{tt} = c^2 \Delta u$ | Wave | 1D, 2D, 3D |
Nonlinear
| Equation | Stepper | Dimensions |
|---|---|---|
| Burgers: $u_t + \frac{1}{2} \nabla \cdot (u \otimes u) = \nu \Delta u$ | Burgers | 1D, 2D, 3D |
| Korteweg-de Vries: $u_t + \frac{1}{2} \nabla \cdot (u \otimes u) - \nabla^3 u = \mu \Delta u$ | KortewegDeVries | 1D, 2D, 3D |
| Kuramoto-Sivashinsky: $u_t + \frac{1}{2} |\nabla u|^2 + \Delta u + \Delta^2 u = 0$ | KuramotoSivashinsky | 1D, 2D, 3D |
| KS (conservative): $u_t + \frac{1}{2} \nabla \cdot (u \otimes u) + \Delta u + \Delta^2 u = 0$ | KuramotoSivashinskyConservative | 1D, 2D, 3D |
| Navier-Stokes (vorticity): $\omega_t + (u \cdot \nabla)\omega = \nu \Delta \omega$ | NavierStokesVorticity | 2D |
| Kolmogorov Flow (vorticity): $\omega_t + (u \cdot \nabla)\omega = \nu \Delta \omega + f$ | KolmogorovFlowVorticity | 2D |
| Navier-Stokes (velocity): $u_t = \nu \Delta u + \mathcal{P}(u \times \omega)$ | NavierStokesVelocity | 3D |
| Kolmogorov Flow (velocity): $u_t = \nu \Delta u + \mathcal{P}(u \times \omega) + f$ | KolmogorovFlowVelocity | 3D |
Reaction-Diffusion
| Equation | Stepper | Dimensions |
|---|---|---|
| Fisher-KPP: $u_t = \nu \Delta u + r, u(1 - u)$ | reaction.FisherKPP | 1D, 2D, 3D |
| Allen-Cahn: $u_t = \nu \Delta u + c_1 u + c_3 u^3$ | reaction.AllenCahn | 1D, 2D, 3D |
| Cahn-Hilliard: $u_t = \nu \Delta(u^3 + c_1 u - \gamma \Delta u)$ | reaction.CahnHilliard | 1D, 2D, 3D |
| Gray-Scott: $u_t = \nu_1 \Delta u + f(1-u) - uv^2, \quad v_t = \nu_2 \Delta v - (f+k)v + uv^2$ | reaction.GrayScott | 1D, 2D, 3D |
| Swift-Hohenberg: $u_t = ru - (k + \Delta)^2 u + g(u)$ | reaction.SwiftHohenberg | 1D, 2D, 3D |
Generic stepper families (for advanced / custom dynamics)
These parametric families generalize the concrete steppers above. Each comes in three flavors: physical coefficients, normalized, and difficulty-based.
| Family | Nonlinearity | Generalizes |
|---|---|---|
GeneralLinearStepper | None | Advection, Diffusion, Dispersion, etc. |
GeneralConvectionStepper | Quadratic convection | Burgers, KdV, KS Conservative |
GeneralGradientNormStepper | Gradient norm | Kuramoto-Sivashinsky |
GeneralVorticityConvectionStepper | Vorticity convection (2D only) | Navier-Stokes, Kolmogorov Flow |
GeneralPolynomialStepper | Arbitrary polynomial | Fisher-KPP, Allen-Cahn, etc. |
GeneralNonlinearStepper | Convection + gradient norm + polynomial | Most of the above |
See the normalized & difficulty interface docs for details.
Features
- Hardware-agnostic — run on CPU, GPU, or TPU in single or double precision.
- Fully differentiable — compute gradients of solutions w.r.t. initial conditions, PDE parameters, or neural network weights when composed with PDE solvers via
jax.grad. - Vectorized batching — advance multiple states or sweep over parameter grids in parallel using
jax.vmap(andeqx.filter_vmap). - Deep-learning native — every stepper is an Equinox Module, composable with neural networks out of the box.
- Lightweight design — no custom grid or state objects; everything is plain
jax.numpyarrays and callable PyTrees. - Initial conditions — library of random IC distributions (truncated Fourier series, Gaussian random fields, etc.).
- Utilities — spectral derivatives, grid creation, autoregressive rollout, interpolation, and more.
- Extensible — add new PDEs by subclassing
BaseStepper.
Documentation
Documentation is available at fkoehler.site/exponax. Key pages:
- 1D Advection Tutorial — learn the basics
- Solver Showcase 1D / 2D / 3D — visual gallery of all dynamics
- Creating Your Own Solvers — extend Exponax with custom PDEs
- Training a Neural Operator — use
Exponaxfor synthetic data generation and training of a neural emulator - Stepper Overview — API reference for all steppers
- Performance Hints — tips for fast simulations
Background
Exponax solves semi-linear PDEs of the form
$$ \partial u / \partial t = Lu + N(u), $$
where $L$ is a linear differential operator and $N$ is a nonlinear differential operator. The linear part is solved exactly via a matrix exponential in Fourier space, while the nonlinear part is integrated using exponential time differencing Runge-Kutta (ETDRK) schemes of order 1 through 4. The complex contour integral method of Kassam & Trefethen is used for numerical stability.
By restricting to periodic domains on scaled hypercubes with uniform Cartesian grids, all transforms reduce to FFTs — yielding blazing-fast simulations. For example, 50 trajectories of the 2D Kuramoto-Sivashinsky equation (200 time steps, 128x128 grid) are generated in under a second on a modern GPU.
References
- Cox, S.M. and Matthews, P.C. "Exponential time differencing for stiff systems." Journal of Computational Physics 176.2 (2002): 430-455. doi:10.1006/jcph.2002.6995
- Kassam, A.K. and Trefethen, L.N. "Fourth-order time-stepping for stiff PDEs." SIAM Journal on Scientific Computing 26.4 (2005): 1214-1233. doi:10.1137/S1064827502410633
- Montanelli, H. and Bootland, N. "Solving periodic semilinear stiff PDEs in 1D, 2D and 3D with exponential integrators." Mathematics and Computers in Simulation 178 (2020): 307-327. doi:10.1016/j.matcom.2020.06.008
Related & Motivation
This package is greatly inspired by the
spinX module of the
ChebFun package in MATLAB. spinX served as a
reliable data generator for early works in physics-based deep learning, e.g.,
DeepHiddenPhysics
and Fourier Neural
Operators.
However, due to the two-language barrier, dynamically calling MATLAB solvers
from Python-based deep learning workflows is hard to impossible. This also
excludes the option to differentiate through them — ruling out
differentiable-physics approaches like solver-in-the-loop correction or
diverted-chain training.
We view Exponax as a spiritual successor of spinX. JAX, as the
computational backend, elevates the power of this solver type with automatic
vectorization (jax.vmap), backend-agnostic execution (CPU/GPU/TPU), and tight
integration for deep learning via its versatile automatic differentiation
engine. With reproducible randomness in JAX, datasets can be re-created in
seconds — no need to ever write them to disk.
Beyond ChebFun, other popular pseudo-spectral implementations include Dedalus in the Python world and FourierFlows.jl in the Julia ecosystem (the latter was especially helpful for verifying our implementation of the contour integral method and dealiasing).
Citation
Exponax was developed as part of the
APEBench benchmark suite for
autoregressive neural emulators of PDEs. The accompanying paper was accepted at
NeurIPS 2024. If you find this package useful for your research, please
consider citing it:
@article{koehler2024apebench,
title={Apebench: A benchmark for autoregressive neural emulators of pdes},
author={Koehler, Felix and Niedermayr, Simon and Westermann, R{\"u}diger and Thuerey, Nils},
journal={Advances in Neural Information Processing Systems},
volume={37},
pages={120252--120310},
year={2024}
}
If you enjoy the project, feel free to give it a star on GitHub!
Funding
The main author (Felix Koehler) is a PhD student in the group of Prof. Thuerey at TUM and his research is funded by the Munich Center for Machine Learning.
License
MIT, see here
fkoehler.site · GitHub @ceyron · X @felix_m_koehler · LinkedIn Felix Köhler
Similar repos
Contributors
Languages
Jupyter Notebook
79.6%
Python
20.4%