ICON-sc: A sympl-Conformant Python Architecture for the ICON Model

Status: design proposal v1.3 (rev.: DaCe-free orchestration; bind-time interface specialization; Tasmania coupling algebra; differentiability contract + functional lowering; project rebrand to ICON-sc — name only, no technical change (work 0053, TD-53.1)) · Scope: global ICON-NWP atmosphere (dycore + tracer transport + full NWP physics suite), ICON-artifact ingestion, Pythonic configuration and output · Compatibility target: scientific equivalence with ICON (same equations, same schemes, same coupling strategy), not bitwise parity · Out of scope for v1 (but not precluded): nesting, ICON-LAM, ocean, AES physics, ART, data assimilation beyond IAU.


0. Executive summary

ICON-sc is a host/orchestration layer that re-expresses ICON-NWP as a composition of sympl-style components over a zero-copy device-field protocol. The design keeps sympl’s three load-bearing ideas — (1) a model is a state dictionary evolved by typed components, (2) components carry self-documenting property contracts (names, dims, units), (3) the run script is the single, legible source of compositional truth — and extends the taxonomy in the four places where ICON’s structure cannot be expressed in vanilla sympl, plus one requirement orthogonal to ICON itself:

  1. A general dynamics–physics coupling algebra, adopted from Ubbiali’s Tasmania framework (ETH thesis no. 28022): federation classes realizing parallel splitting (PS), sequential-tendency splitting (STS), sequential-update splitting (SUS), and symmetrized/Strang splitting (SSUS); ConcurrentCoupling realizing (lazy) full coupling (FC/LFC); and a DynamicalCore base class with slow/fast/super-fast cadence tiers. ICON’s operational coupling — SUS fast physics plus lazily evaluated, piecewise-constant slow tendencies fed to the core — becomes one validated preset in this algebra rather than hard-wired structure. (Vanilla sympl’s TendencyStepper expresses only tendency-summing coupling; the sympl paper concedes sequential splitting needs custom machinery, and Tasmania built precisely that machinery.)

  2. Sub-cycled steppers for the dynamics substepping (ndyn_substeps) and, dually, calling-frequency policies for slow physics (a generalization of sympl’s UpdateFrequencyWrapper, which already has exactly the right piecewise-constant-tendency semantics ICON assumes).

  3. Explicit communication components and halo metadata — sympl has no distributed-memory story at all; ICON-sc makes halo validity part of the field contract and halo exchange a first-class component kind, with GHEX as the substrate.

  4. Location typing (cell/edge/vertex) in the property contracts, because ICON’s C-grid staggering on the icosahedral mesh is not expressible in sympl’s dims-plus-units vocabulary.

  5. A differentiability contract and a functional lowering. Components declare native / custom / none differentiability and expose tunable scheme constants as a ParamTree; the same negotiated composition that feeds the imperative execution tiers also compiles to a pure JAX step function (the F-tier) with explicit-carry PyTrees, lax.scan windows, and a transpose-rule differentiable halo exchange — making jax.jvp/jax.vjp of the composed model a supported operation rather than a rewrite (§8.5–8.7).

The compute substrate is the hybrid agreed on: the state boundary is arrays-with-metadata (xarray DataArray over NumPy/CuPy buffers, exchanged zero-copy via DLPack); GT4Py fields, JAX arrays, or Numba device arrays exist only inside components as ingress views of those buffers. ICON4Py granules are the intended implementation of the dycore, transport, and most physics components — ICON-sc is a host layer for icon4py, not a fork of it. This is the same principle as the ICON4Py↔anemoi decision: shared abstraction lives at the state/data boundary; no attempt is made to unify compute-layer field types across frameworks.

The deliberate trade at the center of the design: composition-time legibility, execution-time fusion. The run script remains a readable, sympl-style description of the model. But at production step rates (milliseconds per step on GPU), a Python dict-of-arrays dispatch loop is a real ceiling, so the architecture splits a one-time negotiation phase (the full sympl contract machinery) from a specialized execution phase: at startup the validated component sequence is compiled into a static execution plan — pre-bound buffers, no dict or xarray traffic — and optionally lowered further to CUDA-graph replay per exchange-free segment, or to a generated native host driver, bypassing the Python boundary under a validated-equivalence contract. The run script then constructs the timestep; it no longer dispatches it. A fourth lowering of the same plan — a pure-functional JAX trace — makes the composed model differentiable without disturbing the production tiers.


1. Where sympl and ICON collide (design tensions)

Being explicit about the collisions up front, because each one drives a taxonomy or protocol extension below.

T1 — Splitting semantics. sympl’s TendencyStepper gathers tendencies from a set of TendencyComponents, sums them, and steps: parallel (process) splitting. ICON-NWP is the opposite on both halves: fast physics is sequential-update splitting (each process consumes the state left by the previous one and updates it in place; tutorial §3.7.2), and slow physics is parallel-split but with piecewise-constant tendencies at per-process cadences injected into the dynamical core mid-substep (vn, exner) and into transport (ρqx) — not summed and stepped by a generic integrator. Consequence: the composition layer adopts Tasmania’s coupling algebra — SequentialUpdateSplitting for the fast suite, a lazy ConcurrentCoupling feeding the dycore’s slow-tendency port for the slow suite — making ICON’s operational arrangement one validated preset in a family that also contains PS, STS, and SSUS (§4.2).

T2 — The dycore is not a component of the time loop; it is a time loop. solve_nonhydro is a two-time-level predictor–corrector, sub-cycled ndyn_substeps times per Δt, with internal halo exchanges at substep boundaries, internal velocity-advection reuse between substeps, and vertically implicit sound-wave treatment. It cannot be decomposed into sympl TendencyComponents without destroying it. Consequence: the dycore is one component — a DynamicalCore in Tasmania’s sense, with the predictor–corrector as its stage structure and ndyn_substeps as its super-fast substep tier — whose interior is GT4Py programs (icon4py’s SolveNonhydro); its double-buffering (nnow/nnew) and substep carry-over state are private, with a restart-serialization protocol (§4.5).

T3 — Staggering and the dyn↔phy interface. Prognostic winds live on edges (vn); physics operates on cell-center (u,v) reconstructions, and turbulence returns increments that are interpolated back to edges and added to vn (tutorial §3.7.3). sympl has no notion of mesh location. Consequence: property contracts gain a location axis; the RBF/edge2cell reconstructions and the increment-return path are explicit DiagnosticComponents / component-internal steps, never hidden unit-conversion-style magic (§3.3).

T4 — Distributed memory. sympl’s paper defers MPI entirely. ICON’s step structure interleaves halo exchanges inside what ICON-sc treats as single components (dycore substeps) and between components (after diffusion, before transport, etc.). Consequence: halo validity becomes field metadata; exchanges between components are auto-inserted or validated by a composition pass; exchanges inside components are declared (communicates_internally=True) and use the same GHEX patterns (§6).

T5 — Units and dims coercion vs. zero-copy. sympl’s ergonomic superpower — silent unit conversion and dim reordering at component ingress — is a performance landmine on GPU. Consequence: canonical-units registry + strict mode: in production, any ingress that would allocate (unit conversion, transpose, host↔device) is an error, not a convenience (§2.4).

T6 — Naming. CF standard names are adopted where CF defines them, but ICON’s internal quantities (exner, vn, θv, ddt_* tendency slots, interface-level variants) are largely outside CF scope — deliberately so, since CF avoids encoding solver-internal temporal/staging distinctions. Consequence: a variable registry with an icon: namespace, CF aliases where they exist, and sympl’s _on_interface_levels suffix convention retained (§2.5).

