Explicit and Implicit Finite Difference Solvers Implemented in JAX for Shock Wave Physics

TL;DR

This paper develops explicit and implicit finite-difference solvers for shock wave modeling using JAX, demonstrating their accuracy, stability, and portability across hardware, and providing a benchmark dataset for CFD and machine learning applications.

Contribution

It introduces a novel JAX-based implementation of finite-difference solvers for shock physics, enabling portable, scalable, and high-performance CFD simulations.

Findings

Explicit scheme captures shocks accurately under strict time steps.

Implicit scheme offers greater stability and accuracy at higher computational cost.

JAX implementation performs well across CPUs, GPUs, and TPUs.

Abstract

Shock dynamics and nonlinear wave propagation are fundamental to computational fluid dynamics (CFD) and high-speed flow modeling. In this study, we developed explicit and implicit finite-difference solvers for the one-dimensional Burgers viscous equation to model shock formation, propagation, and dissipation. The governing equation, which incorporates convective and diffusive effects, serves as a simplified analogue of the Navier-Stokes equations. Using the Finite-JAX framework, each solver is implemented with upwind and central finite-difference schemes for the convective and diffusive terms, respectively. Time integration is performed using explicit forward Euler and implicit backward-time central space (BTCS) schemes under periodic and Dirichlet boundary conditions. Stability is ensured by the Courant-Friedrichs-Lewy (CFL) criteria for the convective and diffusive components. Numerical experiments quantify the accuracy, convergence, and real-time performance of JAX across CPUs, GPUs, and TPUs, demonstrating that JAX maintains fidelity while achieving portability. The results show that the explicit scheme captures impact accurately under strict time-step constraints, while the implicit formulation provides greater stability and accuracy at a higher computational cost. Taken together, these results establish a reproducible dataset for benchmarking CFD solvers and training machine learning models for nonlinear transport and impact-driven phenomena. Our new implementation of FiniteJAX enhances the portability, scalability, and performance of solvers based on the JAX framework developed by Google DeepMind.