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

Abstract

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 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.

53 pages, 23 figures

References in corpus (6)

• Array Programming with NumPy
• The NumPy array: a structure for efficient numerical computation
• Machine learning accelerated computational fluid dynamics
• Learned discretizations for passive scalar advection in a 2-D turbulent flow
• Enhancement of shock-capturing methods via machine learning
• Consistent and symmetry preserving data-driven interface reconstruction for the level-set method

Cited by in corpus (16)

• JAX-FEM: A differentiable GPU-accelerated 3D finite element solver for automatic inverse design and mechanistic data science
• XLB: A differentiable massively parallel lattice Boltzmann library in Python
• Differentiable Programming for Earth System Modeling
• A framework for structural shape optimization based on automatic differentiation, the adjoint method and accelerated linear algebra
• Godunov Loss Functions for Modelling of Hyperbolic Conservation Laws
• Optimized Flow Control based on Automatic Differentiation in Compressible Turbulent Channel Flows
• Scientific Machine Learning of Flow Resistance Using Universal Shallow Water Equations with Differentiable Programming
• Testing and benchmarking emerging supercomputers via the MFC flow solver
• A differentiable programming framework for spin models
• Solving continuum and rarefied flows using differentiable programming
• Temperature transformation recovering the compressible law of the wall for turbulent channel flow
• PICT -- A Differentiable, GPU-Accelerated Multi-Block PISO Solver for Simulation-Coupled Learning Tasks in Fluid Dynamics
• A Comparison of Mesh-Free Differentiable Programming and Data-Driven Strategies for Optimal Control under PDE Constraints
• Neural network-based Godunov corrections for approximate Riemann solvers using bi-fidelity learning
• Lattice Discrete Particle Model (LDPM): Comparison of Various Time Integration Solvers and Implementations
• Flux-controlled wall model for large eddy simulation integrating the compressible law of the wall