google/jax-cfd

JAX‑CFD: Differentiable Computational Fluid Dynamics in JAX

What it is – An open‑source research library that implements core CFD solvers (finite‑volume and pseudospectral) using JAX, so the code is automatically differentiable and can run on GPUs/TPUs. It also ships a few machine‑learning‑augmented models built with Haiku.

Why it matters – By leveraging JAX’s just‑in‑time compilation and automatic differentiation, researchers can:

Train neural networks that interact with fluid simulations (e.g., surrogate models, closure models).
Perform gradient‑based optimisation or data‑assimilation directly on the CFD solver.
Run the same code on CPUs, GPUs, or TPUs without changes.

Key components

jax_cfd.base – Finite‑volume/difference solvers on a staggered (Arakawa‑C/MAC) grid.
jax_cfd.spectral – Pseudospectral vorticity solvers with anti‑aliasing.
jax_cfd.ml – ML‑augmented CFD models (Haiku neural nets) and example pre‑trained models.
jax_cfd.data – Utilities for loading, processing, and visualising simulation data (Xarray + Pillow).

Numerical focus

Unsteady turbulent flows.
Spatial discretisation: staggered‑grid FV/Diff or pseudospectral vorticity.
Temporal discretisation: first‑order explicit advection, implicit/explicit diffusion.
Pressure solve: Conjugate‑gradient or FFT‑based diagonalisation (periodic BCs only).
Advection scheme: 2nd‑order Van Leer.
Turbulence closure: Smagorinsky eddy‑viscosity.

How to get started

Install the core package:
```
pip install jax-cfd
```

Add optional extras for ML or data utilities:

pip install jax-cfd[ml]   # Haiku + ML helpers
pip install jax-cfd[data] # Xarray, Pillow, etc.

Run the ready‑made Colab notebooks in the notebooks/ folder to see:
- 2‑D staggered‑grid FV demo
- 2‑D pseudospectral turbulence demo
- Channel‑flow example
- Reproducing the results from the PNAS 2021 paper (data analysis & model inference)

Project status – Marked no longer maintained (June 2024). The README points to actively maintained alternatives such as JAX‑Fluids, PhiFlow, and Exponax.

Typical use cases

Academic research on differentiable CFD, e.g., learning turbulence closures or inverse observation operators.
Prototyping ML‑CFD hybrid models that require gradients through the fluid solver.
Benchmarking JAX performance on fluid‑dynamics workloads.

Citations – If you use the finite‑volume code or the ML models, cite the PNAS 2021 paper. For the spectral code, cite the arXiv 2022 pre‑print.

Related projects – The README lists other differentiable CFD libraries (PhiFlow, Autograd fluid examples) and broader JAX‑for‑science tools (JAX‑MD, JAX‑DFT, jax‑cosmo, Veros).

All details above are taken directly from the repository’s README.