diff --git a/README.md b/README.md index 8fce5dee..4d187aab 100644 --- a/README.md +++ b/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 diff --git a/src/evalml/config.py b/src/evalml/config.py index 2a8a22f8..6260db95 100644 --- a/src/evalml/config.py +++ b/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 diff --git a/src/evalml/variable_checks.py b/src/evalml/variable_checks.py new file mode 100644 index 00000000..86648d9f --- /dev/null +++ b/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)) diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index 00befe23..103ee07d 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -225,3 +225,67 @@ def test_spectra_init_hours_must_not_be_empty(): with pytest.raises(ValidationError): SpectraConfig(init_hours=[]) assert SpectraConfig(init_hours=[0, 12]).init_hours == [0, 12] + + +def test_output_variables(example_config): + """output_variables is optional, takes a list of names and rejects an empty list.""" + + downscaler = example_config["runs"][0]["temporal_downscaler"] + cfg = ConfigModel.model_validate(example_config) + assert cfg.runs[0].temporal_downscaler.output_variables is None + + downscaler["output_variables"] = ["2t", "tp", "z", "*_850"] + cfg = ConfigModel.model_validate(example_config) + assert cfg.runs[0].temporal_downscaler.output_variables == [ + "2t", + "tp", + "z", + "*_850", + ] + + downscaler["output_variables"] = [] + with pytest.raises(ValueError, match="at least 1 item"): + _ = ConfigModel.model_validate(example_config) + + +def test_output_variables_on_inline_forecaster(example_config): + """The inline forecaster of a temporal downscaler accepts output_variables.""" + + forecaster = example_config["runs"][0]["temporal_downscaler"]["forecaster"] + forecaster["output_variables"] = ["2t", "tp"] + cfg = ConfigModel.model_validate(example_config) + assert cfg.runs[0].temporal_downscaler.forecaster.output_variables == ["2t", "tp"] + + forecaster["output_variables"] = [] + with pytest.raises(ValueError, match="at least 1 item"): + _ = ConfigModel.model_validate(example_config) + + +def test_output_variables_keep_orography(example_config): + """With the lapse-rate correction, a verified run's output_variables needs z.""" + + downscaler = example_config["runs"][0]["temporal_downscaler"] + downscaler["output_variables"] = ["2t", "tp"] + with pytest.raises(ValueError, match="lacks z"): + _ = ConfigModel.model_validate(example_config) + + for names in (["2t", "tp", "z"], ["2t", "*"], ["T_2M", "FIS"]): + downscaler["output_variables"] = names + _ = ConfigModel.model_validate(example_config) + + downscaler["output_variables"] = ["2t", "tp"] + example_config["lapse_rate_correction"] = False + _ = ConfigModel.model_validate(example_config) + + # without an experiment there is no verification, so no lapse-rate correction + example_config["lapse_rate_correction"] = True + del example_config["experiment"] + _ = ConfigModel.model_validate(example_config) + + +def test_output_variables_inline_forecaster_needs_no_orography(example_config): + """The inline forecaster of a downscaler is not verified as such: no z needed.""" + + forecaster = example_config["runs"][0]["temporal_downscaler"]["forecaster"] + forecaster["output_variables"] = ["2t", "tp"] + _ = ConfigModel.model_validate(example_config) diff --git a/tests/unit/test_inference_configs.py b/tests/unit/test_inference_configs.py index edd81947..bc8244ee 100644 --- a/tests/unit/test_inference_configs.py +++ b/tests/unit/test_inference_configs.py @@ -1,3 +1,6 @@ +import copy +import importlib.util +import logging from pathlib import Path import pytest @@ -7,6 +10,17 @@ (Path(__file__).resolve().parents[2] / "resources/inference/configs").glob("*.yaml") ) +SCRIPT = Path(__file__).resolve().parents[2] / "workflow/scripts/inference_prepare.py" +# inference_prepare.py loaded by path (workflow/scripts is not a package). +_spec = importlib.util.spec_from_file_location("inference_prepare", SCRIPT) +_script = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(_script) +_LOG = logging.getLogger(__name__) + + +def set_output_variables(config, variables): + return _script.set_output_variables(config, variables, _LOG) + def _rescales(node, out): """Collect (filter key, param, scale, offset) of every rescale transform filter.""" @@ -47,3 +61,102 @@ def test_rescale_has_grib_name_twin(path): path.name, param, ) + + +VARIABLES = ["2t", "10u", "tp", "*_850"] + + +def _grib_outputs(output, out): + """Collect every grib output spec of an output block (bare grib or nested tee).""" + for kind, spec in output.items(): + if kind == "grib": + out.append(spec) + elif kind == "tee": + for sub in spec.get("outputs", []) if isinstance(spec, dict) else spec: + _grib_outputs(sub, out) + return out + + +def _without_variables(config): + """Copy of a config with `variables` removed from every grib output.""" + config = copy.deepcopy(config) + for spec in _grib_outputs(config["output"], []): + spec.pop("variables", None) + return config + + +@pytest.mark.parametrize("path", CONFIGS, ids=lambda p: p.name) +def test_set_output_variables_on_templates(path): + """Every grib output of every bundled template gets the list; nothing else changes.""" + original = yaml.safe_load(path.read_text()) + config = set_output_variables(copy.deepcopy(original), VARIABLES) + + grib_outputs = _grib_outputs(config["output"], []) + assert grib_outputs + for spec in grib_outputs: + assert spec["variables"] == VARIABLES + assert _without_variables(config) == _without_variables(original) + + +def test_set_output_variables_bare_grib(): + config = {"output": {"grib": {"path": "grib/{step:03}.grib"}}, "lead_time": 6} + set_output_variables(config, ["2t"]) + assert config == { + "output": {"grib": {"path": "grib/{step:03}.grib", "variables": ["2t"]}}, + "lead_time": 6, + } + + +def test_set_output_variables_short_form_grib(): + config = {"output": {"grib": "grib/{step:03}.grib"}} + set_output_variables(config, ["2t"]) + assert config == { + "output": {"grib": {"path": "grib/{step:03}.grib", "variables": ["2t"]}} + } + + +def test_set_output_variables_nested_tee(): + config = { + "output": { + "tee": [ + {"grib": {"path": "a.grib"}}, + {"tee": [{"grib": {"path": "b.grib"}}, {"netcdf": {"path": "c.nc"}}]}, + ] + } + } + set_output_variables(config, ["2t"]) + assert config == { + "output": { + "tee": [ + {"grib": {"path": "a.grib", "variables": ["2t"]}}, + { + "tee": [ + {"grib": {"path": "b.grib", "variables": ["2t"]}}, + {"netcdf": {"path": "c.nc"}}, + ] + }, + ] + } + } + + +def test_set_output_variables_no_grib_output(caplog): + config = {"output": {"netcdf": {"path": "out.nc"}}} + original = copy.deepcopy(config) + with caplog.at_level(logging.WARNING): + set_output_variables(config, ["2t"]) + assert config == original + assert "no GRIB output" in caplog.text + + +def test_set_output_variables_overrides_template(): + config = {"output": {"grib": {"path": "a.grib", "variables": ["sp", "msl"]}}} + set_output_variables(config, ["2t"]) + assert config["output"]["grib"]["variables"] == ["2t"] + + +def test_set_output_variables_one_list_per_output(): + """Each grib output gets its own list, so the yaml dump has no anchors.""" + config = {"output": {"tee": [{"grib": {"path": "a"}}, {"grib": {"path": "b"}}]}} + set_output_variables(config, ["2t"]) + assert "&id" not in yaml.safe_dump(config) diff --git a/tests/unit/test_variable_checks.py b/tests/unit/test_variable_checks.py new file mode 100644 index 00000000..9429a9da --- /dev/null +++ b/tests/unit/test_variable_checks.py @@ -0,0 +1,194 @@ +import json +import logging +import re +import shutil +import zipfile +from pathlib import Path + +import pytest +import yaml + +from evalml.variable_checks import ( + check_output_variables, + checkpoint_variable_names, + inference_config_variable_names, + verified_params, +) + +LOG = logging.getLogger(__name__) +RESOURCES = Path(__file__).resolve().parents[2] / "resources/inference" + +VARIABLES = {"2t", "2d", "10u", "10v", "sp", "msl", "tp", "z", "t_850", "t_500"} +DEFAULT_PARAMS = ["T_2M", "TD_2M", "SP_10M", "PS", "PMSL", "TOT_PREC6"] + + +def _checkpoint(tmp_path, metadata, folder="inference-last"): + path = tmp_path / "inference-last.ckpt" + with zipfile.ZipFile(path, "w") as z: + z.writestr(f"{folder}/data.pkl", b"") + z.writestr(f"{folder}/anemoi-metadata/anemoi.json", json.dumps(metadata)) + return path + + +def test_checkpoint_variable_names(tmp_path): + """Variable names come from the inference metadata, GRIB names from the mars keys.""" + metadata = { + "metadata_inference": { + "seed": 0, + "dataset_names": ["data"], + "data": { + "variable_types": { + "forcing": ["z", "cos_latitude"], + "prognostic": ["2t", "t_850"], + "diagnostic": ["tp"], + } + }, + }, + "dataset": { + "data": { + "variables": ["2t", "cos_latitude", "t_850", "tp", "z"], + # the source datasets before selection list more variables + "specific": {"variables": ["2t", "not_selected"]}, + "variables_metadata": { + "2t": {"mars": {"param": "2t", "levtype": "sfc"}}, + "t_850": {"mars": {"param": "t", "levelist": 850}}, + "tp": {"mars": {"param": "TOT_PREC", "levelist": None}}, + }, + } + }, + } + variables, grib_names = checkpoint_variable_names(_checkpoint(tmp_path, metadata)) + assert variables == {"z", "cos_latitude", "2t", "t_850", "tp"} + # GRIB names equal to a variable name (2t, t_850) are not repeated + assert grib_names == {"TOT_PREC"} + + +def test_checkpoint_variable_names_without_inference_metadata(tmp_path): + """Older checkpoints: the dataset's variables, at the top level of `dataset`.""" + metadata = {"dataset": {"variables": ["2t", "tp"], "variables_metadata": {}}} + variables, grib_names = checkpoint_variable_names(_checkpoint(tmp_path, metadata)) + assert variables == {"2t", "tp"} + assert grib_names == set() + + +def test_checkpoint_without_metadata(tmp_path): + path = tmp_path / "plain.ckpt" + with zipfile.ZipFile(path, "w") as z: + z.writestr("archive/data.pkl", b"") + assert checkpoint_variable_names(path) == (set(), set()) + + +def test_inference_config_variable_names_inline(): + config = { + "typed_variables": {"SP_10M": {"param": "SP_10M"}}, + "patch_metadata": { + "dataset": { + "data": { + "variables_metadata": { + "PTYPE": {"mars": {"param": "PTYPE"}}, + "slor": {"mars": {"param": "SSO_SIGMA"}}, + } + } + } + }, + } + variables, grib_names = inference_config_variable_names(config, Path(".")) + assert variables == {"SP_10M", "PTYPE", "slor"} + assert grib_names == {"PTYPE", "SSO_SIGMA"} + + +def test_inference_config_variable_names_patch_file(tmp_path): + """The bundled downscaler template declares the derive-diagnostic fields in its patch.""" + template = RESOURCES / "configs/sgm-temporal-downscaler-global_trimedge_multi.yaml" + config = yaml.safe_load(template.read_text()) + # prepare_workdir copies the metadata patches to the working directory's resources/ + shutil.copytree(RESOURCES / "metadata", tmp_path / "resources") + variables, _ = inference_config_variable_names(config, tmp_path) + assert {"SP_10M", "DD_10M", "PTYPE", "TOT_SNOW"} <= variables + + +def test_inference_config_without_declarations(): + assert inference_config_variable_names({"lead_time": 6}, Path(".")) == ( + set(), + set(), + ) + + +def test_check_accepts_model_names_and_globs(): + check_output_variables( + ["2t", "2d", "10u", "10v", "sp", "msl", "tp", "z", "t_*"], + VARIABLES, + set(), + DEFAULT_PARAMS, + LOG, + ) + + +@pytest.mark.parametrize( + "names, unknown", + [ + (["T_2M", "tp"], "['T_2M']"), # GRIB short name + (["2t", "2x"], "['2x']"), # typo + (["2t", "q_*"], "['q_*']"), # glob matching nothing + ], +) +def test_check_rejects_names_matching_nothing(names, unknown): + with pytest.raises(ValueError, match=rf"entries {re.escape(unknown)} match no"): + check_output_variables(names, VARIABLES, set(), [], LOG) + + +def test_check_error_lists_the_checkpoint_names(): + with pytest.raises(ValueError, match="10u, 10v, 2d, 2t, msl, sp, t_500"): + check_output_variables(["T_2M"], VARIABLES, set(), [], LOG) + + +def test_check_grib_name_only_warns(caplog): + """A name matching only the GRIB param name of a variable is accepted with a warning.""" + with caplog.at_level(logging.WARNING): + check_output_variables(["TOT_PREC"], {"tp"}, {"TOT_PREC"}, ["TOT_PREC6"], LOG) + assert "GRIB param name" in caplog.text + + +@pytest.mark.parametrize( + "names, param, needs", + [ + (["2t", "z"], "TOT_PREC6", "tp"), + (["2t", "10u", "z"], "SP_10M", "10u, 10v"), + (["10u", "10v", "z"], "T_2M", "2t"), + ], +) +def test_check_rejects_missing_verified_params(names, param, needs): + message = f"['{param}'] are not in output_variables ({param} needs {needs})" + with pytest.raises(ValueError, match=re.escape(message)): + check_output_variables(names, VARIABLES, set(), [param], LOG) + + +def test_check_skips_params_the_checkpoint_lacks(): + """Parameters the model cannot provide are left to the verification (skipped there).""" + check_output_variables(["2t", "z"], {"2t", "z"}, set(), ["T_2M", "HZEROCL"], LOG) + + +def test_check_accepts_derived_params_written_natively(): + """A derived parameter written as such (e.g. by a post-processor) needs no components.""" + check_output_variables( + ["SP_10M", "z"], {"SP_10M", "10u", "10v", "z"}, set(), ["SP_10M"], LOG + ) + + +def test_check_reports_both_errors(): + with pytest.raises(ValueError, match=r"match no.*Verified parameters"): + check_output_variables(["T_2M", "z"], VARIABLES, set(), ["T_2M"], LOG) + + +def test_verified_params(): + assert verified_params({}) == [] + assert verified_params({"experiment": None}) == [] + base = {"stratification": {"regions": ["icon"]}} + assert verified_params({"experiment": base}) == DEFAULT_PARAMS + experiment = { + **base, + "params": ["T_2M"], + "sal": {"enabled": True, "params": ["TOT_PREC1"]}, + "scoremaps": {"enabled": False, "params": ["PMSL"]}, + } + assert verified_params({"experiment": experiment}) == ["T_2M", "TOT_PREC1"] diff --git a/workflow/rules/common.smk b/workflow/rules/common.smk index dc8a4324..34dbc74c 100644 --- a/workflow/rules/common.smk +++ b/workflow/rules/common.smk @@ -398,9 +398,15 @@ def run_specific_hash(run_config: dict, model_type: str) -> str: Changes to these fields create a new run_id (new outputs) but reuse the environment: - steps (lead times) - config YAML file contents (inference parameters) + - output_variables (variables written to the GRIB outputs), only when set, so + that run_ids of configs without it are unchanged; sorted and deduplicated, + so that the order of the list does not matter - For temporal downscalers: the forecaster's run_id (which run's outputs to read) """ - configs_to_hash = [{"steps": run_config["steps"]}] + run_fields = {"steps": run_config["steps"]} + if run_config.get("output_variables") is not None: + run_fields["output_variables"] = sorted(set(run_config["output_variables"])) + configs_to_hash = [run_fields] with open(run_config["config"], "r") as f: configs_to_hash.append(yaml.safe_load(f)) if model_type == "temporal_downscaler" and run_config.get("forecaster"): diff --git a/workflow/rules/inference.smk b/workflow/rules/inference.smk index da517e2c..935e37f9 100644 --- a/workflow/rules/inference.smk +++ b/workflow/rules/inference.smk @@ -193,6 +193,8 @@ rule inference_prepare_forecaster: reftime_to_iso=lambda wc: datetime.strptime( wc.init_time, "%Y%m%d%H%M" ).strftime("%Y-%m-%dT%H:%M"), + output_variables=lambda wc: RUN_CONFIGS[wc.run_id].get("output_variables"), + is_candidate=lambda wc: RUN_CONFIGS[wc.run_id].get("_is_candidate", False), script: "../scripts/inference_prepare.py" @@ -237,6 +239,8 @@ rule inference_prepare_temporal_downscaler: reftime_to_iso=lambda wc: datetime.strptime( wc.init_time, "%Y%m%d%H%M" ).strftime("%Y-%m-%dT%H:%M"), + output_variables=lambda wc: RUN_CONFIGS[wc.run_id].get("output_variables"), + is_candidate=lambda wc: RUN_CONFIGS[wc.run_id].get("_is_candidate", False), forecaster_run_id=lambda wc: ( "null" if RUN_CONFIGS[wc.run_id].get("forecaster") is None diff --git a/workflow/scripts/inference_prepare.py b/workflow/scripts/inference_prepare.py index 2a88f65b..5629f1c8 100644 --- a/workflow/scripts/inference_prepare.py +++ b/workflow/scripts/inference_prepare.py @@ -6,12 +6,20 @@ from pathlib import Path from evalml.helpers import setup_logger +from evalml.variable_checks import ( + check_output_variables, + checkpoint_variable_names, + inference_config_variable_names, + verified_params, +) def prepare_config( default_config_path: str, output_config_path: str, params: dict, + output_variables: list[str] | None = None, + logger: logging.Logger | None = None, ): """Prepare the configuration file for the inference run. @@ -26,17 +34,107 @@ def prepare_config( Path where the updated configuration file will be written. params : dict Dictionary of parameters to override in the default configuration. + output_variables : list of str, optional + Variables to write to every GRIB output (see `set_output_variables`). + If None, the outputs of the default configuration are left as they are. + logger : logging.Logger, optional + Logger for warnings; defaults to a logger named after this module. """ with open(default_config_path, "r") as f: config = yaml.safe_load(f) config = _override_recursive(config, params) + if output_variables: + logger = logger or logging.getLogger(__name__) + config = set_output_variables(config, output_variables, logger) with open(output_config_path, "w") as f: yaml.safe_dump(config, f, sort_keys=False) +def set_output_variables( + config: dict, variables: list[str], logger: logging.Logger +) -> dict: + """Set the variables written by every GRIB output of an anemoi-inference config. + + Sets `variables` on each `grib` output under `config["output"]`, either a bare + `grib` output or any depth of `tee` outputs, overriding a `variables` already + set in the template. If the config has no GRIB output, a warning is logged and + the config is left unchanged. + + Parameters + ---------- + config : dict + The anemoi-inference configuration, modified in place. + variables : list of str + Variable names as the GRIB output sees them; shell-style globs are allowed. + logger : logging.Logger + Logger for the warning. + + Returns + ------- + dict + The updated configuration. + """ + n_grib = _set_grib_variables(config.get("output"), variables) + if n_grib == 0: + logger.warning( + "output_variables is set, but the inference config has no GRIB output; " + "the config is left unchanged." + ) + return config + + +def _set_grib_variables(output, variables: list[str]) -> int: + """Set `variables` on every grib output in an output block; return how many.""" + if not isinstance(output, dict): + return 0 + n = 0 + for kind, spec in output.items(): + if kind == "grib": + if isinstance(spec, str): # short form: the output path only + spec = output[kind] = {"path": spec} + # one list per output, so that yaml does not write anchors and aliases + spec["variables"] = list(variables) + n += 1 + elif kind == "tee": + outputs = spec.get("outputs", []) if isinstance(spec, dict) else spec + for sub in outputs or []: + n += _set_grib_variables(sub, variables) + return n + + +def check_run_output_variables(smk, LOG: logging.Logger): + """Check the run's output_variables against its checkpoint before inference. + + Raises if an entry matches no variable of the checkpoint or of the inference + config (fields added by post-processors) or, for a run that is verified (a + candidate), if a verified parameter the model provides is not written. The + verified parameters are read from the workflow config, not from the rule's + params, so that changing them does not rerun inference. Runs after + prepare_workdir, which stages the metadata patches the config refers to. + """ + output_variables = smk.params.output_variables + if not output_variables: + return + variables, grib_names = checkpoint_variable_names(smk.input.checkpoint) + if not variables: + LOG.warning( + "No variable names found in the checkpoint metadata; " + "output_variables is not checked." + ) + return + with open(smk.input.config) as f: + config = yaml.safe_load(f) + declared, declared_grib = inference_config_variable_names(config, _get_workdir(smk)) + variables |= declared + grib_names = (grib_names | declared_grib) - variables + params = verified_params(smk.config) if smk.params.is_candidate else [] + check_output_variables(output_variables, variables, grib_names, params, LOG) + LOG.info("output_variables checked against the checkpoint and %s", params) + + def prepare_workdir(workdir: Path, resources_root: Path): """Prepare the working directory for the inference run. @@ -111,8 +209,15 @@ def prepare_temporal_downscaler(smk): ) # prepare config + check_run_output_variables(smk, LOG) overrides = _overrides_from_params(smk) - prepare_config(smk.input.config, smk.output.config, overrides) + prepare_config( + smk.input.config, + smk.output.config, + overrides, + output_variables=smk.params.output_variables, + logger=LOG, + ) LOG.info("Wrote config file at %s", smk.output.config) with open(smk.output.config, "r") as f: @@ -141,8 +246,15 @@ def prepare_forecaster(smk): LOG.info("Resources: \n%s", res_list) stage_holdout_stations(smk, workdir, LOG) + check_run_output_variables(smk, LOG) overrides = _overrides_from_params(smk) - prepare_config(smk.input.config, smk.output.config, overrides) + prepare_config( + smk.input.config, + smk.output.config, + overrides, + output_variables=smk.params.output_variables, + logger=LOG, + ) LOG.info("Wrote config file at %s", smk.output.config) with open(smk.output.config, "r") as f: config_content = f.read() diff --git a/workflow/tools/config.schema.json b/workflow/tools/config.schema.json index e6e33d60..59de9990 100644 --- a/workflow/tools/config.schema.json +++ b/workflow/tools/config.schema.json @@ -521,6 +521,23 @@ "title": "Disable Local Eccodes Definitions", "type": "boolean" }, + "output_variables": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "minItems": 1, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "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.", + "title": "Output Variables" + }, "config": { "anyOf": [ { @@ -1129,6 +1146,23 @@ "title": "Disable Local Eccodes Definitions", "type": "boolean" }, + "output_variables": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "minItems": 1, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "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.", + "title": "Output Variables" + }, "config": { "anyOf": [ {