T7 — Mutation and privacy vs. functional purity. ICON’s sequential-update physics, the vault’s in-place/slot-swap execution model, and §4.5’s component-private state are all mutation-shaped; jax.jvp/jax.vjp require a pure function over explicit PyTrees. Consequence: differentiability is a lowering, not a rewrite — the same negotiated composition compiles to a pure step function (F-tier, §8.5) in which in-place updates become functional updates recovered by XLA buffer donation, and component-private state is surfaced as explicit carry; privacy is demoted to an imperative-tier convenience.


2. Layer 0 — State, fields, and the zero-copy protocol

2.1 The state dictionary

Exactly sympl’s model: state: dict[str, DataArray] plus the required time key (datetime/cftime). One logical time level is exposed; multi-time-level schemes own their history privately (T2). The state is the only channel between components; components never share Python object references behind the state’s back.

2.2 The field boundary type

State values are xarray.DataArrays whose .data is any object satisfying a minimal buffer protocol:

@runtime_checkable
class FieldBuffer(Protocol):
    def __dlpack__(self, *, stream: int | None = None): ...
    def __dlpack_device__(self) -> tuple[int, int]: ...
    shape: tuple[int, ...]
    dtype: Any

In practice: numpy.ndarray (CPU) or cupy.ndarray (CUDA/HIP). JAX arrays are permitted as boundary buffers only in single-framework configurations (their immutability breaks the in-place-update contract that ICON’s sequential physics relies on; see §8.4 — the functional lowering of §8.5 instead replaces that contract wholesale). The DataArray layer carries, in dims/attrs:

metadata

values

consumed by

dims

('cell'|'edge'|'vertex', 'height'|'height_interface', …)

ingress adapters, contract checks

attrs['units']

canonical unit string

contract checks (no-op in production)

attrs['location']

cell / edge / vertex (redundant with dims, kept explicit for scalars/sparse fields)

location typing

attrs['halo']

VALID / DIRTY

halo validator pass

attrs['grid_uuid']

uuidOfHGrid

provenance check: refuse to mix fields from different grids

Horizontal layout is flat unstructured (horizontal, vertical) in row-major order, matching icon4py — ICON-Fortran’s nproma/jb-jc blocking is deliberately not adopted; it is a CPU-cache artifact that GT4Py backends re-derive per target as a schedule decision, which is the whole point of separating algorithm from schedule.

2.3 Ingress adapters (zero-copy views)

Each compute framework gets a cached adapter that turns a boundary buffer into its native field type without copying:

  • GT4Py: gtx.as_field(domain, buffer) with the domain derived from dims + the grid object’s sizes; offset providers come from the grid (§3), never from the state.

  • JAX: jax.dlpack.from_dlpack(buffer) for inputs; outputs are written back via a preallocated output buffer (donation is not relied on).

  • Numba/CuPy: direct via __cuda_array_interface__.

Adapters are constructed once per (component, state-schema) pair at bind time — under the execution plan (§8.2) ingress is fully pre-resolved into bound argument packs, so per-step ingress cost is zero; in the interpreted tier it degrades gracefully to one dictionary lookup.

2.4 Contracts, canonical units, strict mode

Components keep sympl’s input_properties / tendency_properties / diagnostic_properties / output_properties, extended with location and halo:

class TurbulentDiffusion(Stepper):
    input_properties = {
        "air_temperature": {
            "dims": ["cell", "height"], "units": "K",
            "halo": "owned",           # computes on owned points only
        },
        "icon:normal_wind": {
            "dims": ["edge", "height"], "units": "m s-1", "halo": "required",
        },
        ...
    }
    output_properties = {
        "icon:normal_wind": {"dims": ["edge", "height"], "units": "m s-1",
                              "halo": "invalidated"},   # dirties the halo
        ...
    }

A canonical-units registry fixes one unit per variable name; components are required to declare canonical units, so sympl’s Pint conversion path exists but compiles to a no-op. ComputeContext(strict=True) (default for production) turns any ingress that would allocate — unit conversion, dim transpose, dtype cast, host↔device transfer — into an exception with a diagnostic naming the offending component and field. strict=False restores full sympl ergonomics for interactive/educational use, which is a first-class use case, not a degraded mode.

2.5 Naming

CF standard names where CF has them (air_temperature, air_pressure, specific_humidity, …); icon:-namespaced names for solver-internal quantities (icon:exner_function, icon:normal_wind, icon:virtual_potential_temperature, icon:ddt_vn_phy, …); sympl’s _on_interface_levels suffix for the height_interface dim variants. The registry is one Python module mapping canonical name ↔ ICON Fortran short name ↔ CF name (if any) ↔ GRIB2 discipline/category/parameter — the last column exists purely to make GRIB2 ingestion (§7.2) table-driven, not to promise GRIB2 output parity.


3. Layer 1 — Grid, geometry, metrics

3.1 IconGrid

One immutable object per horizontal domain, constructed by reading an ICON grid NetCDF file (the same files Zonda / the DWD grid generator produce), keyed by uuidOfHGrid:

grid = icon_grid.from_file("icon_grid_0026_R03B07_G.nc", ctx)
grid.n_cells, grid.n_edges, grid.n_vertices        # local sizes (owned + halo when distributed)
grid.connectivities["C2E"], ["E2V"], ["C2E2C"], ["E2C2V"], ["V2E"], ...
grid.refin_ctrl                                     # retained even without nesting (future LAM/nests)
grid.geometry                                       # edge lengths, dual lengths, areas, orientation, coriolis

Connectivities are exposed in two forms from one storage: raw index arrays (for JAX segment-sum patterns and NumPy reference components) and GT4Py offset providers (for icon4py granules). The grid, not the state, owns topology — components receive the grid at construction time; the state stays pure data. This mirrors the Atlas-vs-GT4Py decomposition finding: mesh location, memory layout, and communication are orthogonal concerns, not bundled into one FunctionSpace-like object.

3.2 Vertical grid and derived fields

VerticalGrid holds vct_a/vct_b, nlev, flat/rayleigh levels, and the SLEVE or Gal-Chen decay parameters. Two factories — MetricsFactory and InterpolationFactory (directly the icon4py metrics/interpolation field factories) — compute the 3D metric terms (ddqz_z_*, wgtfac_*, reference-state fields) and interpolation coefficients (RBF vectors, cell→edge/edge→cell weights, nudging coefficients) once at startup. Their outputs are ordinary DataArrays stored in a read-only static_state dict that components may request in input_properties like any other field; being static, they are exempt from halo tracking.

3.3 Staggering as API, not convention

Because location is contract-checked, the dyn↔phy reconstruction becomes visible architecture instead of buried convention: a WindReconstruction DiagnosticComponent (edges→cells, RBF) runs before the physics suite; the turbulence component’s edge-located output is the increment path (cell-center increments interpolated back to edges inside the component, per tutorial §3.7.3), declared in its output properties. Nothing in the framework silently regrids between locations.


4. Layer 2 — Component taxonomy

4.1 Preserved sympl kinds

DiagnosticComponent, TendencyComponent, ImplicitTendencyComponent, Stepper, Monitor, and the wrapper concept survive unchanged in semantics (property contracts, array_call-style core methods, first-class usability outside any model — the ability to call SaturationAdjustment()(state) interactively is retained and is a genuine research affordance, as the sympl paper argues).

4.2 The coupling algebra (after Tasmania)

Rather than inventing ad hoc containers for ICON’s two coupling styles, ICON-sc adopts — with attribution — the composition layer of Ubbiali’s Tasmania framework, which extended sympl with exactly the general machinery this layer needs: a family of federation classes, each realizing one of the six coupling strategies whose truncation-error and stability properties the thesis analyzes under weak assumptions (independent of the governing equations, the number of parameterizations, and the internals of the time-steppers). Adopting the algebra buys two things at once: ICON’s operational coupling stops being hard-wired structure and becomes one validated preset, and every other point of the algebra becomes available for the coupling-scheme experimentation that production codes structurally cannot do — which was Tasmania’s raison d’être.

