JAX-FVM: A differentiable, entropy-stable finite volume solver on unstructured meshes for compressible flows

Abstract: We present JAX-FVM, an open-source, fully differentiable finite volume method (FVM) for the two-dimensional compressible Euler and Navier-Stokes equations on unstructured triangular meshes. The solver is written entirely in JAX, so that every operation : mesh connectivity, flux evaluation, slope limiting, and time integration is just-in-time compiled, vectorised, and end-to-end differentiable through automatic differentiation (AD), and runs transparently on CPU or GPU. On the numerical side, JAX-FVM is built around an entropy-conservative Tadmor/Ismail-Roe two-point flux supplemented with entropy-variable Rusanov or Roe dissipation, second-order MUSCL reconstruction of primitive variables with least-squares gradients and Venkatakrishnan limiting, and a family of explicit (RK2-4) and matrix-free implicit (Newton, SDIRK2) time integrators whose Jacobian actions are obtained by AD. The combination of an unstructured-mesh compressible FVM with end-to-end differentiability fills a gap left by existing differentiable CFD frameworks, which are almost exclusively restricted to structured grids or spectral discretisations. We describe the governing equations, the discretisation, the software architecture, and a set of standard verification cases. The code is openly available at https://github.com/guigzair/jax_fvm.

• The paper presents an open-source JAX solver that combines end-to-end automatic differentiation, entropy-stable finite-volume discretisation, and native unstructured triangular meshes for 2D compressible Euler and Navier–Stokes flows.
• The solver uses MUSCL reconstruction, entropy-conservative and entropy-stable fluxes, explicit and implicit time integration, and matrix-free JAX derivatives to support inverse design, optimisation, and physics–machine-learning workflows.
• The verification suite reproduces standard shock, vortex, instability, and flow-past-cylinder cases, but the current release remains limited by qualitative validation, two-dimensional meshes, restricted mesh input, no turbulence models, and absent performance benchmarks.

JAX-FVM is an open-source, fully differentiable finite volume solver for the two-dimensional compressible Euler and Navier–Stokes equations on unstructured triangular meshes, written entirely in JAX (2607.07385). Its distinguishing feature is the combination of three properties that, according to the authors, are not jointly available in any existing publicly released code: end-to-end automatic differentiability (AD), entropy-stable discretisation, and native support for unstructured meshes. The code is distributed at github.com/guigzair/jax_fvm.

Motivation and statement of need

The solver targets two audiences simultaneously: CFD practitioners seeking a compact entropy-stable compressible solver, and scientific machine-learning researchers requiring differentiable access to gradients of any scalar flow functional with respect to initial conditions, boundary data, physical parameters, or embedded neural models. The authors argue that existing differentiable CFD platforms — JAX-CFD (2607.07385), JAX-Fluids, ΦFlow, WaterLily, PICT, and Diff-FlowFSI — operate almost exclusively on structured Cartesian grids, whose regular data layout maps naturally onto SIMD/SIMT execution. Unstructured meshes, the workhorse of industrial CFD for body-conforming discretisation of complex geometry, are largely absent from the differentiable-simulation ecosystem because irregular connectivity and variable neighbour counts are hostile to accelerator execution models. The concurrent DiFVM effort addresses unstructured incompressible flow by recasting finite-volume operators as graph message-passing primitives, but no public code is available and the compressible, entropy-stable regime is outside its scope. JAX-FVM therefore claims to be, to the authors' knowledge, among the first openly available differentiable finite volume solvers for compressible flow on unstructured meshes.

The paper also motivates the AD-based design on numerical grounds: the entropy-variable Jacobians underlying the dissipation operators and implicit solvers can be formed exactly by differentiation rather than hand derivation.

Discretisation

The governing equations are the two-dimensional compressible Navier–Stokes equations in conservative form with state $\mathbf{w} = (\rho, \rho u, \rho v, E)^\top$, closed by an ideal-gas equation of state written in non-dimensional form carrying a reference Mach number $M$, with $\gamma = 1.4$ by default. The viscous flux uses a Newtonian stress tensor under the Stokes hypothesis with Fourier heat conduction.

The spatial discretisation is cell-centred finite volume. Second-order accuracy is achieved through MUSCL reconstruction of primitive variables $(\rho, u, v, p)$ with weighted least-squares cell gradients and either a smooth Venkatakrishnan limiter or minmod limiting. The viscous fluxes are evaluated from face-averaged LSQ gradients with a directional correction, which the authors note also avoids odd-even decoupling.

The default convective flux is the entropy-conservative Tadmor two-point flux in the affordable Ismail–Roe parameterisation, augmented with entropy-variable dissipation:

$$\hat{\mathbf{H}} = \hat{\mathbf{H}}^{\mathrm{EC}} - \tfrac{1}{2}\,\alpha\,|\lambda_{\max}|\, \frac{\partial \mathbf{w}}{\partial \boldsymbol{\eta}}\,(\boldsymbol{\eta}_R - \boldsymbol{\eta}_L)$$

