"""Explicit nonlinear time-step policies.
This module keeps the RK/SSP one-step formulas outside the public nonlinear
facade. The integrator still injects the RHS and projection functions so the
step policy remains pure, small, and directly testable.
"""
from __future__ import annotations
from functools import partial
import hashlib
import math
from typing import Any, Callable
import jax
import jax.numpy as jnp
from gkx.config import method_error
import numpy as np
from solvax.ode import RK4_2, PIController, erk_step, step_factor
RhsFn = Callable[..., tuple[jnp.ndarray, object]]
ProjectFn = Callable[[jnp.ndarray], jnp.ndarray]
ScanFn = Callable[..., tuple[jnp.ndarray, Any]]
DiagnosticFn = Callable[..., Any]
FieldSolveFn = Callable[..., Any]
CollisionSplitFn = Callable[[jnp.ndarray, Any, jnp.ndarray, str], jnp.ndarray]
DiagnosticStepFn = Callable[
[tuple[Any, Any, Any, Any, Any, Any], Any],
tuple[tuple[Any, Any, Any, Any, Any, Any], tuple[Any, Any, Any]],
]
SampledDiagnosticScanFn = Callable[
...,
tuple[tuple[Any, Any, Any, Any, Any, Any], tuple[Any, Any, Any]],
]
_SSPX3_ADT = float((1.0 / 6.0) ** (1.0 / 3.0))
_SSPX3_WGTFAC = float((9.0 - 2.0 * (6.0 ** (2.0 / 3.0))) ** 0.5)
_SSPX3_W1 = 0.5 * (_SSPX3_WGTFAC - 1.0)
_SSPX3_W2 = 0.5 * ((6.0 ** (2.0 / 3.0)) - 1.0 - _SSPX3_WGTFAC)
_SSPX3_W3 = (1.0 / _SSPX3_ADT) - 1.0 - _SSPX3_W2 * (_SSPX3_W1 + 1.0)
# rk4-pi: RK4 with the embedded explicit-midpoint error estimate (no extra RHS)
# and an I controller on the step. The estimate is relative to ||G||; the step is
# capped at ``_PI_CAP`` times the CFL step and a step with estimate > ``_PI_REDO``
# is redone once at the controller's smaller step.
_PI_RTOL = 1.0e-2
_PI_CAP = 2.0
_PI_REDO = 3.0
_PI_CONTROL = PIController(k_i=1.0, k_p=0.0, safety=0.9, fac_min=0.5, fac_max=1.3)
_NONLINEAR_METHODS = frozenset(
{"rk2", "rk3", "rk4", "rk4-pi", "sspx3", "imex-ars3", "imex-ark3"}
)
def _rhs_value(rhs_fn: RhsFn, state: jnp.ndarray) -> jnp.ndarray:
dG, _fields = rhs_fn(state)
return dG
def _explicit_rk3_heun_state(
G: jnp.ndarray,
dG: jnp.ndarray,
dt_local: jnp.ndarray,
*,
rhs_fn: RhsFn,
project_state: ProjectFn,
) -> jnp.ndarray:
G1 = project_state(G + (dt_local / 3.0) * dG)
k2 = _rhs_value(rhs_fn, G1)
G2 = project_state(G + (2.0 * dt_local / 3.0) * k2)
k3 = _rhs_value(rhs_fn, G2)
G3 = project_state(G + 0.75 * dt_local * k3)
return G3 + 0.25 * dt_local * dG
def _explicit_rk4_state(
G: jnp.ndarray,
dG: jnp.ndarray,
dt_local: jnp.ndarray,
*,
rhs_fn: RhsFn,
project_state: ProjectFn,
) -> jnp.ndarray:
# Running sum of the slopes: only the latest stage slope stays alive.
k2 = _rhs_value(rhs_fn, project_state(G + 0.5 * dt_local * dG))
total = dG + 2.0 * k2
k3 = _rhs_value(rhs_fn, project_state(G + 0.5 * dt_local * k2))
total = total + 2.0 * k3
k4 = _rhs_value(rhs_fn, project_state(G + dt_local * k3))
return G + (dt_local / 6.0) * (total + k4)
def _explicit_sspx3_state(
G: jnp.ndarray,
dt_local: jnp.ndarray,
*,
rhs_fn: RhsFn,
project_state: ProjectFn,
) -> jnp.ndarray:
def euler_step(G_state: jnp.ndarray) -> jnp.ndarray:
dG_state = _rhs_value(rhs_fn, G_state)
return project_state(G_state + (_SSPX3_ADT * dt_local) * dG_state)
G1 = euler_step(G)
G2_euler = euler_step(G1)
G2 = project_state((1.0 - _SSPX3_W1) * G + (_SSPX3_W1 - 1.0) * G1 + G2_euler)
G3 = euler_step(G2)
return (
(1.0 - _SSPX3_W2 - _SSPX3_W3) * G + _SSPX3_W3 * G1 + (_SSPX3_W2 - 1.0) * G2 + G3
)
def _explicit_stage_update(
G: jnp.ndarray,
dG: jnp.ndarray,
dt_local: jnp.ndarray,
*,
method: str,
rhs_fn: RhsFn,
project_state: ProjectFn,
) -> jnp.ndarray:
"""Advance one nonlinear explicit method before final projection/dtype cast."""
if method in {"imex-ars3", "imex-ark3"}:
chain_linear = getattr(rhs_fn, "chain_implicit_linear", None)
if chain_linear is None:
raise ValueError(
f"method '{method}' needs a chain implicit linear operator"
)
return chain_linear.ars_step(
G, dG, lambda state: _rhs_value(rhs_fn, state), project_state
)
if method == "rk2":
G_half = project_state(G + 0.5 * dt_local * dG)
return G + dt_local * _rhs_value(rhs_fn, G_half)
if method == "rk3":
return _explicit_rk3_heun_state(
G, dG, dt_local, rhs_fn=rhs_fn, project_state=project_state
)
if method in {"rk4", "rk4-pi"}:
return _explicit_rk4_state(
G, dG, dt_local, rhs_fn=rhs_fn, project_state=project_state
)
if method == "sspx3":
return _explicit_sspx3_state(
G, dt_local, rhs_fn=rhs_fn, project_state=project_state
)
raise method_error(method, _NONLINEAR_METHODS)
def _materialize(value: jnp.ndarray) -> jnp.ndarray:
"""Materialize ``value`` on CPU; identity elsewhere (it cost ~20% on GPU)."""
return jax.lax.platform_dependent(
value, cpu=jax.lax.optimization_barrier, default=lambda x: x
)
def _materialized(rhs_fn: RhsFn, project_state: ProjectFn) -> tuple[RhsFn, ProjectFn]:
def materialized_rhs(state: jnp.ndarray) -> tuple[jnp.ndarray, object]:
derivative, fields = rhs_fn(state)
return _materialize(derivative), fields
def materialized_projection(state: jnp.ndarray) -> jnp.ndarray:
return _materialize(project_state(state))
materialized_rhs.chain_implicit_linear = getattr( # type: ignore[attr-defined]
rhs_fn, "chain_implicit_linear", None
)
return materialized_rhs, materialized_projection
[docs]
def advance_rk4_pi_state(
G: jnp.ndarray,
dG: jnp.ndarray,
dt_local: jnp.ndarray,
*,
rhs_fn: RhsFn,
project_state: ProjectFn,
state_dtype: jnp.dtype,
) -> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]:
"""One ``rk4-pi`` step: returns ``(G_new, dt_used, dt_next)``."""
rhs, project = _materialized(rhs_fn, project_state)
def attempt(dt: jnp.ndarray, k1: jnp.ndarray):
G_new, err, _ = erk_step(
lambda _t, y: _rhs_value(rhs, y), 0.0, G, dt, RK4_2, k1=k1, project=project
)
G_new = jnp.asarray(project_state(G_new), dtype=state_dtype)
ratio = jnp.sqrt(jnp.sum(jnp.abs(err) ** 2) / jnp.sum(jnp.abs(G_new) ** 2))
en = ratio / _PI_RTOL
return G_new, en, step_factor(_PI_CONTROL, en, 1.0, RK4_2.embedded_order)
G_new, en, fac = attempt(dt_local, _materialize(dG))
def redo(_):
dt = dt_local * fac
G_redo, _en, fac_redo = attempt(dt, _materialize(dG))
return G_redo, dt, dt * fac_redo
return jax.lax.cond(
en > _PI_REDO,
redo,
lambda _: (G_new, dt_local, dt_local * fac),
None,
)
[docs]
def advance_explicit_nonlinear_state(
G: jnp.ndarray,
dG: jnp.ndarray,
dt_local: jnp.ndarray,
*,
method: str,
rhs_fn: RhsFn,
project_state: ProjectFn,
state_dtype: jnp.dtype,
) -> jnp.ndarray:
"""Advance one explicit nonlinear step with a static method string.
Barriers materialize every stage and stage derivative once. Unbarriered,
XLA:CPU fuses the whole linear RHS into the transpose feeding the
bracket's ``irfft2`` and recomputes it ~40 times per stage whenever the
projector is the identity (``ky >= 0`` layout): 30x on an RK3 step.
jaxlib < 0.11 drops the barrier before CPU fusion.
"""
materialized_rhs, materialized_projection = _materialized(rhs_fn, project_state)
G_new = _explicit_stage_update(
G,
_materialize(dG),
dt_local,
method=method,
rhs_fn=materialized_rhs,
project_state=materialized_projection,
)
G_new = project_state(G_new)
return jnp.asarray(G_new, dtype=state_dtype)
def _checkpoint_block_size(steps: int) -> int:
"""Balance retained block boundaries and states within one reverse block."""
return max(1, math.isqrt(max(int(steps) - 1, 0)) + 1)
def _concatenate_scan_outputs(blocks: Any, tail: Any) -> Any:
"""Flatten block-major scan outputs and append a possibly empty tail."""
return jax.tree.map(
lambda block, remainder: jnp.concatenate(
(block.reshape((-1, *block.shape[2:])), remainder), axis=0
),
blocks,
tail,
)
def _tree_bytes(tree: Any) -> int:
"""Bytes of every array or shape-struct leaf of ``tree``."""
return sum(math.prod(x.shape) * x.dtype.itemsize for x in jax.tree.leaves(tree))
[docs]
def block_checkpoint_plan(
steps: int, carry_bytes: int, residual_bytes: int, memory_budget_bytes: int
) -> int | None:
"""Block length of the block-only schedule, or ``None`` if over budget.
It stores ``steps/B`` carries plus ``B`` step residuals, minimized by
``B = sqrt(steps * carry / residual)``.
"""
count = max(int(steps), 1)
carry = max(int(carry_bytes), 1)
residual = max(int(residual_bytes), 1)
block = int(round(math.sqrt(count * carry / residual)))
block = min(max(block, 1), count)
storage = block * residual + -(-count // block) * carry
return block if storage <= int(memory_budget_bytes) else None
_COMPILED: dict[Any, Callable[..., Any]] = {}
def cached_jit(fn: Callable[..., Any]) -> Callable[..., Any]:
"""``jax.jit(fn)`` shared by every closure with the same graph and constants.
Keys one executable on the jaxpr text, argument avals and a digest of the
closed-over constants, which stay baked: the bits are ``jax.jit(fn)``'s.
As operands they also reuse across geometry but drift 6.5e-8 in float32.
The key is computed once per argument signature of this wrapper, not per
call: tracing, printing and hashing cost ~1 s for an imex RHS, and the
imex probes call one wrapper hundreds of times. ``call.key(*args)`` is that
digest, or ``None`` when ``fn`` closes over an outer trace.
"""
memo: dict[Any, tuple[str, Callable[..., Any], Any]] = {}
def resolve(args: tuple[Any, ...]) -> tuple[str, Callable[..., Any], Any] | None:
leaves, tree = jax.tree.flatten(args)
signature = (tree, tuple(jax.typeof(x) for x in leaves))
hit = memo.get(signature)
if hit is not None:
return hit
closed, out_shape = jax.make_jaxpr(fn, return_shape=True)(*args)
if any(isinstance(c, jax.core.Tracer) for c in closed.consts):
return None # closes over an outer trace: nothing to key
digest = hashlib.sha256(str((closed.jaxpr, closed.in_avals)).encode())
for c in closed.consts:
digest.update(str(jax.typeof(c)).encode() + np.asarray(c).tobytes())
key = digest.hexdigest()
run = _COMPILED.get(key)
if run is None:
if len(_COMPILED) >= 64: # bound the baked constants kept alive
_COMPILED.pop(next(iter(_COMPILED)))
run = _COMPILED[key] = jax.jit(
partial(jax.core.eval_jaxpr, closed.jaxpr, closed.consts)
)
hit = memo[signature] = (key, run, jax.tree.structure(out_shape))
return hit
def call(*args: Any) -> Any:
hit = resolve(args)
if hit is None:
return jax.jit(fn)(*args)
_key, run, out_tree = hit
return jax.tree.unflatten(out_tree, run(*jax.tree.leaves(args)))
def key(*args: Any) -> str | None:
hit = resolve(args)
return None if hit is None else hit[0]
call.key = key # type: ignore[attr-defined]
return call
[docs]
def checkpointed_explicit_scan(
step: Callable[..., Any],
initial_carry: Any,
indices: jnp.ndarray,
*,
checkpoint: bool,
memory_budget_bytes: int | None = None,
) -> tuple[Any, Any]:
"""Scan with a two-level discrete-adjoint checkpoint schedule.
A rematerialized step alone still leaves every scan carry on the reverse
tape. Grouping steps into rematerialized blocks retains only block
boundaries, while the inner rematerialized scan retains the states needed
to reverse one block. A block length near ``sqrt(steps)`` therefore lowers
state storage from ``O(steps)`` to ``O(sqrt(steps))`` without changing the
discrete time integrator. Requested scan outputs retain their inherent
``O(steps)`` result storage.
That nested schedule runs each forward step three times. Under
``memory_budget_bytes`` the per-step remat is dropped when the block-only
schedule fits (:func:`block_checkpoint_plan`, residuals sized from the
abstract pullback, which errs high), so it runs twice. Same discrete map.
"""
steps = int(indices.shape[0])
scan_step = jax.checkpoint(step) if checkpoint else step
if not checkpoint or steps <= 1:
return jax.lax.scan(scan_step, initial_carry, indices, length=steps)
block_size = _checkpoint_block_size(steps)
if memory_budget_bytes is not None:
pullback = jax.eval_shape(
lambda c: jax.vjp(lambda v: step(v, indices[0]), c)[1], initial_carry
)
block_only = block_checkpoint_plan(
steps,
_tree_bytes(initial_carry),
_tree_bytes(pullback),
memory_budget_bytes,
)
if block_only is not None:
block_size = block_only
scan_step = step
block_count, tail_size = divmod(steps, block_size)
blocked_size = block_count * block_size
block_indices = indices[:blocked_size].reshape((block_count, block_size))
def scan_block(carry: Any, local_indices: jnp.ndarray) -> tuple[Any, Any]:
return jax.lax.scan(
scan_step,
carry,
local_indices,
length=block_size,
)
carry, block_outputs = jax.lax.scan(
jax.checkpoint(scan_block),
initial_carry,
block_indices,
length=block_count,
)
carry, tail_outputs = jax.lax.scan(
scan_step,
carry,
indices[blocked_size:],
length=tail_size,
)
return carry, _concatenate_scan_outputs(block_outputs, tail_outputs)
def _maybe_emit_progress(
G: jnp.ndarray,
idx: jnp.ndarray,
*,
steps: int,
dt_val: jnp.ndarray,
real_dtype: jnp.dtype,
show_progress: bool,
) -> jnp.ndarray:
"""Emit executable progress without changing the traced state."""
if not show_progress:
return G
from gkx.callbacks import ( # type: ignore[import-untyped]
print_callback,
should_emit_progress,
)
sim_time = (idx + 1) * dt_val
sim_total = jnp.asarray(steps, dtype=real_dtype) * dt_val
return jax.lax.cond(
should_emit_progress(idx, steps),
lambda state: print_callback(
state, idx, steps, 0.0, 0.0, 0.0, 0.0, sim_time, sim_total
),
lambda state: state,
G,
)
[docs]
@partial(
jax.jit,
static_argnames=(
"rhs_fn",
"rhs_static_args",
"steps",
"method",
"checkpoint",
"project_state",
"show_progress",
"return_fields",
),
donate_argnums=(1,),
)
def integrate_nonlinear_scan(
rhs_fn: RhsFn,
G0: jnp.ndarray,
dt: float,
steps: int,
*,
method: str = "rk4",
checkpoint: bool = False,
project_state: ProjectFn | None = None,
show_progress: bool = False,
return_fields: bool = True,
rhs_args: tuple[Any, ...] = (),
rhs_static_args: tuple[Any, ...] = (),
) -> tuple[jnp.ndarray, Any] | jnp.ndarray:
"""Integrate a cached nonlinear RHS using the explicit solver scan policy."""
state_dtype = jnp.result_type(G0, jnp.complex64)
G0 = jnp.asarray(G0, dtype=state_dtype)
real_dtype = jnp.real(jnp.empty((), dtype=state_dtype)).dtype
dt_val = jnp.asarray(dt, dtype=real_dtype)
projector = project_state if project_state is not None else (lambda G: G)
def bound_rhs(state: jnp.ndarray) -> tuple[jnp.ndarray, object]:
return rhs_fn(state, *rhs_args, *rhs_static_args)
def step(G: jnp.ndarray, idx: jnp.ndarray) -> tuple[jnp.ndarray, Any]:
G = _maybe_emit_progress(
G,
idx,
steps=steps,
dt_val=dt_val,
real_dtype=real_dtype,
show_progress=show_progress,
)
G = jnp.asarray(projector(G), dtype=state_dtype)
dG, _fields = bound_rhs(G)
dG = jnp.asarray(dG, dtype=state_dtype)
G_new = advance_explicit_nonlinear_state(
G,
dG,
dt_val,
method=method,
rhs_fn=bound_rhs,
project_state=projector,
state_dtype=state_dtype,
)
if not return_fields:
return G_new, None
_dG_new, fields_new = bound_rhs(G_new)
return G_new, fields_new
G_final, fields_t = checkpointed_explicit_scan(
step,
G0,
jnp.arange(steps),
checkpoint=checkpoint,
)
if return_fields:
return G_final, fields_t
return G_final
def _apply_explicit_collision_split(
G_new: jnp.ndarray,
dt_local: jnp.ndarray,
*,
use_collision_split: bool,
damping: Any | None,
collision_scheme: str,
apply_collision_split_fn: CollisionSplitFn | None,
) -> jnp.ndarray:
"""Apply the optional post-step collision split used by diagnostics."""
if not use_collision_split or damping is None:
return G_new
if apply_collision_split_fn is None:
raise ValueError(
"apply_collision_split_fn is required when collision split is active"
)
return apply_collision_split_fn(G_new, damping, dt_local, collision_scheme)
def _advance_explicit_diagnostic_state(
G: jnp.ndarray,
dt_prev: jnp.ndarray,
*,
rhs_fn: RhsFn,
method: str,
project_state: ProjectFn,
state_dtype: jnp.dtype,
real_dtype: jnp.dtype,
time_step_policy: Any,
compute_fields_fn: FieldSolveFn,
cache: Any,
params: Any,
term_cfg: Any,
external_phi: jnp.ndarray | float | None,
use_collision_split: bool,
damping: Any | None,
collision_scheme: str,
apply_collision_split_fn: CollisionSplitFn | None,
max_dt: jnp.ndarray | None = None,
) -> tuple[jnp.ndarray, Any, jnp.ndarray, jnp.ndarray]:
"""Advance one explicit diagnostic state and compute its new fields."""
dG, fields = rhs_fn(G)
dt_local = jnp.asarray(
time_step_policy.update_dt(fields, dt_prev), dtype=real_dtype
)
chain = getattr(rhs_fn, "chain_implicit_linear", None)
if chain is not None: # the factor's dt, which the chunk controller sets
dt_local = jnp.asarray(chain.dt, dtype=real_dtype)
scales = time_step_policy.cfl_scales if method == "rk4-pi" else None
if scales is not None:
dt_local = jnp.minimum(
jnp.asarray(dt_prev, real_dtype), jnp.minimum(_PI_CAP * dt_local, scales[5])
)
if max_dt is not None:
dt_local = jnp.minimum(dt_local, jnp.maximum(max_dt, 0.0))
dt_next = dt_local
if scales is not None:
G_new, dt_local, dt_next = advance_rk4_pi_state(
G,
dG,
dt_local,
rhs_fn=rhs_fn,
project_state=project_state,
state_dtype=state_dtype,
)
else:
G_new = advance_explicit_nonlinear_state(
G,
dG,
dt_local,
method=method,
rhs_fn=rhs_fn,
project_state=project_state,
state_dtype=state_dtype,
)
G_new = _apply_explicit_collision_split(
G_new,
dt_local,
use_collision_split=use_collision_split,
damping=damping,
collision_scheme=collision_scheme,
apply_collision_split_fn=apply_collision_split_fn,
)
G_new = jnp.asarray(project_state(G_new), dtype=state_dtype)
fields_new = compute_fields_fn(
G_new, cache, params, terms=term_cfg, external_phi=external_phi
)
return G_new, fields_new, dt_local, dt_next
def _select_explicit_diagnostic(
idx: Any,
*,
diagnostics_stride: int,
diag_prev: Any,
G_new: jnp.ndarray,
fields_new: Any,
G_prev_step: jnp.ndarray,
fields_prev_step: Any,
dt_local: jnp.ndarray,
compute_diag_from_state: DiagnosticFn,
select_diagnostics_fn: Callable[..., Any],
steps: int,
) -> Any:
"""Select or reuse the diagnostic payload for one explicit scan sample."""
def compute_diag():
return compute_diag_from_state(
G_new, fields_new, G_prev_step, fields_prev_step, dt_local
)
return select_diagnostics_fn(
idx,
diagnostics_stride=diagnostics_stride,
diag_prev=diag_prev,
compute_diag_fn=compute_diag,
steps=steps,
)
[docs]
def make_explicit_diagnostic_step(
*,
rhs_fn: RhsFn,
method: str,
project_state: ProjectFn,
state_dtype: jnp.dtype,
real_dtype: jnp.dtype,
time_step_policy: Any,
compute_fields_fn: FieldSolveFn,
cache: Any,
params: Any,
term_cfg: Any,
external_phi: jnp.ndarray | float | None,
compute_diag_from_state: DiagnosticFn,
diagnostics_stride: int,
select_diagnostics_fn: Callable[..., Any],
show_progress: bool,
steps: int,
emit_progress_fn: Callable[..., jnp.ndarray],
use_collision_split: bool = False,
damping: Any | None = None,
collision_scheme: str = "implicit",
apply_collision_split_fn: CollisionSplitFn | None = None,
time_horizon: float | None = None,
) -> DiagnosticStepFn:
"""Build one explicit diagnostic scan step with injected runtime seams."""
def step(
carry: tuple[Any, Any, Any, Any, Any, Any],
idx: Any,
) -> tuple[tuple[Any, Any, Any, Any, Any, Any], tuple[Any, Any, Any]]:
G, G_prev_step, fields_prev_step, diag_prev, t_prev, dt_prev = carry
def advance(max_dt: jnp.ndarray | None):
G_new, fields_new, dt_local, dt_next = _advance_explicit_diagnostic_state(
G,
dt_prev,
rhs_fn=rhs_fn,
method=method,
project_state=project_state,
state_dtype=state_dtype,
real_dtype=real_dtype,
time_step_policy=time_step_policy,
compute_fields_fn=compute_fields_fn,
cache=cache,
params=params,
term_cfg=term_cfg,
external_phi=external_phi,
use_collision_split=use_collision_split,
damping=damping,
collision_scheme=collision_scheme,
apply_collision_split_fn=apply_collision_split_fn,
max_dt=max_dt,
)
diag = _select_explicit_diagnostic(
idx,
diagnostics_stride=diagnostics_stride,
diag_prev=diag_prev,
G_new=G_new,
fields_new=fields_new,
G_prev_step=G_prev_step,
fields_prev_step=fields_prev_step,
dt_local=dt_local,
compute_diag_from_state=compute_diag_from_state,
select_diagnostics_fn=select_diagnostics_fn,
steps=steps,
)
t_new = jnp.asarray(t_prev + dt_local, dtype=real_dtype)
G_new = emit_progress_fn(
G_new,
show_progress=show_progress,
diag=diag,
idx=idx,
steps=steps,
t_new=t_new,
progress_total=time_step_policy.progress_total,
)
return (G_new, G_new, fields_new, diag, t_new, dt_next), (
diag,
t_new,
dt_local,
)
if time_horizon is None:
return advance(None)
remaining = jnp.maximum(
jnp.asarray(time_horizon, dtype=t_prev.dtype) - t_prev, 0.0
)
return jax.lax.cond(
remaining > 0.0,
advance,
lambda _remaining: (carry, (diag_prev, t_prev, jnp.zeros_like(dt_prev))),
remaining,
)
return step
[docs]
def run_explicit_diagnostic_scan(
step_fn: DiagnosticStepFn,
initial_carry: tuple[Any, Any, Any, Any, Any, Any],
*,
steps: int,
stride: int,
sampled_scan: bool,
checkpoint: bool,
sampled_scan_fn: SampledDiagnosticScanFn,
) -> tuple[jnp.ndarray, tuple[Any, Any, Any]]:
"""Run the explicit diagnostic scan using sampled or dense retention."""
scan_step = jax.checkpoint(step_fn) if checkpoint else step_fn
if sampled_scan:
(
(
G_final,
_G_prev_last,
_fields_prev_last,
_diag_last,
_t_last,
_dt_last,
),
scan_diag_out,
) = sampled_scan_fn(
scan_step,
initial_carry,
steps=steps,
stride=stride,
)
return G_final, scan_diag_out
idx = jnp.arange(steps, dtype=jnp.int32)
(
(G_final, _G_prev_last, _fields_prev_last, _diag_last, _t_last, _dt_last),
scan_diag_out,
) = checkpointed_explicit_scan(
step_fn,
initial_carry,
idx,
checkpoint=checkpoint,
)
return G_final, scan_diag_out
__all__ = [
"advance_explicit_nonlinear_state",
"advance_rk4_pi_state",
"block_checkpoint_plan",
"checkpointed_explicit_scan",
"integrate_nonlinear_scan",
"make_explicit_diagnostic_step",
"run_explicit_diagnostic_scan",
]