← all records

Design State and Component

task-eb9e84 · in State and Component

What is a Component?

it is anything that interacts with the State.

Requirements

  • Generic enough to fit current "components" (physics, dycore, diffusion). Does not have to take e.g. IO into account, that can be fitted into components later. Bonus if it works for that as well.
  • Must be more "checkable" (at runtime and/or compile time) than plain sympl components with stringly dicts.

Notes from 2026-09-01 discussion

  • The synthesis proposal is forward looking and contains more than we need for this cycle. Focus mostly on component protocol and a bit on the interaction with model state.
  • Open question: what shape should inputs and outputs be:
    • Dataclasses with fields?
    • Plain python arguments splatted from a dictionary, output a tuple?
    • Strong types for each type of field?
    • Different approaches may have (small?) differences in readability and toolability/type-checking
  • Metadata for fields should always come from one source of truth, not duplicated in input/output definitions
  • Do we need to care about model state?
    • If component protocol is flexible enough, it'll take anything that comes from the model state.
    • We may, however, need adapters for making it more convenient to extract the right fields from the model state and for writing them back to the model state.
    • We may also need to extend the model state to have the same/similar metadata as the input/outputs to components so that they work well together.
  • The synthesis proposal is definitely both over and underspecified. It's not very well considered in terms of the specifics of metadata, roles, kinds, etc.

Jacopo's notes

  1. Look at CF names and see if we can derive one for each weird quantity or if that's futile and we use internal namese for I/O
  2. figure out how to access the metadata "hidden" in Quantity (see Rico).

quantities.py

from

TEMPERATURE: Final[Quantity] = register(
    "air_temperature", units="K", dims=(dims.CellDim, dims.KDim),
    cf_key="temperature",
)
type TemperatureField = Annotated[fa.CellKField[ta.wpfloat], TEMPERATURE]

TENDENCY_TEMPERATURE: Final[Quantity] = register_tendency(
    "ddt_temp", of=TEMPERATURE,
)
type TemperatureTendencyField = Annotated[fa.CellKField[ta.wpfloat], TENDENCY_TEMPERATURE]

to something like:

@field_type(dims, "K", "air_temperature", cf_key="temperature")
type TemperatureField

@field_type(dims, "K", "ddt_temp", tend_of=TemperatureField)
type TemperatureTendencyField


SCRATCHPAD

Quantity

import pint
 
@dataclasses.dataclass(frozen=True)
class Quantity:
    type: Field | Scalar                   # 
    name: str                              # canonical identity
    units: pint.Unit                       # parsed once at registration
    dims: tuple[gtx.Dimension, ...]        # ONE source (full vs half level)
    cf_key: str | None = None              # IO name, required when IO-eligible
    #
    tendency_of: "Quantity | None" = None  # set on tendency quantities only
    #
    metadata: dict = {  # e.g.
        type: Prognostic | Diagnostic
    }
    
def register(name, *, units, dims, cf_key=None) -> Quantity
def register_tendency(name, *, of, cf_key=None) -> Quantity # dims = of.dims; units = of.units / s (checked at registration)

# or

def register(name, *, units, dims, cf_key=None, with_tendency=False) -> Quantity | (Quantity, Quantity)

Component

class Component:
    Input: type = EmptyInput       # the interface slot (a state class)
    Output: type = EmptyOutput
 
    # -- the one abstract member: the computation --
    def run(self, input: Input) -> Output:
        raise NotImplementedError
 
    # -- default methods, driven by the slots; overridable --
    def gather(self, *states) -> Input:
        ...   # sketch under "Input views"; reads self.Input
 
    # -- output routing: a method on the component, symmetric with gather --
    def apply(self, output: Output, *states) -> None:
        ...   # sketch under "Output routing"; reads self.Output

Computing derived quantities

class Muphys(Component)
    class Input:
        c: SomeField
    
    def gather(self, state: State):
        if c in state:
            input.c = state.c
        else:  # option 1: gather is responsible for the compute logic
            input.c = foo(state.a, state.b)

# driver
model_state.a, model_state.b = something()

if muphys is not None:
    # option 2: the driver is responsible for the compute logic
    driver_state.c = foo(model_state.a, model_state.b)
    input  = muphys.gather(model_state, driver_state)
    output = muphys.run(input)
    
if muphys2 is not None:
    driver_state.c = bar(model_state.a, model_state.b)
    input  = muphys.gather(model_state, driver_state)
    output = muphys.run(input)

TODOs

  • Look at CF names and see if we can derive one for each weird quantity or if that's futile and we use internal namese for I/O instead, and cf_names are only labels
  • Remember that outputs can be cached for different Δt and the component returns the cached outputs