JAX-CFD Differentiable Computational Fluid DynamicsNew

Concept

JAX-CFD implements standard CFD solvers (pseudo-spectral NS, finite volume) in JAX. Because JAX is differentiable, you can compute gradients of a simulation trajectory with respect to initial conditions, forcing terms, or model parameters - the foundation of learned turbulence closures and control.

What you learn

• Pseudo-spectral discretization of incompressible NS on periodic domains
• Semi-implicit time stepping and dealiasing practices
• Gradient-based optimization through unrolled simulation trajectories
• Context: Kochkov et al. used this stack for ML-accelerated turbulence (PNAS 2023)

Suggested path

1. Run the decaying turbulence example; compare energy spectra
2. Differentiate a trajectory with respect to the initial field
3. Reproduce a simple learned correction to the sub-grid stress
4. Compare with finite-volume mode for boundary-layer problems

Prerequisites

Comfortable with NumPy and basic fluid mechanics; JAX basics (jit/grad/vmap) are picked up along the way.