JAX-Fluids

JAX-Fluids is a fully auto-differentiable CFD solver built on Google JAX, focused on end-to-end differentiable fluid simulation. Unlike traditional solvers (e.g. OpenFOAM), it back-propagates gradients through simulation parameters (geometry, boundary conditions, properties), naturally fitting Physics-Informed ML, optimization and neural-network-coupled simulation.

Core features

• Fully differentiable: derivatives w.r.t. all simulation parameters, no manual adjoint equations
• Multi-GPU: JAX's SPMD abstraction enables multi-GPU/TPU with zero code changes
• High order: WENO-JS scheme, 3rd~5th order
• Multiphase: two-phase flow (VOF)
• NN coupling: serves as the physics forward solver for Neural ODE / PINN

Use cases: Shape optimization (wing, pipe, bluff body) with gradients back to design params; PINN/PIAN forward solver; data-driven turbulence closure learning; end-to-end simulation-control co-learning.

Note: Marked MIT, but the original CPC paper part is GPLv3 - use as algorithm reference, not embedded in closed-source products. Best for research; cross-validate with OpenFOAM/ANSYS for industry.