arXiv Science⌕ Search
Home/Archive/arXiv

arXiv · 2406.19494

JAX-based differentiable fluid dynamics on GPU and end-to-end optimization

Wenkang Wang·Xuanwei Zhang·Deniz Bezgin·Aaron Buhendwa·Xu Chu·Bernhard Weigand
View original source ↗Read source document ↗Export metadata ↓

Abstract

This project aims to advance differentiable fluid dynamics for hypersonic coupled flow over porous media, demonstrating the potential of automatic differentiation (AD)-based optimization for end-to-end solutions. Leveraging AD efficiently handles high-dimensional optimization problems, offering a flexible alternative to traditional methods. We utilized JAX-Fluids, a newly developed solver based on the JAX framework, which combines autograd and TensorFlow's XLA. Compiled on a HAWK-AI node with NVIDIA A100 GPU, JAX-Fluids showed computational performance comparable to other high-order codes like FLEXI. Validation with a compressible turbulent channel flow DNS case showed excellent agreement, and a new boundary condition for modeling porous media was successfully tested on a laminar boundary layer case. Future steps in our research are anticipated.

Explore related subjects

physics.flu-dyn

Keep this discovery

BibTeXRIS

Wenkang Wang, Xuanwei Zhang, Deniz Bezgin, Aaron Buhendwa, Xu Chu, Bernhard Weigand. 2024-06-27. JAX-based differentiable fluid dynamics on GPU and end-to-end optimization. https://arxiv.org/abs/2406.19494

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

KEEP EXPLORING

Related papers

arXiv2026

Correlative effects of induced magnetic field-buoyancy on reactive solute dispersion dynamics in couple-stress fluids

We investigate the dispersion of a reactive solute in a couple-stress fluid flowing between two parallel plates under the combined effects of pressure-driven flow, buoyancy, and an induced magnetic field. The model incorporates first-order heterogeneous reactions at both channel walls alongside a bulk reaction. Using Mei's multiscale homogenization technique accurate to third order, we develop a higher-order asymptotic formulation to determine the effective longitudinal dispersion coefficient and concentration field. Analytical predictions are complemented by Brownian dynamics simulations and finite-difference solutions, while the Aris method of moments quantifies transient mean displacement, spatial variance, and effective dispersivity. The hydrodynamic analysis reveals a singular branch in the velocity solution when the Hartmann number equals half the couple-stress parameter and identifies a characteristic quarter-power scaling between the Hartmann number and couple-stress parameter, separating couple-stress- and magnetically dominated regimes. The model recovers classical Taylor-dispersion behavior in the non-reactive Newtonian limit and agrees well with experimental measurements. Couple-stress rheology and magnetic damping suppress shear-induced dispersion, whereas buoyancy enhances dispersion through additional transverse velocity gradients. A distinct saturation regime of the dispersion coefficient emerges with an increasing couple-stress parameter, while unequal wall absorption induces persistent transverse asymmetry, and stronger absorption enhances solute removal near the source. Numerical and stochastic results validate the analytical framework while resolving higher-order concentration structures and particle-scale wall adsorption.

physics.flu-dyn↗
arXiv2026

DiffSWE2d: a differentiable Shallow Water Equations solver for end-to-end flood and tsunami modelling

Solving inverse and optimisation problems with traditional shallow water equations (SWE) solvers can be computationally expensive, particularly when gradients with respect to model inputs or parameters must be estimated through repeated forward simulations. In this paper, we introduce DiffSWE2d, an open-source differentiable shallow water equations solver for end-to-end flood and tsunami modelling implemented in PyTorch. By leveraging automatic differentiation, DiffSWE2d represents the time-marching physics as a differentiable computational graph, enabling gradients to be propagated directly through the numerical solver. We validate the solver against two established benchmark cases and demonstrate its application to tsunami waveform inversion, showing its ability to infer model inputs through gradient-based optimisation. DiffSWE2d provides a flexible framework for integrating physics-based hydrodynamic modelling with modern optimisation and machine learning methods. The source code and reproducible examples are publicly available at: https://github.com/ZhonghouXu/DiffSWE2d

physics.flu-dyn↗
arXiv2026

Low inertia limit of elasto-inertial turbulence

Pipe and channel flows of viscoelastic fluids display chaotic dynamics at unusually low speeds, a phenomenon referred to as elasto-inertial turbulence, EIT. First reported in experiments a century ago, recent theoretical studies and model computations predict a variety of scenarios for the phenomenon's origin, ranging from hoop stress modes to center modes and to Tollmien-Schlichting waves. Lacking experimental confirmation, the relevant scenario in actual flows of polymer solutions remains unknown. We here determine the transition threshold of EIT in pipe experiments, covering three decades in elasticity number. Across this entire parameter range, the transition features center mode structures at onset. Eventually the instability diverges at a lower inertia (upper elasticity) limit, which is a robust signature of this center mode scenario. Finally, we report the first experimental observation of a traveling wave in viscoelastic pipe flow, and the sequences of localized structures found, are in excellent agreement with a center mode traveling wave, the "arrowhead" solution, discovered in model simulations.

physics.flu-dyn↗