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

TL;DR

JAX-Shock is a GPU-accelerated, fully differentiable solver for compressible flow that combines high-order shock capturing with neural augmentation, enabling efficient simulation, optimization, and inverse modeling of shock interactions.

Contribution

The paper introduces JAX-Shock, a novel differentiable, GPU-accelerated solver that integrates high-order shock capturing with neural flux corrections for improved accuracy and flexibility in compressible flow simulations.

Findings

Achieves high-fidelity shock resolution with WENO and HLLC flux.

Enhances accuracy through neural flux augmentation.

Supports gradient-based optimization and inverse modeling.

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.