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

Fuente: arXiv
Salvato in:
Dettagli Bibliografici
Autori principali: Bezgin, Deniz A., Buhendwa, Aaron B., Adams, Nikolaus A.
Natura: Preprint
Pubblicazione: 2024
Soggetti:
Accesso online:
Tags: Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
_version_ 1866909098792976384
author Bezgin, Deniz A.
Buhendwa, Aaron B.
Adams, Nikolaus A.
author_facet Bezgin, Deniz A.
Buhendwa, Aaron B.
Adams, Nikolaus A.
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2402_05193
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle JAX-Fluids 2.0: Towards HPC for Differentiable CFD of Compressible Two-phase Flows
Bezgin, Deniz A.
Buhendwa, Aaron B.
Adams, Nikolaus A.
Fluid Dynamics
Computational Engineering, Finance, and Science
Machine Learning
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.
title JAX-Fluids 2.0: Towards HPC for Differentiable CFD of Compressible Two-phase Flows
topic Fluid Dynamics
Computational Engineering, Finance, and Science
Machine Learning
url https://arxiv.org/abs/2402.05193