Scan rollouts and checkpointing

veris.step repeats a fixed number of model timesteps using jax.lax.scan. A timestep advances the model by one time interval. veris.advance performs one complete physics timestep. The shared rollout driver repeats it; a rollout is this sequence of steps. Use functools.partial to fix Configuration, PhysicalConstants and the physics mode for the run. The default result is the final State, the collection of model arrays, without a saved State history:

from functools import partial
import jax
jax.config.update("jax_enable_x64", True)

from veris import Experiment, advance, initialize_experiment, step

experiment = Experiment.builtin("island")
initial, settings, constants = initialize_experiment(experiment)
transition = partial(
    advance, configuration=settings, constants=constants,
    forcing={"Qnet": 100.0, "Qsw": 0.0}, physics=experiment.physics,
)
final = step(initial, transition, 3, has_aux=True)

steps is a nonnegative integer fixed during compilation. The transition must be pure (no changes to outside resources) and traceable by JAX (JAX can inspect its operations). The carry, or data passed from step to step, must keep the same structure, array shapes and dtypes (number types). Reuse the same functions with fixed arguments across rollouts to reuse compiled scans. Settings and physical constants keep their existing fixed treatment. State and forcing arrays are differentiable: JAX can calculate how results change when these inputs change. Keep changes to Python objects, NumPy conversions and file output outside the transition and observer. advance returns (State, Diagnostics). has_aux=True passes the updated State to the next step as the rollout carry.

Forcing fields are prescribed inputs, such as wind or heat, and are reapplied before each timestep. You can supply scalars or physical (nx, ny) JAX arrays. Each field gets its own halo preparation; halos are extra cells that hold boundary or neighboring data. Supplying uWind or vWind calculates wSpeed unless it is explicitly prescribed. Missing fields keep their current State values, including growth Qnet/Qsw carried from the previous step. Host Experiment definitions check finite values and units outside JAX operations. Checks for initial ice and prescribed heat/wind sensitivities cover JIT (compilation), gradients, JVP (forward directional derivatives), VJP (reverse derivatives) and finite differences (small input changes). Fixed settings, constants and host definitions stay outside differentiation. Physical threshold choices remain those of Automatic differentiation at model boundaries.

Physics modes

advance selects the complete physics sequence through physics:

  • "coupled" advances dynamics (ice motion), transport (movement between cells) and growth (freezing and melting). Prescribed heat supplied through forcing is refreshed before each timestep.

  • "dynamics" advances momentum and transport with prescribed ocean and wind fields.

  • "thermodynamics" advances growth. Missing Qnet and Qsw keep the values returned by the preceding step, as in the reference column.

The core veris.physics.dynamics.driver.dynamics_transport stage advances momentum, calculates ocean stress before transport, then applies advection (movement with ice velocity), cleanup and ridging (ice piling up). It returns (state, stress_u, stress_v) with arrays that include halos. Specialized combinations add forcing, optional growth, the context for splitting arrays across devices, and final State/Diagnostics halos. When arrays are split this way, this local stage runs inside the caller’s jax.shard_map.

For example, a growth rollout uses the same public transition:

experiment = Experiment.builtin("growth")
initial, settings, constants = initialize_experiment(experiment)
transition = partial(
    advance, configuration=settings, constants=constants,
    physics=experiment.physics,
)
final = step(initial, transition, 3, has_aux=True)

Selected observations and diagnostics

A pure observer selects an array PyTree, a nested collection of arrays, after each State update. For the advance transition, it accepts State and Diagnostics. The result is (final_state, history). Each array in the history has a first axis for time, with length steps. The history excludes the initial State:

final, history = step(
    initial, transition, 3, has_aux=True,
    observe=lambda state, diagnostic: {"ice": state.hIceMean[2:-2, 2:-2]},
)

For output-only coupling fields, use has_aux=True with a transition returning (state, diagnostics) and an observer accepting both arguments:

final, history = step(
    initial, transition, 3, has_aux=True,
    observe=lambda state, diagnostic: {
        "ice": state.hIceMean,
        "freshwater": diagnostic.EmPmR,
    },
)

Only State passes from step to step. Diagnostics are discarded when no observer is supplied. An observer can keep diagnostics calculated partway through the physics sequence, such as ocean stress before transport and growth.

Zero steps keep the initial State and return empty histories with the expected shapes. JAX still inspects the transition and observer to determine those shapes, including in auxiliary-output mode. A callback that JAX cannot trace is still invalid when the step count is zero. Observe the initial State separately if needed.

Time-varying forcing and AD

