Integration

Differentiable rollouts

The package-level driver repeats a function that updates the model by one step. The function is pure: it returns results without changing Python objects or writing files. JAX can differentiate the repeated calculation to measure how results change with inputs. See Scan rollouts and checkpointing for external inputs that change over time, observations stored as nested collections called PyTrees, step diagnostics, checkpointing to limit memory use, and examples that split arrays across devices.

veris.step(initial: _Carry, advance: Callable[[...], Any], steps: int, *, checkpoint: bool = True, inputs: Any = None, observe: Callable[[...], Any] | None = None, has_aux: bool = False) → Any

Advance a fixed number of timesteps with jax.lax.scan.

advance(state) returns the next State. When inputs is supplied, its array PyTree must have leading length steps on every leaf, and the transition receives advance(state, input_slice) instead. Bind settings using a closure or functools.partial. Physical parameter PyTrees and forcing may be captured as traced closure values or supplied through inputs.

By default return only the final State. With a pure observe(state) callable, return (final_state, observations); observation leaves stack along a new leading time axis and exclude the initial State. For transitions returning (state, diagnostics), pass has_aux=True and use observe(state, diagnostics). Unobserved diagnostics are discarded.

steps, callables and switches are static; carry and forcing arrays stay differentiable. Carry leaf shapes/dtypes and PyTree structure must remain constant. The default checkpoint wraps the scan body with jax.checkpoint. Multi-process execution saves halo exchange results across recomputation. Zero steps preserve the initial State; JAX still traces the pure transition and observer to infer empty observation shapes. Host I/O and mutable callbacks belong outside this API. For sharded physics, retain the setup’s active mesh.

Dynamics and transport stage

Custom step functions can reuse the shared dynamics stage, which computes ice motion and transport. It returns (state, stress_u, stress_v). The stresses, or forces per unit area, are computed before transport moves ice across cells. The caller supplies atmospheric inputs, optional freezing and melting, and the final refresh of State and Diagnostics halos. Halos are extra cells that hold boundary or neighboring values. When splitting arrays across devices, call this local stage inside the caller’s jax.shard_map.

Sea-ice dynamics and transport composed from the core numerical kernels.

Compose the existing Hibler/Kimmritz momentum kernels, ocean stress, directional advection, cleanup and ridging in reference order. Inputs and returned State and stress arrays retain their halo-inclusive (x, y) storage. Callers own atmospheric forcing, optional thermodynamics and the final State/Diagnostics halo refresh. Under sharding this stage must run inside the caller’s shard_map.

veris.physics.dynamics.driver.dynamics_transport(vs: State, conf: Configuration, phys: PhysicalConstants) → tuple[State, Array, Array]

Advance momentum and transport, retaining pre-transport ocean stresses.

Host timing and scheduled output

These helpers compile and time model runs and connect them to an OutputManager. Keep them outside JAX transformations such as compilation and automatic differentiation. Use veris.step() directly to compute derivatives. See Scan rollouts and checkpointing for timing and Scheduled integration for scheduling rules.

veris.execution.rollout_output.run_timed(state: _T, advance: Callable[[...], _T], steps: int, *, observe: Callable[[Any, int], None] | None = None, select: Callable[[_T], Any] | None = None, chunk_size: int | None = None, checkpoint: bool = True, inputs: Any = None, history_select: Callable[[_T], Any] | None = None, history_sink: Callable[[Any], None] | None = None, validate: Callable[[_T], None] | None = None) → tuple[_T, float, float]

Compile scan shapes, then time evolution with optional host observations.

advance and select must be pure and JAX traceable. Callbacks returned by output_callbacks use scheduled device reductions, with no default chunk cap. An explicit positive chunk_size caps scheduled segments. Arbitrary observers receive each step via histories capped at eight steps by default; with no selector they receive the whole carry. Every used chunk length is lowered and compiled without executing physics. Return final State, compilation time and actual integration time (including host observation and output). Optional host validation runs before selected histories or scheduled samples are committed.

veris.execution.rollout_output.output_callbacks(manager: OutputManager, step_seconds: float, *, continue_runs: bool = False, start_step: int = 0, mean_coverage_start_us: int = 0, emit_initial_instant: bool = True) → tuple[Callable[[Any, int], None] | None, Callable[[Any], Any] | None]

Select the union of configured fields and retain halos for the collector.

Disabled output returns no callbacks, so the host runner allocates no history. The callable carries scheduling metadata for device reductions in run_timed. It also supports direct host calls with the manager’s ordinary sample API.