zoomy_jax.fvm.reconstruction_jax module

Contents

zoomy_jax.fvm.reconstruction_jax module#

JIT-compatible FVM face reconstruction and diffusion for JAX.

Mirrors the NumPy reconstruction classes in zoomy_core.fvm.reconstruction with the same interface: recon(Q) → (Q_L, Q_R).

All operations use JAX primitives (jnp.at[].add, jnp.where) for full JIT and autodiff compatibility.

Classes#

  • ConstantReconstruction: 1st-order piecewise-constant.

  • MUSCLReconstruction: 2nd-order MUSCL with slope limiting.

  • FreeSurfaceMUSCL: MUSCL with wet-dry fallback and h-positivity.

  • DiffusionOperatorJAX: Sparse discrete Laplacian with Crank-Nicolson implicit solve.

class zoomy_jax.fvm.reconstruction_jax.ConstantReconstruction(mesh, dim)#

Bases: object

First-order piecewise-constant reconstruction (JAX).

Takes inner-cells-only Q and BC-provided bf_face_values so boundary faces get Q_R = bf_values (not a stale ghost cell).

class zoomy_jax.fvm.reconstruction_jax.MUSCLReconstruction(mesh, dim, limiter='venkatakrishnan')#

Bases: object

Second-order MUSCL reconstruction with slope limiting (JAX, JIT-compatible).

Parameters:

  • mesh (MeshJAX) –

  • dim (int) –

  • limiter ("venkatakrishnan", "barth_jespersen", or "minmod") –

class zoomy_jax.fvm.reconstruction_jax.LSQMUSCLReconstructionJAX(mesh, dim, limiter='venkatakrishnan', unlimited_indices=None)#

Bases: object

LSQ-stencil MUSCL reconstruction (JAX) — mirrors the NumPy zoomy_core.fvm.reconstruction.LSQMUSCLReconstruction.

The LSQ stencil + boundary-face augmentation come from the mesh (lsq_gradQ, lsq_neighbors, lsq_boundary_face_neighbors, lsq_scale_factors). Q has shape (n_var, n_inner_cells) (no ghost cells). Boundary face values are passed in by the caller (typically the flux operator evaluates them via the BC kernel before calling reconstruct).

The class returns (Q_L, Q_R) of shape (n_var, n_faces):

  • Interior faces: both sides reconstructed from inner cells.

  • Boundary faces: Q_L = inner cell reconstructed at the face, Q_R = the BC-provided face value.

Limiter coefficients come from neighbor min/max (interior + boundary-face values). Currently supports Venkatakrishnan, Barth-Jespersen, and minmod.

class zoomy_jax.fvm.reconstruction_jax.FreeSurfaceLSQMUSCLJAX(mesh, dim, h_index, eps_wet=0.001, limiter='venkatakrishnan', unlimited_indices=None, momentum_indices=None)#

Bases: LSQMUSCLReconstructionJAX

LSQ-MUSCL with wet-dry fallback for free-surface flows (JAX) — mirrors NumPy FreeSurfaceLSQMUSCL.

In dry cells (h < eps_wet) drops to first-order (φ = 0). Clamps h ≥ 0 at face states and zeros the corresponding momentum components so hu/h doesn’t blow the flux to NaN at the dry front.

class zoomy_jax.fvm.reconstruction_jax.FreeSurfaceMUSCL(mesh, dim, h_index, eps_wet=0.001, limiter='venkatakrishnan', momentum_indices=None)#

Bases: MUSCLReconstruction

[LEGACY] Green-Gauss MUSCL with wet-dry fallback (JAX).

Superseded by FreeSurfaceLSQMUSCLJAX for the new NSM-routed JAX HyperbolicSolver — kept temporarily for tests that pin against the old reconstruction.

In dry cells (h < eps_wet), falls back to 1st order (φ = 0). Clamps h ≥ 0 at face states after reconstruction, and zeros the corresponding momentum components so dry-face velocities never become infinite (without this, hu/h blows the HLLC flux to NaN at the dry front of a Ritter dam-break).

class zoomy_jax.fvm.reconstruction_jax.DiffusionOperatorJAX(mesh, dim, nu=1.0)#

Bases: object

Sparse discrete diffusion operator for JAX: L(u) = nabla . (nu nabla u).

Assembled once per mesh + viscosity as a dense (nc, nc) matrix. Provides: - explicit(u): L @ u (for explicit stepping) - implicit_solve(u_star, dt): Crank-Nicolson solve

The dense matrix approach is JIT-compatible and works inside jax.lax.while_loop. For typical 1D/2D FVM grids (up to ~1000 cells) this is efficient; for larger grids a matrix-free GMRES variant is also available via implicit_solve_gmres.

explicit(u)#

Compute L @ u[:nc] (for explicit stepping). Returns shape (nc,).

implicit_solve(u_star, dt)#

Crank-Nicolson: (I - dt/2 * L) u^{n+1} = (I + dt/2 * L) u*.

Second-order in time for diffusion. Uses dense linear solve (jnp.linalg.solve) which is fully JIT-compatible.

Parameters:

  • u_star (jnp.ndarray, shape (n_cells,)) – State after explicit advection step (includes ghost cells).

  • dt (scalar) – Time step.

Returns:

u_new – Updated state with ghost cells copied from inner neighbors.

Return type:

jnp.ndarray, shape (n_cells,)

implicit_solve_gmres(u_star, dt, tol=1e-08, maxiter=100)#

Crank-Nicolson via matrix-free GMRES. JIT-compatible.

Use this for larger grids where the dense solve becomes expensive.

Contents