Automatic differentiation (AD) calculates how a result changes with its inputs. inputs accepts a nonempty array PyTree where every array has a first axis of length steps. The transition then receives transition(state, input_slice). For example, calculate how a thickness objective changes with each prescribed cooling value:

import jax
import jax.numpy as jnp
jax.config.update("jax_enable_x64", True)
from veris import Experiment, advance as physics_timestep, initialize_experiment, step

experiment = Experiment.builtin("island")
initial, settings, constants = initialize_experiment(experiment)
def transition(state, prescribed):
    return physics_timestep(
        state, settings, constants, prescribed, physics=experiment.physics,
    )
cooling = jnp.asarray([80.0, 100.0, 120.0], dtype=initial.hIceMean.dtype)

def objective(forcing):
    final = step(
        initial, transition, 3,
        inputs={"Qnet": forcing, "Qsw": jnp.zeros_like(forcing)},
        has_aux=True, checkpoint=True,
    )
    return final.hIceMean[2:-2, 2:-2].mean()

sensitivity = jax.grad(objective)(cooling)

checkpoint=True is the default. It wraps the repeated step calculation, including the observer, with jax.checkpoint. Reverse AD recalculates intermediate step values so it needs to save fewer values, at the cost of more computation. This does not guarantee constant memory as a rollout grows longer. Saved carry values and requested observation histories can still grow with the number of timesteps. Use checkpoint=False to compare the ordinary scan. Multi-process rollouts retain halo exchange results for reverse AD rather than exchanging those values again. Arithmetic is still recalculated; saved halo values add storage at each timestep and inner solver iteration. See JAX’s gradient checkpointing guide and Automatic differentiation at model boundaries for physical derivative conventions.

Sharded rollouts

A mesh arranges devices along the grid axes. Initialize it and keep its context active around the rollout and its JAX operations. Define global physical dimensions and matching partitions. Regional initialization creates the packed State, which includes each partition’s halos:

import jax
from veris import Execution, Experiment, Grid, advance, initialize_experiment, step

layout = (1, jax.device_count())
mesh = jax.make_mesh(layout, ("x", "y"))
experiment = Experiment.builtin(
    "dynamics", grid=Grid(8, 8 * layout[1]),
    execution=Execution(partitions=layout),
)
with jax.set_mesh(mesh):
    initial, settings, constants = initialize_experiment(experiment, mesh=mesh)
    transition = partial(
        advance, configuration=settings, constants=constants,
        physics=experiment.physics,
    )
    final = step(initial, transition, 3, has_aux=True)

Each partition keeps its own halos. Selected history arrays gain a first axis for time that is not split across devices. Select one time slice before using the spatial output collector, which gathers grid fields. Coupled presets and custom OceanGeometry use the same rules for global physical inputs. Keep jax.set_mesh(mesh) active around pure advance/step and their JAX operations. Public advance accepts scalars or forcing arrays on the global physical grid. Forcing series keep a first axis for time that is not split across devices. See Experiment definitions.

Host output and timing

The dynamics, growth and parallel command-line runners use Simulation, which calls veris.execution.rollout_output.run_timed. You can also use the low-level timing API directly. It prepares and compiles each required scan shape with jax.jit(...).lower(...).compile() before running the model. Compilation can inspect pure functions, but it executes no model transitions, device callbacks, sample collection or file writes. The returned tuple contains final State, compilation seconds and integration seconds. Integration timing includes execution, waiting for device results and host output.

Use output_callbacks(manager, step_seconds) to connect an OutputManager:

from veris.execution.rollout_output import output_callbacks, run_timed

def state_transition(state):
    return transition(state)[0]

observe, select = output_callbacks(manager, experiment.step_duration.total_seconds())
final, compilation_seconds, integration_seconds = run_timed(
    initial, state_transition, 100, observe=observe, select=select,
)

The Scheduled integration reference defines where chunks of steps end, how means are accumulated, what collectors must do, and how to use a manager dedicated to writing reduced fields. It also describes calendar and sampling rules alongside direct OutputManager.sample usage.

Custom host observers keep the existing observe(sample, iteration) interface. They receive selected initial and later samples from histories of limited length. The default chunk cap is eight; chunk_size changes that cap. With no observer, including when OutputManager output is disabled, the runner compiles and executes one full-length scan, saves no observation histories and ignores the chunk cap. Zero-step host runs neither compile nor execute physics, but can write a requested initial instantaneous record. Final snapshots remain separate host operations. Use veris.step directly for AD. The timing/output wrapper runs only on the host.

See Integration for Python signatures and docstrings.