← all records

Prototype implementation

task-083973 · in State and Component

Based on lots of work by lots of people (Mikael, Yilu, Chia Rui, Enrique, Magda, ...)


Intro

  • This is a v1, plenty of space for improvements
  • Likely already contains too much, even for a v1; can be split
  • Everything is in mwe/components on branch components-mwe of the knowledge base
  • We had four or five design docs since July bewtween Chia Rui, Enrique, Mikael, Yilu and me. There are still gaps even in Mikael's v4 (e.g. no home for dt, apply cannot run the physics-driver, no way to address now and next). So Mikael and I we decided to move to implementation.

Problems we are trying to solve

This is not a list of criticisms, it serves to provide a direction.

  1. Current States are rigid boxes which are costing us duplication of code and memory (PrepAdvection and AdvectionPrepAdvState held the same 3 quantities, allocated separately, with nothing copying one into the other; T/Tv/p derived in >3 places with different inputs: IO uses permanently-zero hydrometeors, physics uses real tracers; the temperature written to output is not the temperature physics computed with. vn -> u,v computed 3x with different domain bounds; one even skips the halo exchange, .... )
  2. Agreement by string. "te" must match in the adapter, in the component and in INPUTS_PROPERTIES. tmx computes ddt_temperature and exposes it as tend_temperature through a port table. A wrong key type-checks and runs.
  3. Intent lives in the composer. Where a tendency is applied is decided in the driver in ApplyToPrognostic, by key name, once, for every configuration. muphys never declares, e.g., that its temperature tendency goes into temperature and then through the equation of state into exner and theta_v. tmx never says that its tend_u becomes a vn increment. A new process emitting tend_w would be summed and dropped without complaints.
  4. Wasted computations: temperature and u are diagnosed twice per step whatever the configuration: at physics entry and again for output. With no physics and no output that is still sixteen calls over four steps.
  5. No introspection: nothing can tell you "who reads theta_v, who writes it", unless you go through the entire code.
  6. Types are not informative: dict[str, DataField] is opaque to mypy. PredictorCorrectorPair's accessors type as Callable[..., Any] -> type-ignore pragmas in the dycore.
  7. One decision, several owners: the driver swaps the dycore's predictor/corrector pair while the dycore decides, from the same flags, whether to recompute the predictor. The condition exists twice.
  8. initialize_granules is all handwritten
  9. store picks fields by name out of a dict

Design

Primary

  1. Every (physical) quantity is one frozen Quantity(name, units, cf_key), registered once by name, using a type alias through Annotated:
type VnField = Annotated[EdgeKField[wpfloat], quantity("normal_velocity", units="m s-1", cf_key="vn")]

Tendencies and Increments are never declared by hand: tendency_of(q) and increment_of(q) derive them and keep the parent in of. For mypy and gt4py the Annotated alias erases to the plain field type (so we can use them directly without changing all the call sites/programs/field_operators/...)

  1. A State is a frozen dataclass whose leaves are those aliases. Generic wrappers add metadata without changing the type: Read, ReadWrite, Tendency, Increment, Now, Next. rho_now: Read[Now[RhoField]] is still a CellKField[wpfloat] to mypy. The declarations can be read through typing.get_type_hints.

  2. A Component is a typed box: Input, Output, run, collect_inputs. A base class, not a Protocol, so a default method can be inherited and overridden. collect_inputs(*states) selects pointers by (quantity, level) from the states it is handed (never computes). Missing or duplicated input keys raise MissingInput, AmbiguousSource.

Secondary

  1. Derived quantities: the component declares, the composer computes. A process that needs temperature declares it like this in the Input:
class Input(State):
  temperature: Read[TemperatureField] = derived_by(recipes.TemperatureFromThetaExner)

Recipes are Components in a shared module, so Components pick a recipe rather than writing one TODO(what if doesn't exist). The composer resolves once at init: collects the recipes, deduplicates, and runs them once per step before the processes. Two wrappers picking different recipes for one quantity is an init-time error, not two temperatures.

  1. Time level TODO(time level/step) belongs to the declaration and the Pair, not the quantity. Double buffering is a driver decision, not a property of air density. TimeStepPair[PrognosticState] is two whole states and one swap, same as currently implemented. When a pair is handed to collect_inputs, its now leaves resolve at level NOW and its next leaves at NEXT. Only the dycore and tracer advection say Now/Next; every other granule declares a level-free quantity and gets whichever state the driver hands over.