zoomy_jax.transformation.jax_runtime module

zoomy_jax.transformation.jax_runtime module#

JaxRuntime — clean JIT-vmapped JAX runtime over a NumericalSystemModel.

This is the new JAX runtime: the documented workflow is

Model → SystemModel → NumericalSystemModel → JaxRuntime

Every operator stored on the underlying SystemModel (flux, source, hydrostatic_pressure, nonconservative_matrix, quasilinear_matrix, mass_matrix, eigenvalues, the indexed BC kernels) is lambdified once with the JAX module, then wrapped in jax.jit(jax.vmap(...)) so a single call evaluates the operator at every cell (or every face) in one shot.

The Riemann Numerics (built from nsm.build_numerics()) is lambdified the same way and exposes per-face numerical_flux / numerical_fluctuations vmapped over the face axis.

Design choices:

  • No source-Model dependency. JaxRuntime consumes only the NSM (and its embedded SystemModel). The legacy path that required Kernel(model) + JaxRuntimeModel(model, kernel=...) is gone.

  • Parameters are live. self.parameters is a property that reads nsm.sm.parameter_values fresh on every access, so user mutations (e.g. nsm.sm.parameter_values.g = 12.0) flow through without a runtime rebuild. jax.grad can be taken against the parameter axis of any operator.

  • No broadcast hack. Operators are vmap’d over the cell/face axis explicitly — no ones_like(anchor) wrapping; each lambdified function takes pure scalar inputs.

Replaces the legacy JaxRuntimeModel (Model-based) and JaxRuntimeSymbolic (Numerics-based, broadcast-tricked).

class zoomy_jax.transformation.jax_runtime.JaxRuntime(nsm)#

Bases: object

JIT-vmapped JAX runtime over a NumericalSystemModel.

Parameters:

  • nsm (NumericalSystemModel) – The NSM whose embedded sm (and the Numerics built from the NSM’s riemann class) provide every operator. A bare SystemModel (or Model) is auto-promoted via NumericalSystemModel.from_system_model().

  • exposed (Operators) –

  • ----------------- –

  • Qaux) (Per-cell (vmap over the cell axis of Q and) –

  • ``flux(Q (*) –

  • Qaux –

  • ``(n_eq (normal)`` →) –

  • n_dim –

  • n_cells)`` –

  • ``hydrostatic_pressure(Q (*) –

  • Qaux –

  • ``(n_eq –

  • n_dim –

  • n_cells)`` –

  • ``source(Q (*) –

  • Qaux –

  • ``(n_eq –

  • n_cells)`` –

  • ``mass_matrix(Q (*) –

  • Qaux –

  • ``(n_eq –

  • n_state –

  • n_cells)`` –

  • ``nonconservative_matrix(Q (*) –

  • Qaux –

  • ``(n_eq –

  • n_state –

  • n_dim –

  • n_cells)`` –

  • ``quasilinear_matrix(Q (*) –

  • Qaux –

  • NCP (parameters)`` → same shape as) –

  • ``eigenvalues(Q (*) –

  • Qaux –

  • parameters –

  • ``(n_eq –

  • n_cells)`` –

  • axis) (Per-face (vmap over the face) –

  • ``numerical_flux(qL (*) –

    → (n_eq, n_faces)

  • qR – → (n_eq, n_faces)

  • qauxL – → (n_eq, n_faces)

  • qauxR – → (n_eq, n_faces)

  • parameters – → (n_eq, n_faces)

  • normal)`` – → (n_eq, n_faces)

  • ``numerical_fluctuations(qL (*) –

    → (2*n_eq, n_faces) (Dp and Dm stacked)

  • qR – → (2*n_eq, n_faces) (Dp and Dm stacked)

  • qauxL – → (2*n_eq, n_faces) (Dp and Dm stacked)

  • qauxR – → (2*n_eq, n_faces) (Dp and Dm stacked)

  • parameters – → (2*n_eq, n_faces) (Dp and Dm stacked)

  • normal)`` – → (2*n_eq, n_faces) (Dp and Dm stacked)

  • (indexed) (Boundary kernel) –

  • ``boundary_conditions(i_bc (*) –

    qaux_cell, parameters, normal)`` → (n_state,) ghost cell value. Invoke this inside the solver’s fori_loop over boundary faces — it is not itself vmap’d because the BC index varies per face.

  • time – qaux_cell, parameters, normal)`` → (n_state,) ghost cell value. Invoke this inside the solver’s fori_loop over boundary faces — it is not itself vmap’d because the BC index varies per face.

  • position – qaux_cell, parameters, normal)`` → (n_state,) ghost cell value. Invoke this inside the solver’s fori_loop over boundary faces — it is not itself vmap’d because the BC index varies per face.

  • distance – qaux_cell, parameters, normal)`` → (n_state,) ghost cell value. Invoke this inside the solver’s fori_loop over boundary faces — it is not itself vmap’d because the BC index varies per face.

  • q_cell – qaux_cell, parameters, normal)`` → (n_state,) ghost cell value. Invoke this inside the solver’s fori_loop over boundary faces — it is not itself vmap’d because the BC index varies per face.

:paramqaux_cell, parameters, normal)`` → (n_state,) ghost cell value.

Invoke this inside the solver’s fori_loop over boundary faces — it is not itself vmap’d because the BC index varies per face.

classmethod from_nsm(nsm)#

Auto-promote a Model / SystemModel to NSM and wrap.

Return type:

JaxRuntime

property parameters: Array#

Live numeric parameter array, read fresh from self.sm.parameter_values on every access. Mutating nsm.sm.parameter_values.g = 12.0 is reflected on the next call — no runtime rebuild needed. Pass this array (or a jax.grad-tracked clone) as the parameter argument of any operator to take AD w.r.t. parameters.

property parameter_names: List[str]#
property parameter_symbols: List[Symbol]#