zoomy_jax.fvm.solver_jax module

zoomy_jax.fvm.solver_jax module#

JAX FVM solvers: JIT-compiled time stepping.

Solver hierarchy (mirrors NumPy):
HyperbolicSolver (explicit flux + source)

setup_simulation — mesh/model → JAX operators (closures) step — single explicit timestep run_simulation — jax.lax.while_loop over step solve — init + setup + run

Inherits param definitions from NumPy HyperbolicSolverNumpy but overrides all computational methods with JAX implementations.

zoomy_jax.fvm.solver_jax.log_callback_hyperbolic(iteration, time, dt, time_stamp, log_every=10)#

Log callback hyperbolic.

zoomy_jax.fvm.solver_jax.log_callback_poisson(iteration, res)#

Log callback poisson.

zoomy_jax.fvm.solver_jax.log_callback_execution_time(time)#

Log callback execution time.

zoomy_jax.fvm.solver_jax.newton_solver(residual)#

Newton solver.

class zoomy_jax.fvm.solver_jax.HyperbolicSolver(**kwargs)#

Bases: HyperbolicSolver

JAX HyperbolicSolver — JIT-compiled explicit time stepping.

Follows the setup_simulation / step / run_simulation pattern. Inherits param definitions from the NumPy base class.

initialize(mesh, model)#

Allocate state arrays of shape (n_var, n_inner_cells) — no ghost cells. Boundary values are computed inline inside the flux operator from the indexed BC kernel; the LSQ-MUSCL reconstruction sees them via the LSQ-augmented stencil (lsq_boundary_face_neighbors) and uses the BC face values as Q_R at boundary faces (NumPy parity).

Matches the NumPy contract: Q.shape[1] == mesh.n_inner_cells.

create_runtime(Q, Qaux, mesh, model)#

Create the JAX runtime over the NSM.

model here is self.nsm.sm — the SystemModel inside the NSM that setup_simulation resolved. Builds a JaxRuntime (jit-vmapped per-cell + per-face operators) and converts the LSQ mesh to its JAX form.

update_q(Q, Qaux, mesh, model, parameters)#

JIT-compatible update_variables via vmap (replaces NumPy cell loop).

get_compute_source(mesh, runtime)#

JIT-compiled source operator using JaxRuntime.source (which is itself jit-vmap’d over the cell axis).

get_apply_boundary_conditions(mesh, runtime)#

JIT-compiled BC operator that fills ghost cells via the indexed kernel on JaxRuntime.boundary_conditions.

get_compute_max_abs_eigenvalue(mesh, runtime)#

Max |eigenvalue| over the faces, using JaxRuntime.eigenvalues (which is jit-vmap’d per cell).

get_flux_operator(mesh, runtime)#

Flux operator — literal NumPy parity (S5b).

  • Q has shape (n_var, n_inner_cells) — no ghost cells.

  • Boundary face values are evaluated inline from the BC kernel (vmap’d over the boundary face axis).

  • The LSQ-MUSCL reconstruction is fed Q + bf_values so its limiter bounds and Q_R at boundary faces are physically consistent.

  • Interior + boundary faces are accumulated in separate accumulators (split loops).

setup_simulation(mesh, model)#

Build all JAX operators from mesh and model.

model may be a Model, a SystemModel, or a NumericalSystemModel. All numerical knobs (reconstruction.order, reconstruction.limiter, regularization.eigenvalue_eps) live on the NSM; the mesh stencil uses nsm.resolved_lsq_degree().

Until the JAX runtime is wired through JaxRuntimeModel.from_system_model (S4b), callers passing a bare SystemModel or an NSM built from one must do so with a source Model — JAX’s Kernel + Model-based JaxRuntimeModel paths still need it.

Returns:

Q, Qaux – Initial state arrays on device.

Return type:

jnp.ndarray

step(dt, time, Q, Qaux)#

Perform a single explicit time step.

Each RK stage starts by refreshing the boundary ghost cells from the current interior state — mirroring the numpy solver, which evaluates the indexed BC kernel inside its flux operator. Without per-stage BCs the second RK stage reads stale ghosts from the previous outer step, contaminating a few cells next to each boundary and stripping the global L2 of its 2nd-order rate even though the interior is fine (interior errs ≈ O(dx²), boundary errs ≈ O(dx)).

post_step(time, dt, Q, Qold, Qaux)#

Post-step processing: BCs, update_q, update_qaux.

Separated from step so that subclasses (e.g. IMEX) can insert implicit solves between the explicit step and the post-processing.

Parameters:

  • time (scalar) – Current simulation time (after dt advance).

  • dt (scalar) – Time step size.

  • Q (jnp.ndarray) – State after explicit step.

  • Qold (jnp.ndarray) – State before the step (for aux updates).

  • Qaux (jnp.ndarray) – Auxiliary state before the step.

Returns:

Q_new, Qaux_new – Fully updated state and auxiliary arrays.

Return type:

jnp.ndarray

compute_timestep(Q, Qaux)#

Compute the adaptive time step using the stored eigenvalue operator.

JIT-compatible. Uses self.compute_dt (from param) with the precomputed eigenvalue operator and min inradius.

Returns:

dt

Return type:

scalar

run_simulation(Q, Qaux, write_output=True)#

JIT-compiled time loop using jax.lax.while_loop.

Calls compute_timestep -> step -> post_step in a while_loop until time >= time_end.

Parameters:

  • Q (jnp.ndarray) – Initial state (from setup_simulation).

  • Qaux (jnp.ndarray) – Initial state (from setup_simulation).

  • write_output (bool) – Whether to write snapshots to HDF5.

Returns:

Q, Qaux – Final state.

Return type:

jnp.ndarray

solve(mesh, model, write_output=True)#

Full solve: initialize -> setup -> run.

This is the main entry point, compatible with the NumPy solver interface. Calls setup_simulation then run_simulation.

name = 'HyperbolicSolver'#
class zoomy_jax.fvm.solver_jax.PoissonSolver(**kwargs)#

Bases: Solver

PoissonSolver. (class).

get_residual(Qaux, Qold, Qauxold, parameters, mesh, model, boundary_operator, time, dt)#

Get residual.

solve(mesh, model, write_output=True)#

Solve.

name = 'PoissonSolver'#