The operator vocabulary. All operators consume the same property-contracted components of §4.1 and are themselves components (they compose):

  • ConcurrentCoupling(comps) — Tasmania’s heterogeneous composite over TendencyComponents and DiagnosticComponents: sums tendency contributions, unions diagnostics in declared order. It is the building block of full coupling (FC: evaluate at every integrator stage) and lazy full coupling (LFC: evaluate once per step and hold constant — COSMO’s operational choice, and, when the “hold” is stretched over N steps by CallingFrequency, ICON’s slow-physics choice).

  • TendencyStepper — kept in its Tasmania form: a registry-based integrator family ("forward_euler", "rk2", "rk3ws", …; name-keyed Factory/MetaFactory registration, so user-defined steppers are first-class citizens of any downstream federation) wrapping a ConcurrentCoupling or bare TendencyComponent. Its sibling SequentialTendencyStepper implements the two-input-state signature E(ψⁿ, Δt; P + (ψ_prov − ψⁿ)/Δt) that makes STS expressible — the accumulated provisional tendency of upstream processes enters as a constant forcing.

  • ParallelSplitting(sections) — each section integrated independently from ψⁿ; the recombination ψⁿ⁺¹ = Σₗ ψₗⁿ⁺¹ − L·ψⁿ (Tasmania’s DataArrayDictOperator step) compiles, in the execution plan, to a single fused multi-axpy vault op.

  • SequentialTendencySplitting(sections) — the STS chain of SequentialTendencySteppers, dynamics first, then physics ideally ordered slowest→fastest (Beljaars’ constant-forcing argument, with the thesis’s caveat that ranking atmospheric time scales is genuinely hard).

  • SequentialUpdateSplitting(sections) — the SUS chain: each process corrects the state left by its predecessor. Sections may be bare Steppers (adjustment-type components like saturation adjustment enter directly) or (TendencyComponent, stepper_name) pairs resolved through the registry. Ordering is explicit and validated against declared must_follow/must_precede constraints; ICON’s fixed fast-physics order ships as NWP_FAST_ORDER, including the surface-transfer-at-end time-level optimization documented in the tutorial.

  • SSUS(sections, core, lam=0.5, pre_steppers=None) — the symmetrized/Strang composition: sections traversed in reverse over λΔt, the core over Δt, sections forward over (1−λ)Δt, with per-side stepper choice (Eₗ* ≠ Eₗ is legal, per thesis §2.3.5). λ = ½ is the only choice with second-order coupling error regardless of problem structure; the operator exists because ICON-sc should be able to measure what that costs and buys on a real dycore.

  • DynamicalCore — Tasmania’s base class for two-time-level, multi-stage cores, with three cadence tiers whose terminology Tasmania itself borrowed from ICON: a slow-tendency input port (tendencies computed outside, consumed inside at the point the core’s numerics dictate); an optional per-stage fast-tendency ConcurrentCoupling (FC within the core, invoked at every stage); and a super-fast substep tier with its own hook. NonhydroSolver derives from it: the predictor–corrector is the stage structure; ndyn_substeps (with the CFL-adaptive ratio provider, escalation bounded at 12) is the super-fast tier; the per-stage fast slot is empty in the ICON preset but is the natural experiment port (e.g., per-substep saturation adjustment or microphysics — precisely the kind of question the coupling literature keeps open). The generic Subcycle combinator survives for uses outside the core (tracer-transport call-frequency reduction); the dycore’s own substepping now lives inside its DynamicalCore structure.

The bus, reframed. The v1.0 SlowTendencyBus survives unchanged in mechanics but is now recognized as the naming-and-checking convention for the DynamicalCore slow-tendency port: slow processes publish piecewise-constant tendencies into icon:ddt_* slots under CallingFrequency (LFC’s lazy evaluation stretched over N steps), the dycore and transport declare those slots as inputs, and the composition-time check that every published tendency has exactly one consumer is retained verbatim.

The ICON-NWP preset, in algebra terms: per Δt — dyn→phy diagnostics; ConcurrentCoupling of CallingFrequency-wrapped slow processes publishing to the bus; NonhydroSolver(DynamicalCore) consuming the bus in its substeps; horizontal diffusion; tracer transport consuming the ρqₓ bus slots; SequentialUpdateSplitting(NWP_FAST_ORDER). Presets carry a validated / experimental label; only validated presets inherit the scientific-equivalence claim of §9.

Adopted sympl-fork engineering. Ubbiali’s fork of sympl contributed several mechanisms that ICON-sc adopts outright — two of which v1.1 of this document had independently reinvented, which is worth flagging as convergent evidence:

  • The object-oriented component ABI: __call__(state, *, out=None) with caller-provided output DataArrays, array_call(inputs, outputs) receiving both input and output raw buffers, and allocate_* hooks for framework-side allocation. Caller-provided outputs are exactly what the execution plan’s preallocated buffers and slot-swap machinery require, so T0 and T1–T3 share one ABI instead of the plan bolting a second calling convention onto sympl components.

  • The Checker/Operator decomposition with static variants (component-definition information only) and dynamic variants (crossing definitions with actual data): this is the internal structure of ICON-sc’s negotiation phase — static checkers at composition time, dynamic checkers and operators at bind time, neither on the step path. Ubbiali’s “disable runtime checks” switch is subsumed by the tier system.

  • Factory/MetaFactory name-keyed registries (subclasses self-register on import): used for tendency steppers, physics-scheme selection, and backends, so configuration strings resolve declaratively through registries rather than if-ladders.

  • Avoidance of axis permutations that coerce duck arrays to numpy.ndarray — subsumed by, and one of the original motivations validated by, strict mode (§2.4).

Not adopted: Tasmania’s Cartesian spatial-domain machinery (HorizontalBoundary, Topography registries) — the ICON grid layer (§3) owns that ground on the icosahedron — and Tasmania’s multi-backend stencil-dispatch infrastructure, which ICON-sc replaces with the ComputeContext/ingress-adapter design of §2 and §5.2.

4.3 The ICON-NWP → ICON-sc mapping

ICON process

ICON-sc class

cadence

implementation source

notes

solve_nonhydro (predictor–corrector, vertically implicit)

NonhydroSolver(DynamicalCore); substep tier = ndyn_substeps (CFL-adaptive)

Δτ

icon4py solve_nonhydro

consumes icon:ddt_vn_phy, icon:ddt_exner_phy via the slow port; communicates_internally=True; private nnow/nnew + velocity-advection carry-over; per-stage fast slot empty in preset

horizontal diffusion

HorizontalDiffusion(Stepper)

Δt

icon4py diffusion

Smagorinsky + background 4th-order

tracer transport (FFSL horiz. + PPM vert., directional split)

TracerTransport(Stepper)

Δt

icon4py advection

applies icon:ddt_qx_phy post-horizontal-transport; per-tracer scheme config; reduced calling frequency (§3.6.4) via Subcycle inverse (substep_ratio)

saturation adjustment (×2)

SaturationAdjustment(Stepper)

Δt

icon4py / port

appears twice in NWP_FAST_ORDER

TERRA (land), FLake, sea ice

each a Stepper

Δt

port from ICON Fortran

tiled surface; old-time-level input semantics honored by suite ordering

turbulent diffusion (turbdiff)

TurbulentDiffusion(Stepper)

Δt

port

implicit vertical solve; increment path back to edges

microphysics (one-moment graupel)

Microphysics(Stepper)

Δt

muphys-cpp bindings or GT4Py port

scheme-selectable

surface transfer

