JAX-Shock: A Differentiable, GPU-Accelerated, Shock-Capturing Neural Solver for Compressible Flow Simulation

Abstract

Understanding shock-solid interactions remains a central challenge in compressiblefluiddynamics. WepresentJAX-Shock: afully-differentiable,GPU-accelerated, high-order shock-capturing solver for efficient simulation of the compressible Navier-Stokes equations. Built entirely in JAX, the framework leverages automatic differentiation to enable gradient-based optimization, parameter inference, and end-to-end training of deep learning-augmented models. The solver integrates fifth-order WENO reconstruction with an HLLC flux to resolve shocks and discontinuities with high fidelity. To handle complex geometries, an immersed boundary method is implemented for accurate representation of solid interfaces within the compressible flow field. In addition, we introduce a neural flux module trained to augment the numerical fluxes with data-driven corrections, significantly improving accuracy and generalization. JAX-Shock also supports sequence-to-sequence learning for shock interaction prediction and reverse-mode inference to identify key physical parameters from data. Compared with purely data-driven approaches, JAX-Shock enhances generalization while preserving physical consistency. The framework establishes a flexible platform for differentiable physics, learning-based modeling, and inverse design in compressible flow regimes dominated by complex shock-solid interactions.