A differentiable, shock-capturing neural solver for compressible flow simulation

Understanding shock–solid interactions remains a central challenge in compressible fluid dynamics. We present JAX-Shock (just-in-time compiled, exact shock-capturing solver with automatic differentiation): a fully-differentiable, graphics processing unit-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 weighted essentially non-oscillatory reconstruction with an Harten-Lax-van Leer-Contact flux to resolve shocks and discontinuities with high fidelity. To handle complex geometries, an immersed boundary method is implemented for the 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 the 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.