SurfaceTransfer(Stepper)

Δt

port

called last (tutorial note), consistent time levels

convection (Tiedtke–Bechtold)

ImplicitTendencyComponent under CallingFrequency(dt_conv)

dt_conv

port (large)

publishes to tendency bus

cloud cover

DiagnosticComponent under CallingFrequency(dt_conv)

dt_conv

port

diagnostic cloud fields for radiation

radiation

TendencyComponent under CallingFrequency(dt_rad) + ReducedGridWrapper

dt_rad

pyRTE-RRTMGP (primary, GPU-capable); ecRad via CFFI as validation reference

fluxes → heating at constant volume (c_v)

non-orographic GWD

TendencyComponent under CallingFrequency(dt_gwd)

dt_gwd

port

tendency bus

SSO drag (Lott–Miller)

TendencyComponent under CallingFrequency(dt_sso)

dt_sso

port

tendency bus

p, T diagnosis (hydrostatic pressure integration)

HydrostaticPressureDiagnostics(DiagnosticComponent)

Δt

port of §3.7.3 recipe

the dyn→phy interface, explicit

wind reconstruction edges→cells

WindReconstruction(DiagnosticComponent)

Δt

icon4py interpolation

RBF

IAU

IncrementalAnalysisUpdate(TendencyComponent)

Δt within window

new

weighted increments; iterative IAU is a driver recipe (re-run loop), not v1

The isochoric (constant-volume) coupling convention — heating rates converted with c_v, hydrostatically integrated pressure for physics — is a stated suite-wide contract documented on the tendency bus slots, since it is exactly the kind of invisible convention that corrupts ported schemes inherited from constant-pressure hosts.

4.4 What is deliberately not decomposed

Inside NonhydroSolver, the ~50 stencil programs of the predictor/corrector are GT4Py programs (icon4py’s existing granule structure, including its hand-fused program variants where launch count matters), not ICON-sc components. The component boundary is drawn where scientific recomposition is plausible (swap a microphysics scheme, reorder physics, disable convection) — not at the stencil level, where recomposition is a performance/correctness hazard and GT4Py already provides the right abstraction. This two-granularity picture (component = scheme at the composition layer; stencil = unit of scheduling below it) matches the cross-model finding that modern rewrites moved the unit of optimization to the stencil while the unit of composition stays the scheme.

4.5 Component-private state and restart

Components may hold private persistent fields (dycore time levels and substep carry-overs, turbulence TKE history if configured, radiation’s stored fluxes between calls, CallingFrequency’s cached tendencies and phase). The contract: implement restart_state() -> dict[str, DataArray] and load_restart_state(d). RestartMonitor serializes the public state + every component’s private state + full config + grid UUIDs; restart reproducibility (same trajectory across a restart, ICON’s standard requirement) is a CI-enforced property.


5. Layer 3 — Composition and the time loop

5.1 The canonical run script

The primary UX is a legible run script, per sympl. The complete global NWP configuration:

from icon_sc import (ComputeContext, ConcurrentCoupling, SequentialUpdateSplitting,
                    CallingFrequency, ReducedGridWrapper)
from icon_sc.grid import icon_grid, vertical_grid, extpar
from icon_sc.init import from_dwd_analysis
from icon_sc.components import *          # the table in §4.3
from icon_sc.io import ZarrMonitor, RestartMonitor

ctx   = ComputeContext(backend="gtfn_gpu", comm=MPI.COMM_WORLD,
                       io_ranks=4, strict=True)
grid  = icon_grid.from_file("icon_grid_0026_R03B07_G.nc", ctx)        # decomposes if comm.size > 1
vgrid = vertical_grid.sleve(nlev=90, ...)
static = extpar.load("extpar_0026_R03B07_G.nc", grid) | metrics(grid, vgrid) | interpolation(grid)

cfg   = NWPConfig(dtime=timedelta(seconds=180), ndyn_substeps=5,
                  dt_rad="1h", dt_conv="9min", dt_gwd="9min", dt_sso="9min")

state = from_dwd_analysis(fg="fc_R03B07+0000.grb", an="an_R03B07.grb",
                          grid=grid, vgrid=vgrid, ctx=ctx, mode="iau")

dycore    = NonhydroSolver(grid, vgrid, static, cfg.dynamics, ctx,
                           substeps=cfl_adaptive(base=cfg.ndyn_substeps, max_ratio=12))
diffusion = HorizontalDiffusion(grid, static, cfg.diffusion, ctx)
transport = TracerTransport(grid, static, cfg.transport, ctx)

fast = SequentialUpdateSplitting(NWP_FAST_ORDER, dict(
          satad=SaturationAdjustment(...), land=Terra(...), lake=FLake(...),
          seaice=SeaIce(...), turb=TurbulentDiffusion(...),
          mphys=Microphysics(scheme="graupel", ...), sfc=SurfaceTransfer(...)), ctx)

slow = ConcurrentCoupling(
        CallingFrequency(Convection(grid, ...), cfg.dt_conv),
        CallingFrequency(CloudCover(grid, ...), cfg.dt_conv),
        CallingFrequency(ReducedGridWrapper(Radiation(rte_rrtmgp_cfg), rrg_pair(grid)), cfg.dt_rad),
        CallingFrequency(NonOrographicGWD(...), cfg.dt_gwd),
        CallingFrequency(SSODrag(...), cfg.dt_sso))

iau     = IncrementalAnalysisUpdate(window="3h", increments=state.pop_increments())
diag_pT = HydrostaticPressureDiagnostics(grid, vgrid)
recon   = WindReconstruction(grid, static)

monitor = ZarrMonitor("run.zarr", every="1h", variables=OUTPUT_SET, ctx=ctx)
restart = RestartMonitor("restart/", every="12h", ctx=ctx)

with ctx.timeloop(state, until=cfg.run_length) as loop:
    for dt in loop:
        state.update(diag_pT(state))                 # dyn→phy diagnostics
        state.update(recon(state))                   # vn → (u,v) at cells
        tend, diag = slow(state, dt)                 # lazy, piecewise-constant → icon:ddt_* bus
        state.update(diag); state.update(tend)
        state.update(iau(state, dt))
        state = dycore(state, dt)                    # consumes vn/exner bus slots; internal halos
        state = diffusion(state, dt)
        state = transport(state, dt)                 # consumes ρqx bus slots post-advection
        state = fast(state, dt)                      # sequential updates
        monitor.store(state); restart.store(state)
        state["time"] += dt

Every claim the sympl paper makes about legibility survives: which schemes, what order, what cadences, where output happens — all in one screen. (Model/Federation-style convenience wrappers exist for education, assembling exactly this loop from cfg; they add nothing architecturally.)

The coupling algebra makes scheme experiments equally legible. Swapping the fast-physics coupling from operational SUS to Strang splitting around the dry step is a three-line change, not a fork:

dry  = chain(dycore, diffusion, transport)                      # E0 in thesis notation
step = SSUS(sections=fast.sections, core=dry, lam=0.5)          # E*_l over λΔt, reversed; E0; E_l over (1-λ)Δt
# loop body: bus publication as before, then `state = step(state, dt)`

Composition-time constraint checking (§4.2, risk 7 of §11) decides whether such a composite is merely structurally legal or also scientifically admissible; experimental composites never inherit the validated-preset label.

5.2 ComputeContext

The one object threaded through construction: array namespace + allocator (numpy/cupy, pooled), GT4Py backend selection (gtfn_cpu, gtfn_gpu, embedded for debugging), MPI communicators (compute / I/O split, §6.4), decomposition handle, strict-mode flag, execution tier (interpret / plan / graph / native, §8.3), and the plan compiler (§8.2). Components allocate their private fields through it; the state initializer allocates public fields through it. Nothing else in the system touches devices or communicators directly.

