JANC: A cost-effective, differentiable compressible reacting flow solver featured with JAX-based adaptive mesh refinement

Abstract

The compressible reacting flow numerical solver is an essential tool in the study of combustion, energy disciplines, as well as in the design of industrial power and propulsion devices. We have established the first JAX-based block-structured adaptive mesh refinement (AMR) framework, called JAX-AMR, and then developed a fully-differentiable solver for compressible reacting flows, named JANC. JANC is implemented in Python and features automatic differentiation capabilities, enabling an efficient integration of the solver with machine learning. Furthermore, benefited by multiple acceleration features such as XLA-powered JIT compilation, GPU/TPU computing, parallel computing, and AMR, the computational efficiency of JANC has been significantly improved. In a comparative test of a two-dimensional detonation tube case, the computational cost of the JANC core solver, running on a single A100 GPU, was reduced to 1% of that of OpenFOAM, which was parallelized across 384 CPU cores. When the AMR method is enabled for both solvers, JANC's computational cost can be reduced to 1-2% of that of OpenFOAM. The core solver of JANC has also been tested for parallel computation on a 4-card A100 setup, demonstrating its convenient and efficient parallel computing capability. JANC also shows strong compatibility with machine learning by combining adjoint optimization to make the whole dynamic trajectory efficiently differentiable. JANC provides a new generation of high-performance, cost-effective, and high-precision solver framework for large-scale numerical simulations of compressible reacting flows and related machine learning research. Now, the source codes have been available under the MIT license at https://github.com/JA4S/JAX-AMR and https://github.com/JA4S/JANC.