JAX-Fluids 2.0: Towards HPC for Differentiable CFD of Compressible Two-phase Flows

TL;DR

JAX-Fluids 2.0 delivers a fully differentiable CFD framework for compressible single- and two-phase flows that scales on HPC hardware using a JAX-based distributed parallelization. It couples a high-order Godunov finite-volume discretization with either a level-set sharp-interface or a five-equation diffuse-interface model, supported by positivity-preserving limiters and stretched Cartesian meshes. The paper validates the solver on canonical benchmarks including turbulent boundary layers and shock–interface interactions, and demonstrates automatic-differentiation gradients computed through distributed simulations, with strong weak scaling on GPUs and TPUs. These capabilities enable end-to-end optimization, data assimilation, and uncertainty quantification in differentiable CFD, with potential extensions to multi-physics and adaptive multiresolution. The work thus advances HPC-ready, differentiable multiphase CFD for data-driven modeling and control.

Abstract

In our effort to facilitate machine learning-assisted computational fluid dynamics (CFD), we introduce the second iteration of JAX-Fluids. JAX-Fluids is a Python-based fully-differentiable CFD solver designed for compressible single- and two-phase flows. In this work, the first version is extended to incorporate high-performance computing (HPC) capabilities. We introduce a parallelization strategy utilizing JAX primitive operations that scales efficiently on GPU (up to 512 NVIDIA A100 graphics cards) and TPU (up to 1024 TPU v3 cores) HPC systems. We further demonstrate the stable parallel computation of automatic differentiation gradients across extended integration trajectories. The new code version offers enhanced two-phase flow modeling capabilities. In particular, a five-equation diffuse-interface model is incorporated which complements the level-set sharp-interface model. Additional algorithmic improvements include positivity-preserving limiters for increased robustness, support for stretched Cartesian meshes, refactored I/O handling, comprehensive post-processing routines, and an updated list of state-of-the-art high-order numerical discretization schemes. We verify newly added numerical models by showcasing simulation results for single- and two-phase flows, including turbulent boundary layer and channel flows, air-helium shock bubble interactions, and air-water shock drop interactions.

Figures (26)

• Figure 1: Schematic of the computational domain. We denote the axis dircetions .
• Figure 2: (a) Schematic of the spatial reconstruction and (b) schematic of the HLLC approximate solution to the resulting Riemann problem. (a): The stencils $S_k$ corresponding to the interpolants $p_{k,i+\frac{1}{2}}^{-}$ for WENO5-Z and TENO6-A reconstruction are shown. (b): The corresponding Riemann problem is solved by the HLLC Riemann solver. The emanating wave structure ($s^L, s^*, s^R$) and the resulting intermediate start states $\mathbf{U}^{*L}$ and $\mathbf{U}^{*R}$ are visualized.
• Figure 3: Schematic finite-volume discretization for cut cell $(i,j,k)$ on a Cartesian grid. The black dots represent the cell centers. The blue line indicates the interface, and the red line depicts the linear approximation of the interface. The figure illustrates a two-dimensional slice in the $(x,y)$-plane. (For interpretation of the references to color in this figure legend, the reader is referred to the web version of this article.)
• Figure 4: Exemplary domain decomposition with four blocks arranged in a $2\times2\times1$ grid.
• Figure 5: Schematic illustration of two adjacent blocks in $x$-direction within the homogeneous domain decomposition during a halo update. Depicted are the face (blue), edge (green), and vertex (red) halo cells of block 1 at the east side of the block. We highlight the corresponding cells in block 2 that are required for the halo update accordingly. This update entails a west permutation, i.e., a permutation of data from block 2 to block 1.