Search papers, labs, and topics across Lattice.
The paper introduces JAX-Shock, a fully differentiable, GPU-accelerated solver for the compressible Navier-Stokes equations, built using JAX to enable gradient-based optimization and learning. It combines a fifth-order WENO scheme with an HLLC flux for shock capturing and an immersed boundary method for complex geometries. A neural flux module is trained to augment numerical fluxes, improving accuracy and generalization compared to purely data-driven methods.
Differentiable fluid dynamics just got a speed boost: JAX-Shock offers a GPU-accelerated, high-order shock-capturing solver that lets you backprop through compressible flow simulations.
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.