yielding an entropy-stable Rusanov scheme; a full Roe matrix-dissipation flux with entropy-variable wave-strength scaling and a plain central/Rusanov flux are provided as alternatives. A notable implementation choice is that the Jacobian–vector product $$\partial \mathbf{w}/\partial\boldsymbol{\eta}\,(\boldsymbol{\eta}_R - \boldsymbol{\eta}_L)$$ is evaluated by a forward-mode AD call (jax.jvp) rather than an analytical derivation. The semi-discrete scheme consequently satisfies a discrete entropy inequality, and the paper states that this property is "difficult to obtain and verify without an AD-enabled implementation" — a strong claim in that it positions differentiability as integral to the numerical scheme's correctness, not merely a convenience.

Boundary conditions (periodic, slip/no-slip wall, supersonic inlet, outlet, subsonic inlet) are encoded through integer face markers with ghost/exterior states.

Time integration

Explicit options comprise two-stage SSP-Heun (RK2, the default), RK3, and RK4. For stiff or steady problems, backward Euler and SDIRK2 (with $x = 1 - 1/\sqrt{2}$) are implemented, each implicit stage solved by Newton iterations in which Jacobian actions $$\mathbf{v} \mapsto \mathbf{v} + \Delta t\,(\partial\mathcal{R}/\partial\mathbf{w})\,\mathbf{v}$$ are computed matrix-free via jax.jvp and passed to GMRES. The explicit step size follows a convective CFL condition, augmented by a diffusive limit $\Delta t \sim \Delta x^2/\nu$ for viscous runs.

Implementation and differentiability

All operations — mesh connectivity, flux evaluation, slope limiting, time integration — are just-in-time compiled with the static mesh passed as a compile-time argument; per-cell operations such as the LSQ solve are expressed as nested vmap maps. Reverse-mode AD propagates gradients of any scalar objective through the full residual, and the package's dependencies on flax and optax allow optimisation and inverse-design loops to be assembled directly. Execution is hardware-agnostic across CPU and GPU, with double precision available for strict conservation checks. A diagnostics suite covering total entropy, kinetic energy, enstrophy, palinstrophy, vorticity, and Mach number supports quantitative verification of structure-preserving behaviour.

Verification

The headline case is the Mach 3 forward-facing step of Woodward and Colella, computed at $t = 4$ on a fully unstructured triangulation with the entropy-conservative flux, Venkatakrishnan-limited MUSCL reconstruction, and explicit RK2. The simulation resolves the canonical shock system — detached bow shock, Mach stem, reflected shocks, and contact/slip line. Beyond this, the distribution ships the full set of nineteen two-dimensional Riemann configurations of Lax and Liu, the double Mach reflection, Kelvin–Helmholtz, the Taylor–Green vortex, several isentropic low-Mach vortex problems initialised in mechanical equilibrium via pressure-Poisson projection to suppress spurious acoustics, and a cylinder-in-channel geometry. Structure preservation is monitored through the entropy, energy, and enstrophy diagnostics.

It should be noted that the verification is qualitative and standard-case-based: the paper reports no convergence-rate studies, no quantitative error norms, no entropy-production measurements against reference solutions, and no performance benchmarks (e.g., wall-clock comparisons against structured-grid differentiable solvers or against hand-coded adjoint codes). The claim of filling a gap in the differentiable-CFD ecosystem is therefore supported by existence and correctness of the capability rather than by quantitative superiority.

Limitations and open questions

The paper is explicit about several restrictions. Mesh input is confined to the built-in MeshPy/Triangle generator; a VTK reader for externally generated meshes is planned but not yet available. The implementation is limited to two-dimensional triangular meshes, with extension to three-dimensional tetrahedra planned but not demonstrated. No turbulence closure is included; Smagorinsky and RANS models are listed as future work. Additionally, the comparison with DiFVM rests on the absence of a public DiFVM release, a claim that is time-sensitive. An open technical question is whether the vmap-based handling of irregular connectivity scales to industrially relevant mesh sizes on GPUs, since no scaling data is presented.

Conclusion

JAX-FVM delivers a compact, entropy-stable, end-to-end differentiable compressible Euler/Navier–Stokes finite volume solver on unstructured triangular meshes, executable on CPU or GPU. Its combination of unstructured geometry, entropy-consistent numerics, and JAX-based AD makes it a candidate platform for gradient-based inverse problems, PDE-constrained optimisation, and hybrid physics/machine-learning modelling in a regime not covered by existing structured-grid differentiable CFD frameworks. The principal caveats are its two-dimensionality, restricted mesh input, absence of turbulence modelling, and the lack of quantitative verification or performance data in the present release.