feat: select the variables written to the inference GRIB output

README.md

@@ -61,6 +61,8 @@ runs:
       label: M-1 forecaster
       steps: 0/120/6
       config: resources/inference/configs/sgm-forecaster-global_trimedge.yaml
+      # Optional: variables written to the GRIB output files (all if omitted), see "Output variables".
+      # output_variables: [2t, 2d, 10u, 10v, msl, tp, z]
   # A temporal_downscaler entry can optionally embed its driving forecaster config inline.
   # - temporal_downscaler:
   #     checkpoint: /path/to/temporal_downscaler.ckpt
@@ -295,6 +297,60 @@ gets a station-group selector; scorecards and the standard plots use `all`.
   directory (e.g. `holdout_station_file: holdout_stations.csv` for the
   `nudge_toward_observation` filter).
 
+### Output variables
+
+By default, the inference step writes every model field to every GRIB output of the
+inference config. A `forecaster` or `temporal_downscaler` entry (including the inline
+`forecaster` of a temporal downscaler) can restrict this with `output_variables`:
+
+```yaml
+runs:
+  - forecaster:
+      checkpoint: https://...
+      steps: 0/120/6
+      output_variables: [2t, 2d, 10u, 10v, sp, tp, z, "*_850"]
+```
+
+The list is written as `variables` to every `grib` output of the inference config
+(including those inside a `tee`), replacing any `variables` set in the template.
+
+- Names are those of the fields that reach the GRIB output: usually the checkpoint's
+  variable names, as listed by `anemoi-inference metadata`. For the checkpoints tested so
+  far these are names like `2t`, `tp` and `t_850`, not the names in the GRIB files
+  (`T_2M`, `TOT_PREC`). Post-processors can rename fields or add new ones, and this
+  depends on the anemoi-inference version: from 0.11.2, a transform filter
+  (`forward_transform_filter`) passes the fields back under a name built from their
+  GRIB param, which a metadata patch can change (e.g. `10u` → `U_10M` in
+  `sgm-multidataset-ich1-oper-patch.yaml`). anemoi-inference ignores a name that matches
+  nothing, so when in doubt, check the GRIB files of a short run.
+- Shell-style globs, e.g. `"*_850"` for all fields on 850 hPa (quote them in YAML), need
+  anemoi-inference 0.13.0 or later. Older versions match names literally, so a pattern
+  matches nothing there.
+- Before inference, the prepare step checks the list against the checkpoint and stops with
+  an error if
+  - an entry matches no variable of the checkpoint (the error lists them), nor a field the
+    inference config declares for a post-processor that adds it (in `typed_variables` or
+    in the `variables_metadata` of `patch_metadata`, as the derive-diagnostic fields of
+    `sgm-temporal-downscaler-global_trimedge_multi.yaml`). An entry that matches only the
+    GRIB param name of a variable is accepted with a warning;
+  - a parameter verified from the run (`experiment.params`, and the params of enabled
+    score maps and SAL) is not written, while the model has it. A derived parameter
+    needs its components: `SP_10M` and `DD_10M` need `10u` and `10v`, `RELHUM_2M` needs
+    `2t` and `2d`, `TW_2M` needs `2t`, `2d` and `sp`. A period such as `TOT_PREC6` needs
+    `tp`. Parameters the model does not have are skipped by the verification, as without
+    `output_variables`.
+
+  The verified parameters are checked when the inference config is prepared: adding a
+  parameter to an experiment whose inference has run is not checked. Showcase parameters
+  are not checked.
+- Keep `z` (the surface geopotential, written as `FIS`): the model orography is read from
+  it for the lapse-rate correction of the verification. With an `experiment` section and
+  `lapse_rate_correction` on (the default), a run without it is rejected.
+- A forecaster that drives a temporal downscaler must keep every field the downscaler
+  reads, from both the LAM files and the `ifs-*` files. This is not checked.
+- Changing the list gives the run a new `run_id`, so inference reruns; the inference
+  environment is reused. The order of the list does not matter.
+
 
 ## Installation
 