5.3 Configuration

Typed, layered, Pythonic: frozen dataclasses per component (DiffusionConfig, TransportConfig(tracer_schemes={"qv": "ffsl_miura", ...}), …) composed into NWPConfig; YAML/TOML ingestion is a thin loader into the same dataclasses; validation at construction (ranges, cross-field constraints like the dt_rad/dt_conv multiple rule). No namelist emulation layer in v1 — but the registry (§2.5) plus per-config icon_namelist_origin annotations document, field by field, which ICON namelist variable each option corresponds to, so a future from_namelist() adapter is a table walk, not archaeology. Full config + grid UUIDs + package versions + git SHAs are stamped into every output artifact (the reproducibility story sympl promises, made enforceable).


6. Layer 4 — Distributed execution

6.1 Decomposition

decompose(grid, comm) partitions cells (METIS or ICON’s geometric decomposition — pluggable), derives edge/vertex ownership by the usual adjacency rules, and builds DecompositionInfo: owned/halo index sets per location type, halo levels (two cell rows, matching ICON’s stencil depth requirements), and global↔local index maps. The result is a local IconGrid (local connectivity tables with halo entries) plus exchange metadata — components are SPMD and see only local grids; nothing in any component’s numerics is decomposition-aware.

6.2 GHEX substrate and CommunicationComponent

One GHEX unstructured domain descriptor and pattern per location type (cell/edge/vertex) and halo depth, built once from DecompositionInfo and cached on the context. HaloExchange(fields=[...]) is a component: input properties = the fields with halo: DIRTY, output = same fields halo: VALID; its array_call posts one fused GHEX exchange for all listed fields (message aggregation matters at scale). GPU-aware MPI is assumed available and verified at startup; the fallback path stages through pinned host buffers with an explicit startup warning.

6.3 The halo validator pass

At composition time (ctx.timeloop(...) entry), ICON-sc walks the loop’s component sequence symbolically: each component’s declared halo: required inputs are checked against the running validity map produced by predecessors’ halo: invalidated/valid outputs. Two modes: halos="auto" inserts minimal HaloExchange components at the latest valid point (and reports the inserted schedule); halos="manual" (production default) only verifies hand-placed exchanges and fails loudly on a gap or a redundant exchange. Components with communicates_internally=True (the dycore) assert their own postconditions. This pass is the single most valuable safety net the architecture adds over hand-rolled MPI hosts: halo staleness bugs become composition-time errors.

6.4 I/O and restart ranks

ICON’s worker/output-server split is reproduced as communicator topology, not process zoo: ComputeContext splits comm into a compute communicator and an I/O communicator (io_ranks). Monitors run a nonblocking gather-or-forward: on store(), compute ranks post sends of their owned slices (with global index maps) to I/O ranks and return immediately; I/O ranks assemble and write NetCDF or Zarr (region-writes with a global chunk map — Zarr makes the parallel-write path far simpler than NetCDF and is the recommended default). RestartMonitor uses the same machinery on its own cadence (the num_restart_procs analog). Asynchrony is bounded by a two-deep buffer; backpressure blocks the loop rather than exhausting memory.


7. Layer 5 — Ingestion, initialization, output

7.1 What is ingested vs. what is emitted

Per the agreed reconciliation: ICON-sc reads ICON’s artifact ecosystem — grid NetCDF (§3.1), ExtPar NetCDF (topography incl. SSO fields, land-use tiles, soil, climatologies), DWD analysis/first-guess and IFS initial data in GRIB2 — and writes Pythonic artifacts (Zarr/NetCDF with CF-compliant metadata where applicable). No GRIB2 output, no namelist reading in v1.

7.2 GRIB2 ingestion

A thin eccodes-based reader driven by the variable registry’s GRIB2 columns: shortName/level-type → canonical name + vertical placement; fields land directly in device-resident DataArrays via the context allocator. Horizontal remapping from IFS grids is out of scope (external preprocessing with iconremap/earthkit-regrid documented as the supported path); ingestion assumes data already on the target ICON grid, exactly like ICON itself after remap.

7.3 Initialization modes

from_dwd_analysis(mode="uninitialized" | "iau" | "initialized") and from_ifs_analysis(...) reproduce tutorial ch. 5’s start modes: construct the prognostic set (vn from u/v via edge projection where needed, exner/θv from p/T, hydrostatic balancing of the reference state), stash IAU increments when mode="iau" for the IncrementalAnalysisUpdate component. Idealized initializers (Jablonowski–Williamson, Straka) are DiagnosticComponent-shaped functions used by the validation ladder (§9).

7.4 Monitors

ZarrMonitor (default; anemoi-conformant layout available as a flag — this is the cheap end of the ICON4Py↔anemoi roadmap and falls out of the design nearly for free), NetCDFMonitor (CF-ish, xarray-written), RestartMonitor, and a PlotMonitor for education. Output-set definitions are named variable lists with optional on-the-fly derived diagnostics (declared as DiagnosticComponents evaluated on I/O ranks when their inputs are cell-located and local, else on compute ranks).


8. Execution architecture: performance tiers and differentiability

8.1 The honest problem statement

The sympl boundary (dict lookups, DataArray attribute traffic, per-call contract matching, wrapper hops, Python control flow) costs microseconds per component call — and inside components, per-stencil Python launches cost more. At climate-toy scale that is invisible; at production ICON scale on GPUs — where a full Δt step is tens of milliseconds and a dycore substep is single-digit milliseconds — interpreted per-step dispatch of ~20 components plus halo logic plus hundreds of kernel launches is a hard ceiling on strong scaling. Two orthogonal observations structure the solution: (a) nothing about the interfaces changes during execution — the state schema, buffer set, component sequence, ingress plans, and halo schedule are all fixed at composition time, so every lookup performed in the loop is recomputing an invariant; (b) once (a) is exploited, the residual cost is pure dispatch — Python function calls and kernel-launch latency — which is a runtime problem (graph replay, native driver), not a compiler problem. Neither requires whole-program dataflow compilation, which is why no DaCe-style SDFG orchestration appears anywhere in this design.

8.2 Bind-time specialization: negotiation vs. execution phases

The full sympl machinery — name resolution, dims/units matching, Pint no-op verification, location and halo checks, xarray wrapping — runs exactly once, at bind time (entry of ctx.timeloop(...)), against the composed sequence and a schema-representative state. Its product is a frozen execution plan; nothing sympl-shaped executes per step.

StateVault. The state’s execution-phase form is a dense, slotted container: a flat list of raw buffers plus an interned name → index map consulted only at bind time. Fields are referenced by handles:

@dataclass(frozen=True, slots=True)
class FieldHandle:
    idx: int
    name: str          # debug/repr only

class StateVault:
    __slots__ = ("buffers", "names", "schema_hash", "epoch")

The public dict-of-DataArrays view still exists — as a lazily materialized façade over the vault (DataArray wrappers cached per slot, reconstructed only on epoch change) — so monitors, interactive inspection, and the interpreted tier see unmodified sympl semantics. xarray is thereby confined to negotiation, monitoring, and debugging; it never appears on the step path.

Bound calls. For each component, bind time produces a specialized entry: the ingress adapters of §2.3 are resolved to concrete zero-copy views, argument order is fixed, and the result is a plain tuple pack — deliberately a tuple/NamedTuple rather than anything richer, because tuple indexing is the cheapest structured access CPython offers and the pack is built once:

class BoundCall(NamedTuple):
    fn: Callable            # the component's raw kernel entry (array_call specialized)
    args: tuple             # pre-resolved buffers / gt4py fields / scalars
    tag: str                # profiling/debug label

