JAX-Fluids: A fully-differentiable high-order computational fluid dynamics solver for compressible two-phase flows

Physical systems are governed by partial differential equations (PDEs). The Navier-Stokes equations describe fluid flows and are representative of nonlinear physical systems with complex spatio-temporal interactions. Fluid flows are omnipresent in nature and engineering applications, and their accurate simulation is essential for providing insights into these processes. While PDEs are typically solved with numerical methods, the recent success of machine learning (ML) has shown that ML methods can provide novel avenues of finding solutions to PDEs. ML is becoming more and more present in computational fluid dynamics (CFD). However, up to this date, there does not exist a general-purpose ML-CFD package which provides 1) powerful state-of-the-art numerical methods, 2) seamless hybridization of ML with CFD, and 3) automatic differentiation (AD) capabilities. AD in particular is essential to ML-CFD research as it provides gradient information and enables optimization of preexisting and novel CFD models. In this work, we propose JAX-Fluids: a comprehensive fully-differentiable CFD Python solver for compressible two-phase flows. JAX-Fluids is intended for ML-supported CFD research. The framework allows the simulation of complex fluid dynamics with phenomena like three-dimensional turbulence, compressibility effects, and two-phase flows. Written entirely in JAX, it is straightforward to include existing ML models into the proposed framework. Furthermore, JAX-Fluids enables end-to-end optimization. I.e., ML models can be optimized with gradients that are backpropagated through the entire CFD algorithm, and therefore contain not only information of the underlying PDE but also of the applied numerical methods. We believe that a Python package like JAX-Fluids is crucial to facilitate research at the intersection of ML and CFD and may pave the way for an era of differentiable fluid dynamics. Program title: JAX-Fluids CPC Library link to program files: https://doi.org/10.17632/pzvkwn5s6p.1 Developer's repository link: https://github.com/tumaer/JAXFLUIDS Code Ocean capsule: https://codeocean.com/capsule/6819679 Licensing provisions: GNU GPLv3 Programming language: Python Supplementary material: Source code; Examples; Videos: Moving solid bodies, Taylor-Green vortex, Rising bubble, Shock-bubble interaction. Nature of problem: The compressible Navier-Stokes equations describe continuum-scale fluid flows. These flows often involve highly complex flow phenomena such as shocks, material interfaces, and turbulence. The intrinsic nonlinear dynamics render the numerical simulation of the these equations challenging. Machine learning provides novel avenues for describing partial differential equations. Machine learning models rely on gradient information provided by automatic differentiation and are often implemented in Python. In contrast, existing high-performance computational fluid dynamics codes are typically written in Fortran or C++ and do not offer inherent automatic differentiation capabilities. These discrepancies hinder the advance of machine-learning-supported computational fluid dynamics. Up to this day, a general-purpose fully-differentiable computational fluid dynamics solver for compressible two-phase flows is missing. Solution method: We introduce JAX-Fluids: a general-purpose three-dimensional fully-differentiable computational fluid dynamics solver for compressible two-phase flows. JAX-Fluids is a simulation framework intended for machine-learning-supported computational fluid dynamics research. Our framework is written entirely in JAX, a high-performance numerical computing library with automatic differentiation capabilities. We have used an object-oriented programming style and a modular design philosophy. This allows the straightforward exchange of numerical methods. We provide a wide variety of state-of-the-art high-order computational methods for compressible flows. The modularity of our framework additionally facilitates the integration of custom subroutines. We use the sharp-interface level-set method to model two-phase flows. The software package can easily be installed as a Python package. We have build the source code around the JAX NumPy API. This makes JAX-Fluids accessible and performant. JAX-Fluids runs on CPUs, GPUs, and TPUs. We use HDF5 in combination with XDMF for writing output quantities. The Python packages Haiku and Optax are used for implementation and training of machine learning methods. Additional comments including restrictions and unusual features: JAX-Fluids relies on open-source third-party Python libraries. These are automatically installed. In the current version, JAX-Fluids only runs on a single accelerator (CPU/GPU/TPU). Future versions will include support for parallel execution. JAX-Fluids has been tested on Linux and macOS operating systems.