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

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.