Search arXiv⌕ Search
Home/Archive/arXiv

arXiv · 2402.05193

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

Deniz A. Bezgin·Aaron B. Buhendwa·Nikolaus A. Adams
View original source ↗Read source document ↗Export metadata ↓

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.

Explore related subjects

physics.flu-dyncs.CEcs.LG

Keep this discovery

Explore connections, maps & timelines

BibTeXRIS

Deniz A. Bezgin, Aaron B. Buhendwa, Nikolaus A. Adams. 2024-02-07. JAX-Fluids 2.0: Towards HPC for Differentiable CFD of Compressible Two-phase Flows. https://arxiv.org/abs/2402.05193

Cite the original work for its findings. Save a collection to share your selection of sources.

KEEP EXPLORING

Related papers

arXiv2026

Data-driven low-dimensional model for the onset of turbulence in plane Couette flow

We construct low-dimensional dynamical systems for low-Reynolds-number turbulence in plane Couette flow using Kernel Quantile Regression. Exploiting the finite-dimensional structure of attractors in dissipative systems, reduced-order models are learned from direct numerical simulation data using a small set of physically meaningful observables. The resulting discrete-time models accurately reproduce periodic and chaotic dynamics near the onset of turbulence. The minimum number of variables required for accurate prediction is found to be consistent with embedding theory based on the attractor dimension. For chaotic regimes, the models capture both short-term trajectory evolution and long-term statistical properties, including probability density functions. By incorporating the Reynolds number as an additional input variable, we further develop a parameter-dependent model that successfully reproduces the bifurcation structure over a range of Reynolds numbers, including bifurcations between training points. These results demonstrate the effectiveness of machine-learning-based reduced-order modeling for capturing the essential dynamics and statistics of weakly turbulent flows.

physics.flu-dyn↗
arXiv2026

Rapidly rotating internally heated convection: bounds on long-time averages

Convection on geophysical and astrophysical scales is subject to rapid rotation and strong heating from within the domain. In studying the long-time behaviour of the solutions for such a system, energy identities fail to capture the effects of rotation because the Coriolis force does no work, and rapid rotation can be prohibitive for direct numerical simulations. Instead, we derive an asymptotically reduced model for rapidly rotating convection driven by uniform internal heating between isothermal stress-free boundaries in a plane periodic layer. The main contribution is the proof of bounds on the mean temperature, and the mean vertical convective heat transport, in terms of the Rayleigh and Ekman numbers, in the limit of infinite Prandtl number. The first quantity represents the mixing of the flow, and the second the asymmetry in heat leaving the bottom and top boundaries due to convection, and unlike Rayleigh-Bénard convection, the two are not a priori related. We employ alternative estimation techniques to those used in previous studies (Grooms \& Whitehead, 2014 \textit{Nonlinearity}, 28, 29) and identify two distinct scaling behaviours for both quantities. Finally, our bounds are optimised, within the methodology, and provide a rigorous constraint for future studies of rotation-dominated internally heated convection.

physics.flu-dyn↗
arXiv2026

Baroclinic wave dynamics in the Ekman-free rotating rectangular annulus with localized forced plume

We report numerical simulations of a rotating rectangular annulus that isolates the Ekman-free bulk of the cylindrical baroclinic annulus, subjected to bi-directional temperature gradients imposed by a uniformly cooled inner wall and a localized forced heated plume at the outer bottom. The finite-volume OpenFOAM solver is employed across combinations of source Richardson number $Ri_0 = 99, 4, 1$ and Rossby number $Ro = 0.3, 0.1, 0.07$. A non-dimensional scaling of the governing equations identifies geostrophic-hydrostatic balance as the leading-order bulk state, a result confirmed a posteriori by the $x$ and $z-$momentum budgets. Baroclinic waves of mode $m=2$ at $Ro=0.3$ transition to $m=3$ as $Ro$ decreases, consistent with the contraction of the Eady deformation radius $L_ρ= NH/f$; Complex Empirical Orthogonal Function (CEOF) analysis characterizes the wave regime and detects a Hopf-bifurcated vacillating state at $Ri_0 = 99,~Ro = 0.1$. The plume morphology, classified through the Morton length scale and source flux-balance parameter, transitions from weak, laterally-swept structures at $Ri_0 = 99$ to sustained columnar plumes traversing the full baroclinic depth at $Ri_0 \leq 4$. The plume entrainment coefficient $Γ(z)$ shows opposite rotational sensitivities at low and high $Ri_0$, which we organize through a local plume Rossby number $Ro_p = w/(2Ωb)$. A mixing-length argument predicts a bulk turbulent heat flux $\overline{u'T'} \propto Ri_0^{-1/2}$, anticipating an order-of-magnitude enhancement from $Ri_0 = 99$ to $Ri_0 = 1$, in agreement with the simulations. A regime map in the $(Ri_0, Ro)$ plane reveals that, within the explored range, the plume-regime and wave-selection problems are approximately separable: $Ri_0$ sets the plume regime while $Ro$ selects the dominant baroclinic wave mode.

physics.flu-dyn↗