src/evalml/config.py

@@ -1,4 +1,5 @@
 import re
+from fnmatch import fnmatchcase
 from math import isfinite
 from pathlib import Path
 from typing import Dict, List, Any, ClassVar, FrozenSet, Literal, Optional, Union
@@ -113,6 +114,19 @@ class RunConfig(BaseModel):
         False,
         description="If true, the ECCODES_DEFINITION_PATH environment variable will not be set to the COSMO local definitions.",
     )
+    output_variables: List[str] | None = Field(
+        None,
+        min_length=1,
+        description="List of variables written to the GRIB output files of the inference run. "
+        "Names are those of the fields that reach the GRIB output: the checkpoint's variable "
+        "names as listed by `anemoi-inference metadata` (e.g. 2t, tp, t_850), unless a "
+        "post-processor renames them. Before inference, the list is checked against the "
+        "checkpoint and the verified parameters. Shell-style globs (e.g. '*_850') need "
+        "anemoi-inference >= 0.13.0; older versions match names literally. The list is "
+        "written to every GRIB output of the inference config. If omitted, all variables "
+        "are written. Keep z (written as FIS): the model orography is read from it for the "
+        "lapse-rate correction.",
+    )
     config: Dict[str, Any] | str
 
     model_config = {"extra": "forbid"}
@@ -1006,6 +1020,24 @@ def validate_unique_labels(self) -> "ConfigModel":
             )
         return self
 
+    @model_validator(mode="after")
+    def validate_output_variables_keep_orography(self) -> "ConfigModel":
+        """With the lapse-rate correction, verified runs must write the orography (z)."""
+        if self.experiment is None or not self.lapse_rate_correction:
+            return self
+        for item in self.runs:
+            run = getattr(item, next(iter(type(item).model_fields)))
+            names = getattr(run, "output_variables", None)
+            if names and not any(
+                fnmatchcase(o, n) for n in names for o in ("z", "FIS")
+            ):
+                raise ValueError(
+                    f"output_variables {names} of run '{run.label or run.checkpoint}' "
+                    "lacks z: the lapse-rate correction of the verification reads the "
+                    "model orography from it. Add z, or set lapse_rate_correction: false."
+                )
+        return self
+
     @model_validator(mode="after")
     def validate_scorecard_references(self) -> "ConfigModel":
         sc = self.experiment.scorecards if self.experiment else None

src/evalml/variable_checks.py