Buffer identity and the “new state” problem. sympl Steppers notionally return fresh dicts; a zero-lookup plan requires stable buffer identities. The plan therefore pre-allocates all component outputs — delivered through the __call__(state, out=...) / array_call(inputs, outputs) ABI adopted from Ubbiali’s sympl fork (§4.2) — and encodes state evolution as either in-place updates (the common case under ICON’s sequential-update physics) or slot swaps. Swaps are not indirected at runtime: for every ping-pong pair the compiler emits even/odd plan variants with references pre-bound, and the interpreter alternates variants — zero per-step indirection, at the cost of 2× plan memory (kilobytes).

Wrapper and federation compilation. Wrappers that are pure control flow dissolve into the plan: CallingFrequency becomes a per-step cadence mask (the plan precomputes, from the lcm structure of Δt/dt_conv/dt_rad/dt_gwd/dt_sso, the small set of distinct step signatures and the op list active under each); Subcycle and the DynamicalCore substep tier unroll or loop with pre-bound dt/n; scaling-type wrappers fold into bound constants. The federation classes of §4.2 dissolve identically, being pure control flow over their sections: SequentialUpdateSplitting flattens into the op list; ParallelSplitting’s recombination ψⁿ⁺¹ = Σψₗ − L·ψⁿ compiles to one fused multi-axpy vault op over the affected slots; SequentialTendencySplitting’s provisional-tendency forcings become dedicated vault slots plus a diff-and-scale op feeding each SequentialTendencyStepper; SSUS simply doubles the op list with reversed section order and λ-scaled dt constants. Wrappers that do real compute (ReducedGridWrapper) remain ordinary bound calls. HaloExchange components bind to pre-built GHEX bulk-exchange handles (pattern + field set frozen), so an exchange is one FFI call.

Why not a dataclass state? The schema is configuration-dependent (the tracer set alone varies), so no static dataclass can be the public state type. Where typed attribute access is genuinely wanted — component-internal views, tests — the bind phase can emit a generated slotted dataclass (dataclasses.make_dataclass(..., slots=True)) over vault slots; but inside the plan, handles and tuples beat slot-descriptor access anyway, so this is an ergonomic option, not the mechanism. For scale: a dict lookup is ~40–60 ns and slotted attribute access ~20–40 ns, but those were never the real cost — the µs-scale per-field DataArray construction, attrs traffic, and sympl match/convert logic are, and phase separation removes them categorically rather than shaving constants.

Staleness guards. The plan is valid for one (composition, schema, decomposition) triple, enforced by schema_hash + an epoch counter bumped by any out-of-band state mutation through the façade; a stale plan raises at the next step rather than silently binding dead buffers. Debug builds re-run negotiation every N steps and diff against the plan.

8.3 Execution tiers (DaCe-free orchestration)

Four tiers over one composition model, ordered by implementation cost; all share the plan from §8.2, and CI enforces cross-tier agreement to validation tolerances on the idealized suite.

T0 — interpreted sympl dispatch. Full per-call contract machinery, dict/DataArray semantics. Reference behavior; education and component development.

T1 — plan interpreter. A Python loop over tuple[BoundCall, ...] per step signature. No dict, no xarray, no contract logic; per-component overhead drops to one Python call + tuple splat (order 1 µs), i.e. tens of µs per Δt step for the component layer. This is the default production tier on CPU and the development tier on GPU.

T2 — graph replay. On CUDA/HIP, the plan interpreter is run once per step signature under stream capture (CuPy exposes begin_capture/end_capture); each exchange-free segment of the plan becomes a captured graph, replayed thereafter with a single launch (µs-scale, independent of interior kernel count). Segments are delimited by GHEX exchanges, Fortran bridge calls, and framework seams, because host-side MPI cannot sit inside a capture; the per-step Python cost collapses to (number of segments) × (one graph launch + one exchange call). The adaptive ndyn_substeps escalation and slow-physics cadences are handled by the signature cache — the signature set is small and bounded (substep ratio ∈ [5..12] × cadence phases), so graphs are captured lazily and reused. Requirements this imposes, checked at bind: static shapes, stable buffer pointers (guaranteed by the vault), no data-dependent host branching inside segments (the CFL decision reads back one scalar per step, between segments). HIP graph support via CuPy is flagged experimental — verified at startup, with T1 fallback.

T3 — generated native host driver. For the last factor — removing Python from the step entirely, and serving CPU targets where graphs don’t apply — the bind phase can emit a small C++ translation unit that sequences the step directly: GTFN-compiled programs are C++ entry points in shared objects (the Python bindings are a convenience layer, not the interface), GHEX is native C++, and Fortran bridges have C ABI. The generated driver is a few hundred lines of calls with baked-in pointers, strides, cadence branches (if (step % k == 0)), and the even/odd swap logic; it is compiled once at bind (cached by plan hash) and exposed as one symbol: step(vault, step_index). Python then makes one FFI call per Δt — or per output interval. ML/JAX components and any Python-only component return control to Python at their seam, segmenting the driver exactly as they segment T2 graphs. T3 is an escalation path, not a prerequisite: T2 is expected to reach within noise of a native driver on GPU, so T3 is built only if profiling of T1/T2 says so, and its scope can stay limited to the hottest region (the sub-cycled dycore + diffusion + transport block).

The fusion question, answered without DaCe. What SDFG orchestration would have bought beyond dispatch removal is inter-kernel fusion across stencil boundaries. That lever is retained at the code level rather than the compiler level: gt4py programs may contain multiple field operators, and icon4py already maintains hand-fused program variants for launch-bound regions of solve_nonhydro; coarsening granule programs is the sanctioned way to cut kernel count. The residual loss — automatic cross-component fusion — is accepted and quantified: with T2/T3 the cost of an unfused boundary is one kernel launch (~5 µs amortized to ~0 under graph replay) plus a possible missed memory-locality win; the latter is real but second-order next to the halo-exchange and dycore-interior costs that dominate ICON’s step, and it can be recovered case-by-case by fusing at the gt4py program level where profiles justify it.

A note on precompiled gt4py entry points. Both T1 (call overhead) and T3 (C++ linkage) lean on gt4py.next’s compiled-program path — pinning offset providers, shapes, and scalar-vs-field argument status at compile time so the per-call marshalling disappears. That API surface has been evolving (compile-time connectivities, static-args variants); the design treats it as a bind-time detail behind the ingress adapter, but the exact pinning mechanism should be re-verified against the current gt4py.next release before implementation — sketch-status, not pinned.

8.4 Multi-framework components

The boundary protocol admits JAX or Numba components (e.g., an ML parameterization replacing convection — the anemoi-adjacent use case). Rules: DLPack ingress is zero-copy but read-only-by-convention; JAX components must be functional at the boundary (inputs → new output buffers written back by the adapter into pre-registered vault slots), never in-place; a framework seam is a segment boundary in T2/T3 (graphs and the native driver stop and resume around it); stream/queue synchronization at seams is the adapter’s job (DLPack stream argument, external-stream interop). Inside its segment a JAX component is free to be one jitted function — its own internal fusion story — which composes cleanly with this tiering. These rules are the state/data-boundary principle applied inside a single model; a JAX component honoring them is automatically differentiable: "native" under the §8.6 contract.

8.5 The functional lowering (F-tier)

ctx.functional(composition, window=...) is a fourth consumer of the §8.2 negotiation. Instead of an imperative op list it emits a pure JAX function over explicit PyTrees:

step_fn: (StateTree, ParamTree, StaticArgs) -> StateTree          # one Δt, traced
window  = scan_window(step_fn, n_steps, remat="per_step")          # lax.scan + jax.checkpoint policy
loss_grads = jax.vjp(lambda p: J(window(state0, p)), params)       # or jax.jvp for tangent-linear

