astronomix._integrators._explicit_rk

Generic explicit Runge-Kutta time-integration schemes.

These drivers are agnostic to the spatial discretisation and to the state layout: the state u is an arbitrary JAX pytree (a single conserved-state array for the hydro solvers, or the (q, bx, by, bz) tuple for the MHD constrained-transport solver) and the model is supplied as a small set of closures. The finite-difference (WENO / CT) and finite-volume (unsplit) solvers build those closures and call into the scheme; the heavy, discretisation-specific work lives there, not here.

Model contract (all closures operate on the state pytree u):

pre_stage(u) -> u

Hooks applied at the start of every stage (positivity flooring, boundary handling, …). Defaults to the identity.

rhs(u, dt_stage) -> du

The stage increment dt_stage * L(u), a pytree matching u.

finalize(u) -> u

Hooks applied once after the last stage (e.g. recomputing the cell-centered MHD fields from the interface fields, final positivity). Defaults to the identity.

lsrk_increment(u, du, a_coef, dt) -> du [optional, low-storage only]

Returns the new low-storage register a_coef * du + dt * L(u). Used by lsrk4 to let a model fuse the a_coef * du accumulate into its flux-divergence kernel (memory optimisation). When omitted, lsrk4 falls back to a_coef * du + rhs(u, dt).

Schemes:

ssprk4 - 5-stage 4th-order Spiteri-Ruuth SSPRK (3-register). lsrk4 - 5-stage 4th-order Carpenter-Kennedy 2N-storage low-storage RK. rk2_ssp - 2-stage 2nd-order SSP RK (Heun).

Module Contents

Functions

ssprk4	5-stage 4th-order SSPRK (Spiteri-Ruuth).
lsrk4	5-stage 4th-order Carpenter-Kennedy 2N-storage low-storage RK4.
rk2_ssp	2-stage 2nd-order SSP RK (Heun).

API

astronomix._integrators._explicit_rk.ssprk4(u0, dt, *, rhs, pre_stage=_identity, finalize=_identity)

5-stage 4th-order SSPRK (Spiteri-Ruuth).

Args:

u0: Initial state pytree at t = t_n. dt: Full time step. rhs: rhs(u, dt_stage) -> du stage increment. pre_stage: per-stage hook, defaults to identity. finalize: post-integration hook, defaults to identity.

Returns:

The state pytree at t = t_n + dt.

astronomix._integrators._explicit_rk.lsrk4(u0, dt, *, pre_stage=_identity, finalize=_identity, rhs=None, lsrk_increment=None)

5-stage 4th-order Carpenter-Kennedy 2N-storage low-storage RK4.

Either rhs or lsrk_increment must be supplied. lsrk_increment lets the model fuse the a_coef * du accumulate into its flux kernel; when only rhs is given the accumulate is done here via tree_map.

astronomix._integrators._explicit_rk.rk2_ssp(u0, dt, *, rhs, pre_stage=_identity, finalize=_identity)

2-stage 2nd-order SSP RK (Heun).