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

Understanding shock-solid interactions remains a central challenge in fluid dynamics. We introduce JAX-Shock: a fully-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 and end-to-end training of deep learning-accelerated models. JAX-Shock integrates fifth-order WENO reconstruction and HLLC flux to resolve shocks and discontinuities with high fidelity. It supports sequence-to-sequence learning for shock interaction prediction and reverse-mode inference to identify key physical parameters from data. Additionally, we introduce a neural flux module trained to augment the numerical fluxes with data-driven corrections, significantly improving accuracy and generalization. Compared to purely data-driven approaches, JAX-Shock improves generalization while preserving physical consistency. The framework provides a flexible platform for differentiable physics, learning-based modeling, and inverse design in compressible flow regimes involving complex shock dynamics.