Development of a ML-enabled high-order Discontinuous Galerkin solver for compressible flow simulations

Modern high-performance computing and advanced algorithms have presented several opportunities for the advancement of engineering. Despite this, fluid dynamics simulations are not yet fully leveraging the emerging opportunities for enabling the integration of CFD with machine-learning on modern hardware architectures such as massively-parallel GPUs. Thus, we present a new Discontinuous Galerkin (DG) solver built using the ML-enabled JAX library. The DG method offers high order of accuracy as well as high arithmetic intensity, which optimally maps on the GPU hardware, thereby maximizing energy efficiency for tightly integrating scientific computing and machine-learning tasks. We demonstrate scalability of the resulting JAX-based DG solver in applications to different flow problems.