StateTree is a generated frozen-dataclass PyTree derived from the vault schema (one leaf per slot), extended with explicit carry: every leaf a component declared in functional_state() — dycore time levels and velocity-advection carry-over, CallingFrequency tendency caches, turbulence history — is surfaced into the tree (tension T7: §4.5’s privacy is an imperative-tier convenience, demoted here by contract). The semantic mapping from the plan is mechanical: in-place updates become functional updates whose memory cost donate_argnums recovers under jit; even/odd slot swaps become tuple rebinding with zero post-trace cost; cadence masks become per-step boolean masks with jnp.where selecting cached-vs-recomputed slow tendencies (the cache being carry), which keeps the trace static across step signatures. Two declared freezes: the CFL-adaptive substep ratio is pinned to a static value in F-tier configurations (differentiating through discrete ratio switching is ill-posed, and the freeze is recorded in provenance, never silent), and monitors/restart are outside the trace (io_callback permitted only outside differentiated regions).

A pleasant inversion: under jit, XLA performs precisely the cross-component fusion the imperative tiers renounce (§8.3) — the differentiable path is incidentally the most aggressively fused one, paid for in XLA compile times and in XLA’s weaker fit to icosahedral indirect addressing, where the required idiom is flatten + segment_ids + padded neighborhoods for the connectivity gathers.

8.6 The differentiability contract

Property contracts gain two axes: differentiable: "native" | "custom" | "none", and a params declaration exposing tunable scheme constants (entrainment coefficients, autoconversion thresholds, …) as a ParamTree distinct from state — parameter estimation and gradient-based calibration never smuggle constants through state fields.

  • native — the component provides a pure functional_call(inputs, params) -> outputs co-located with its imperative kernel, both drawing on one shared scheme-constants module (single source of numerical truth; CI holds the pair together with the same ladder-level-2 tolerances as the kernel itself). Traced directly; jvp/vjp for free. A JAX component under the §8.4 seam rules is automatically native.

  • custom — the primal stays GT4Py: the gtfn kernel is registered as an XLA FFI target (jax.ffi.ffi_call; young API, pinned under §8.3’s discipline) and paired via jax.custom_vjp/jax.custom_jvp with hand-written tangent/adjoint GT4Py programs. This is the classical adjoint-stencil route of operational 4D-Var, made per-component and composable. Implicit structure gets implicit-function-theorem treatment rather than unrolling: the adjoint of the turbulence vertical implicit solve is the transposed tridiagonal solve; saturation adjustment’s fixed point differentiates through lax.custom_root-style rules, not through recorded Newton iterations — which also sidesteps while_loop’s reverse-mode prohibition.

  • none — bridges and unported schemes. Per differentiated region, composition-time policy is error (default) or stop_gradient — explicit, warned, and stamped into provenance. Gradient truncation is never silent.

Numerical guidance ships as documentation-with-tests: where-branching and thresholds yield subgradients (legitimate, but characterized in L8); hard positivity clips truncate sensitivities and differentiable configurations should prefer smooth clamps; fp64 is the default for gradient work.

8.7 Differentiable distributed execution

Halo exchange is linear in the field values, which settles both directions at once: the JVP of an exchange is the same exchange applied to tangents; the VJP is the transpose, not the exchange run backward — cotangents accumulated on ghost points scatter-add onto the owning ranks’ interior points, after which the ghost region of the cotangent is zeroed. Both rules derive mechanically from DecompositionInfo (swap the send/recv index sets; switch the write mode to accumulate) and are registered via custom_vjp/custom_jvp on a DifferentiableHaloExchange, which the F-tier substitutes at every point the halo validator would place a HaloExchange. Implementation path: mpi4jax point-to-point first (existing jit/token integration), GHEX behind jax.ffi with stream-correct custom calls when performance demands it — closing the known gap that GHEX has no JAX/XLA autodiff integration. Execution stays SPMD: one jit per rank over the local domain, MPI underneath; no single-controller shard_map construction is attempted for the unstructured pattern. Reverse-mode memory over a window is managed by the lax.scan + per-step jax.checkpoint policy (≈ one extra forward), and the adjoint test in L8 runs through the exchange, since the transpose rule is exactly the part hand-rolled adjoint models historically get wrong.

8.8 What is not attempted

No cross-framework compute unification (no GT4Py-Field-in-JAX-trace heroics — the custom route crosses an FFI boundary, it does not trace GT4Py); no source-transformation AD of gtfn-generated kernels (Enzyme-JAX over the generated LLVM/MLIR is noted as exploratory, nothing more); no whole-program dataflow compilation in the imperative tiers (§8.3 accepts and bounds that loss; the F-tier gets XLA fusion as a side effect); no nproma emulation; no attempt to make T0/T1 competitive with Fortran ICON — T2/T3 vs. icon4py-standalone remains the meaningful forward-performance benchmark, while F-tier performance is measured against NeuralGCM-class JAX models, not against the imperative tiers.


9. Validation ladder (what “scientific equivalence” means operationally)

  1. Stencil parity — inherited from icon4py: granule outputs vs. serialized ICON reference data (Serialbox datasets), tolerance-based.

  2. Component parity — single-column and single-process tests against ICON offline drivers/SCM: satad, graupel, turbdiff, TERRA on fixed forcing; radiation vs. ecRad reference (the CFFI-wrapped ecRad exists for exactly this).

  3. Suite behavior — fast-physics SequentialUpdateSplitting vs. ICON column output over multi-day SCM integrations; slow-tendency cadence semantics (piecewise constancy, phase after restart) verified analytically.

  4. Idealized dynamics — Jablonowski–Williamson baroclinic wave and Straka density current per tutorial ch. 4, compared against published/ICON-generated references (norms, growth rates, structure).

  5. Global real-data — R2B4→R2B6 forecasts from DWD analysis; probtest-style ensemble-tolerance comparison against ICON (perturbed-IC ensembles define the equivalence band; ICON-sc must fall inside it), plus standard skill scores vs. analysis.

  6. Invariants in CI — restart reproducibility, run-to-run determinism per backend, tracer mass conservation, dry-mass conservation drift bounds, cross-tier agreement (T0 interpreted vs. T1 plan vs. T2/T3 lowered execution).

  7. Coupling-algebra self-convergence — the federation classes themselves are validated by the thesis’s grid-refinement methodology: self-convergence tests (a Burgers-type problem, then the level-4 idealized cases) verifying that each operator attains its formal order around ICON-sc components — with the documented caveat, directly relevant to the ICON preset where saturation adjustment appears twice per step, that satad caps coupling-error convergence at first order unless a prognostic condensation formulation is substituted (thesis §2.5/App. 2.B).

  8. Gradient verification — for every native/custom component and for composed windows: Taylor-remainder tests for jvp (correct first- and second-order decay), adjoint consistency ⟨Jv, w⟩ = ⟨v, Jᵀw⟩ to fp64 tolerances (the dot-product test, per component and through multi-step scan windows including the halo-exchange transpose), and finite-difference cross-checks on scalar functionals. Long-window gradient growth is characterized, not gated — sensitivity blow-up over chaotic horizons is physics, not a defect.

Explicitly not claimed: bitwise cross-platform reproducibility (that is the grid-generator problem’s territory — libm/correct-rounding — and out of scope for a model targeting scientific equivalence).


10. Future-proofing and extension points

Coupling-scheme research (new headline payoff). Because the ICON preset is one point in the federation algebra, ICON-sc supports exactly the class of experiments that motivated Tasmania and that production codes structurally cannot host: swapping SUS ↔ STS ↔ SSUS around an operational-grade dycore and physics suite, λ-sweeps, ordering studies, per-stage (FC) physics through the DynamicalCore fast-tendency port, and per-substep physics through its super-fast port — all under composition-time constraint validation and the cross-tier equivalence machinery, at kilometer-scale-capable performance rather than on idealized slice models. This closes the gap the thesis names explicitly: coupling research done on bespoke toys because flexible production code didn’t exist.