@@ -0,0 +1,229 @@
+"""Checks of a run's `output_variables` against its checkpoint and the verified params.
+
+anemoi-inference writes a field to a GRIB output only if its name is in the output's
+`variables` list (exact match, shell-style globs from anemoi-inference 0.13.0), and
+silently ignores names that match nothing. The names are those of the fields that reach
+the output: the checkpoint's variable names, unless a post-processor renames them (from
+anemoi-inference 0.11.2, transform filters rename fields to their GRIB param names) or
+adds new fields (declared in the inference config, see `inference_config_variable_names`). These checks run before inference so that a typo, a GRIB short name
+or a forgotten component of a verified parameter fails early instead of costing an
+inference run.
+"""
+
+import json
+import logging
+import zipfile
+from fnmatch import fnmatchcase
+from pathlib import Path
+from typing import Iterable
+
+import yaml
+
+from data_input import _ICON_TO_IFS, _IFS_TO_ICON, DERIVED_PARAMS
+from data_input import parse_aggregated_param
+
+_GLOB_CHARS = set("*?[")
+
+
+def is_glob(name: str) -> bool:
+    return bool(_GLOB_CHARS & set(name))
+
+
+def checkpoint_variable_names(checkpoint: str | Path) -> tuple[set[str], set[str]]:
+    """Return the variable names of a checkpoint and the GRIB param names stored with them.
+
+    The variable names are read from the `variable_types` of each dataset in the
+    inference metadata (forcings, prognostic and diagnostic variables), or else from the
+    `variables` lists of the dataset metadata. The GRIB param names are those of the
+    `mars` request of each variable (`param`, with `_{levelist}` for a level), as
+    anemoi-inference names the fields a transform filter passes back; names equal to a
+    variable name are left out.
+
+    Returns two empty sets if the checkpoint has no anemoi metadata.
+    """
+    with zipfile.ZipFile(checkpoint) as z:
+        name = next(
+            (n for n in z.namelist() if n.endswith("anemoi-metadata/anemoi.json")), None
+        )
+        if name is None:
+            return set(), set()
+        metadata = json.loads(z.read(name))
+
+    variables: set[str] = set()
+    for dataset in (metadata.get("metadata_inference") or {}).values():
+        if isinstance(dataset, dict):
+            for names in (dataset.get("variable_types") or {}).values():
+                variables.update(names)
+    datasets = _datasets(metadata)
+    if not variables:
+        for dataset in datasets:
+            variables.update(dataset.get("variables") or [])
+    return variables, _grib_names(datasets) - variables
+
+
+def inference_config_variable_names(
+    config: dict, workdir: Path
+) -> tuple[set[str], set[str]]:
+    """Return the field names an inference config declares, and their GRIB param names.
+
+    Fields a post-processor adds (e.g. derive-diagnostic) are not in the checkpoint, but
+    the config declares them in `typed_variables` or in the `variables_metadata` of its
+    `patch_metadata` (inline or a file, relative to the working directory). The patch
+    can also change the GRIB param names of checkpoint variables.
+    """
+    variables = set(config.get("typed_variables") or {})
+    patch = config.get("patch_metadata") or {}
+    if isinstance(patch, str):
+        with open(Path(workdir) / patch) as f:
+            patch = yaml.safe_load(f) or {}
+    datasets = _datasets(patch)
+    for dataset in datasets:
+        variables.update(dataset.get("variables") or [])
+        variables.update(dataset.get("variables_metadata") or {})
+    return variables, _grib_names(datasets)
+
+
+def _datasets(metadata: dict) -> list[dict]:
+    """The per-dataset metadata entries ("data" for single-dataset checkpoints).
+
+    Older checkpoints have the dataset metadata at the top level. Only the top-level
+    `variables` of a dataset are the model's: the nested `specific` entries list the
+    source datasets before the variables were selected.
+    """
+    dataset_md = metadata.get("dataset") or {}
+    if "variables" in dataset_md or "variables_metadata" in dataset_md:
+        dataset_md = {"data": dataset_md}
+    return [d for d in dataset_md.values() if isinstance(d, dict)]
+
+
+def _grib_names(datasets: list[dict]) -> set[str]:
+    names = set()
+    for dataset in datasets:
+        for entry in (dataset.get("variables_metadata") or {}).values():
+            mars = entry.get("mars") if isinstance(entry, dict) else None
+            if isinstance(mars, dict) and mars.get("param"):
+                level = mars.get("levelist")
+                names.add(f"{mars['param']}_{level}" if level else mars["param"])
+    return names
+
+
+def _matches(entry: str, names: Iterable[str]) -> set[str]:
+    if is_glob(entry):
+        return {n for n in names if fnmatchcase(n, entry)}
+    return {entry} & set(names)
+
+
+def _icon_names(names: Iterable[str]) -> set[str]:
+    """The evalml (ICON) parameter names a set of model field names is loaded as."""
+    return {_IFS_TO_ICON.get(n, n) for n in names}
+
+
+def _provided_by(param: str, icon_names: set[str]) -> bool:
+    """Whether a verified param (e.g. TOT_PREC6, SP_10M) can be loaded from these fields."""
+    base, _ = parse_aggregated_param(param)
+    if base in icon_names:
+        return True
+    components = DERIVED_PARAMS.get(base)
+    return components is not None and set(components) <= icon_names
+
+
+def _requirement(param: str) -> str:
+    """Model names a user would write to provide a param, for error messages."""
+    base, _ = parse_aggregated_param(param)
+    names = DERIVED_PARAMS.get(base, (base,))
+    return ", ".join(_ICON_TO_IFS.get(n, n) for n in names)
+
+
+def check_output_variables(
+    output_variables: list[str],
+    variables: set[str],
+    grib_names: set[str],
+    verified_params: Iterable[str],
+    logger: logging.Logger,
+) -> None:
+    """Raise ValueError if `output_variables` would silently write less than intended.
+
+    Parameters
+    ----------
+    output_variables : list of str
+        The run's `output_variables` (names and shell-style globs).
+    variables : set of str
+        The checkpoint's variable names (see `checkpoint_variable_names`).
+    grib_names : set of str
+        GRIB param names stored with the checkpoint's variables that differ from the
+        variable names. An entry matching only one of these is accepted with a warning:
+        it is written only if a post-processor renames the field.
+    verified_params : iterable of str
+        Parameters verified from this run's output (e.g. T_2M, SP_10M, TOT_PREC6). A
+        parameter the checkpoint can provide must be written, or its components for a
+        derived parameter. Parameters the checkpoint cannot provide are not checked.
+    logger : logging.Logger
+        Logger for the warnings.
+
+    Raises
+    ------
+    ValueError
+        If an entry matches no checkpoint name, or a verified parameter the checkpoint
+        can provide is not written.
+    """
+    unknown, written = [], set()
+    for entry in output_variables:
+        matched = _matches(entry, variables)
+        if not matched and (matched := _matches(entry, grib_names)):
+            logger.warning(
+                "output_variables entry %r matches the GRIB param name of a checkpoint "
+                "variable, not a variable name: it is written only if a post-processor "
+                "renames the field to its GRIB param name (e.g. a transform filter with "
+                "anemoi-inference >= 0.11.2).",
+                entry,
+            )
+        if not matched:
+            unknown.append(entry)
+        written |= matched
+
+    errors = []
+    if unknown:
+        errors.append(
+            f"output_variables entries {unknown} match no variable of the checkpoint "
+            "and would write nothing. Use the checkpoint's variable names: "
+            f"{', '.join(sorted(variables))}."
+        )
+
+    provided = _icon_names(variables | grib_names)
+    written_icon = _icon_names(written)
+    missing = [
+        p
+        for p in dict.fromkeys(verified_params)
+        if _provided_by(p, provided) and not _provided_by(p, written_icon)
+    ]
+    if missing:
+        needs = "; ".join(f"{p} needs {_requirement(p)}" for p in missing)
+        errors.append(
+            f"Verified parameters {missing} are not in output_variables ({needs})."
+        )
+
+    if errors:
+        raise ValueError(" ".join(errors))
+
+
+def verified_params(config: dict) -> list[str]:
+    """Parameters the experiment verifies from each candidate run's GRIB output.
+
+    experiment.params (with its default) and the params of the enabled score maps and
+    SAL verification. Empty without an experiment section. The showcase params are not
+    included: a showcase section always exists (with defaults), so they cannot tell
+    whether a showcase is produced.
+    """
+    from evalml.config import ExperimentConfig, SalConfig, ScoreMapsConfig
+
+    experiment = config.get("experiment")
+    if not experiment:
+        return []
+    params = list(
+        experiment.get("params") or ExperimentConfig.model_fields["params"].default
+    )
+    for key, model in (("scoremaps", ScoreMapsConfig), ("sal", SalConfig)):
+        section = experiment.get(key) or {}
+        if section.get("enabled", False):
+            params += section.get("params") or model.model_fields["params"].default
+    return list(dict.fromkeys(params))