zoomy_jax.fvm.solver_chorin_vam_jax module

Contents

zoomy_jax.fvm.solver_chorin_vam_jax module#

ChorinSplitVAMSolverJax — JAX port of the NumPy ChorinSplitVAMSolver.

Three-step Chorin projection split for the VAM(1,2,2) chain:

predictor → pressure → corrector

Each substep is driven by its own JaxRuntime over the corresponding sub-SystemModel (sm_pred / sm_press / sm_corr). The pressure step is matrix-free Newton-GMRES via jax.scipy.sparse.linalg.gmres(); the corrector is one state_update call on sm_corr’s runtime.

Reuses the pure-symbolic helpers from the NumPy implementation (_pad_to_square, _substitute_dt) — those are backend-agnostic.

Status#

  • Predictor: JAX HyperbolicSolver flux loop (symbolic Riemann via JaxRuntime).

  • Pressure: matrix-free GMRES (jax.scipy.sparse.linalg.gmres).

  • Corrector: lambdified state_update from sm_corr.

  • Aux pools: per-sub-system, refreshed via LSQ stencils on the mesh (mirrors NumPy update_aux_variables).

class zoomy_jax.fvm.solver_chorin_vam_jax.ChorinSplitVAMSolverJax(sm_pred, sm_press, sm_corr, *, reconstruction=None, **kwargs)#

Bases: HyperbolicSolver

JAX port of ChorinSplitVAMSolver.

Parameters:

  • sm_pred (SystemModel) – The three sub-systems from split_simple(sm, pressure_vars, dt_symbol).

  • sm_press (SystemModel) – The three sub-systems from split_simple(sm, pressure_vars, dt_symbol).

  • sm_corr (SystemModel) – The three sub-systems from split_simple(sm, pressure_vars, dt_symbol).

  • reconstruction (ReconstructionSpec, optional) – Predictor reconstruction spec.

  • pressure_tol (float) –

  • pressure_maxit (int) –

  • time_order (int) – 1 = single Chorin cycle per step; 2 = SSPRK2 wrap.

pressure_tol = 1e-06#
pressure_maxit = 100#
time_order = 1#
setup_simulation(mesh, write_output=False)#

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

property nc#
update_aux_variables()#

Host-side wrapper that mutates self — convenience for notebooks / tests that drive the solver imperatively. The pure-functional equivalent (used by the JIT-able step) is _refresh_aux_for_sm() called per sub-system.

chorin_cycle(dt, time, Q, Qaux_pred, Qaux_press, Qaux_corr)#

One predictor → pressure → corrector cycle as a PURE function of (Q, Qaux_pred, Qaux_press, Qaux_corr) — no self._sim_* mutations during the call. This is the JIT-able entry point. Returns the updated 4-tuple.

step(dt, time, Q, Qaux)#

Host-side wrapper that drives one Chorin cycle and mutates self._sim_Q / self.Qaux_press / self.Qaux_corr. Calls the pure-functional chorin_cycle() internally. Tests and notebooks that want a JIT’d run loop should call chorin_cycle() (or run_jit_steps()) directly.

run_jit_steps(dt, n_steps, Q, Qaux_pred, Qaux_press, Qaux_corr, t_start=0.0)#

JIT’d time loop via jax.lax.scan. Calls chorin_cycle() n_steps times. Returns final (Q, Qaux_pred, Qaux_press, Qaux_corr, time).

name = 'ChorinSplitVAMSolverJax'#
Contents