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

Fuente: arXiv
Gespeichert in:
Bibliographische Detailangaben
1. Verfasser: Zhang, Bo
Format: Preprint
Veröffentlicht: 2026
Schlagworte:
Online-Zugang:
Tags: Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
_version_ 1866917190508216320
author Zhang, Bo
author_facet Zhang, Bo
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2601_04400
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle JAX-Shock: A Differentiable, GPU-Accelerated, Shock-Capturing Neural Solver for Compressible Flow Simulation
Zhang, Bo
Fluid Dynamics
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.
title JAX-Shock: A Differentiable, GPU-Accelerated, Shock-Capturing Neural Solver for Compressible Flow Simulation
topic Fluid Dynamics
url https://arxiv.org/abs/2601.04400