Data assimilation and hybrid-ML training. The F-tier turns the composed model into what variational DA and differentiable-hybrid research need: jax.vjp of a forecast window for 4D-Var-style adjoint experiments; ParamTree gradients for scheme calibration against observations or reference simulations; and end-to-end training of an ML component (§8.4) inside the physical model — which upgrades the anemoi coupling below from offline adaptation to online differentiable coupling, the NeuralGCM pattern with an ICON dycore. The differentiability contract keeps this honest: any none-labeled component in the window is a composition-time decision (error or explicit stop_gradient), so a truncated gradient is always a documented choice.

Nesting / LAM (the known v1 exclusions). Three provisions keep them cheap later: (i) refin_ctrl and grid-hierarchy metadata are read and retained; (ii) the driver holds state_per_domain: dict[DomainId, State] internally even when its length is 1 — no flat-single-state assumption is baked in; (iii) parent↔child feedback and lateral-boundary nudging are shaped as paired CouplingComponents (a communication-component subtype) plus a per-domain loop nesting in the driver — the tutorial’s processing sequence (§3.9.2) maps onto nested Subcycles. The LAM prefetch PE becomes another async role on the I/O communicator.

anemoi / ML coupling. ZarrMonitor’s anemoi layout + a PrognosticState↔anemoi.State adapter component realize the previously scoped roadmap stages 1–3 inside this architecture; an ML emulator replacing a physics component is just a JAX/Torch component under §8.4 rules.

Other physics suites. The AES package’s dyn2phy/phy2dyn transposition layer maps naturally onto explicit DiagnosticComponent pairs; nothing in the taxonomy is NWP-specific except NWP_FAST_ORDER.

Ocean. Out of scope, but the taxonomy is location-typed and grid-parameterized, and YAC-style coupling would enter as CouplingComponents between two ICON-sc models with distinct grids.


11. Risks and open questions

  1. Porting mass. Tiedtke–Bechtold convection, TERRA (tiles), and turbdiff are large, gnarly Fortran bodies with decades of accumulated fixes. Mitigation: CFFI-wrapped Fortran as bridge components (CPU-only, segmenting T2/T3) so the model is whole early, with GT4Py ports replacing bridges scheme by scheme against ladder-level-2 tests. The architecture explicitly tolerates permanently-bridged components at the cost of T2/T3 segmentation.

  2. Radiation choice. pyRTE-RRTMGP (GPU, maintained, Python-first) is not ecRad; scientific equivalence at suite level requires either an ecRad port (expensive) or accepting RRTMGP-vs-ecRad as a documented scheme substitution validated at ladder level 5. Decision needed early; the architecture is agnostic but the equivalence claim is not.

  3. Two-time-level privacy vs. introspection. Hiding nnow/nnew inside the dycore is clean but makes some ICON diagnostics (which peek at both levels) awkward; escape hatch is the restart_state() surface, which doubles as an introspection API — acceptable, slightly ugly.

  4. Graph-capture and native-driver constraints. T2 requires static shapes, stable buffer pointers, and no host-side data-dependent branching inside segments; the adaptive ndyn_substeps and slow-physics cadences are absorbed by the bounded step-signature cache, but any future component with in-segment host control flow forces a split. HIP graph support through CuPy is markedly less mature than CUDA’s. T3’s generated driver trades Python overhead for a build-and-debug surface (a compiler at bind time, symbol/ABI hygiene across gtfn shared objects, GHEX linkage). Both T1 and T3 lean on gt4py.next’s still-evolving precompiled-program / static-args API — verify and pin before implementation. Build T2 incrementally (dycore segment first, widen outward); treat T3 strictly as profiling-driven escalation.

  5. Plan staleness and dual-semantics drift. Bind-time specialization creates two representations of one model (T0 semantics vs. the frozen plan). The schema-hash/epoch guards catch out-of-band structural mutation, but semantic drift — a wrapper that turns out not to be pure control flow, a component mutating an input it declared read-only, an aliasing assumption violated by a swap — is the bug class to design tests around. CI’s cross-tier agreement checks and the debug-mode periodic renegotiation-and-diff exist for exactly this; treat any T0/T1 divergence as a release blocker, never a tolerance to widen.

  6. GHEX Python bindings + GPU-aware MPI on target machines (Alps-class) need early smoke-testing; the pinned-host fallback protects correctness, not scalability.

  7. Coupling swaps are structurally easy and scientifically treacherous. ICON’s physics carries implicit contracts that hold only under the operational arrangement: the surface-transfer-at-end time-level optimization, satad’s placement (twice per step) and its convergence-capping behavior, turbulence’s old-time-level input semantics, the isochoric heating convention on the bus. A federation swap that violates any of these is legal code and wrong science — and the thesis’s own results (ordering sensitivity, scheme-dependent convergence of moist variables) show the differences are not academic. Mitigation: components declare coupling_constraints (must_follow/must_precede, required time-level semantics, admissible operator families) checked at composition; presets carry the validated/experimental label; experimental composites never inherit the scientific-equivalence claim; and the level-7 self-convergence tests exist to characterize, not merely permit, non-preset composites.

  8. Dual-implementation and adjoint drift. native components carry two implementations of one scheme (gtfn kernel + JAX functional core); custom components carry a primal and a hand-written adjoint. Both pairs drift silently under maintenance unless held together mechanically: shared scheme-constants modules (one source of numerical truth), ladder-level-2 tolerances applied to functional cores against the same serialized references as the kernels, and level-8 adjoint tests in CI so an un-updated adjoint fails the dot-product test on the next run. Two further accepted realities: jax.ffi and mpi4jax are young surfaces, pinned under the constraints/ discipline like gt4py’s; and reverse-mode memory at production resolution is real even with per-step remat — long differentiation windows belong to coarse configurations and SCM columns, which is where the actual use cases (parameter estimation, hybrid training, 4D-Var research) live.


12. Summary of the extension delta over sympl

Everything sympl provides is kept. The complete list of additions: location + halo axes in property contracts; canonical-units registry + strict mode; FieldBuffer DLPack boundary + cached ingress adapters; IconGrid/VerticalGrid/static-state; the Tasmania coupling algebra (ConcurrentCoupling, registry-based TendencyStepper/SequentialTendencyStepper, ParallelSplitting/SequentialTendencySplitting/SequentialUpdateSplitting/SSUS federations, DynamicalCore with slow/fast/super-fast tiers, the out=/array_call(inputs, outputs) ABI, Checker/Operator-structured negotiation, Factory registries) with coupling_constraints validation added on top; Subcycle, CallingFrequency (generalized), ReducedGridWrapper, SlowTendencyBus convention + checker, CommunicationComponent; halo validator pass; ComputeContext; component-private restart protocol; the negotiation/execution phase split (StateVault, execution plans with even/odd swap variants and cadence masks) with its optional graph-replay and generated-native-host-driver lowerings; and the differentiability layer — the differentiable contract axis with ParamTree parameters, the F-tier functional lowering with explicit-carry PyTrees and scan/remat windows, and the transpose-rule DifferentiableHaloExchange. Each addition traces to a tension in §1 or to a Tasmania mechanism adopted in §4.2 — nothing else about sympl needed to change to host a production-structured NWP model, which is a nontrivial validation of the sympl thesis and of Ubbiali’s extension of it: between them, the two lineages had already built most of the composition layer a production model needs; what ICON-sc adds is the device-field boundary, the distributed-memory story, and the execution tiers that make the composition affordable at scale.