diff --git a/.gitignore b/.gitignore index 29bc353f..f8a0f12e 100644 --- a/.gitignore +++ b/.gitignore @@ -57,3 +57,6 @@ uv.lock # evalml .evalml_snakemake_cmd.txt + +# Core dumps +core_* diff --git a/config/varda-single-1.0-inference-only.yaml b/config/varda-single-1.0-inference-only.yaml new file mode 100644 index 00000000..fd311e82 --- /dev/null +++ b/config/varda-single-1.0-inference-only.yaml @@ -0,0 +1,43 @@ +# yaml-language-server: $schema=../workflow/tools/config.schema.json +description: | + Run inference only for Varda-single-1.0 (no evaluation against observations). + +config_label: varda-single-1.0-inference-only + +dates: + start: 2025-03-01T00:00 + end: 2025-03-03T00:00 + frequency: 24h + +runs: + - temporal_downscaler: + checkpoint: https://service.meteoswiss.ch/mlstore#/models/sruc-m-2-interpolator/versions/3 + label: Varda-single-1.0 + steps: 0/120/1 + config: resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi.yaml + output: fdb + extra_requirements: + - anemoi-datasets==0.5.35 + - /scratch/mch/cosuna/evalml_fdb/anemoi-inference + forecaster: + checkpoint: https://service.meteoswiss.ch/mlstore#/models/sruc-m-1-forecaster/versions/4 + config: resources/inference/configs/sgm-multidataset-forecaster-global-ich1-oper.yaml + steps: 0/120/6 + extra_requirements: + - anemoi-inference==0.11.1 + +locations: + output_root: output/ + +profile: + executor: slurm + global_resources: + gpus: 16 + default_resources: + slurm_partition: "postproc" + cpus_per_task: 1 + mem_mb_per_cpu: 1800 + runtime: "1h" + gpus: 0 + slurm_account: s83 + jobs: 50 diff --git a/pyproject.toml b/pyproject.toml index a2587ec1..463ed699 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,6 +26,7 @@ dependencies = [ "scores>=2.0.0", "eccodes>=2.44,<2.48", "eccodes-cosmo-resources-python==2.44.0.1", + "pyfdb", # for reading VARDA-SINGLE forecast output archived directly to FDB "earthkit-data>=1.0,<1.1", "earthkit-utils>=1.0,<1.1", "earthkit-plots>=1.0,<1.1", diff --git a/resources/fdb/patch_metkit_language.py b/resources/fdb/patch_metkit_language.py new file mode 100644 index 00000000..c8a71259 --- /dev/null +++ b/resources/fdb/patch_metkit_language.py @@ -0,0 +1,75 @@ +"""Patch metkitlib's bundled language.yaml to recognize VARDA-SINGLE's custom +`model` values for FDB *retrieve* requests. + +`model` is not a standard MARS/metkit key -- it was added purely for FDB archival, +via the `marsModel` concepts in eccodes-cosmo-mars's local.215.def/local.98.def +(see resources/fdb/realtime-varda.schema and the VARDA-SINGLE FDB archival notes). +`fdb.list()` doesn't validate key values against any vocabulary, so archival and +completeness checks work fine without this patch. But `fdb.retrieve()` parses +requests through metkit's MarsLanguage, which validates enum-typed keys like +`model` against language.yaml's fixed value list -- and rejects anything not +already registered there, even though the value is a perfectly valid archived key. + +Run once per venv build, after `metkitlib` is installed (it ships its own copy of +language.yaml under its package data; there's no config-based override for this, +unlike ECCODES_DEFINITION_PATH/FDB5_CONFIG_FILE, so the file has to be edited in place). +""" + +import argparse +import re +import sys +from pathlib import Path + +# Custom VARDA-SINGLE model names, per marsModel concepts in eccodes-cosmo-mars's +# local.215.def (ICON, GPI=180) / local.98.def (IFS, GPI=180). "varda-single-g" is +# the global-domain variant seen from the temporal downscaler's IFS-grid output. +_CUSTOM_MODELS = ["varda-single", "varda-single-g"] + +# Structural anchor, not a specific value: metkitlib/share/metkit/language.yaml's +# top-level ` model:` key ends with the default (no-context) `type: enum` block, +# listing recognized model names -- its last entry (and thus the exact line to +# anchor on) shifts between metkitlib releases as upstream adds new models, so +# this matches the *end of the model: block* (a blank line before the next +# top-level ` key:`) instead of any one entry. +_BLOCK_END_RE = re.compile(r"(\n model:\n(?:.*\n)*?)(\n [a-zA-Z_]+:)") + + +def patch(language_yaml: Path) -> None: + text = language_yaml.read_text() + if all(f"- [{m}]" in text for m in _CUSTOM_MODELS): + print(f"[patch_metkit_language] {language_yaml} already patched, skipping") + return + m = _BLOCK_END_RE.search(text) + if m is None: + raise RuntimeError( + f"Could not locate the end of the 'model:' block in {language_yaml} -- " + "metkitlib's language.yaml layout may have changed; update " + "_BLOCK_END_RE in this script." + ) + addition = "\n".join(f" - [{model}]" for model in _CUSTOM_MODELS) + text = text[: m.end(1)] + addition + text[m.end(1) :] + language_yaml.write_text(text) + print(f"[patch_metkit_language] added {_CUSTOM_MODELS} to {language_yaml}") + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "venv", + type=Path, + help="Path to the venv root (contains lib/python*/site-packages)", + ) + args = parser.parse_args() + + matches = list( + args.venv.glob("lib/python*/site-packages/metkitlib/share/metkit/language.yaml") + ) + if not matches: + raise RuntimeError(f"metkitlib language.yaml not found under {args.venv}") + for language_yaml in matches: + patch(language_yaml) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/resources/fdb/realtime-varda.schema b/resources/fdb/realtime-varda.schema new file mode 100644 index 00000000..c64fd7e4 --- /dev/null +++ b/resources/fdb/realtime-varda.schema @@ -0,0 +1,107 @@ + +# * Format of the rules is: + +# [a1, a2, a3 ...[b1, b2, b3... [c1, c2, c3...]]] + +# - The first level (a) defines which attributes are used to name the top level directory +# - The second level (b) defines which attributes are used to name the data files +# - The third level (c) defines which attributes are used as index keys + +# * Rules can be grouped + +# [a1, a2, a3 ... +# [b1, b2, b3... [c1, c2, c3...]] +# [B1, B2, B3... [C1, C2, C3...]] +# ] + +# * A list of values can be given for an attribute +# [ ..., stream=enfo/efov, ... ] +# This will be used when matching rules. + +# * Attributes can be typed +# Globally, at the begining of this file: + +# refdate: Date; + +# or in the context of a rule: +# [type=cl, ... [date:ClimateMonth, ...]] + +# Typing attributes is done when the user's requests or the GRIB values need to be modified before directories, files and indexes are created. For example, ClimateMonth will transform 2010-04-01 to 'may' internally. + +# * Attributes can be optional +# [ step, levelist?, param ] +# They will be replaced internally by an empty value. It is also posiible to provide a default subtitution value: e.g. [domain?g] will consider the domain to be 'g' if missing. + +# * Attributes can be removed: +# [grid-] +# This is useful to remove attributes present in the GRIB that should not be ignored + +# * Rules are matched: + +# - If the attributes are present in the GRIB/Request, or marked optional or ignored +# - If a list of possible value is provided, one of them must match, for example +# [ class, expver, stream=enfo/efov, date, time, domain ] +# will match either stream=enfo or stream=efov, all other attributes will be matched if they exist in the GRIB or user's request + +# * On archive: +# - Attributes are extracted from the GRIB (namespace 'mars'), possibly modified by the attribute type +# - Only the first rule is used, so order is important +# - All GRIB attributes must be used by the rules, otherwise an error is raised + +# * On retrieve: +# - Attributes are extracted from the user's request, possibly modified by the attribute type (e.g. for handling of U/V) +# - All the matching rules are considered +# - Only attributes listed in the rules are used to extract values from the user's request + + +# Default types + +param: Param; +step: Step; +date: Date; +hdate: Date; +refdate: Date; +latitude: Double; +longitude: Double; +level: Double; +levelist: Double; +grid: Grid; +expver: Expver; + +time: Time; +fcmonth: Integer; + +number: Integer; +frequency: Integer; +direction: Integer; +channel: Integer; + +instrument: Integer; +ident: Integer; + +diagnostic: Integer; +iteration: Integer; +system: Integer; +method: Integer; + +# ??????? + +# reference: Integer; +# fcperiod: Integer; + +# opttime: Integer; +# leadtime: Integer; + +# quantile: ?????? +# range: ?????? + +# band: Integer; + + +######################################################## +# These rules must be first, otherwise fields of These +# classes will be index with the default rule for oper #expver?, model?, type? + +[ date, time, stream, class, expver, model, type, domain- + [ levtype, number? + [ step, param, levelist?, timespan?none ]]] diff --git a/resources/inference/configs/sgm-multidataset-forecaster-global-ich1-oper.yaml b/resources/inference/configs/sgm-multidataset-forecaster-global-ich1-oper.yaml index b6a0c4cb..6fc349c2 100644 --- a/resources/inference/configs/sgm-multidataset-forecaster-global-ich1-oper.yaml +++ b/resources/inference/configs/sgm-multidataset-forecaster-global-ich1-oper.yaml @@ -32,6 +32,7 @@ output: path: grib/{date}{time:04}_{step:03}.grib encoding: typeOfGeneratingProcess: 2 + generatingProcessIdentifier: 180 centre: lssw templates: samples: resources/templates_index_icon.yaml @@ -84,7 +85,17 @@ output: path: grib/ifs-{date}{time:04}_{step:03}.grib encoding: typeOfGeneratingProcess: 2 + generatingProcessIdentifier: 180 centre: ecmf + # localNumberOfExperiment is not available in IFS/ECMWF GRIB templates; + # omitting it means marsClass=od (GPI=180 fallback) for the global grid + # Set native marsStream/marsType codetable keys so g2_mars_labeling returns + # stream=enfo and type=fc. Cannot use a concept override in local.98.def + # for marsStream because it conflicts with the codetable key that + # g2_mars_labeling expects as input. Setting them natively is portable + # across eccodes builds. + marsStream: enfo + marsType: fc templates: samples: resources/templates_index_ifs.yaml post_processors: diff --git a/resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi.yaml b/resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi.yaml index cc1be1bb..dd1003de 100644 --- a/resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi.yaml +++ b/resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi.yaml @@ -134,6 +134,7 @@ output: path: grib/{dateTime}_{step:03}.grib encoding: typeOfGeneratingProcess: 2 + generatingProcessIdentifier: 180 templates: samples: resources/templates_index_icon.yaml post_processors: @@ -182,19 +183,33 @@ output: offset: 0 param: CLCH # here, the trimedge mask can be specified when available - # at overlap steps (multiples of the forecaster stride) replace the - # re-predicted prognostics with the forecaster's boundary values - - forward_transform_filter: - copy-prognostic-from-forecaster: - forecaster_path: forecaster/20* - common_leadtime: 6h - params_to_keep: [tp] # NOTE, if more diagnostics are added, update accordingly. - namer: *namer + # DISABLED: copy-prognostic-from-forecaster is incompatible with an + # FDB-backed forecaster (it globs the local forecaster/ dir). Waiting + # on the design decision in ecmwf/anemoi-inference#468 for a native, + # backend-agnostic replacement before re-enabling. Until then, + # overlap steps (multiples of the forecaster stride) carry the + # downscaler's own re-predicted prognostics, not the forecaster's. + # - forward_transform_filter: + # copy-prognostic-from-forecaster: + # forecaster_path: forecaster/20* + # common_leadtime: 6h + # params_to_keep: [tp] # NOTE, if more diagnostics are added, update accordingly. + # namer: *namer - grib: path: grib/ifs-{dateTime}_{step:03}.grib encoding: typeOfGeneratingProcess: 2 + generatingProcessIdentifier: 180 + # localNumberOfExperiment is not available in IFS/ECMWF GRIB templates; + # omitting it means marsClass=od (GPI=180 fallback) for the global grid + # Set native marsStream/marsType codetable keys so g2_mars_labeling returns + # stream=enfo and type=fc. Cannot use a concept override in local.98.def + # for marsStream because it conflicts with the codetable key that + # g2_mars_labeling expects as input. Setting them natively is portable + # across eccodes builds. + marsStream: enfo + marsType: fc templates: samples: resources/templates_index_ifs.yaml post_processors: @@ -204,12 +219,13 @@ output: inverse: true - assign_mask: # fill local/global overlapping points with nan mask: "global/cutout_mask" - - forward_transform_filter: - copy-prognostic-from-forecaster: - forecaster_path: forecaster/ifs* - common_leadtime: 6h - params_to_keep: [tp] # NOTE, if more diagnostics are added, update accordingly. - namer: *namer + # DISABLED: see note above the lam_0 tee's copy-prognostic-from-forecaster block. + # - forward_transform_filter: + # copy-prognostic-from-forecaster: + # forecaster_path: forecaster/ifs* + # common_leadtime: 6h + # params_to_keep: [tp] # NOTE, if more diagnostics are added, update accordingly. + # namer: *namer modifiers: - patches: - variable: diff --git a/resources/inference/templates/templates.md b/resources/inference/templates/templates.md new file mode 100644 index 00000000..22067fee --- /dev/null +++ b/resources/inference/templates/templates.md @@ -0,0 +1,27 @@ +# GRIB Template Notes + +## VARDA templates (`varda-*.grib`) + +Created for VARDA-SINGLE (generatingProcessIdentifier=180), a deterministic AI forecast model. +Derived from the `icon-ch1-*.grib` templates with the following changes: + +| Key | icon-ch1 value | varda value | Reason | +|-----|---------------|-------------|--------| +| `productDefinitionTemplateNumber` | 1 (ensemble instantaneous) / 11 (ensemble accumulation) | 0 (deterministic instantaneous) / 8 (deterministic accumulation) | VARDA-SINGLE is deterministic, not an ensemble | +| `typeOfGeneratingProcess` | 4 (ensemble forecast) | 2 (forecast) | Same reason | +| `typeOfProcessedData` | 5 (control forecast member) | 1 (forecast products) | Same reason | + +### Files + +| File | levtype | PDT | Notes | +|------|---------|-----|-------| +| `varda-typeOfLevel=isobaricInhPa.grib` | pl | 0 | Pressure level instantaneous fields | +| `varda-typeOfLevel=surface.grib` | sfc | 0 | Surface instantaneous fields (FR_LAND, PS, FIS, SKT, SSO_SIGMA, SSO_STDH, CLCT) | +| `varda-typeOfLevel=heightAboveGround.grib` | sfc | 0 | Height-above-ground fields (T_2M, TD_2M, U_10M, V_10M) | +| `varda-typeOfLevel=meanSea.grib` | sfc | 0 | Mean sea level fields (PMSL) | +| `varda-shortName=TOT_PREC.grib` | sfc | 8 | Total precipitation (temporal accumulation) | +| `varda-shortName=CLCH.grib` | pl | 0 | High cloud cover | +| `varda-shortName=CLCM.grib` | pl | 0 | Medium cloud cover | + +Low cloud cover (CLCL) uses `icon-ch1-shortName=CLCL.grib` (not a `varda-*` variant) -- +merged in from main's own template batch, kept for consistency with its siblings there. diff --git a/resources/inference/templates/templates_index_icon.yaml b/resources/inference/templates/templates_index_icon.yaml index 4a5f9198..73953b84 100644 --- a/resources/inference/templates/templates_index_icon.yaml +++ b/resources/inference/templates/templates_index_icon.yaml @@ -1,25 +1,25 @@ -# ICON-CH1 templates +# VARDA templates (PDT=0/8, typeOfGeneratingProcess=2 — deterministic forecast) - - {levtype: pl, param: [T, QV, U, V, OMEGA, FI]} - - resources/icon-ch1-typeOfLevel=isobaricInhPa.grib + - resources/varda-typeOfLevel=isobaricInhPa.grib - - {levtype: sfc, param: [T_2M, TD_2M, U_10M, V_10M]} - - resources/icon-ch1-typeOfLevel=heightAboveGround.grib + - resources/varda-typeOfLevel=heightAboveGround.grib - - {levtype: sfc, param: [FR_LAND, PS, FIS, SKT, SSO_SIGMA,SSO_STDH, CLCT]} - - resources/icon-ch1-typeOfLevel=surface.grib + - resources/varda-typeOfLevel=surface.grib - - {levtype: sfc, param: [PMSL]} - - resources/icon-ch1-typeOfLevel=meanSea.grib + - resources/varda-typeOfLevel=meanSea.grib - - {levtype: sfc, param: [TOT_PREC]} - - resources/icon-ch1-shortName=TOT_PREC.grib + - resources/varda-shortName=TOT_PREC.grib - - {levtype: sfc, param: [CLCH]} - - resources/icon-ch1-shortName=CLCH.grib + - resources/varda-shortName=CLCH.grib - - {levtype: sfc, param: [CLCM]} - - resources/icon-ch1-shortName=CLCM.grib + - resources/varda-shortName=CLCM.grib - - {levtype: sfc, param: [CLCL]} - resources/icon-ch1-shortName=CLCL.grib diff --git a/resources/inference/templates/varda-shortName=CLCH.grib b/resources/inference/templates/varda-shortName=CLCH.grib new file mode 100644 index 00000000..5da38897 Binary files /dev/null and b/resources/inference/templates/varda-shortName=CLCH.grib differ diff --git a/resources/inference/templates/varda-shortName=CLCM.grib b/resources/inference/templates/varda-shortName=CLCM.grib new file mode 100644 index 00000000..9a9c0670 Binary files /dev/null and b/resources/inference/templates/varda-shortName=CLCM.grib differ diff --git a/resources/inference/templates/varda-shortName=TOT_PREC.grib b/resources/inference/templates/varda-shortName=TOT_PREC.grib new file mode 100644 index 00000000..47deaf2d Binary files /dev/null and b/resources/inference/templates/varda-shortName=TOT_PREC.grib differ diff --git a/resources/inference/templates/varda-typeOfLevel=heightAboveGround.grib b/resources/inference/templates/varda-typeOfLevel=heightAboveGround.grib new file mode 100644 index 00000000..2aaa9491 Binary files /dev/null and b/resources/inference/templates/varda-typeOfLevel=heightAboveGround.grib differ diff --git a/resources/inference/templates/varda-typeOfLevel=isobaricInhPa.grib b/resources/inference/templates/varda-typeOfLevel=isobaricInhPa.grib new file mode 100644 index 00000000..334f7796 Binary files /dev/null and b/resources/inference/templates/varda-typeOfLevel=isobaricInhPa.grib differ diff --git a/resources/inference/templates/varda-typeOfLevel=meanSea.grib b/resources/inference/templates/varda-typeOfLevel=meanSea.grib new file mode 100644 index 00000000..f43f0748 Binary files /dev/null and b/resources/inference/templates/varda-typeOfLevel=meanSea.grib differ diff --git a/resources/inference/templates/varda-typeOfLevel=surface.grib b/resources/inference/templates/varda-typeOfLevel=surface.grib new file mode 100644 index 00000000..71053eca Binary files /dev/null and b/resources/inference/templates/varda-typeOfLevel=surface.grib differ diff --git a/src/data_input/__init__.py b/src/data_input/__init__.py index 3c8c29ce..e29b2649 100644 --- a/src/data_input/__init__.py +++ b/src/data_input/__init__.py @@ -72,9 +72,134 @@ def rename_ifs_to_icon(ds: xr.Dataset, mapping: dict[str, str]) -> xr.Dataset: return ds +# FDB schema shared by archival (workflow/rules/inference.smk) and reads (this module). +# Root marker prefix: load_forecast_data(root=...) treats an "fdb:" string the +# same way it already treats "jretrievedwh:..." for truth data -- a marker selecting a +# non-filesystem source, here an FDB root, rather than a real local directory. +_FDB_SCHEMA = ( + Path(__file__).resolve().parents[2] / "resources/fdb/realtime-varda.schema" +) +_FDB_ROOT_MARKER = "fdb:" +# VARDA-SINGLE archives two grids under distinct `model` values: "varda-single" +# (ICON-origin, regional/LAM domain) and "varda-single-g" (IFS-origin, global +# domain) -- they have different point counts and can't be merged into one +# dataset via xarray. _collect_ml_grib_files' local-file equivalent only ever +# reads the regional grid (its "20*.grib" glob excludes the "ifs-"-prefixed +# global files), so load_forecast_data's FDB path filters to the same model for +# parity. plotting.compat.load_state_from_fdb uses both (see _FDB_GLOBAL_MODEL), +# mirroring load_state_from_grib's local ifs-*.grib companion-file merge, since +# that one works with flat point lists rather than an aligned xarray grid. +_FDB_MODEL = "varda-single" +_FDB_GLOBAL_MODEL = "varda-single-g" + +# Numeric FDB "param" codes for VARDA-SINGLE's custom MARS namespace (GPI=180 +# local table, eccodes-cosmo-mars's local.215.def). FDB's retrieve() `param` +# key takes these numeric codes, not GRIB shortNames -- there's no lookup +# table for this shipped anywhere, so these were discovered empirically by +# decoding real archived messages under the correct ECCODES_DEFINITION_PATH +# (eccodes-cosmo-mars/definitions) and reading each field's `paramId`. To add +# a new one: fetch any archived field for that param, print +# `eccodes.codes_get(gid, "paramId")` with that definitions path active. +_FDB_PARAM_IDS: dict[str, str] = { + # levtype=sfc + "PS": "500000", + "PMSL": "500002", + "T_2M": "500011", + "TD_2M": "500017", + "U_10M": "500027", + "V_10M": "500029", + "TOT_PREC": "500041", + # levtype=pl + "FI": "500006", + "T": "500014", + "U": "500028", + "V": "500030", + "QV": "500035", +} +_FDB_PL_PARAMS = {"FI", "T", "U", "V", "QV"} + +# The schema's `timespan` key defaults to "none" when omitted from a request +# -- fine for instantaneous fields, but accumulated-from-start fields like +# TOT_PREC are archived with timespan="fs" and simply don't match (retrieve() +# silently returns nothing at all, for any step) if timespan isn't set +# explicitly. Confirmed empirically: TOT_PREC values this dict doesn't list +# default to "none". +_FDB_PARAM_TIMESPAN: dict[str, str] = { + "TOT_PREC": "fs", +} + +# _FDB_GLOBAL_MODEL ("varda-single-g", IFS-origin global grid) is archived +# under standard WMO/ECMWF param ids, not VARDA's custom local table -- see +# the `modifiers.patches` block in +# resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi.yaml, +# which explicitly pairs each ICON name with its archived (param, shortName). +# Pressure-level ids (T/U/V/QV/FI/OMEGA) match that same file's regex patches +# (e.g. "^t_(\\d+)$" -> param 130) -- needed when the downscaler's own model +# input reads the global grid's pressure levels directly from FDB (see +# inference_prepare.py's _convert_cutout_to_fdb), not just this module's own +# _fetch_fdb_fieldlist reads (which so far have only ever requested surface +# params from the global grid). +_FDB_GLOBAL_PARAM_IDS: dict[str, str] = { + "PS": "134", + "PMSL": "151", + "T_2M": "167", + "TD_2M": "168", + "U_10M": "165", + "V_10M": "166", + "TOT_PREC": "228", + "FR_LAND": "172", + "T": "130", + "U": "131", + "V": "132", + "QV": "133", + "FI": "129", + "OMEGA": "135", +} +_FDB_GLOBAL_PL_PARAMS: set[str] = set() +_FDB_GLOBAL_PARAM_TIMESPAN: dict[str, str] = { + "TOT_PREC": "fs", +} + +# MARS identity (class/stream/type/expver) that a GPI=180 archive resolves to for +# each model -- confirmed empirically by decoding a real archived message's `mars` +# namespace under the correct ECCODES_DEFINITION_PATH (eccodes-cosmo-mars's +# local.215.def maps GPI=180 to class=ai for the local/ICON table; local.98.def has +# no such concept for the ECMWF/IFS table, so the global grid stays class=od and +# needs marsStream/marsType set natively in its own output encoding -- see +# resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi.yaml). +# Used both when archiving (inference.smk's fdb5_config.yaml) and when a +# downstream run reads another run's FDB output directly as its own model input +# (inference_prepare.py's _convert_cutout_to_fdb). +_FDB_IDENTITY: dict[str, dict[str, str]] = { + _FDB_MODEL: { + "class": "ai", + "stream": "oper", + "type": "fc", + "expver": "0001", + "model": _FDB_MODEL, + }, + _FDB_GLOBAL_MODEL: { + "class": "od", + "stream": "enfo", + "type": "fc", + "expver": "0001", + "model": _FDB_GLOBAL_MODEL, + }, +} + +# Per-model (param_ids, pl_params, param_timespan) registries, keyed by the +# `model` value used in _fetch_fdb_fieldlist's requests. +_FDB_MODEL_REGISTRY: dict[str, tuple[dict[str, str], set[str], dict[str, str]]] = { + _FDB_MODEL: (_FDB_PARAM_IDS, _FDB_PL_PARAMS, _FDB_PARAM_TIMESPAN), + _FDB_GLOBAL_MODEL: ( + _FDB_GLOBAL_PARAM_IDS, + _FDB_GLOBAL_PL_PARAMS, + _FDB_GLOBAL_PARAM_TIMESPAN, + ), +} + XARRAY_ENGINE_PROFILE = { "ensure_dims": ["z", "number", "step", "forecast_reference_time"], - "add_valid_time_coord": True, "global_attrs": [{"institution": "MeteoSwiss"}, {"Conventions": "CF-1.8"}], } @@ -601,7 +726,11 @@ def _discover_icon_member_ids( ) -def load_from_grib_file(file: str | list[str], sel_kwargs): +def load_fieldlist_from_grib_file(file: str | list[str], sel_kwargs): + """Load an earthkit FieldList from one or more local GRIB files, filtered + by `sel_kwargs` -- the part of :func:`load_from_grib_file` that's reusable + by callers (e.g. plotting.compat) that need the fieldlist itself rather + than an already-decoded xarray Dataset.""" # Coerce Path objects to str: earthkit-data unwraps a single-element list # into one File source without converting, and then fails on non-str paths. if isinstance(file, (list, tuple)): @@ -609,14 +738,22 @@ def load_from_grib_file(file: str | list[str], sel_kwargs): else: file = str(file) fieldlist = ekd.from_source("file", file, lazily=True).to_fieldlist() - return fieldlist_to_xarray(fieldlist.sel(**sel_kwargs)) + return fieldlist.sel(**sel_kwargs) + + +def load_from_grib_file(file: str | list[str], sel_kwargs): + return fieldlist_to_xarray(load_fieldlist_from_grib_file(file, sel_kwargs)) def variable_name_profile( level_type: Literal[ - "height_above_ground_level", - "mean_sea", + "heightAboveGround", + "meanSea", "surface", + "isobaricInhPa", + "atmosphere", + "isobaricLayer", + "unknown", "pressure", "entire_atmosphere", "isothermZero", @@ -625,11 +762,29 @@ def variable_name_profile( "pressure_layer", ], ) -> dict[str, Any]: - """Resolve variable name profile based on the level type.""" + """Resolve variable name profile based on the raw GRIB `typeOfLevel` value. + + Note: earlier versions of this profile used earthkit-data's composite + "vertical.level_type"/"parameter.variable" metadata keys (semantic names + like "pressure", template "{parameter.variable}_{vertical.level}"). Those + keys don't resolve at all with the currently pinned earthkit-data (they + raise KeyError / silently match nothing) -- this now uses the raw, + always-available `typeOfLevel`/`shortName`/`levelist` GRIB keys instead. + + "isobaricLayer" is the raw-key equivalent of the semantic "pressure_layer" + level type (CLCH/CLCM: a pressure layer between two isobaric surfaces). + "unknown" covers CLCL specifically: ICON-CH1 defines its low-cloud layer + with an unset second fixed surface, so eccodes can't name a typeOfLevel + for it at all -- confirmed on the real GRIB template, not just a loading + gap (see resources/inference/templates/templates.md). + """ if level_type in [ - "height_above_ground_level", - "mean_sea", + "heightAboveGround", + "meanSea", "surface", + "atmosphere", + "isobaricLayer", + "unknown", "entire_atmosphere", # single-level ICON fields: HZEROCL, SNOWLMT, CAPE_*, CIN_* "isothermZero", @@ -638,10 +793,10 @@ def variable_name_profile( "pressure_layer", ]: return {} - elif level_type == "pressure": + elif level_type == "isobaricInhPa": return { "variable_key": "p_l", - "remapping": {"p_l": "{parameter.variable}_{vertical.level}"}, + "remapping": {"p_l": "{shortName}_{levelist}"}, } else: raise ValueError(f"Unsupported level type: {level_type}") @@ -651,11 +806,27 @@ def fieldlist_to_xarray(fieldlist) -> xr.Dataset: ds = xr.Dataset() if len(fieldlist) == 0: return ds - for level_type_group in fieldlist.group_by("vertical.level_type"): - # earthkit-data should return the group key...TODO: open issue? - level_type = level_type_group.get("vertical.level_type")[0] + # FieldList.group_by("typeOfLevel") is broken on the pinned earthkit-data + # (>=1.0,<1.1): it always yields a single group containing every field, + # regardless of how many distinct typeOfLevel values are actually present + # (reproduced in BUG_REPORT.md). Group by hand instead, using per-field + # .metadata() reads (reliable) rather than .group_by()/.sel() with the raw + # key (also unreliable on this version -- only the composite semantic keys + # like "parameter.variable" work via .sel()). + fields_by_level_type: dict[str, list] = {} + for fld in fieldlist: + fields_by_level_type.setdefault(fld.metadata("typeOfLevel"), []).append(fld) + for level_type, fields in fields_by_level_type.items(): + level_type_group = ekd.FieldList.from_fields(fields) profile = XARRAY_ENGINE_PROFILE | variable_name_profile(level_type) _ds = level_type_group.to_xarray(**profile, allow_holes=True) + # add_valid_time_coord=True (earthkit's own convenience for this) isn't + # supported together with allow_holes=True on the currently pinned + # earthkit-data (raises NotImplementedError) -- compute it by hand instead. + if "forecast_reference_time" in _ds.coords and "step" in _ds.coords: + _ds = _ds.assign_coords( + valid_time=_ds["forecast_reference_time"] + _ds["step"] + ) ds = ds.merge( _ds, compat="no_conflicts", combine_attrs="drop_conflicts", join="outer" ) @@ -759,6 +930,149 @@ def _mask_icon_missing_values(ds: xr.Dataset) -> xr.Dataset: return ds +def _fdb_config(fdb_root: str) -> dict: + return { + "type": "local", + "engine": "toc", + "schema": str(_FDB_SCHEMA.resolve()), + "spaces": [{"handler": "Default", "roots": [{"path": fdb_root}]}], + } + + +_fdb_identity_cache: dict[tuple, dict[str, tuple]] = {} + + +def _discover_fdb_identity( + fdb_root: str, model: str, date: str, time_: str +) -> dict[str, tuple]: + """Return ``{levtype: (class, stream, type, expver)}`` actually archived + for this date/time/model -- these four keys vary by experiment/deployment + (e.g. class=od for a real operational archive vs ai/rd here), unlike + param/step, which callers always specify explicitly and so never need + discovering. `list()` is metadata-only (no data transfer) and fast, but + this is still cached per (fdb_root, model, date, time) since a single + `load_forecast_data` call looks it up twice (the caller's own params, then + the elevation fallback's FIS/FI lookup). + """ + import pyfdb + + cache_key = (fdb_root, model, date, time_) + if cache_key in _fdb_identity_cache: + return _fdb_identity_cache[cache_key] + + with pyfdb.FDB(_fdb_config(fdb_root)) as fdb: + entries = list(fdb.list({"date": date, "time": time_, "model": model}, level=2)) + identities = {} + for entry in entries: + key = entry.combined_key() + identities[key["levtype"]] = ( + key["class"], + key["stream"], + key["type"], + key["expver"], + ) + _fdb_identity_cache[cache_key] = identities + return identities + + +def _fetch_fdb_fieldlist( + fdb_root: str, + reftime: datetime, + steps: list[int], + params: list[str], + model: str = _FDB_MODEL, +): + """Fetch exactly the requested fields from FDB -- one `retrieve()` per + (levtype, timespan) group, each requesting exactly the caller's params via + FDB's own `param` request key (mapped through `_FDB_MODEL_REGISTRY`), + never "every param this block happens to have" discovered via `list()` + first. + + Uses earthkit-data's native `"fdb"` source (`stream=False`, which + retrieves via pyfdb and caches the result to a local temp file) rather + than calling pyfdb directly and hand-rolling our own tempfile/byte + concatenation -- earthkit already does exactly that, correctly (verified + directly against this archive; an earlier version of this function + claimed earthkit's `"fdb"` source was buggy, which doesn't reproduce on + the currently pinned earthkit-data). + + Returns `None` if nothing is archived for this date/time/model at all. + Raises `ValueError` for any param without a known FDB code rather than + silently fetching nothing for it, or for an unregistered `model` -- + see `_FDB_MODEL_REGISTRY`. + """ + if model not in _FDB_MODEL_REGISTRY: + raise ValueError( + f"No FDB param registry for model={model!r} -- add one to _FDB_MODEL_REGISTRY." + ) + param_ids_table, pl_params, timespan_table = _FDB_MODEL_REGISTRY[model] + + unknown = [p for p in params if p not in param_ids_table] + if unknown: + raise ValueError( + f"No known FDB param id for {unknown} (model={model!r}) -- add it " + "to the relevant table in _FDB_MODEL_REGISTRY." + ) + + date = reftime.strftime("%Y%m%d") + time_ = reftime.strftime("%H%M") + identities = _discover_fdb_identity(fdb_root, model, date, time_) + if not identities: + return None + + by_group: dict[tuple[str, str], list[str]] = {} + for p in params: + levtype = "pl" if p in pl_params else "sfc" + timespan = timespan_table.get(p, "none") + by_group.setdefault((levtype, timespan), []).append(param_ids_table[p]) + + config = _fdb_config(fdb_root) + fields = [] + for (levtype, timespan), param_ids in by_group.items(): + if levtype not in identities: + continue + cls_, stream, type_, expver = identities[levtype] + request = { + "class": cls_, + "stream": stream, + "type": type_, + "expver": expver, + "model": model, + "date": date, + "time": time_, + "levtype": levtype, + "step": [str(s) for s in steps], + "param": sorted(set(param_ids)), + "timespan": timespan, + } + src = ekd.from_source("fdb", request, config=config, stream=False) + fl = src.to_fieldlist() if hasattr(src, "to_fieldlist") else src + fields.extend(fl) + + return ekd.FieldList.from_fields(fields) if fields else None + + +def _load_forecast_data_from_fdb( + fdb_root: str, reftime: datetime, steps: list[int], params: list[str] +) -> xr.Dataset: + """Load forecast data directly from an FDB store (no local GRIB files). + + External callers should use :func:`load_forecast_data` (root prefixed with + ``"fdb:"``), which handles param expansion, step expansion, IC synthesis, + and disaggregation, same as the GRIB-file path. Restricted to `_FDB_MODEL` + (the regional grid, which uses VARDA's own ICON-convention param names + directly -- no IFS-alias expansion needed here, unlike the local-grib + path, which also has to handle the global forecaster's IFS-native output). + """ + fieldlist = _fetch_fdb_fieldlist(fdb_root, reftime, steps, params, model=_FDB_MODEL) + if fieldlist is None: + raise FileNotFoundError( + f"No FDB data found for date={reftime:%Y%m%d} time={reftime:%H%M} " + f"model={_FDB_MODEL} in {fdb_root}" + ) + return fieldlist_to_xarray(fieldlist) + + def _jretrieve_df_to_xarray(df, short_names, catalog) -> xr.Dataset: """Pivot long-form jretrieve obs into a (time, values) cube aligned to the catalog, NaN-filled for missing cells.""" @@ -1556,14 +1870,45 @@ def load_forecast_data( - Plain ``TOT_PREC`` is returned as cumulative-from-start without disaggregation Routing (in order): - 1. ``*.grib`` files present in *root* → :func:`_load_forecast_data_from_grib` + 1. ``root`` starts with ``"fdb:"`` → :func:`_load_forecast_data_from_fdb` + (inference output archived directly to FDB, no local GRIB at all) + 2. ``*.grib`` files present in *root* → :func:`_load_forecast_data_from_grib` (ML inference output) - 2. ``INCA`` in path parts → :func:`_load_INCA_baseline_from_netcdf` - 3. Otherwise → ICON operational archive (via :func:`_load_icon_baseline_from_grib`) + 3. ``INCA`` in path parts → :func:`_load_INCA_baseline_from_netcdf` + 4. Otherwise → ICON operational archive (via :func:`_load_icon_baseline_from_grib`) """ + # FDB-backed runs only ever archive the narrower get_base_params set (see + # _FDB_PARAM_IDS/_FDB_GLOBAL_PARAM_IDS) -- the derive-diagnostic post-processor's + # extra fields (SP_10M, DD_10M, ...) aren't in those tables, so the richer + # get_grib_load_params set is only applied below, for the grib/icon-archive + # branches main's own routing never shared with the (main-unaware) FDB branch. + load_params = get_base_params(params) + load_steps = get_steps(steps, params) + if str(root).startswith(_FDB_ROOT_MARKER): + LOG.info("Loading forecasts from FDB...") + fdb_root = str(root)[len(_FDB_ROOT_MARKER) :] + ds = _load_forecast_data_from_fdb(fdb_root, reftime, load_steps, load_params) + # Try to derive elevation from surface geopotential (FIS/z at step 0) + # before falling back to ICON grid constants lookup. + try: + fis_ds = _load_forecast_data_from_fdb(fdb_root, reftime, [0], ["FIS"]) + except (FileNotFoundError, ValueError): + # ValueError: "FIS" (surface geopotential) has no known FDB param + # id -- it's simply never archived for this model, confirmed by + # direct inspection; this probe is expected to come up empty here + # and fall through to the ICON-constants lookup below. + fis_ds = None + if fis_ds is not None and "FIS" in fis_ds: + elevation = ekdv.geopotential_height_from_geopotential( + fis_ds["FIS"].isel(step=0, drop=True).squeeze(drop=True) + ) + ds = ds.assign_coords(elevation=elevation) + if "elevation" not in ds.coords: + ds = _try_assign_elevation(ds) + return _disaggregated_and_derived_params(ds, steps, params) + root = Path(root) load_params = get_grib_load_params(params) - load_steps = get_steps(steps, params) if any(root.glob("*.grib")): LOG.info("Loading forecasts from GRIB files...") ds = _load_forecast_data_from_grib( diff --git a/src/evalml/cli.py b/src/evalml/cli.py index bd2e96b1..3f172496 100644 --- a/src/evalml/cli.py +++ b/src/evalml/cli.py @@ -176,6 +176,32 @@ def cli(): pass +@cli.command( + help="Run inference only (no verification) as defined by a config YAML file." +) +@click.argument( + "configfile", type=click.Path(exists=True, dir_okay=False, path_type=Path) +) +@workflow_options +def inference( + configfile, cores, verbose, dry_run, unlock, report, dag, rulegraph, extra_smk_args +): + raise SystemExit( + execute_workflow( + configfile, + "inference_all", + cores, + verbose, + dry_run, + unlock, + report, + dag, + rulegraph, + extra_smk_args, + ) + ) + + @cli.command(help="Launch an experiment defined by a config YAML file.") @click.argument( "configfile", type=click.Path(exists=True, dir_okay=False, path_type=Path) diff --git a/src/evalml/config.py b/src/evalml/config.py index 919d62cd..38ce8fbb 100644 --- a/src/evalml/config.py +++ b/src/evalml/config.py @@ -71,6 +71,10 @@ class InferenceResources(BaseModel): None, description="Number of tasks per submission.", ) + srun_prefix: str | None = Field( + None, + description="Optional command to prepend to the srun invocation, e.g. 'uenv run --view=realtime fdb/5.19:v2 --'.", + ) class RunConfig(BaseModel): @@ -104,19 +108,42 @@ class RunConfig(BaseModel): description="List of extra dependencies to install for this model. " "These will be added to the requirements.txt file in the run directory.", ) - inference_resources: InferenceResources | None = Field( - None, - description="Resource requirements for inference jobs (optional; defaults handled externally).", - ) disable_local_eccodes_definitions: bool = Field( False, description="If true, the ECCODES_DEFINITION_PATH environment variable will not be set to the COSMO local definitions.", ) + output: Literal["grib", "fdb"] = Field( + "grib", + description="Where this run's forecast output is stored, and therefore where " + "downstream reads (verification, meteograms, frames, scoremaps) should look for " + "it. Must match what the referenced inference `config`'s own output.tee actually " + "writes to (an 'fdb' tee block requires output: fdb here, and vice versa) -- this " + "is the single declared source of truth for backend routing; nothing probes for " + "data presence or parses the inference config to guess it.", + ) + root: str | None = Field( + None, + description="Only meaningful when output: fdb. Root directory of a shared/global " + "FDB store this run reads from and writes into, used as-is (same pattern as " + "baseline.root) -- the caller is responsible for pointing different checkpoints " + "at different roots if they shouldn't share one. When unset (the default), this " + "run's FDB output lives in a local per-checkpoint directory under the output root.", + ) config: Dict[str, Any] | str model_config = {"extra": "forbid"} + @model_validator(mode="after") + def _root_requires_fdb_output(self) -> "RunConfig": + if self.root and self.output != "fdb": + raise ValueError( + "root is set but output is not 'fdb' -- a custom root only applies to " + "FDB-backed runs (grib output always lives under the local run " + "directory); set output: fdb or remove root." + ) + return self + @field_validator("checkpoint") def resolve_registered_model(cls, v: str) -> str: """Resolve a registered model name (e.g. 'amber-ridge') to its checkpoint path.""" @@ -185,7 +212,9 @@ class TemporalDownscalerConfig(RunConfig): forecaster: ForecasterConfig | None = Field( None, - description="Configuration for the forecaster run that this temporal downscaler is based on.", + description="Configuration for the forecaster run that this temporal downscaler is based on. " + "If the forecaster declares output: fdb, the downscaler reads its output directly from FDB " + "(inference_prepare.py's _convert_cutout_to_fdb) instead of the local grib symlink.", ) @@ -822,6 +851,10 @@ class Profile(BaseModel): """Workflow execution profile.""" executor: str = Field(..., description="Job executor, e.g. 'slurm'.") + fdb: InferenceResources | None = Field( + None, + description="FDB and inference resource settings (optional; overrides defaults for all inference jobs).", + ) global_resources: GlobalResources default_resources: DefaultResources jobs: int = Field(..., ge=1, description="Maximum number of parallel jobs.") @@ -943,7 +976,7 @@ class ConfigModel(BaseModel): ..., description="List of experiment participants, including forecaster/temporal downscaler ML runs and baselines.", ) - truth: TruthConfig | None + truth: TruthConfig | None = None lapse_rate_correction: bool = Field( default=True, description="Apply lapse-rate correction to T_2M (0.0065 K/m) and PS (11.5 Pa/m).", diff --git a/src/evalml/leadtimes.py b/src/evalml/leadtimes.py new file mode 100644 index 00000000..3dc96de3 --- /dev/null +++ b/src/evalml/leadtimes.py @@ -0,0 +1,69 @@ +import logging +import re + +# Params with no value at lead time 0. Two distinct reasons land here: +# - period-accumulated params (TOT_PREC/tp) verify a [lead - period, lead] +# window, so they have no value at lead times shorter than one step spacing +# (e.g. no 0h precip map). +# - diagnostic params (e.g. CLCT/tcc in AIFS) aren't part of the model's input state, +# so they're simply absent from the initial-state GRIB file written at step 0. +# Cloud cover is prognostic in some models (Varda-single RC); listing it here +# only drops their 0h animation frame. +# Short and canonical names both appear across the workflow (showcases vs maps). +PARAMS_WITHOUT_STEP_ZERO_VALUE = { + "TOT_PREC", + "tp", + "CLCT", + "tcc", + "CLCL", + "lcc", + "CLCM", + "mcc", + "CLCH", + "hcc", + "SSRD", + "ssrd", +} + + +def resolve_leadtimes(steps_spec, requested="all", param=None): + """Lead times to compute for a single participant. + + A run or baseline produces only the lead times in its own ``steps`` spec + (``start/stop/step``, hours). This returns those of the ``requested`` + selection that the participant actually produces — the literal ``"all"`` + (every produced lead time) or an explicit list of ints — so a 36h lead is + never requested of an ICON-CH1 baseline (steps ``0/33/6``), nor a >120h + lead of ICON-CH2. Explicitly requested lead times the participant cannot + produce are skipped with a warning. For accumulated ``param``s, lead times + shorter than one step spacing are dropped (no accumulation window). + """ + start, end, step = map(int, steps_spec.split("/")) + supported = set(range(start, end + 1, step)) + wanted = supported if requested == "all" else set(requested) + + unsupported = sorted(wanted - supported) + if unsupported: + logging.getLogger("snakemake").warning( + "Skipping lead time(s) %sh: not produced by forecast steps '%s'.", + unsupported, + steps_spec, + ) + + valid = wanted & supported + # Showcase-facing names encode an accumulation window as a numeric suffix + # (e.g. "TOT_PREC1", "TOT_PREC6"); strip it before checking membership so + # those aggregated variants are recognized the same as their base param. + base_param = re.sub(r"\d+$", "", param) if param else param + if ( + param in PARAMS_WITHOUT_STEP_ZERO_VALUE + or base_param in PARAMS_WITHOUT_STEP_ZERO_VALUE + ): + # For aggregated params (e.g. TOT_PREC6), the accumulation window encoded + # in the suffix is the minimum valid lead time — not the run's step spacing. + # At shorter leads the de-accumulation window doesn't fit (prev_lt < 0), + # so the frame would fall back to showing a raw cumulative total. + agg_match = re.search(r"\d+$", param) if param and base_param != param else None + min_lt = int(agg_match.group()) if agg_match else step + valid = {lt for lt in valid if lt >= min_lt} + return sorted(valid) diff --git a/src/plotting/__init__.py b/src/plotting/__init__.py index 6c78cfdb..4d5c8d0e 100644 --- a/src/plotting/__init__.py +++ b/src/plotting/__init__.py @@ -1,323 +1,25 @@ -from contextlib import contextmanager -from pathlib import Path +"""Lightweight package root. Deliberately avoids importing anything heavy +(cartopy, earthkit.plots, matplotlib) at module level: Python always runs a +package's __init__.py before any of its submodules, so an eager import here +would force that whole visualization stack to load even for a submodule like +plotting.compat, which only loads data and never renders anything. + +DOMAINS/get_projection/StatePlotter (plotting.viz's actual content) are +re-exported lazily via module __getattr__ (PEP 562), so `from plotting import +DOMAINS` still works exactly as before for callers that do plot -- they just +pay the cartopy/earthkit.plots import cost at that point instead of always. +""" -import cartopy.crs as ccrs -import earthkit.plots as ekp import numpy as np -from matplotlib.tri import Triangulation State = dict[str, np.ndarray | dict[str, np.ndarray]] -_PROJECTIONS: dict[str, ccrs.Projection] = { - "platecarree": ccrs.PlateCarree(), - "orthographic": ccrs.Orthographic(central_longitude=5.0, central_latitude=45.0), - # added some pojections to test the behaviour, can be deleted later - "rotatedlatlon": ccrs.RotatedPole(pole_longitude=-170.0, pole_latitude=43.0), - "azimuthalequidist": ccrs.AzimuthalEquidistant(), -} +_LAZY_VIZ_NAMES = ("DOMAINS", "get_projection", "StatePlotter", "resolve_region") -# Mapping of region names to their geographic extent and projection -# extent [lon_min, lon_max, lat_min, lat_max] in PlateCarree coordinates -def get_projection(name: str) -> "ccrs.Projection": - """Look up a projection by name.""" - if name not in _PROJECTIONS: - raise ValueError( - f"Unknown projection {name!r}. Available: {list(_PROJECTIONS)}" - ) - return _PROJECTIONS[name] +def __getattr__(name): + if name in _LAZY_VIZ_NAMES: + from . import viz - -DOMAINS = { - "globe": { - "extent": None, # full globe view - "projection": _PROJECTIONS["orthographic"], - }, - "europe": { - "extent": [-16.0, 25.0, 30.0, 65.0], - "projection": _PROJECTIONS["orthographic"], - }, - "centraleurope": { - "extent": [-1.5, 18, 41.5, 51], - "projection": _PROJECTIONS["orthographic"], - }, - "icon-ch": { - "extent": [0, 17.5, 40.5, 53.0], - "projection": _PROJECTIONS["orthographic"], - }, - "alps": { - "extent": [4.5, 16.0, 43.5, 48.5], - "projection": _PROJECTIONS["orthographic"], - }, - "radar": { - "extent": [3.2, 12.5, 43.6, 49.4], - "projection": _PROJECTIONS["orthographic"], - }, - "switzerland": { - "extent": [5.6, 10.8, 45.0, 48.6], - "projection": _PROJECTIONS["orthographic"], - }, -} - - -def resolve_region( - region_name: str, region_cfg: dict, lead_time: int = 0 -) -> tuple["ccrs.Projection", list[float] | None]: - """Return (projection, extent) for a showcase region. - - ``region_cfg`` is one entry of the dict built by ``parse_showcase_regions``. - A custom ``extent`` takes precedence; otherwise the named domain in - ``DOMAINS`` is used, rotated by ``lead_time`` when ``rotate`` is set. - """ - if region_cfg.get("extent") is not None: - projection = get_projection(region_cfg.get("projection") or "orthographic") - return projection, region_cfg["extent"] - domain = DOMAINS[region_name] - if region_cfg.get("rotate"): - base = domain["projection"].proj4_params - central_longitude = ( - base["lon_0"] + 360.0 * lead_time / region_cfg["hours_per_revolution"] - ) % 360 - # central_longitude=central_longitude, central_latitude=0.0 for zero-centered - projection = ccrs.Orthographic( - central_longitude=central_longitude, central_latitude=base["lat_0"] - ) - return projection, domain["extent"] - return domain["projection"], domain["extent"] - - -class StatePlotter: - """A class to plot state fields on various DOMAINS.""" - - def __init__( - self, - lon: np.ndarray, - lat: np.ndarray, - out_dir: Path, - ): - """Initialize the StatePlotter object. - - Latitudes and longitudes are passed during initialization so that - the triangulation can be computed once and reused for all plots. - The reference coordinate system of lon/lat is assumed to be PlateCarree() - currently. - - Parameters - ---------- - lon : np.ndarray - The longitudes of the grid. - lat : np.ndarray - The latitudes of the grid. - out_dir : Path - The output directory to save the plots. - """ - - self.lon = lon - self.lat = lat - out_dir.mkdir(exist_ok=True, parents=True) - self.out_dir = out_dir - - self.tri = Triangulation(self.lon, self.lat) - - def init_geoaxes( - self, - projection: ccrs.Projection, - bbox: list[float] | None, - nrows: int = 1, - ncols: int = 1, - size: tuple[float] | None = None, - name: str | None = None, - ) -> ekp.Figure: - """Initialize a figure with GeoAxes for plotting fields. - - Parameters - ---------- - projection : cartopy.crs.Projection - The projection used for the region. - bbox : list[float] or None - The bounding box [lon_min, lon_max, lat_min, lat_max] in PlateCarree coordinates. - If None, the full projection extent is used. - name : str, optional - The name of the region. - nrows : int, optional - The number of rows in the figure, by default 1. - ncols : int, optional - The number of columns in the figure, by default 1. - size : tuple - size of the figure in inches - - Returns - ------- - earthkit.plots.Figure - The figure object. - """ - if bbox is not None and len(bbox) != 4: - raise ValueError( - "bbox must be a list of four floats [lon_min, lon_max, lat_min, lat_max]" - ) - - ekp_fig = ekp.Figure( - crs=projection, - domain=bbox, - rows=nrows, - columns=ncols, - size=size, - ) - self.fig = ekp_fig - return ekp_fig - - def plot_field( - self, - subplot: ekp.Map, - field: np.ndarray, - style: ekp.styles.Style | None = None, - colorbar: bool = True, - title: str | None = None, - gridline_labels: bool = True, - **kwargs, - ): - """Plot a field on a Map object. - - Parameters - ---------- - subplot : earthkit.plots.Map - The Map subplot object to plot on. - field : np.ndarray - The field to plot. - style : ekp.styles.Style, optional - Earthkit.plots style for the map plot. - colorbar : bool - Whether to plot a colorbar, by default True. - title: str, optional - Map subplot title. - gridline_labels : bool - Whether to draw lat/lon degree labels on the gridlines, by - default True. Set to False for views whose center rotates - between frames (e.g. a rotating globe animation), where the - labels' varying width would otherwise make the map's position - shift from frame to frame. - kwargs : dict - Additional keyword arguments to pass to ax.tripcolor, including cmap, - vmin, vmax, etc. - """ - - proj = subplot._crs - # transform data coordinates to map coordinate reference system outside - # of the plotting function is a lot faster than letting tricontourf or - # tripcolor handle it in general, but not sure if using earthkit - # removed for now to simplify the workflow - if isinstance(proj, ccrs.Orthographic): - triang, mask = self._orthographic_tri(proj) - else: - triang, mask = self.tri, slice(None, None) - x, y = triang.x, triang.y - # TODO: this is hardcoded for when the initial state has two timesteps - # need to ditch this later - field = field[mask] - field = field[-1] if field.ndim == 2 else field.squeeze() - finite = np.isfinite(field) - # TODO: clip data to domain would make plotting faster (especially tripcolor) - # tried using Map.domain.extract() but too memory heavy (probably uses - # meshgrid in the background), implement clipping with e.g. - # points = np.column_stack((lon,lat)); mask = domain_polygon.contains_points(points) - # would work but needs handling domain crossing dateline etc. - - # TODO: tricontourf/tripcolor can handle a Triangulation when used directly, - # for some reason this doesn not work when using it with earthkit-plot, - # guess is that the earthkit-plots check throwing: - # "ValueError: x and y arrays must have the same length" is incorrect, - # therefore using x,y here - # subplot.tripcolor( # also works but is slower - # have to overwrite _plot_kwargs to avoid earthkit-plots trying to pass transform - # PlateCarree based on NumpySource - - # Normalize style and color-related kwargs - style_to_use, plot_kwargs = self._prepare_plot_kwargs(style, kwargs) - - # An all-NaN field (e.g. accumulations at lead 0) leaves nothing to - # contour: only the map layers below are drawn. - if finite.any(): - # Temporarily suppress earthkit-plots internal source-based kwargs - with self._temporary_plot_kwargs_override(subplot): - subplot.tricontourf( - x=x[finite], - y=y[finite], - z=field[finite], - style=style_to_use, - transform=proj, - **plot_kwargs, - ) # for earthkit.plots to work properly cmap and norm are needed here - # TODO: gridlines etc would be nicer to have in the init, but I didn't get - # them to overlay the plot layer - - if gridline_labels: - subplot.standard_layers() - else: - subplot.land() - subplot.coastlines() - subplot.borders() - subplot.gridlines(draw_labels=False) - - if colorbar: - subplot.legend() - if title: - subplot.title(title) - - def _prepare_plot_kwargs( - self, - style: ekp.styles.Style | None, - kwargs: dict, - ) -> tuple[ekp.styles.Style | None, dict]: - """Return a cleaned style and plot kwargs without mutating the input.""" - plot_kwargs = dict(kwargs) - - # Discrete colors mode: if explicit 'colors' provided, drop cmap - colors = plot_kwargs.get("colors", None) - if colors is not None: - plot_kwargs.pop("cmap", None) - plot_kwargs.setdefault( - "no_style", True - ) # avoid interpolation being performed by earthkit-plots resulting in an error - return style, plot_kwargs - - # Continuous mode: remove None entries so configure_style doesn't call - # style.with_overrides(colors=None, vmin=None, ...) and wipe the style's - # pre-configured values. - for key in ("colors", "levels", "cmap", "norm", "vmin", "vmax"): - if plot_kwargs.get(key) is None: - plot_kwargs.pop(key, None) - - return style, plot_kwargs - - @contextmanager - def _temporary_plot_kwargs_override(self, subplot: ekp.Map): - """Temporarily override internal _plot_kwargs to avoid transform issues.""" - has_attr = hasattr(subplot, "_plot_kwargs") - old = getattr(subplot, "_plot_kwargs", None) - subplot._plot_kwargs = lambda source: {} - try: - yield - finally: - if has_attr: - subplot._plot_kwargs = old - else: - try: - delattr(subplot, "_plot_kwargs") - except Exception: - pass - - def _orthographic_tri( - self, proj: ccrs.Projection - ) -> tuple[Triangulation, np.ndarray]: - """Compute the triangulation for an orthographic-family projection.""" - x, y, _ = proj.transform_points(ccrs.PlateCarree(), self.lon, self.lat).T - # Drop limb points falling just outside the projection limits (cartopy - # shrinks them slightly below the globe radius). Contour vertices beyond - # the limits make cartopy re-project the paths via shapely instead of - # passing them through, which can turn a band's holes inside out and - # paint it over the whole globe. - x_min, x_max = proj.x_limits - y_min, y_max = proj.y_limits - with np.errstate(invalid="ignore"): - mask = (x >= x_min) & (x <= x_max) & (y >= y_min) & (y <= y_max) - return Triangulation(x[mask], y[mask]), mask + return getattr(viz, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/src/plotting/compat.py b/src/plotting/compat.py index b07b7014..5556259c 100644 --- a/src/plotting/compat.py +++ b/src/plotting/compat.py @@ -4,7 +4,14 @@ import geopandas as gpd import numpy as np from shapely.geometry import MultiPoint -from data_input import load_from_grib_file, rename_ifs_to_icon +from data_input import ( + load_fieldlist_from_grib_file, + fieldlist_to_xarray, + rename_ifs_to_icon, + _fetch_fdb_fieldlist, + _FDB_MODEL, + _FDB_GLOBAL_MODEL, +) PARAMS_MAP = { @@ -26,28 +33,31 @@ PARAMS_MAP_INV = {v: k for k, v in PARAMS_MAP.items()} -def load_state_from_grib( - file: Path, paramlist: list[str] | None = None +def _decode_state( + regional_fieldlist, + global_fieldlist, + paramlist: list[str] | None, + source: str, ) -> dict[str, np.ndarray | dict[str, np.ndarray] | gpd.GeoSeries]: - # Also request IFS shortname aliases - paramlist_extended = list( - {p for p in (paramlist or [])} - | {PARAMS_MAP[p] for p in (paramlist or []) if p in PARAMS_MAP} - ) - ds = load_from_grib_file(file, {"parameter.variable": paramlist_extended}) + """Build a plotting state dict from the regional (lam_0) and optional + global (IFS) fieldlists for one lead time. + + This is the common decoding step shared by :func:`load_state_from_grib` + and :func:`load_state_from_fdb`, which differ only in how they fetch the + two fieldlists (local files vs FDB retrieve) -- once both have an + earthkit FieldList in hand, the xarray conversion, renaming, and + flat-point-list state assembly are identical. `source` is only used to + make the "nothing found" error message useful. + """ + ds = fieldlist_to_xarray(regional_fieldlist) # Rename any IFS shortnames back to COSMO names ds = rename_ifs_to_icon(ds, PARAMS_MAP_INV) - # TODO check if needed - # if "tp" in ifs_rename and "TOT_PREC" in ds: - # # IFS/ECMWF convention: "tp" is accumulated precip in meters. - # # Convert to kg m-2 (mm) to match the ICON-native convention. - # ds["TOT_PREC"] = ds["TOT_PREC"] * 1000 state = {} ref_param = next((p for p in (paramlist or []) if p in ds), None) if ref_param is None: raise ValueError( - f"None of the requested params {paramlist} found in {file}. " - "The GRIB file may not contain these fields at this lead time " + f"None of the requested params {paramlist} found in {source}. " + "The source may not contain these fields at this lead time " "(e.g. accumulated fields like TOT_PREC are undefined at step 0)." ) # Cast datetime64 directly: datetime.fromtimestamp() would read the epoch @@ -77,24 +87,24 @@ def load_state_from_grib( state["fields"][param] = np.full( ds["longitude"].values.size, np.nan, dtype=float ) - global_file = str(file.parent / f"ifs-{file.stem}.grib") - if Path(global_file).exists(): + + if global_fieldlist is not None: _paramlist_ecmwf = [PARAMS_MAP[p] for p in paramlist] - ds = load_from_grib_file(global_file, {"parameter.variable": _paramlist_ecmwf}) - mask = ~np.isnan(ds[_paramlist_ecmwf[0]].values.squeeze()) - global_lons = ds["longitude"].values.flatten() + gds = fieldlist_to_xarray(global_fieldlist) + mask = ~np.isnan(gds[_paramlist_ecmwf[0]].values.squeeze()) + global_lons = gds["longitude"].values.flatten() if np.max(global_lons) > 180: global_lons = ((global_lons + 180) % 360) - 180 state["longitudes"] = np.concatenate([state["longitudes"], global_lons[mask]]) state["latitudes"] = np.concatenate( - [state["latitudes"], ds["latitude"].values.flatten()[mask]] + [state["latitudes"], gds["latitude"].values.flatten()[mask]] ) for param in _paramlist_ecmwf: - if param in ds: + if param in gds: state["fields"][PARAMS_MAP_INV[param]] = np.concatenate( [ state["fields"][PARAMS_MAP_INV[param]], - ds[param].values.flatten()[mask], + gds[param].values.flatten()[mask], ] ) else: @@ -107,6 +117,70 @@ def load_state_from_grib( return state +def load_state_from_grib( + file: Path, paramlist: list[str] | None = None +) -> dict[str, np.ndarray | dict[str, np.ndarray] | gpd.GeoSeries]: + # Also request IFS shortname aliases + paramlist_extended = list( + {p for p in (paramlist or [])} + | {PARAMS_MAP[p] for p in (paramlist or []) if p in PARAMS_MAP} + ) + regional_fieldlist = load_fieldlist_from_grib_file( + file, {"parameter.variable": paramlist_extended} + ) + + global_fieldlist = None + global_file = file.parent / f"ifs-{file.stem}.grib" + if global_file.exists(): + _paramlist_ecmwf = [PARAMS_MAP[p] for p in paramlist] + global_fieldlist = load_fieldlist_from_grib_file( + global_file, {"parameter.variable": _paramlist_ecmwf} + ) + + return _decode_state( + regional_fieldlist, global_fieldlist, paramlist, source=str(file) + ) + + +def load_state_from_fdb( + fdb_root: str, reftime: datetime, step: int, paramlist: list[str] | None = None +) -> dict[str, np.ndarray | dict[str, np.ndarray] | gpd.GeoSeries]: + """FDB equivalent of :func:`load_state_from_grib` for one lead time. + + Fetches the regional (`_FDB_MODEL`) and global (`_FDB_GLOBAL_MODEL`) grids + from FDB -- exactly the requested params, via `_fetch_fdb_fieldlist` -- + and decodes them directly through the same :func:`_decode_state` shared + with the local-file path (it works with flat point lists, not an aligned + xarray grid, so the regional/global point-count mismatch that + :func:`data_input.load_forecast_data` sidesteps by dropping the global + grid entirely isn't a problem here). + """ + regional_fieldlist = _fetch_fdb_fieldlist( + fdb_root, reftime, [step], paramlist, model=_FDB_MODEL + ) + if regional_fieldlist is None: + raise FileNotFoundError( + f"No FDB data found for step={step} model={_FDB_MODEL} in {fdb_root}" + ) + try: + global_fieldlist = _fetch_fdb_fieldlist( + fdb_root, reftime, [step], paramlist, model=_FDB_GLOBAL_MODEL + ) + except ValueError: + # Some requested params may have no known id for the global grid + # (_FDB_GLOBAL_PARAM_IDS only covers what's actually been needed from + # it) -- tolerate that the same way as "nothing archived" below, + # rather than failing the whole frame over a companion grid. + global_fieldlist = None + + return _decode_state( + regional_fieldlist, + global_fieldlist, + paramlist, + source=f"{fdb_root} (step={step})", + ) + + def load_state_from_raw( file: Path, paramlist: list[str] | None = None ) -> dict[str, np.ndarray | dict[str, np.ndarray]]: diff --git a/src/plotting/viz.py b/src/plotting/viz.py new file mode 100644 index 00000000..fb2e5a64 --- /dev/null +++ b/src/plotting/viz.py @@ -0,0 +1,330 @@ +"""Visualization helpers (projections, domains, map rendering). + +Split out of plotting/__init__.py so that importing it (cartopy, +earthkit.plots, matplotlib) is only paid by callers that actually plot -- +plotting/__init__.py re-exports DOMAINS/get_projection/StatePlotter lazily +(see its module __getattr__) so that importing a sibling submodule like +plotting.compat (pure data loading, no rendering) never triggers this. +""" + +from contextlib import contextmanager +from pathlib import Path + +import cartopy.crs as ccrs +import earthkit.plots as ekp +import numpy as np +from matplotlib.tri import Triangulation + +_PROJECTIONS: dict[str, ccrs.Projection] = { + "platecarree": ccrs.PlateCarree(), + "orthographic": ccrs.Orthographic(central_longitude=5.0, central_latitude=45.0), + # added some pojections to test the behaviour, can be deleted later + "rotatedlatlon": ccrs.RotatedPole(pole_longitude=-170.0, pole_latitude=43.0), + "azimuthalequidist": ccrs.AzimuthalEquidistant(), +} + + +# Mapping of region names to their geographic extent and projection +# extent [lon_min, lon_max, lat_min, lat_max] in PlateCarree coordinates +def get_projection(name: str) -> "ccrs.Projection": + """Look up a projection by name.""" + if name not in _PROJECTIONS: + raise ValueError( + f"Unknown projection {name!r}. Available: {list(_PROJECTIONS)}" + ) + return _PROJECTIONS[name] + + +DOMAINS = { + "globe": { + "extent": None, # full globe view + "projection": _PROJECTIONS["orthographic"], + }, + "europe": { + "extent": [-16.0, 25.0, 30.0, 65.0], + "projection": _PROJECTIONS["orthographic"], + }, + "centraleurope": { + "extent": [-1.5, 18, 41.5, 51], + "projection": _PROJECTIONS["orthographic"], + }, + "icon-ch": { + "extent": [0, 17.5, 40.5, 53.0], + "projection": _PROJECTIONS["orthographic"], + }, + "alps": { + "extent": [4.5, 16.0, 43.5, 48.5], + "projection": _PROJECTIONS["orthographic"], + }, + "radar": { + "extent": [3.2, 12.5, 43.6, 49.4], + "projection": _PROJECTIONS["orthographic"], + }, + "switzerland": { + "extent": [5.6, 10.8, 45.0, 48.6], + "projection": _PROJECTIONS["orthographic"], + }, +} + + +def resolve_region( + region_name: str, region_cfg: dict, lead_time: int = 0 +) -> tuple["ccrs.Projection", list[float] | None]: + """Return (projection, extent) for a showcase region. + + ``region_cfg`` is one entry of the dict built by ``parse_showcase_regions``. + A custom ``extent`` takes precedence; otherwise the named domain in + ``DOMAINS`` is used, rotated by ``lead_time`` when ``rotate`` is set. + """ + if region_cfg.get("extent") is not None: + projection = get_projection(region_cfg.get("projection") or "orthographic") + return projection, region_cfg["extent"] + domain = DOMAINS[region_name] + if region_cfg.get("rotate"): + base = domain["projection"].proj4_params + central_longitude = ( + base["lon_0"] + 360.0 * lead_time / region_cfg["hours_per_revolution"] + ) % 360 + # central_longitude=central_longitude, central_latitude=0.0 for zero-centered + projection = ccrs.Orthographic( + central_longitude=central_longitude, central_latitude=base["lat_0"] + ) + return projection, domain["extent"] + return domain["projection"], domain["extent"] + + +class StatePlotter: + """A class to plot state fields on various DOMAINS.""" + + def __init__( + self, + lon: np.ndarray, + lat: np.ndarray, + out_dir: Path, + ): + """Initialize the StatePlotter object. + + Latitudes and longitudes are passed during initialization so that + the triangulation can be computed once and reused for all plots. + The reference coordinate system of lon/lat is assumed to be PlateCarree() + currently. + + Parameters + ---------- + lon : np.ndarray + The longitudes of the grid. + lat : np.ndarray + The latitudes of the grid. + out_dir : Path + The output directory to save the plots. + """ + + self.lon = lon + self.lat = lat + out_dir.mkdir(exist_ok=True, parents=True) + self.out_dir = out_dir + + self.tri = Triangulation(self.lon, self.lat) + + def init_geoaxes( + self, + projection: ccrs.Projection, + bbox: list[float] | None, + nrows: int = 1, + ncols: int = 1, + size: tuple[float] | None = None, + name: str | None = None, + ) -> ekp.Figure: + """Initialize a figure with GeoAxes for plotting fields. + + Parameters + ---------- + projection : cartopy.crs.Projection + The projection used for the region. + bbox : list[float] or None + The bounding box [lon_min, lon_max, lat_min, lat_max] in PlateCarree coordinates. + If None, the full projection extent is used. + name : str, optional + The name of the region. + nrows : int, optional + The number of rows in the figure, by default 1. + ncols : int, optional + The number of columns in the figure, by default 1. + size : tuple + size of the figure in inches + + Returns + ------- + earthkit.plots.Figure + The figure object. + """ + if bbox is not None and len(bbox) != 4: + raise ValueError( + "bbox must be a list of four floats [lon_min, lon_max, lat_min, lat_max]" + ) + + ekp_fig = ekp.Figure( + crs=projection, + domain=bbox, + rows=nrows, + columns=ncols, + size=size, + ) + self.fig = ekp_fig + return ekp_fig + + def plot_field( + self, + subplot: ekp.Map, + field: np.ndarray, + style: ekp.styles.Style | None = None, + colorbar: bool = True, + title: str | None = None, + gridline_labels: bool = True, + **kwargs, + ): + """Plot a field on a Map object. + + Parameters + ---------- + subplot : earthkit.plots.Map + The Map subplot object to plot on. + field : np.ndarray + The field to plot. + style : ekp.styles.Style, optional + Earthkit.plots style for the map plot. + colorbar : bool + Whether to plot a colorbar, by default True. + title: str, optional + Map subplot title. + gridline_labels : bool + Whether to draw lat/lon degree labels on the gridlines, by + default True. Set to False for views whose center rotates + between frames (e.g. a rotating globe animation), where the + labels' varying width would otherwise make the map's position + shift from frame to frame. + kwargs : dict + Additional keyword arguments to pass to ax.tripcolor, including cmap, + vmin, vmax, etc. + """ + + proj = subplot._crs + # transform data coordinates to map coordinate reference system outside + # of the plotting function is a lot faster than letting tricontourf or + # tripcolor handle it in general, but not sure if using earthkit + # removed for now to simplify the workflow + if isinstance(proj, ccrs.Orthographic): + triang, mask = self._orthographic_tri(proj) + else: + triang, mask = self.tri, slice(None, None) + x, y = triang.x, triang.y + # TODO: this is hardcoded for when the initial state has two timesteps + # need to ditch this later + field = field[mask] + field = field[-1] if field.ndim == 2 else field.squeeze() + finite = np.isfinite(field) + # TODO: clip data to domain would make plotting faster (especially tripcolor) + # tried using Map.domain.extract() but too memory heavy (probably uses + # meshgrid in the background), implement clipping with e.g. + # points = np.column_stack((lon,lat)); mask = domain_polygon.contains_points(points) + # would work but needs handling domain crossing dateline etc. + + # TODO: tricontourf/tripcolor can handle a Triangulation when used directly, + # for some reason this doesn not work when using it with earthkit-plot, + # guess is that the earthkit-plots check throwing: + # "ValueError: x and y arrays must have the same length" is incorrect, + # therefore using x,y here + # subplot.tripcolor( # also works but is slower + # have to overwrite _plot_kwargs to avoid earthkit-plots trying to pass transform + # PlateCarree based on NumpySource + + # Normalize style and color-related kwargs + style_to_use, plot_kwargs = self._prepare_plot_kwargs(style, kwargs) + + # An all-NaN field (e.g. accumulations at lead 0) leaves nothing to + # contour: only the map layers below are drawn. + if finite.any(): + # Temporarily suppress earthkit-plots internal source-based kwargs + with self._temporary_plot_kwargs_override(subplot): + subplot.tricontourf( + x=x[finite], + y=y[finite], + z=field[finite], + style=style_to_use, + transform=proj, + **plot_kwargs, + ) # for earthkit.plots to work properly cmap and norm are needed here + # TODO: gridlines etc would be nicer to have in the init, but I didn't get + # them to overlay the plot layer + + if gridline_labels: + subplot.standard_layers() + else: + subplot.land() + subplot.coastlines() + subplot.borders() + subplot.gridlines(draw_labels=False) + + if colorbar: + subplot.legend() + if title: + subplot.title(title) + + def _prepare_plot_kwargs( + self, + style: ekp.styles.Style | None, + kwargs: dict, + ) -> tuple[ekp.styles.Style | None, dict]: + """Return a cleaned style and plot kwargs without mutating the input.""" + plot_kwargs = dict(kwargs) + + # Discrete colors mode: if explicit 'colors' provided, drop cmap + colors = plot_kwargs.get("colors", None) + if colors is not None: + plot_kwargs.pop("cmap", None) + plot_kwargs.setdefault( + "no_style", True + ) # avoid interpolation being performed by earthkit-plots resulting in an error + return style, plot_kwargs + + # Continuous mode: remove None entries so configure_style doesn't call + # style.with_overrides(colors=None, vmin=None, ...) and wipe the style's + # pre-configured values. + for key in ("colors", "levels", "cmap", "norm", "vmin", "vmax"): + if plot_kwargs.get(key) is None: + plot_kwargs.pop(key, None) + + return style, plot_kwargs + + @contextmanager + def _temporary_plot_kwargs_override(self, subplot: ekp.Map): + """Temporarily override internal _plot_kwargs to avoid transform issues.""" + has_attr = hasattr(subplot, "_plot_kwargs") + old = getattr(subplot, "_plot_kwargs", None) + subplot._plot_kwargs = lambda source: {} + try: + yield + finally: + if has_attr: + subplot._plot_kwargs = old + else: + try: + delattr(subplot, "_plot_kwargs") + except Exception: + pass + + def _orthographic_tri( + self, proj: ccrs.Projection + ) -> tuple[Triangulation, np.ndarray]: + """Compute the triangulation for an orthographic-family projection.""" + x, y, _ = proj.transform_points(ccrs.PlateCarree(), self.lon, self.lat).T + # Drop limb points falling just outside the projection limits (cartopy + # shrinks them slightly below the globe radius). Contour vertices beyond + # the limits make cartopy re-project the paths via shapely instead of + # passing them through, which can turn a band's holes inside out and + # paint it over the whole globe. + x_min, x_max = proj.x_limits + y_min, y_max = proj.y_limits + with np.errstate(invalid="ignore"): + mask = (x >= x_min) & (x <= x_max) & (y >= y_min) & (y <= y_max) + return Triangulation(x[mask], y[mask]), mask diff --git a/uv.lock b/uv.lock index 589f78fc..b71c5d42 100644 --- a/uv.lock +++ b/uv.lock @@ -1272,6 +1272,7 @@ dependencies = [ { name = "mlflow" }, { name = "netcdf4" }, { name = "pydantic" }, + { name = "pyfdb" }, { name = "pyproj" }, { name = "pysteps" }, { name = "pyzmq" }, @@ -1306,6 +1307,7 @@ requires-dist = [ { name = "mlflow", specifier = ">=3.1.1" }, { name = "netcdf4", specifier = ">=1.7.2" }, { name = "pydantic", specifier = ">=2.11.7" }, + { name = "pyfdb" }, { name = "pyproj", specifier = ">=3.7.2" }, { name = "pysteps", specifier = ">=1.12" }, { name = "pyzmq", specifier = ">=27.1.0" }, @@ -1378,6 +1380,34 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cb/a8/20d0723294217e47de6d9e2e40fd4a9d2f7c4b6ef974babd482a59743694/fastjsonschema-2.21.2-py3-none-any.whl", hash = "sha256:1c797122d0a86c5cace2e54bf4e819c36223b552017172f32c5c024a6b77e463", size = 24024, upload-time = "2025-08-14T18:49:34.776Z" }, ] +[[package]] +name = "fdb5lib" +version = "5.21.4.21" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "eccodeslib" }, + { name = "eckitlib" }, + { name = "metkitlib" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/21/42/ad862302f4252ee81045a1d3b7dfd18c91afe71a565c2009ed3ead4e2828/fdb5lib-5.21.4.21-cp311-cp311-macosx_13_0_arm64.whl", hash = "sha256:47457a9e524adc1e6386af29fcfdd8134e9ce302918eb88dcd40687c9bbef52b", size = 1383780, upload-time = "2026-06-25T14:50:12.38Z" }, + { url = "https://files.pythonhosted.org/packages/1b/0d/225080639df23bbc168afbe8a315713127946581650c38104f392dd4369b/fdb5lib-5.21.4.21-cp311-cp311-macosx_13_0_x86_64.whl", hash = "sha256:c8382a26ddb024750e0ef2079900c37bbb85e0a714577930d0f6eb964a6c9792", size = 1413311, upload-time = "2026-06-25T14:39:14.826Z" }, + { url = "https://files.pythonhosted.org/packages/0c/d1/08c3a597684b40901c39cd932c7e3a11590f85f2d9c8545472ac70754223/fdb5lib-5.21.4.21-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:62a31a0f9d65ff97d87c2c9e2dd966a68f68fd34f048ee4c8c906f69e942ff89", size = 1687866, upload-time = "2026-06-25T14:35:12.348Z" }, + { url = "https://files.pythonhosted.org/packages/f6/67/3f6fb9e06f5434e3c001881eae57ea1e1fe33602cf1ed0375fdce343de44/fdb5lib-5.21.4.21-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:3b1f97c93c02d1eac2a0f7c80207942fc7ee1687117c90bf9c31b75957f38f4f", size = 1667418, upload-time = "2026-06-25T14:35:24.423Z" }, + { url = "https://files.pythonhosted.org/packages/30/b6/7cf7763e593c920b54e7be47b0a676a819b5b2aa7416e083444585eae04f/fdb5lib-5.21.4.21-cp312-cp312-macosx_13_0_arm64.whl", hash = "sha256:4328703222058fe67dc590c7957aa0b09d92887455212c0bfd851fed2edad4bb", size = 1383779, upload-time = "2026-06-25T15:07:59.25Z" }, + { url = "https://files.pythonhosted.org/packages/ce/5f/a36d1c694859bce2511b548fea4fb42175689fb44e96d999ba9f6e1c5af7/fdb5lib-5.21.4.21-cp312-cp312-macosx_13_0_x86_64.whl", hash = "sha256:fe343e453db17ae80adad14d13f169feb99a7a81baa0446cb9dc6addba4ec5eb", size = 1413310, upload-time = "2026-06-25T15:20:02.611Z" }, + { url = "https://files.pythonhosted.org/packages/b9/2e/6168244bf0da56887fb7a212276326f18b8a7553378e0d20e7f4c29f6fc0/fdb5lib-5.21.4.21-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:6d072697b93ebfda93791993a68188baee8137ea7a5fbd540333a85d80a0b9ec", size = 1687865, upload-time = "2026-06-25T14:36:11.968Z" }, + { url = "https://files.pythonhosted.org/packages/13/9a/c716cc83f3557d678de74acd7c9dbad3eeacbde227fda9d7db8bea5fbe0a/fdb5lib-5.21.4.21-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:1bd16c7515840aff8fcdf4fdb260a7c23e6d1e8c9ea54e0c7cea3a5974eba16a", size = 1667418, upload-time = "2026-06-25T14:35:51.007Z" }, + { url = "https://files.pythonhosted.org/packages/66/48/149ef9a854187423dbdf57133e2ba3655d6a784acf1bf51b9d56903c1dd8/fdb5lib-5.21.4.21-cp313-cp313-macosx_13_0_arm64.whl", hash = "sha256:afd95eab549646717c5258757661c7b30854b2b79fd6f54f3dfe7f70be0e971b", size = 1383782, upload-time = "2026-06-25T15:04:00.29Z" }, + { url = "https://files.pythonhosted.org/packages/b8/bd/27a70984ee0a14013d046108d971327632bd5a65f0b092c66b55314dab03/fdb5lib-5.21.4.21-cp313-cp313-macosx_13_0_x86_64.whl", hash = "sha256:0a1e4349c6c1e89cd8dfbd0c87964a493e2906e25f17c967cbf3e2468b44c4a8", size = 1413311, upload-time = "2026-06-25T14:40:14.193Z" }, + { url = "https://files.pythonhosted.org/packages/b0/95/2d2a435d1476c346e66950f11b4c181bca2e7ce06fbfa2ec74f3049e01af/fdb5lib-5.21.4.21-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:7fc98326abfa122103d760138e0baaa63fd509ca23d5bb266701cc722d7e9a09", size = 1687866, upload-time = "2026-06-25T14:33:09.171Z" }, + { url = "https://files.pythonhosted.org/packages/87/9a/f5645fe4c7d7e141a3f107cdd1c946f2ee8a120b7c861527f8b41ae5ceab/fdb5lib-5.21.4.21-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:f43fd9d0aa1961905a0c973471ed98c7f33867647e5178816a36d8068fab5029", size = 1667421, upload-time = "2026-06-25T14:35:40.179Z" }, + { url = "https://files.pythonhosted.org/packages/37/80/fa2b8dde32cdf73b23deaa5c7716b9fa24e652a7664903ac4a25f9cd4451/fdb5lib-5.21.4.21-cp314-cp314-macosx_13_0_arm64.whl", hash = "sha256:59715fdb92026f6ff6aa9e450da2d61fbdbd43e4b273fa92993580a9928fa43f", size = 1383781, upload-time = "2026-06-25T14:35:24.469Z" }, + { url = "https://files.pythonhosted.org/packages/c4/b1/be84089e4e930367e8ef9ef122ba483bc5ce08f486447bc2b161ad0cb8f6/fdb5lib-5.21.4.21-cp314-cp314-macosx_13_0_x86_64.whl", hash = "sha256:792f1c9e42e2ff9133c2a495b601c412c9fd9abf7afee6ca343b0b97f41b4e0b", size = 1413312, upload-time = "2026-06-25T14:59:19.479Z" }, + { url = "https://files.pythonhosted.org/packages/48/5c/e9858d63d3d69d8539cedb3f9dfaff0eda4c7a972271eb2a4c0be39f1b5c/fdb5lib-5.21.4.21-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:ac5f25d58b80c674b3b4186a06709898bb06d05195292498820d861b1237c872", size = 1687866, upload-time = "2026-06-25T14:34:53.662Z" }, + { url = "https://files.pythonhosted.org/packages/b9/20/5c5be93b09a59d5eb546fd67f723bd6f3585784d7bd081099cb2ea864a7c/fdb5lib-5.21.4.21-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:e649d51bfc7797f3989481932b7889a72b5255d9b46722ba83505a035ba5c764", size = 1667418, upload-time = "2026-06-25T14:35:26.413Z" }, +] + [[package]] name = "filelock" version = "3.29.0" @@ -2511,6 +2541,33 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979, upload-time = "2022-08-14T12:40:09.779Z" }, ] +[[package]] +name = "metkitlib" +version = "1.18.3.21" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "eccodeslib" }, + { name = "eckitlib" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/83/39/22796962d3fa1f9c1933ec435cae770a6e9c12cd84d2b60c36049ce73a6c/metkitlib-1.18.3.21-cp311-cp311-macosx_13_0_arm64.whl", hash = "sha256:2744b70ba96c0e575473147a7629d3bda32bcfb3680ee09fe80d0e874328b766", size = 948486, upload-time = "2026-06-25T14:50:14.901Z" }, + { url = "https://files.pythonhosted.org/packages/be/5f/788970e80ac0460c0ea0567c49df546d2b46f2c816be5f4f7d3b49d2d836/metkitlib-1.18.3.21-cp311-cp311-macosx_13_0_x86_64.whl", hash = "sha256:1204ce430a0a103a7a0b54ede4320643736ec9eaefb4f4a3a3184f799aed4bbd", size = 993570, upload-time = "2026-06-25T14:39:18.166Z" }, + { url = "https://files.pythonhosted.org/packages/00/75/f00cc2c7b36471c396acbaf70dfc424d94ec6e3f7ee22632d8f82714b00c/metkitlib-1.18.3.21-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:7898fca26401fe091872a3214cbdd46731b5304c3dd7b08f78c635d0c74b0f6c", size = 1038267, upload-time = "2026-06-25T14:35:14.733Z" }, + { url = "https://files.pythonhosted.org/packages/25/ce/a528262530c8c39aa8b7cbe35b6abfbc8db32dc3e75cf2aa86d000a6dde6/metkitlib-1.18.3.21-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:f2047b67f95a498d315a00a6e24cade20cc20479f7c3ca4650ad30a083d78c83", size = 1073555, upload-time = "2026-06-25T14:35:27.688Z" }, + { url = "https://files.pythonhosted.org/packages/e7/a2/775eecc029b413f019844bd8fdc6c268fdad99958e7e4246091bc2b26d2a/metkitlib-1.18.3.21-cp312-cp312-macosx_13_0_arm64.whl", hash = "sha256:6daf0c0695b3d6ee84508059e9112c7b5befe5376b55468fd3ccb87bdc0f2106", size = 948484, upload-time = "2026-06-25T15:08:01.746Z" }, + { url = "https://files.pythonhosted.org/packages/f7/98/2c28bdbc4300022418c5b05e72208434bcb608b852f1fa151d0d2cf5dbd3/metkitlib-1.18.3.21-cp312-cp312-macosx_13_0_x86_64.whl", hash = "sha256:8e1309b9dfd02b9611628f38695b3991c7c05a678566da3baab828878d004f66", size = 993570, upload-time = "2026-06-25T15:20:05.454Z" }, + { url = "https://files.pythonhosted.org/packages/22/4a/91df7945c9673e9db26c0a9fd04085ccb6c0b6287213ad9fabbda3d4241a/metkitlib-1.18.3.21-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:f93c8046368d49534844c2aaa7b45a8810ea14f291ea9241e345ddecbda22463", size = 1038267, upload-time = "2026-06-25T14:36:14.432Z" }, + { url = "https://files.pythonhosted.org/packages/22/2c/6b65b295e689ef40605c0da4e58f6042cc56061132cacefb9ba668da144d/metkitlib-1.18.3.21-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:42560798e800dec6f251c8d947d4560422dbb33cc139436954cbee674ebf8b18", size = 1073556, upload-time = "2026-06-25T14:35:54.441Z" }, + { url = "https://files.pythonhosted.org/packages/50/3a/aa7cba991330a143334f3a25958790df045a425e8f62fb6608811216b1b4/metkitlib-1.18.3.21-cp313-cp313-macosx_13_0_arm64.whl", hash = "sha256:ed70c936f9ef1ec898c7d2d4450548da9c51a289bd15ee5abbab4ac1a5b84436", size = 948486, upload-time = "2026-06-25T15:04:02.689Z" }, + { url = "https://files.pythonhosted.org/packages/b0/4e/7c28c315d4be71270120631af02b87beca7674f3b1e4e9be404a7db62a11/metkitlib-1.18.3.21-cp313-cp313-macosx_13_0_x86_64.whl", hash = "sha256:308e3e2084da51650b09657731b911f18e4edef26f16548c50c2125b61666bba", size = 993575, upload-time = "2026-06-25T14:40:17.046Z" }, + { url = "https://files.pythonhosted.org/packages/cc/67/095a563914a43250430f3c69fd235d0a0d5e8bc405e3ce2c97f2874e6a83/metkitlib-1.18.3.21-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:0155b398de2bf16759495718d9ebf43463865c2b1444f0c837ce158d67111c4c", size = 1038265, upload-time = "2026-06-25T14:33:11.495Z" }, + { url = "https://files.pythonhosted.org/packages/1f/41/e65cb9cb3949440a5b691adc2035f8897863b802ee747bddf49a2e8eb811/metkitlib-1.18.3.21-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:1d47fcd78529e67a72fad30ee7df85f2a4455245ae58cd9c38a2d3b3902311d0", size = 1073557, upload-time = "2026-06-25T14:35:43.43Z" }, + { url = "https://files.pythonhosted.org/packages/8d/b4/d99fb6f49937958ce1f2bfaa98c407b7d38e2be7dacbde51a0173cdd2504/metkitlib-1.18.3.21-cp314-cp314-macosx_13_0_arm64.whl", hash = "sha256:b2a793f597766ec1f1dfc510706ec6c7c061b5150c0f2d9f0f9914d28ab3442b", size = 948485, upload-time = "2026-06-25T14:35:28.261Z" }, + { url = "https://files.pythonhosted.org/packages/34/47/ec16aafb2e3ceea1725a5aaf33362b378c2e327c30fe96c8347ff1d94b61/metkitlib-1.18.3.21-cp314-cp314-macosx_13_0_x86_64.whl", hash = "sha256:a41b4d7a723237f2b76adc566f589d4f6f8ddaa42b8d83aaa513fdfd405f0557", size = 993573, upload-time = "2026-06-25T14:59:22.827Z" }, + { url = "https://files.pythonhosted.org/packages/4d/d0/b19d70b208731dbb88be918c066f369e5eda6d4f056e9f1dd212b134264f/metkitlib-1.18.3.21-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:13873e85ee6a0c561e5c1d0421bc80540f0ef06be3b710f53a6be6019c84397e", size = 1038265, upload-time = "2026-06-25T14:34:56.306Z" }, + { url = "https://files.pythonhosted.org/packages/fd/78/c5235fe5cd54fad3e96aa2b9e5f616dd121fbfa68283c2744d208a080606/metkitlib-1.18.3.21-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:a20603e0c9726ad5fa4f5f2a5633c6cdf6d3651d26807d8cfa812aa904c5d001", size = 1073557, upload-time = "2026-06-25T14:35:30.473Z" }, +] + [[package]] name = "mir-python" version = "1.29.0.21" @@ -3695,6 +3752,34 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/4b/2d/69abac8f838090bbecd5df894befb2c2619e7996a98ddb949db9f3b93225/pydantic_core-2.46.4-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:d51026d73fcfd93610abc7b27789c26b313920fcfb20e27462d74a7f8b06e983", size = 2193071, upload-time = "2026-05-06T13:38:08.682Z" }, ] +[[package]] +name = "pyfdb" +version = "5.21.4.21" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "fdb5lib" }, + { name = "findlibs" }, + { name = "pyyaml" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/fb/da/8a15cef3fd46a36621b51916b17b57bdbd94464b66efd93ef9118afa57b2/pyfdb-5.21.4.21-cp311-cp311-macosx_13_0_arm64.whl", hash = "sha256:13ac8d8d550d13d3780616beef413d8508f345e20c40f9d782be193be0b139d5", size = 149245, upload-time = "2026-06-25T14:50:24.918Z" }, + { url = "https://files.pythonhosted.org/packages/9e/d0/878a420309deb9f5fe3fce79322cd2f83f203715d3adce8ac95c774de949/pyfdb-5.21.4.21-cp311-cp311-macosx_13_0_x86_64.whl", hash = "sha256:2c4f12a493153771df25ea36b898f0785c23af8c5a80c6551ba2df0c4d46a996", size = 156545, upload-time = "2026-06-25T14:39:28.737Z" }, + { url = "https://files.pythonhosted.org/packages/ab/7d/f3a912f9cf22b4aab9828d720346c81635e304d925cc35c852a508f86297/pyfdb-5.21.4.21-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:de41c16f3c0a3b4c4d517fe48597a145b475df413bfcb23ce121d19bf09fc6e4", size = 143177, upload-time = "2026-06-25T14:35:23.566Z" }, + { url = "https://files.pythonhosted.org/packages/f1/ef/f11938857d4947e079661d03ec2814ccfa7f0cd9cd5e0455bffe798ed450/pyfdb-5.21.4.21-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:ea0eb7f2410c8e56a6f2493de535c8d7e323bf97ec656627ba1b145f80b8e97b", size = 149461, upload-time = "2026-06-25T14:35:40.975Z" }, + { url = "https://files.pythonhosted.org/packages/5b/8f/29a20431c99615dcac06627a9a9b9560e12a9a88751fe95d5f97696a34d3/pyfdb-5.21.4.21-cp312-cp312-macosx_13_0_arm64.whl", hash = "sha256:f4c386a219bb39952e890f550eeb52d5ed2cbcd92b14fcc9102a3ccca21311d2", size = 166134, upload-time = "2026-06-25T15:08:10.709Z" }, + { url = "https://files.pythonhosted.org/packages/d7/d6/27569905374b1c927066b742a3e22f76276cfee55d5594fa1ab8a440ff05/pyfdb-5.21.4.21-cp312-cp312-macosx_13_0_x86_64.whl", hash = "sha256:111db0977abd0ac822246ff06002b504e715f3debc1ae1f103a61b639f851e0c", size = 175417, upload-time = "2026-06-25T15:20:16.483Z" }, + { url = "https://files.pythonhosted.org/packages/28/29/66a73a0d052d94e3764fd338846515fef9bc80cd7a647054ed7606712554/pyfdb-5.21.4.21-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:6fa072926285b0b97bb097c1a8317c21ef0176887ba85957245d7cfb315b379f", size = 143569, upload-time = "2026-06-25T14:36:23.404Z" }, + { url = "https://files.pythonhosted.org/packages/83/73/3159c403f5e36b39ae90848924df531ba8aae067e7ad7bab4c3d5d885984/pyfdb-5.21.4.21-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:01be2c0f5722746c689f5a905b7100dc0dbcfa5baf52fc57385b48080d248563", size = 150025, upload-time = "2026-06-25T14:36:06.596Z" }, + { url = "https://files.pythonhosted.org/packages/a4/a5/d7ede9ce50b84f38251225b7aa9d4b718738f1ec38bbee304b482a75125c/pyfdb-5.21.4.21-cp313-cp313-macosx_13_0_arm64.whl", hash = "sha256:90d0fced85646b3813cbf02c67dad3e459d8e881403c62d5507dc25a20769e2b", size = 166136, upload-time = "2026-06-25T15:04:11.847Z" }, + { url = "https://files.pythonhosted.org/packages/6e/cc/26717b54d031fc5b954e9ccae4b1a16e580cd9bef39825498381fe4ff9b9/pyfdb-5.21.4.21-cp313-cp313-macosx_13_0_x86_64.whl", hash = "sha256:17c408cc9d9f2d63ebb6dda993336b572d124a9a34e0edc83410c244a9161038", size = 175435, upload-time = "2026-06-25T14:40:28.179Z" }, + { url = "https://files.pythonhosted.org/packages/b8/82/a26a36ff544620e2f27a529d872d71b5b925083934baab228a4df1b1d790/pyfdb-5.21.4.21-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:1e0b85ec7ee5963abacda33e8c08a71415ff3d73d03ca9ab9edf6a1919b522b6", size = 143606, upload-time = "2026-06-25T14:33:20.076Z" }, + { url = "https://files.pythonhosted.org/packages/cf/0b/21be146d12fc8b4e42a62ef2f64b8d3e97c68c54998a7d53c85028358361/pyfdb-5.21.4.21-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:83d2425d33a2bacdf0b05136a374c05be0dd34f1bf34e4b4996f1d8ab6e38645", size = 149988, upload-time = "2026-06-25T14:35:55.227Z" }, + { url = "https://files.pythonhosted.org/packages/7b/27/ce41bef04b9aa5ecce4914a10cd52e0e2d2e117fa32646b833861d20674f/pyfdb-5.21.4.21-cp314-cp314-macosx_13_0_arm64.whl", hash = "sha256:7924fd9176597fc51cd74e6037632af938795c25d26fd32b9a0a8ab1673e6744", size = 165426, upload-time = "2026-06-25T14:35:38.449Z" }, + { url = "https://files.pythonhosted.org/packages/e6/42/4ee89d919b03732cdb4cfda7572794cc1766a52ab7827cf0e8a44cb9803e/pyfdb-5.21.4.21-cp314-cp314-macosx_13_0_x86_64.whl", hash = "sha256:21c19111739778c1381cba57dd0e3e03a4b3991a4f884e1887fa7552853205d2", size = 175474, upload-time = "2026-06-25T14:59:33.672Z" }, + { url = "https://files.pythonhosted.org/packages/f6/9c/936cdaad1e92753b92f7cc8acc5cb77fc470c15f4be748c0dfb0cc0228da/pyfdb-5.21.4.21-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:4fbb88afaae3f3f7152f2deebe962ae42a93baa659b341c91eade500ac2a4f9c", size = 143661, upload-time = "2026-06-25T14:35:05.139Z" }, + { url = "https://files.pythonhosted.org/packages/2c/c5/1d11a305af088097983401130c78eb611f5dacc2318e8db2971065364a53/pyfdb-5.21.4.21-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:49d4ec5137482101f207468e08cb0a7c7755485630ad682633fe04fc459edc1d", size = 150001, upload-time = "2026-06-25T14:35:43.604Z" }, +] + [[package]] name = "pygments" version = "2.20.0" diff --git a/workflow/Snakefile b/workflow/Snakefile index 7a3809bc..8d461fff 100644 --- a/workflow/Snakefile +++ b/workflow/Snakefile @@ -3,6 +3,7 @@ # https://snakemake.readthedocs.io/en/stable/snakefiles/best_practices.html, import time +from types import SimpleNamespace from evalml.config import ConfigModel from evalml.failures import summarize_failures @@ -38,7 +39,8 @@ CONFIG_LABEL = config["config_label"] or CONFIG_FILE.stem EXPERIMENT_NAME = f"{WHEN}_{CONFIG_LABEL}_{CONFIG_HASH}" CANDIDATE_LABELS = ", ".join(v.get("label", k) for k, v in CANDIDATES.items()) BASELINE_LABELS = ", ".join(v.get("label", k) for k, v in BASELINES.items()) -GROUNDTRUTH_LABEL = config["truth"].get("label", config["truth"]["root"]) +_truth_cfg = config.get("truth") +GROUNDTRUTH_LABEL = _truth_cfg.get("label", _truth_cfg["root"]) if _truth_cfg else None DATA_DIR = OUT_ROOT / "data" LOGS_DIR = OUT_ROOT / "logs" @@ -305,14 +307,20 @@ rule sandbox_all: ), +def _inference_all_okfiles(wc): + return [ + _get_inference_okfile( + SimpleNamespace(run_id=run_id, init_time=t.strftime("%Y%m%d%H%M")) + ) + for run_id in CANDIDATES + for t in REFTIMES + ] + + rule inference_all: """Run inference for all reference times as defined in the configuration.""" input: - expand( - rules.inference_execute.output.okfile, - init_time=[t.strftime("%Y%m%d%H%M") for t in REFTIMES], - run_id=CANDIDATES, - ), + _inference_all_okfiles, output: run_all_ok=touch(OUT_ROOT / "logs/run_inference_all.ok"), shell: diff --git a/workflow/rules/common.smk b/workflow/rules/common.smk index 2ad240ec..3789c1ee 100644 --- a/workflow/rules/common.smk +++ b/workflow/rules/common.smk @@ -27,7 +27,7 @@ ENV_HASH_FIELDS = { "disable_local_eccodes_definitions", } # Fields excluded from ALL hashing (display/resource metadata only). -RUN_HASH_EXCLUDE = {"label", "inference_resources", "_is_candidate", "model_type"} +RUN_HASH_EXCLUDE = {"label", "_is_candidate", "model_type", "fdb_root"} # Fields excluded from baseline hashing (display metadata only). BASELINE_HASH_EXCLUDE = {"label"} TRUTH_HASH_EXCLUDE = {"label"} @@ -398,11 +398,17 @@ 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 (grib/fdb), only when non-default -- see below - For temporal downscalers: the forecaster's run_id (which run's outputs to read) """ configs_to_hash = [{"steps": run_config["steps"]}] with open(run_config["config"], "r") as f: configs_to_hash.append(yaml.safe_load(f)) + # Hashed only when set to the non-default "fdb": hashing it unconditionally would + # change run_id (and thus the output directory) for every existing run in the repo + # the moment this field shipped, making every already-computed grib run look stale. + if run_config.get("output", "grib") != "grib": + configs_to_hash.append({"output": run_config["output"]}) if model_type == "temporal_downscaler" and run_config.get("forecaster"): # run output depends on which forecaster RUN was used (not just the env) configs_to_hash.append(run_config["forecaster"].get("run_id")) @@ -428,8 +434,10 @@ def master_hash() -> str: return generate_json_hash(configs_to_hash) -def truth_hash(truth_config: dict) -> str: +def truth_hash(truth_config: dict | None) -> str: """Generate a short hash of the configs for the truth data.""" + if truth_config is None: + return generate_json_hash({}) cfg = {k: v for k, v in truth_config.items() if k not in TRUTH_HASH_EXCLUDE} return generate_json_hash(cfg) @@ -442,9 +450,11 @@ def verif_hash(full_config: dict) -> str: changing the station_holdout selection) produces new output paths and unconditionally triggers a rerun. """ - truth_cfg = { - k: v for k, v in full_config["truth"].items() if k not in TRUTH_HASH_EXCLUDE - } + truth_cfg = ( + {k: v for k, v in full_config["truth"].items() if k not in TRUTH_HASH_EXCLUDE} + if full_config.get("truth") is not None + else {} + ) experiment_verif_cfg = { "lapse_rate_correction": full_config.get("lapse_rate_correction", True), "station_holdout": HOLDOUT_STATIONS, @@ -515,7 +525,7 @@ elif TRUTH_TYPE == "inca": inca.check_prerequisites(config["truth"]["root"]) -TRUTH_HASH = truth_hash(config["truth"]) +TRUTH_HASH = truth_hash(config.get("truth")) HOLDOUT_STATIONS = sorted(EXPERIMENT_CONFIG.get("station_holdout") or []) HOLDOUT_FILE = ( OUT_ROOT / f"data/station_holdout/{generate_json_hash(HOLDOUT_STATIONS)}.csv" diff --git a/workflow/rules/inference.smk b/workflow/rules/inference.smk index da517e2c..b844daef 100644 --- a/workflow/rules/inference.smk +++ b/workflow/rules/inference.smk @@ -5,6 +5,7 @@ import os from pathlib import Path from datetime import datetime +from types import SimpleNamespace rule inference_get_checkpoint: @@ -83,8 +84,11 @@ rule inference_extract_requirements: # Prepare the inference environment for a specific checkpoint. The venv is built # in /dev/shm (RAM) to avoid heavy IO on the parallel filesystem, then squashed -# to a .squashfs image tracked directly as the rule output. -# See https://docs.cscs.ch/guides/storage/#python-virtual-environments-with-uenv. +# to a .squashfs image tracked directly as the rule output. FDB native libraries +# are bundled straight into this image via pip (pyfdb pulls in fdb5lib/ +# eccodeslib/eckitlib/metkitlib, all manylinux wheels with vendored .so's) -- +# no separate FDB uenv squashfs needs mounting at run time, matching how the +# main evalml venv already reads FDB in verification.smk/plot.smk. rule inference_prepare_env: input: metadata=OUT_ROOT / "data/runs/{env_id}/anemoi.json", @@ -94,6 +98,8 @@ rule inference_prepare_env: log: OUT_ROOT / "logs/inference_prepare_env/{env_id}.log", localrule: True + params: + needs_fdb=lambda wc: _env_needs_fdb(wc.env_id), shell: """ ( @@ -109,6 +115,35 @@ rule inference_prepare_env: echo "[$(date)] Installing requirements from {input.requirements}..." uv pip install -r {input.requirements} + # pyfdb's own wheel dependencies (fdb5lib etc.) bundle the native FDB + # stack, so this is just another pip install -- no uenv involved. + if [ "{params.needs_fdb}" = "True" ]; then + echo "[$(date)] Installing pyfdb (pure-Python FDB binding)..." + uv pip install pyfdb + + echo "[$(date)] Cloning eccodes-cosmo-mars (varda-ext branch) for MARS definitions..." + mkdir -p "$VENV/data/share" + # Pinned to d043635 (pre "add a grid Name") rather than a bare + # --depth 1 branch clone: f9b4603 adds a mars.gridName concept + # that the current realtime-varda.schema doesn't account for, + # breaking FDB archival with a "Keywords not used" error. + # Un-pin once the schema is updated to match. + git clone --branch varda-ext \ + git@github.com:meteoswiss/eccodes-cosmo-mars.git \ + "$VENV/data/share/eccodes-cosmo-mars" + git -C "$VENV/data/share/eccodes-cosmo-mars" checkout d04363540bb24b1fd432dce2b1e65ba0f27ba591 + echo "[$(date)] Cloned eccodes-cosmo-mars definitions (pinned to d043635)" + + # metkitlib (pulled in by pyfdb) ships its own language.yaml, used to + # validate fdb.retrieve() requests. Our custom `model` values (added + # purely for archival, via eccodes-cosmo-mars's marsModel concepts) + # were never registered there, so retrieve() rejects them even though + # they're valid archived keys -- fdb.list() has no such validation, so + # this only bites downstream reads, not archival or completeness checks. + echo "[$(date)] Patching metkitlib's language.yaml for custom FDB model names..." + python resources/fdb/patch_metkit_language.py "$VENV" + fi + echo "[$(date)] Compiling Python bytecode..." python -m compileall -j 8 -o 0 -o 1 -o 2 $VENV/lib/python*/site-packages echo "[$(date)] Testing that eccodes is working..." @@ -155,14 +190,179 @@ rule inference_create_sandbox: def get_resource(wc, field: str, default): - """Fetch a resource field from the run config, or return the default.""" - rc = RUN_CONFIGS[wc.run_id] - if rc["inference_resources"] is None: - return default - if isinstance(rc["inference_resources"], dict): - return rc["inference_resources"].get(field, default) or default + """Fetch a resource field from profile.fdb, or return the default.""" + fdb = (config.get("profile") or {}).get("fdb") or {} + if isinstance(fdb, dict): + return fdb.get(field, default) or default + else: + return getattr(fdb, field, default) or default + + +def _env_needs_fdb(env_id): + """Whether any run sharing this env_id declares output: fdb -- i.e. whether + this shared venv needs the FDB native stack installed. pyfdb bundles its own + native libraries (fdb5lib/eccodeslib/eckitlib/metkitlib wheels), so this is + a plain pip install -- no uenv/squashfs mount of a separate native image + needed.""" + return any( + run_cfg["env_id"] == env_id and run_cfg.get("output") == "fdb" + for run_cfg in RUN_CONFIGS.values() + ) + + +def _count_fdb_blocks(run_id): + """Count how many output.tee targets this run will write into FDB. + + Zero whenever this run doesn't declare output: fdb -- regardless of what + its config's output.tee contains, a grib-output run never writes to FDB. + + For output: fdb runs, inference configs are maintained once, in grib: form, + as the single source of truth for both backends -- inference_prepare.py's + _convert_tee_to_fdb rewrites every 'grib:' tee entry to 'fdb:' at + config-prep time (see resources/inference/configs/ + sgm-temporal-downscaler-global_trimedge_multi.yaml). So each on-disk + 'grib:' entry becomes one FDB target; a literal 'fdb:' entry (a config + authored directly in fdb: form, bypassing the conversion) is also counted, + for robustness. + + A single init_time can write several output blocks (e.g. an ICON-grid target + and an IFS-grid target) into the same env_id-keyed root; the check needs to + know how many to expect so a run that silently dropped one entirely doesn't + read as complete just because the other is. + """ + if RUN_CONFIGS[run_id]["output"] != "fdb": + return 0 + with open(RUN_CONFIGS[run_id]["config"]) as f: + run_cfg = yaml.safe_load(f) + tee = ((run_cfg.get("output") or {}).get("tee")) or [] + return sum( + 1 + for entry in tee + if isinstance(entry, dict) and ({"fdb", "grib"} & entry.keys()) + ) + + +checkpoint inference_check_fdb: + """Check whether (run_id, init_time) output already exists, complete, in FDB. + +Checks this run's one configured FDB root (see _resolve_fdb_root_for_run -- +the run's own `root:` field if set, otherwise its local per-checkpoint +directory). A 'hit' requires every requested lead time to be present in +that root -- see inference_check_fdb.py for why no further MARS-key +disambiguation (class/expver/model) is needed. Writes 'hit\\t' or +'miss\\t' to the status file. Runs inside the squashfs venv so that pyfdb (and +the native FDB C library it bundles) is available. Skipped entirely (writes +'miss') unless this run's own config declares `output: fdb`. +""" + input: + image=lambda wc: OUT_ROOT + / f"data/runs/{RUN_CONFIGS[wc.run_id]['env_id']}/venv.squashfs", + script="workflow/scripts/inference_check_fdb.py", + output: + status=OUT_ROOT / "logs/inference_check_fdb/{run_id}-{init_time}.status", + log: + OUT_ROOT / "logs/inference_check_fdb/{run_id}-{init_time}.log", + localrule: True + params: + image_path=lambda wc, input: str(Path(input.image).resolve()), + fdb_configured=lambda wc: RUN_CONFIGS[wc.run_id]["output"] == "fdb", + fdb_root=lambda wc: _resolve_fdb_root_for_run(wc.run_id), + fdb_schema=str(Path("resources/fdb/realtime-varda.schema").resolve()), + date=lambda wc: datetime.strptime(wc.init_time, "%Y%m%d%H%M").strftime("%Y%m%d"), + time=lambda wc: datetime.strptime(wc.init_time, "%Y%m%d%H%M").strftime("%H%M"), + steps=lambda wc: RUN_CONFIGS[wc.run_id]["steps"], + expected_blocks=lambda wc: _count_fdb_blocks(wc.run_id), + # fmt: off + shell: + """ + ( + set -euo pipefail + if [ "{params.fdb_configured}" != "True" ]; then + printf 'miss\t' >{output.status} + echo "FDB check skipped: FDB not configured for this environment" + else + # Mount just the venv -- pyfdb bundles its own native FDB libraries + # (fdb5lib etc.), no separate uenv image needed. + VENV_MOUNT=$(mktemp -d) + trap "rmdir '$VENV_MOUNT' 2>/dev/null || true" EXIT + squashfs-mount -s "{params.image_path}:$VENV_MOUNT" -- bash -c ' + source '"$VENV_MOUNT"'/bin/activate + python {input.script} \ + --fdb-root {params.fdb_root} \ + --fdb-schema {params.fdb_schema} \ + --date {params.date} \ + --time {params.time} \ + --steps {params.steps} \ + --expected-blocks {params.expected_blocks} \ + --output {output.status} + ' + fi + ) >{log} 2>&1 + """ +# fmt: on + + +def _resolve_fdb_root_for_run(run_id): + """This run's one FDB root: its own `root:` config value if set (a + shared/global store, used as-is -- same pattern as baseline.root), + otherwise a local directory dedicated to this checkpoint (env_id-keyed) + under the output root.""" + custom_root = RUN_CONFIGS[run_id].get("root") + if custom_root: + return custom_root + env_id = RUN_CONFIGS[run_id]["env_id"] + return str((OUT_ROOT / f"data/fdb/{env_id}").resolve()) + + +def _get_inference_status(wc): + """Return (okfile_path, fdb_root_or_None) for (run_id, init_time). + + fdb_root is non-None whenever this run's own config writes to FDB at all + (its output.tee has fdb: targets) -- whether the data was already there + (a check 'hit') or is about to be freshly written by inference_execute + (a 'miss', or FDB-check not configured but still an FDB-writing run). + It's None only when this run's config has no FDB targets at all (plain + local-grib output), in which case downstream code reads local grib as + before. + + For output: fdb runs that actually execute (a 'miss'), the returned okfile + is inference_check_fdb_duplicates's own -- not inference_execute's -- + so nothing downstream (verification, plotting, a chained downscaler's + forecaster input) consumes this run's FDB data until it's confirmed free + of MARS-key collisions. A cache 'hit' skips this: that data was already + checked when it was first produced. + """ + if RUN_CONFIGS[wc.run_id]["output"] == "fdb": + status_file = checkpoints.inference_check_fdb.get( + run_id=wc.run_id, init_time=wc.init_time + ).output.status + status, _, root = Path(status_file).read_text().strip().partition("\t") + if status == "hit": + return str(status_file), root + okfile = str( + OUT_ROOT + / f"logs/inference_check_fdb_duplicates/{wc.run_id}-{wc.init_time}.ok" + ) else: - return getattr(rc["inference_resources"], field) or default + okfile = str(OUT_ROOT / f"logs/inference_execute/{wc.run_id}-{wc.init_time}.ok") + fdb_root = ( + _resolve_fdb_root_for_run(wc.run_id) + if _count_fdb_blocks(wc.run_id) > 0 + else None + ) + return okfile, fdb_root + + +def _get_inference_okfile(wc): + """Dependency file marking (run_id, init_time) data as ready, whether freshly + executed or already complete in FDB.""" + return _get_inference_status(wc)[0] + + +def _get_inference_fdb_root(wc): + """FDB root already holding this (run_id, init_time)'s data (a check 'hit'), + or None if the data is local grib / about to be freshly executed.""" + return _get_inference_status(wc)[1] def get_leadtime(wc): @@ -193,6 +393,7 @@ 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_backend=lambda wc: RUN_CONFIGS[wc.run_id]["output"], script: "../scripts/inference_prepare.py" @@ -202,6 +403,22 @@ def _get_forecaster_run_id(run_id): return RUN_CONFIGS[run_id]["forecaster"]["run_id"] +def _resolve_forecaster_fdb_root(run_id): + """FDB root to read this run's forecaster output from, or None. + + None whenever there's no forecaster, or the forecaster writes local grib + (the default) -- in which case the existing local-grib-symlink input + chaining applies instead. inference_prepare.py's _convert_cutout_to_fdb + only runs when this is non-None. + """ + if RUN_CONFIGS[run_id].get("forecaster") is None: + return None + fct_run_id = _get_forecaster_run_id(run_id) + if RUN_CONFIGS[fct_run_id]["output"] != "fdb": + return None + return _resolve_fdb_root_for_run(fct_run_id) + + # Prepare temporal downscaler for a specific run ID rule inference_prepare_temporal_downscaler: input: @@ -211,8 +428,12 @@ rule inference_prepare_temporal_downscaler: holdout=lambda wc: holdout_dep(wc) if run_uses_holdout(wc.run_id) else [], forecasts=lambda wc: ( [ - OUT_ROOT - / f"logs/inference_execute/{_get_forecaster_run_id(wc.run_id)}-{wc.init_time}.ok" + _get_inference_okfile( + SimpleNamespace( + run_id=_get_forecaster_run_id(wc.run_id), + init_time=wc.init_time, + ) + ) ] if RUN_CONFIGS[wc.run_id].get("forecaster") is not None else [] @@ -220,7 +441,6 @@ rule inference_prepare_temporal_downscaler: output: config=Path(OUT_ROOT / "data/runs/{run_id}/{init_time}/config.yaml"), resources=directory(OUT_ROOT / "data/runs/{run_id}/{init_time}/resources"), - grib_out_dir=directory(OUT_ROOT / "data/runs/{run_id}/{init_time}/grib"), forecaster=directory(OUT_ROOT / "data/runs/{run_id}/{init_time}/forecaster"), okfile=touch( OUT_ROOT @@ -242,6 +462,8 @@ rule inference_prepare_temporal_downscaler: if RUN_CONFIGS[wc.run_id].get("forecaster") is None else _get_forecaster_run_id(wc.run_id) ), + forecaster_fdb_root=lambda wc: _resolve_forecaster_fdb_root(wc.run_id), + output_backend=lambda wc: RUN_CONFIGS[wc.run_id]["output"], script: "../scripts/inference_prepare.py" @@ -288,6 +510,10 @@ rule inference_execute: disable_local_definitions=lambda wc: RUN_CONFIGS[wc.run_id].get( "disable_local_eccodes_definitions", False ), + srun_prefix=lambda wc: get_resource(wc, "srun_prefix", "") or "", + fdb_configured=lambda wc: RUN_CONFIGS[wc.run_id]["output"] == "fdb", + fdb_root=lambda wc: _resolve_fdb_root_for_run(wc.run_id), + fdb_schema=str(Path("resources/fdb/realtime-varda.schema").resolve()), account_flag=lambda wc: ( f"--account={config['profile']['default_resources']['slurm_account']}" if config["profile"]["default_resources"].get("slurm_account") @@ -301,8 +527,29 @@ rule inference_execute: cd {params.workdir} + # Write into this run's one FDB root (see _resolve_fdb_root_for_run): its + # own `root:` config value if set, otherwise its local per-checkpoint root. + # Exported so _run_inference (executed by squashfs-mount in a fresh bash) can see it. + export FDB5_CONFIG_OVERRIDE="" + if [ "{params.fdb_configured}" = "True" ]; then + mkdir -p "{params.fdb_root}" + # Use a patched schema that strips 'domain' (present in IFS GRIB MARS + # namespace via mars_labeling.def but absent from MeteoSwiss FDB schema). + cat > fdb5_config.yaml < 1, add parallel flag to CMD_ARGS and override automatic cluster detection @@ -322,17 +580,56 @@ rule inference_execute: CMD_ARGS+=(runner.parallel.cluster=slurm) fi + set +e # Equivalent to `anemoi-inference run`, but imports pyproj before eccodes: # the eckitlib wheel bundles its own PROJ, which aborts at exit (double free) # if loaded before pyproj. See https://github.com/ecmwf/eckit/issues/354 python -c 'import sys, importlib.util; importlib.util.find_spec("pyproj") and __import__("pyproj"); from anemoi.inference.__main__ import main; sys.argv[0] = "anemoi-inference"; sys.exit(main())' \ - run config.yaml "${{CMD_ARGS[@]}}" + run config.yaml "${{CMD_ARGS[@]}}" 2>&1 | tee inference_stdout.log + local rc=${{PIPESTATUS[0]}} + set -e + if [ "$rc" -ne 0 ]; then + # anemoi-inference, when using the copy-prognostic-from-forecaster + # post-processor (anemoi-plugins-meteoswiss), is known to crash + # during native library teardown strictly AFTER all forecast windows + # have been computed and written out (free(): invalid pointer, + # munmap_chunk(): invalid pointer, or similar glibc heap-corruption + # diagnostics observed across runs -- likely an eccodes/earthkit-data + # handle cleanup ordering bug tied to + # that filter's cached forecaster GRIB index being garbage-collected + # at interpreter shutdown). The exact signal raised depends on memory + # layout at the time: observed as both SIGABRT (134) and SIGBUS (135) + # across otherwise-identical runs. Verified empirically (including + # "Flushing FDB" x2 immediately preceding the crash, for FDB-backed + # runs): all expected output steps are present and bit-correct on + # disk/FDB despite the abort. Treat this exact signature as success + # instead of discarding valid output and re-running (expensive) GPU + # inference for nothing. + local last_window + last_window=$(grep -oE 'window [0-9]+/[0-9]+' inference_stdout.log | tail -1 | awk '{{print $2}}') + local win_num=${{last_window%%/*}} + local win_total=${{last_window##*/}} + if [ "$rc" = "134" -o "$rc" = "135" ] \ + && grep -q -e "free(): invalid pointer" -e "munmap_chunk(): invalid pointer" -e "double free or corruption" inference_stdout.log \ + && ! grep -q "Traceback (most recent call last)" inference_stdout.log \ + && [ -n "$win_num" ] && [ "$win_num" = "$win_total" ]; then + echo "WARNING: anemoi-inference aborted during process teardown (exit $rc) after completing all $win_total forecast windows -- known benign native cleanup crash (see copy-prognostic-from-forecaster + eccodes handle GC), treating as success." >&2 + else + return "$rc" + fi + fi }} export -f _run_inference - srun \ + # squashfs-mount needs the mountpoint dir to already exist, and -- since srun + # may dispatch to a different physical node -- it must be on the shared + # filesystem, not /tmp (node-local). + VENV_MOUNT="{params.workdir}/.venv-mount" + mkdir -p "$VENV_MOUNT" + trap "rmdir '$VENV_MOUNT' 2>/dev/null || true" EXIT + + {params.srun_prefix} srun \ --job-name=anemoi-inference \ - --uenv={params.env_path}:/user-environment \ --unbuffered \ {params.account_flag} \ --partition={resources.slurm_partition} \ @@ -341,8 +638,74 @@ rule inference_execute: --time={resources.runtime} \ --gres={resources.gres} \ --ntasks={resources.ntasks} \ - bash -c '_run_inference' + squashfs-mount -s "{params.env_path}:$VENV_MOUNT" -- bash -c "_run_inference $VENV_MOUNT" + ) >{log} 2>&1 && touch {output.okfile} + """ +# fmt: on + + +rule inference_check_fdb_duplicates: + """Fail if this run's own FDB archival collided with an already-present + record -- i.e. a MARS key with more than one entry (live + masked) in + FDB's toc index for this init time (see inference_check_fdb_duplicates.py + for why this can only be caught by inspecting the archive after the fact, + not at write time). Only runs for output: fdb runs, right after + inference_execute actually writes -- not on a cache hit (status == "hit" + in _get_inference_status), whose data was already checked when it was + first produced. + """ + input: + image=lambda wc: OUT_ROOT + / f"data/runs/{RUN_CONFIGS[wc.run_id]['env_id']}/venv.squashfs", + okfile=OUT_ROOT / "logs/inference_execute/{run_id}-{init_time}.ok", + script="workflow/scripts/inference_check_fdb_duplicates.py", + output: + okfile=OUT_ROOT / "logs/inference_check_fdb_duplicates/{run_id}-{init_time}.ok", + log: + OUT_ROOT / "logs/inference_check_fdb_duplicates/{run_id}-{init_time}.log", + localrule: True + params: + image_path=lambda wc, input: str(Path(input.image).resolve()), + fdb_root=lambda wc: _resolve_fdb_root_for_run(wc.run_id), + fdb_schema=str(Path("resources/fdb/realtime-varda.schema").resolve()), + date=lambda wc: datetime.strptime(wc.init_time, "%Y%m%d%H%M").strftime("%Y%m%d"), + time=lambda wc: datetime.strptime(wc.init_time, "%Y%m%d%H%M").strftime("%H%M"), + # fmt: off + shell: + """ + ( + set -euo pipefail + VENV_MOUNT=$(mktemp -d) + trap "rmdir '$VENV_MOUNT' 2>/dev/null || true" EXIT + squashfs-mount -s "{params.image_path}:$VENV_MOUNT" -- bash -c ' + source '"$VENV_MOUNT"'/bin/activate + python {input.script} \ + --fdb-root {params.fdb_root} \ + --fdb-schema {params.fdb_schema} \ + --date {params.date} \ + --time {params.time} \ + --output {output.okfile} + ' ) >{log} 2>&1 - touch {output.okfile} """ # fmt: on + + +# Fail fast: a run declaring `output: fdb` is relying on its own referenced +# inference config to actually deliver FDB output -- catch the one way that +# silently doesn't happen, at workflow-build time, before any (expensive) +# inference job runs: +# - the referenced inference config has zero `fdb:` output.tee blocks -> +# inference_check_fdb.py's completeness check is vacuously true when +# expected_blocks == 0 (0 complete blocks found >= 0 expected), reporting a false +# "hit" against whatever (possibly unrelated, possibly empty) data already sits in +# that FDB root. +for _run_id, _run_cfg in RUN_CONFIGS.items(): + if _run_cfg.get("output") != "fdb": + continue + if _count_fdb_blocks(_run_id) == 0: + raise ValueError( + f"Run '{_run_id}' declares output: fdb but its inference config " + f"({_run_cfg['config']}) has no 'fdb:' blocks under output.tee -- either " + "add one, or set output: grib to match what it actually writes." + ) diff --git a/workflow/rules/plot.smk b/workflow/rules/plot.smk index 7b30124e..26badd4b 100644 --- a/workflow/rules/plot.smk +++ b/workflow/rules/plot.smk @@ -28,7 +28,7 @@ def _get_available_baselines(wc) -> list[dict[str, str]]: rule plot_meteogram: input: script="workflow/scripts/plot_meteogram.py", - inference_okfile=rules.inference_execute.output.okfile, + inference_okfile=_get_inference_okfile, truth_dep=truth_file_dep, output: expand( @@ -44,7 +44,7 @@ rule plot_meteogram: runtime="60m", params: ana_label=lambda wc: config["truth"]["label"], - truth_root=config["truth"]["root"], + truth_root=(config.get("truth") or {}).get("root", ""), fcst_grib=lambda wc: ( Path(OUT_ROOT) / f"data/runs/{wc.run_id}/{wc.init_time}/grib" ).resolve(), @@ -64,17 +64,38 @@ rule plot_meteogram: if config.get("lapse_rate_correction", True) else "" ), + fdb_configured=lambda wc: RUN_CONFIGS[wc.run_id]["output"] == "fdb", + fdb_forecast=lambda wc: ( + "fdb:" + _resolve_fdb_root_for_run(wc.run_id) + if RUN_CONFIGS[wc.run_id]["output"] == "fdb" + else "" + ), shell: """ set -euo pipefail - export ECCODES_DEFINITION_PATH=$(realpath .venv/share/eccodes-cosmo-resources/definitions) BASELINE_ROOTS=({params.baseline_roots:q}) BASELINE_STEPS=({params.baseline_steps:q}) BASELINE_LABELS=({params.baseline_labels:q}) + # If this run declares output: fdb, its forecast output lives there + # (never in local grib -- see workflow/rules/inference.smk), so always + # prefer the FDB source over the (nonexistent) local grib dir. + # Reading from FDB needs pyfdb (an evalml dependency, see pyproject.toml) + # and eccodes-cosmo-mars's MARS-namespace GRIB concepts (a repo-level + # clone, resources/fdb/patch_metkit_language.py's docstring has details) + # -- no squashfs-mount of the per-checkpoint venv/FDB uenv needed, unlike + # inference_execute: pyfdb bundles its own native FDB libraries. + FORECAST_SRC={params.fcst_grib:q} + if [ "{params.fdb_configured}" = "True" ]; then + FORECAST_SRC={params.fdb_forecast:q} + export ECCODES_DEFINITION_PATH="$(realpath eccodes-cosmo-mars/definitions):$(realpath .venv/share/eccodes-cosmo-resources/definitions)" + else + export ECCODES_DEFINITION_PATH=$(realpath .venv/share/eccodes-cosmo-resources/definitions) + fi + CMD_ARGS=( - --forecast {params.fcst_grib:q} + --forecast "$FORECAST_SRC" --forecast_steps {params.fcst_steps:q} --forecast_label {params.fcst_label:q} --analysis {params.truth_root:q} @@ -99,7 +120,7 @@ rule plot_meteogram: rule plot_forecast_frame: input: script="workflow/scripts/plot_forecast_frame.py", - inference_okfile=rules.inference_execute.output.okfile, + inference_okfile=_get_inference_okfile, output: expand( OUT_ROOT @@ -114,7 +135,11 @@ rule plot_forecast_frame: resources: slurm_partition="postproc", cpus_per_task=4, - runtime="20m", + # Bumped from 20m: under heavy concurrent load on the shared postproc + # partition (many plot_forecast_frame jobs at once), a single region's + # render has been observed taking ~16m vs ~1.5m uncontended -- 20m was + # killing otherwise-healthy jobs via TIMEOUT, not a correctness bug. + runtime="45m", params: grib_out_dir=lambda wc: str( (Path(OUT_ROOT) / f"data/runs/{wc.run_id}/{wc.init_time}/grib").resolve() @@ -124,15 +149,36 @@ rule plot_forecast_frame: (Path(OUT_ROOT) / f"data/runs/{wc.run_id}/{wc.init_time}/frames").resolve() ), accu=lambda wc: int(RUN_CONFIGS[wc.run_id]["steps"].split("/")[2]), + fdb_configured=lambda wc: RUN_CONFIGS[wc.run_id]["output"] == "fdb", + fdb_forecast=lambda wc: ( + "fdb:" + _resolve_fdb_root_for_run(wc.run_id) + if RUN_CONFIGS[wc.run_id]["output"] == "fdb" + else "" + ), shell: """ - export ECCODES_DEFINITION_PATH=$(realpath .venv/share/eccodes-cosmo-resources/definitions) - python {input.script} \ - --input {params.grib_out_dir:q} --date {wildcards.init_time:q} \ - --param {wildcards.param:q} --leadtime {wildcards.leadtime:q} \ - --regions_json {params.regions_json:q} \ - --outdir {params.outdir:q} \ - --accu {params.accu} >{log} 2>&1 + set -euo pipefail + + # See plot_meteogram for why this always wins over the (nonexistent) + # local grib dir when the run declares output: fdb, and why no + # squashfs-mount is needed to read it (pyfdb bundles its own native libs). + INPUT_SRC={params.grib_out_dir:q} + if [ "{params.fdb_configured}" = "True" ]; then + INPUT_SRC={params.fdb_forecast:q} + export ECCODES_DEFINITION_PATH="$(realpath eccodes-cosmo-mars/definitions):$(realpath .venv/share/eccodes-cosmo-resources/definitions)" + else + export ECCODES_DEFINITION_PATH=$(realpath .venv/share/eccodes-cosmo-resources/definitions) + fi + + CMD_ARGS=( + --input "$INPUT_SRC" --date {wildcards.init_time:q} + --param {wildcards.param:q} --leadtime {wildcards.leadtime:q} + --regions_json {params.regions_json:q} + --outdir {params.outdir:q} + --accu {params.accu} + ) + + python {input.script} "${{CMD_ARGS[@]}}" >{log} 2>&1 """ diff --git a/workflow/rules/report.smk b/workflow/rules/report.smk index 7dea3a93..a06adbd1 100644 --- a/workflow/rules/report.smk +++ b/workflow/rules/report.smk @@ -9,7 +9,7 @@ include: "common.smk" def make_header_text(): dates = config["dates"] - truth = config["truth"]["label"] + truth = (config.get("truth") or {}).get("label", "") if isinstance(dates, list): return f"Explicit initializations from {len(dates)} runs have been used." text = f"Verification against {truth} with initializations from {dates.get('start')} to {dates.get('end')} by {dates.get('frequency')}" @@ -53,7 +53,9 @@ rule report_experiment_dashboard: ) for sid in EXPERIMENT_PARTICIPANTS ) - + ",truth-{}:{}".format(TRUTH_HASH, config["truth"]["label"]), + + ",truth-{}:{}".format( + TRUTH_HASH, (config.get("truth") or {}).get("label", "") + ), shell: """ python {input.script} \ diff --git a/workflow/rules/verif_obs.smk b/workflow/rules/verif_obs.smk index c603870b..795b27c0 100644 --- a/workflow/rules/verif_obs.smk +++ b/workflow/rules/verif_obs.smk @@ -1,6 +1,7 @@ import json from pathlib import Path from datetime import timedelta +from types import SimpleNamespace def _parse_steps(steps: str) -> list[int]: @@ -38,11 +39,14 @@ if config["mec"] is not None: rule prepare_mec_input: """Collect EKF SYNOP, monSYNOP, and reference verSYNOP observation files into the MEC input_obs directory.""" input: - inference_ok=lambda wc: expand( - rules.inference_execute.output.okfile, - run_id=wc.run_id, - init_time=[t.strftime("%Y%m%d%H%M") for t in REFTIMES], - ), + inference_ok=lambda wc: [ + _get_inference_okfile( + SimpleNamespace( + run_id=wc.run_id, init_time=t.strftime("%Y%m%d%H%M") + ) + ) + for t in REFTIMES + ], output: obs=directory(OUT_ROOT / "data/runs/{run_id}/mec/{init_time}/input_obs"), ekf_file=OUT_ROOT diff --git a/workflow/rules/verification.smk b/workflow/rules/verification.smk index 4dc25bfa..39d8ee25 100644 --- a/workflow/rules/verification.smk +++ b/workflow/rules/verification.smk @@ -2,6 +2,7 @@ # VERIFICATION WORKFLOW # # ----------------------------------------------------- # from datetime import datetime +from types import SimpleNamespace import pandas as pd @@ -108,7 +109,7 @@ rule verification_metrics_baseline: params: baseline_steps=lambda wc: BASELINE_CONFIGS[wc.baseline_id]["steps"], member=lambda wc: BASELINE_CONFIGS[wc.baseline_id].get("member", "000"), - truth=config["truth"]["root"], + truth=(config.get("truth") or {}).get("root", ""), truth_source_id=f"truth-{TRUTH_HASH}", regions=REGIONS, experiment_params=",".join(EXPERIMENT_PARAMS), @@ -156,7 +157,7 @@ rule verification_metrics: "src/verification/spatial.py", "src/data_input/__init__.py", script="workflow/scripts/verification_metrics.py", - inference_okfile=rules.inference_execute.output.okfile, + inference_okfile=_get_inference_okfile, truth_dep=truth_file_dep, holdout=holdout_dep, valid_masks=participant_valid_mask_files, @@ -173,7 +174,7 @@ rule verification_metrics: # TODO: implement logic to use experiment name instead of run_id as wildcard params: fcst_steps=lambda wc: RUN_CONFIGS[wc.run_id]["steps"], - truth=config["truth"]["root"], + truth=(config.get("truth") or {}).get("root", ""), truth_source_id=f"truth-{TRUTH_HASH}", regions=REGIONS, grib_out_dir=lambda wc: ( @@ -189,24 +190,46 @@ rule verification_metrics: if config.get("lapse_rate_correction", True) else "" ), + fdb_configured=lambda wc: RUN_CONFIGS[wc.run_id]["output"] == "fdb", + fdb_forecast=lambda wc: ( + "fdb:" + _resolve_fdb_root_for_run(wc.run_id) + if RUN_CONFIGS[wc.run_id]["output"] == "fdb" + else "" + ), valid_masks_flag=valid_masks_flag, shell: """ - export ECCODES_DEFINITION_PATH=$(realpath .venv/share/eccodes-cosmo-resources/definitions) - uv run {input.script} \ - --forecast {params.grib_out_dir} \ - --truth "{params.truth}" \ - --reftime {wildcards.init_time} \ - --steps "{params.fcst_steps}" \ - --source_id "{wildcards.run_id}" \ - --truth_source_id "{params.truth_source_id}" \ - --regions '{params.regions}' \ - --params "{params.experiment_params}" \ - --threshold_dict "{params.threshold_dict}" \ - {params.holdout_flag} \ - {params.lapse_rate_flag} \ - {params.valid_masks_flag} \ - --output {output} >{log} 2>&1 + set -euo pipefail + + # See plot_meteogram (plot.smk) for why this always wins over the + # (nonexistent) local grib dir when the run declares output: fdb, and + # why no squashfs-mount is needed to read it (pyfdb bundles its own + # native libs). + FORECAST_SRC={params.grib_out_dir:q} + if [ "{params.fdb_configured}" = "True" ]; then + FORECAST_SRC={params.fdb_forecast:q} + export ECCODES_DEFINITION_PATH="$(realpath eccodes-cosmo-mars/definitions):$(realpath .venv/share/eccodes-cosmo-resources/definitions)" + else + export ECCODES_DEFINITION_PATH=$(realpath .venv/share/eccodes-cosmo-resources/definitions) + fi + + CMD_ARGS=( + --forecast "$FORECAST_SRC" + --truth {params.truth:q} + --reftime {wildcards.init_time:q} + --steps {params.fcst_steps:q} + --source_id {wildcards.run_id:q} + --truth_source_id {params.truth_source_id:q} + --regions {params.regions:q} + --params {params.experiment_params:q} + --threshold_dict {params.threshold_dict:q} + {params.holdout_flag} + {params.lapse_rate_flag} + {params.valid_masks_flag} + --output {output:q} + ) + + uv run python {input.script} "${{CMD_ARGS[@]}}" >{log} 2>&1 """ @@ -281,7 +304,9 @@ rule verification_metrics_plot: ) for sid in EXPERIMENT_PARTICIPANTS ) - + ",truth-{}:{}".format(TRUTH_HASH, config["truth"]["label"]), + + ",truth-{}:{}".format( + TRUTH_HASH, (config.get("truth") or {}).get("label", "") + ), shell: """ uv run {input.script} {input.verif} --output_dir {output} --labels "{params.label_map}" >{log} 2>&1 @@ -293,11 +318,10 @@ rule verification_scoremaps: "src/verification/__init__.py", "src/data_input/__init__.py", script="workflow/scripts/verification_scoremaps.py", - inference_okfiles=lambda wc: expand( - rules.inference_execute.output.okfile, - init_time=_restrict_reftimes_to_hours(REFTIMES), - allow_missing=True, - ), + inference_okfiles=lambda wc: [ + _get_inference_okfile(SimpleNamespace(run_id=wc.run_id, init_time=it)) + for it in _restrict_reftimes_to_hours(REFTIMES) + ], truth_dep=truth_file_dep, output: OUT_ROOT @@ -317,19 +341,52 @@ rule verification_scoremaps: fcst_steps=lambda wc: RUN_CONFIGS[wc.run_id]["steps"], truth=config["truth"]["root"], truth_label=config["truth"]["label"], - reftimes=" ".join(t.strftime("%Y%m%d%H%M") for t in REFTIMES), + reftimes=[t.strftime("%Y%m%d%H%M") for t in REFTIMES], run_root=lambda wc: (Path(OUT_ROOT) / f"data/runs/{wc.run_id}").resolve(), + fdb_configured=lambda wc: RUN_CONFIGS[wc.run_id]["output"] == "fdb", + # Bare root (no "fdb:" marker prefix): --fdb_root is consumed by + # iter_init_dirs, which adds the marker itself when constructing + # load_forecast_data's root argument for each reftime -- unlike + # plot.smk/verification_metrics' single --forecast argument, which IS + # load_forecast_data's root directly. A run's declared `output` is the + # same for every reftime under its run_root (it's a property of the run + # config, not of any one reftime), so this is the one root iter_init_dirs + # uses for all of them -- it never needs to guess per-directory. + fdb_root=lambda wc: ( + _resolve_fdb_root_for_run(wc.run_id) + if RUN_CONFIGS[wc.run_id]["output"] == "fdb" + else "" + ), shell: """ - export ECCODES_DEFINITION_PATH=$(realpath .venv/share/eccodes-cosmo-resources/definitions) - uv run {input.script} \ - --run_root {params.run_root} \ - --reftimes {params.reftimes} \ - --truth "{params.truth}" \ - --step {wildcards.leadtime} \ - --steps "{params.fcst_steps}" \ - --param {wildcards.param} \ - --output {output} >{log} 2>&1 + set -euo pipefail + + # See plot_meteogram (plot.smk) for why no squashfs-mount is needed to + # read from FDB (pyfdb bundles its own native libs). + if [ "{params.fdb_configured}" = "True" ]; then + export ECCODES_DEFINITION_PATH="$(realpath eccodes-cosmo-mars/definitions):$(realpath .venv/share/eccodes-cosmo-resources/definitions)" + else + export ECCODES_DEFINITION_PATH=$(realpath .venv/share/eccodes-cosmo-resources/definitions) + fi + + # --reftimes is nargs="+" in argparse: pass each timestamp as its own + # array element (not one space-joined string) so it fans out correctly. + REFTIMES_ARR=({params.reftimes:q}) + + CMD_ARGS=( + --run_root {params.run_root:q} + --reftimes "${{REFTIMES_ARR[@]}}" + --truth {params.truth:q} + --step {wildcards.leadtime:q} + --steps {params.fcst_steps:q} + --param {wildcards.param:q} + --output {output:q} + ) + if [ -n "{params.fdb_root}" ]; then + CMD_ARGS+=(--fdb_root {params.fdb_root:q}) + fi + + uv run python {input.script} "${{CMD_ARGS[@]}}" >{log} 2>&1 """ diff --git a/workflow/scripts/inference_check_fdb.py b/workflow/scripts/inference_check_fdb.py new file mode 100644 index 00000000..5adf5a0b --- /dev/null +++ b/workflow/scripts/inference_check_fdb.py @@ -0,0 +1,207 @@ +"""Check whether inference output for a given init_time already exists in FDB. + +Each FDB root passed on the command line is assumed dedicated to a single +checkpoint (the caller passes this run's own resolved root -- see +`_resolve_fdb_root_for_run` in workflow/rules/inference.smk: either the run's +explicit `root:` config value, or else its local per-checkpoint directory), +so no cross-checkpoint MARS-key disambiguation (class/expver/model) is +needed here. + +A single init_time's output can still be split into several distinct +*blocks* within that one root -- e.g. an ICON-grid target and an IFS-grid +target from the same run's `output.tee`, which go through different +post-processing and so don't necessarily share the same param set. A block +is identified by its combined key minus {step, param, levelist, timespan} +(i.e. by whatever class/stream/type/etc. distinguishes it). A "hit" requires: + + - at least `expected_blocks` distinct, complete blocks are present; + - each block holds every requested lead time (from --steps); + - each block's params are consistent across steps: every step other than + the first has exactly the same param set (the union across those steps), + and the first step's params are a subset of it. This tolerates fields + that are legitimately absent only at the initial step (e.g. accumulated + params with no meaningful value at lead time 0) without needing to name + them, while still catching a run that crashed partway through writing a + step's fields (which would otherwise still show up as "present"). + +Roots are checked in the order given; the first one holding a complete set +wins. Writes "hit\t" or "miss\t" to the output file. + +Intended to run inside a squashfs-mounted venv that provides pyfdb and the +native FDB C library (libfdb5.so). The squashfs is the same one built by +inference_create_venv / inference_make_squashfs_image. +""" + +import argparse +import sys +from pathlib import Path + + +def _make_fdb_config(fdb_root: str, fdb_schema: str) -> dict: + return { + "type": "local", + "engine": "toc", + "schema": fdb_schema, + "spaces": [{"handler": "Default", "roots": [{"path": fdb_root}]}], + } + + +def _parse_steps(steps_spec: str) -> set[int]: + """Parse a 'start/end/step' lead-time spec (hours) into the expected step set.""" + start, end, step = map(int, steps_spec.split("/")) + return set(range(start, end + 1, step)) + + +def _group_by_block(entries) -> dict[tuple, dict[int, set]]: + """Group list() entries into {block_key: {step: {(levtype, param), ...}}}. + + block_key is the combined key with step/param/levelist/timespan/levtype + removed -- i.e. whatever identifies a distinct *output.tee target* + (class/stream/type/...) sharing this date+time. levtype is deliberately + NOT part of block_key: one target (e.g. an ICON-grid `fdb:` entry) writes + both 'sfc' and 'pl' levtypes, and those must be checked together as a + single block -- otherwise a run that crashed after writing only 'sfc' + would still count as its own separate, "complete" block, undermining the + --expected-blocks check. levtype is instead folded into each step's + param-identity tuple so sfc/pl params are never conflated with each other. + """ + blocks: dict[tuple, dict[int, set]] = {} + for el in entries: + key = el.combined_key() + if "step" not in key or "param" not in key: + continue + block_key = tuple( + sorted( + (k, v) + for k, v in key.items() + if k not in ("step", "param", "levelist", "timespan", "levtype") + ) + ) + item = (key.get("levtype"), key["param"]) + blocks.setdefault(block_key, {}).setdefault(int(key["step"]), set()).add(item) + return blocks + + +def _block_is_complete(step_params: dict[int, set], expected_steps: set[int]) -> bool: + found_steps = set(step_params) + if expected_steps - found_steps: + return False + + first_step = min(expected_steps) + other_steps = expected_steps - {first_step} + if not other_steps: + return True # single-step run: nothing to cross-check against + + reference = set() + for s in other_steps: + reference |= step_params[s] + if not reference: + return False # no params at all on the non-first steps + + for s in other_steps: + if step_params[s] != reference: + return False + return step_params[first_step] <= reference + + +def check_fdb( + fdb_root: str, + fdb_schema: str, + date: str, + time: str, + steps_spec: str, + expected_blocks: int, +) -> bool: + """Return True if `fdb_root` holds a complete output for this init time.""" + import pyfdb + + if not Path(fdb_root).is_dir(): + return False + + expected_steps = _parse_steps(steps_spec) + + # Pass the config as a dict directly. pyfdb.FDB(...) treats a `str` argument as + # inline YAML *content* to parse, not a file path (only a `Path` object is read + # as a file) -- passing str(config_path) here silently mis-parses and falls back + # to FDB5's built-in default config, which doesn't exist in this venv and fails + # with "Cannot open .../fdb5lib/etc/fdb/schema". A dict sidesteps the ambiguity + # entirely and needs no temp file. + with pyfdb.FDB(_make_fdb_config(fdb_root, fdb_schema)) as fdb: + # level=3 reaches the index level, where 'step'/'param' are part of the + # combined key (schema: [date,time,...,type[levtype,number? + # [step,param,levelist?,timespan?]]]). + entries = fdb.list({"date": date, "time": time}, level=3) + blocks = _group_by_block(entries) + + complete_blocks = [ + step_params + for step_params in blocks.values() + if _block_is_complete(step_params, expected_steps) + ] + if len(complete_blocks) < expected_blocks: + print( + f"FDB check: {fdb_root} has {len(complete_blocks)}/{expected_blocks} " + f"complete block(s) for date={date} time={time}" + ) + return False + return True + + +def main() -> int: + parser = argparse.ArgumentParser( + description="Check if FDB already holds a complete output for a given init_time." + ) + parser.add_argument( + "--fdb-root", + required=True, + action="append", + dest="fdb_roots", + help="Root path of an FDB store to check. Repeatable; checked in order, " + "first complete match wins.", + ) + parser.add_argument( + "--fdb-schema", required=True, help="Path to the FDB5 schema file." + ) + parser.add_argument("--date", required=True, help="Initialisation date (YYYYMMDD).") + parser.add_argument("--time", required=True, help="Initialisation time (HHMM).") + parser.add_argument( + "--steps", + required=True, + help="Expected lead times as 'start/end/step' (hours).", + ) + parser.add_argument( + "--expected-blocks", + required=True, + type=int, + help="Number of distinct fdb: output targets this run's config declares.", + ) + parser.add_argument("--output", required=True, help="File to write the result to.") + args = parser.parse_args() + + hit_root = "" + for root in args.fdb_roots: + if check_fdb( + root, + args.fdb_schema, + args.date, + args.time, + args.steps, + args.expected_blocks, + ): + hit_root = root + break + status = "hit" if hit_root else "miss" + + output = Path(args.output) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(f"{status}\t{hit_root}") + + print( + f"FDB check: date={args.date} time={args.time} steps={args.steps} " + f"-> {status}" + (f" ({hit_root})" if hit_root else "") + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/workflow/scripts/inference_extract_requirements.py b/workflow/scripts/inference_extract_requirements.py index dc3441ae..2f46c817 100644 --- a/workflow/scripts/inference_extract_requirements.py +++ b/workflow/scripts/inference_extract_requirements.py @@ -180,6 +180,29 @@ def extract_git_requirements( return requirements +def _parse_local_package_name(path: str) -> str: + """Derive a package name from a local filesystem path. + + Reads the ``[project] name`` field from ``pyproject.toml`` if present; + falls back to the last path component. + """ + import os + import re + + pyproject = os.path.join(path, "pyproject.toml") + if os.path.exists(pyproject): + try: + with open(pyproject) as f: + content = f.read() + m = re.search(r'^\s*name\s*=\s*["\']([^"\']+)["\']', content, re.MULTILINE) + if m: + return m.group(1) + except Exception: + pass + + return os.path.basename(os.path.normpath(path)) + + def _parse_url_package_name(url: str) -> str: """Derive a package name from a VCS/HTTP URL. @@ -215,8 +238,11 @@ def parse_overrides(overrides: list[str]) -> dict[str, str | None]: - ``name==version`` → pinned PyPI package - a URL (``git+``, ``http://``, ``https://``) → VCS / direct install + - a local filesystem path (``/``, ``./``, ``../``) → local directory install - bare ``name`` → unpinned PyPI package (value is ``None``) """ + import os + result: dict[str, str | None] = {} for item in [*BASE_DEPENDENCIES, *(overrides or [])]: @@ -230,6 +256,11 @@ def parse_overrides(overrides: list[str]) -> dict[str, str | None]: elif any(item.startswith(prefix) for prefix in ("git+", "http://", "https://")): name = _parse_url_package_name(item) result[name] = item + elif item.startswith(("/", "./", "../")) or ( + os.sep in item and os.path.exists(item) + ): + name = _parse_local_package_name(item) + result[name] = item else: result[item] = None @@ -285,7 +316,9 @@ def format_requirements( pypi_requirements.pop(name, None) git_requirements.pop(name, None) - if isinstance(value, str) and value.startswith(("http://", "https://", "git+")): + if isinstance(value, str) and value.startswith( + ("http://", "https://", "git+", "/", "./", "../") + ): git_requirements[name] = value else: pypi_requirements[name] = value # type: ignore[assignment] # may be None diff --git a/workflow/scripts/inference_prepare.py b/workflow/scripts/inference_prepare.py index 2a88f65b..e3682bd1 100644 --- a/workflow/scripts/inference_prepare.py +++ b/workflow/scripts/inference_prepare.py @@ -6,12 +6,22 @@ from pathlib import Path from evalml.helpers import setup_logger +from data_input import ( + _FDB_MODEL, + _FDB_GLOBAL_MODEL, + _FDB_PARAM_IDS, + _FDB_GLOBAL_PARAM_IDS, + _FDB_IDENTITY, + _FDB_SCHEMA, +) def prepare_config( default_config_path: str, output_config_path: str, params: dict, + output_backend: str = "grib", + forecaster_fdb_root: str | None = None, ): """Prepare the configuration file for the inference run. @@ -26,11 +36,29 @@ def prepare_config( Path where the updated configuration file will be written. params : dict Dictionary of parameters to override in the default configuration. + output_backend : str + "grib" (default) leaves output.tee untouched -- the inference config is + itself the single source of truth, always written in grib: form. "fdb" + rewrites every `output.tee` entry from `grib:` to `fdb:` (dropping the + grib-only `path` key) -- see `_convert_tee_to_fdb`. + forecaster_fdb_root : str | None + If this run's forecaster declares `output: fdb`, the FDB root its output + was archived into -- rewrites this downscaler's own input.cutout entries + from `grib:` to `fdb:` so it reads the forecaster's output directly from + FDB instead of a local grib symlink. None (default) leaves input.cutout + untouched (the forecaster writes local grib, read via the symlink + prepare_temporal_downscaler already creates). """ with open(default_config_path, "r") as f: config = yaml.safe_load(f) + if output_backend == "fdb": + config = _convert_tee_to_fdb(config) + + if forecaster_fdb_root is not None: + config = _convert_cutout_to_fdb(config, forecaster_fdb_root) + config = _override_recursive(config, params) with open(output_config_path, "w") as f: @@ -50,7 +78,6 @@ def prepare_workdir(workdir: Path, resources_root: Path): Path to the root directory containing resource files. """ workdir.mkdir(parents=True, exist_ok=True) - (workdir / "grib").mkdir(parents=True, exist_ok=True) (workdir / "resources").mkdir(parents=True, exist_ok=True) shutil.copytree( resources_root / "templates", workdir / "resources", dirs_exist_ok=True @@ -94,14 +121,28 @@ def prepare_temporal_downscaler(smk): # prepare forecaster directory fct_run_id = smk.params.forecaster_run_id + forecaster_fdb_root = smk.params.forecaster_fdb_root if fct_run_id != "null": - fct_workdir = ( - smk.params.output_root / "runs" / fct_run_id / smk.wildcards.init_time - ) - (workdir / "forecaster").symlink_to(fct_workdir / "grib") - LOG.info( - "Created symlink to forecaster grib directory at %s", workdir / "forecaster" - ) + if forecaster_fdb_root is None: + fct_workdir = ( + smk.params.output_root / "runs" / fct_run_id / smk.wildcards.init_time + ) + (workdir / "forecaster").symlink_to(fct_workdir / "grib") + LOG.info( + "Created symlink to forecaster grib directory at %s", + workdir / "forecaster", + ) + else: + # Still an output directory of this rule (see inference.smk's + # inference_prepare_temporal_downscaler), so it must exist, but + # unused -- input.cutout reads the forecaster's output straight + # from FDB instead. + (workdir / "forecaster").mkdir(parents=True, exist_ok=True) + LOG.info( + "Forecaster output is in FDB (%s) -- downscaler input.cutout will be " + "rewritten to read it directly, no local grib symlink needed.", + forecaster_fdb_root, + ) else: (workdir / "forecaster").mkdir(parents=True, exist_ok=True) (workdir / "forecaster/.dataset").touch() @@ -112,7 +153,13 @@ def prepare_temporal_downscaler(smk): # prepare config overrides = _overrides_from_params(smk) - prepare_config(smk.input.config, smk.output.config, overrides) + prepare_config( + smk.input.config, + smk.output.config, + overrides, + output_backend=smk.params.output_backend, + forecaster_fdb_root=forecaster_fdb_root, + ) LOG.info("Wrote config file at %s", smk.output.config) with open(smk.output.config, "r") as f: @@ -136,13 +183,19 @@ def prepare_forecaster(smk): workdir = _get_workdir(smk) prepare_workdir(workdir, smk.params.resources_root) + (workdir / "grib").mkdir(parents=True, exist_ok=True) LOG.info("Prepared working directory at %s", workdir) res_list = "\n".join([str(fn) for fn in Path(workdir / "resources").rglob("*")]) LOG.info("Resources: \n%s", res_list) stage_holdout_stations(smk, workdir, 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_backend=smk.params.output_backend, + ) LOG.info("Wrote config file at %s", smk.output.config) with open(smk.output.config, "r") as f: config_content = f.read() @@ -177,6 +230,106 @@ def _setup_logger(smk) -> logging.Logger: return LOG +def _convert_tee_to_fdb(config: dict) -> dict: + """Rewrite every `output.tee` entry from `{grib: {...}}` to `{fdb: {...}}`. + + The inference config is maintained once, in grib: form, as the single + source of truth for both backends -- see + resources/inference/configs/sgm-temporal-downscaler-global_trimedge_multi.yaml, + whose output.tee encoding/templates/post_processors/modifiers are already + FDB-compatible (generatingProcessIdentifier, marsStream/marsType, etc.) and + were kept in parity with the (now-deleted) *-fdb-parity.yaml by hand before + this. Only the grib-specific `path` key doesn't apply to fdb: and is + dropped; everything else under the tee entry carries over unchanged. + """ + tee = config.get("output", {}).get("tee") + if not tee: + raise ValueError( + "output_backend='fdb' but this inference config has no output.tee " + "entries to convert -- check the referenced config file." + ) + new_tee = [] + for entry in tee: + if set(entry.keys()) != {"grib"}: + raise ValueError( + f"output_backend='fdb' expects every output.tee entry to be a " + f"single 'grib:' block, got keys {sorted(entry.keys())} -- " + "this config isn't written in the expected grib: source-of-truth form." + ) + inner = dict(entry["grib"]) + inner.pop("path", None) + new_tee.append({"fdb": inner}) + config["output"]["tee"] = new_tee + return config + + +# Which FDB model/param-id table a downscaler input.cutout entry reads from, +# keyed by its grid name -- "lam_0" is the regional/ICON grid (VARDA's own local +# GPI=180 table), "global" is the IFS/global grid (standard WMO/ECMWF param ids). +# See _FDB_IDENTITY (data_input/__init__.py) for the class/stream/type/expver +# each model resolves to. +_CUTOUT_FDB_PARAM_IDS = { + "lam_0": (_FDB_MODEL, _FDB_PARAM_IDS), + "global": (_FDB_GLOBAL_MODEL, _FDB_GLOBAL_PARAM_IDS), +} + + +def _convert_cutout_to_fdb(config: dict, fdb_root: str) -> dict: + """Rewrite the downscaler's input.cutout entries from `{grib: {...}}` to + `{fdb: {...}}`, so it reads its forecaster's output directly from FDB + instead of a local grib symlink. + + `pre_processors`/`namer` under the grib: block carry over unchanged -- they + operate on the already-decoded field regardless of how the raw bytes were + fetched. The grib-only `path` key is dropped. In its place we add + `fdb_config` (pointing at the forecaster's FDB root), `param_id_map` + (translating the checkpoint's own param names -- already rewritten to ICON + shortnames by this config's `patch_metadata`, see anemoi-inference's + Metadata.patch -- to this archive's real param ids), and the static + class/stream/type/expver/model identity, via anemoi-inference's FDBInput + (registered as `fdb`, inputs/fdb.py in the anemoi-inference source): + these keys all pass straight through as constructor kwargs. + """ + cutout = config.get("input", {}).get("cutout") + if not cutout: + raise ValueError( + "forecaster declares output: fdb but this downscaler's input has no " + "input.cutout entries to convert -- check the referenced inference config." + ) + new_cutout = [] + for entry in cutout: + if len(entry) != 1: + raise ValueError( + f"Expected a single-key cutout entry, got {sorted(entry.keys())}" + ) + ((grid_name, block),) = entry.items() + if grid_name not in _CUTOUT_FDB_PARAM_IDS: + raise ValueError( + f"Unknown cutout grid '{grid_name}' -- add its FDB identity to " + "_CUTOUT_FDB_PARAM_IDS before converting it to an fdb: input." + ) + if set(block.keys()) != {"grib"}: + raise ValueError( + f"Expected cutout entry '{grid_name}' to be a single 'grib:' block, " + f"got keys {sorted(block.keys())} -- this config isn't written in " + "the expected grib: source-of-truth form." + ) + inner = dict(block["grib"]) + inner.pop("path", None) + model, param_id_map = _CUTOUT_FDB_PARAM_IDS[grid_name] + inner["fdb_config"] = { + "type": "local", + "engine": "toc", + "schema": str(_FDB_SCHEMA.resolve()), + "spaces": [{"handler": "Default", "roots": [{"path": fdb_root}]}], + } + inner["param_id_map"] = dict(param_id_map) + inner.update(_FDB_IDENTITY[model]) + new_cutout.append({grid_name: {"fdb": inner}}) + config["input"]["cutout"] = new_cutout + return config + + def _override_recursive(original: dict, updates: dict) -> dict: """Recursively override values in the original dictionary with those from the updates dictionary.""" for key, value in updates.items(): diff --git a/workflow/scripts/plot_forecast_frame.py b/workflow/scripts/plot_forecast_frame.py index 5f895a1b..ef72648e 100644 --- a/workflow/scripts/plot_forecast_frame.py +++ b/workflow/scripts/plot_forecast_frame.py @@ -1,6 +1,7 @@ import json import logging from argparse import ArgumentParser +from datetime import datetime from pathlib import Path import cartopy.crs as ccrs @@ -8,7 +9,7 @@ from data_input import parse_aggregated_param from plotting import resolve_region from plotting import StatePlotter -from plotting.compat import load_state_from_grib +from plotting.compat import load_state_from_grib, load_state_from_fdb from plotting.styles import get_style from plotting.styles import preprocess_field @@ -38,7 +39,7 @@ def main(): ) args = parser.parse_args() - grib_dir = Path(args.input) + is_fdb = args.input.startswith("fdb:") init_time = args.date lead_time = args.leadtime param = args.param @@ -67,62 +68,101 @@ def main(): else: paramlist = [base_param] - # Load grib once — shared across all region plots - # TODO: fix file pattern & globbing - grib_file = list(grib_dir.glob(f"2*_{lead_time}.grib")) - if not grib_file: - grib_file = list(grib_dir.glob(f"2*_{lead_time:03}.grib")) - if not grib_file: - LOG.warning( - "No GRIB file found for lead_time=%s in %s (model may not write initial state)." - " Creating empty placeholder frames.", - lead_time, - grib_dir, - ) - for region_name in regions: - outfn = outdir / f"frame_{lead_time}_{param}_{region_name}.png" - outfn.touch() - return - grib_file = Path(grib_file[0]) - LOG.info("Loading grib file %s", grib_file) - try: - state = load_state_from_grib(grib_file, paramlist=paramlist) - except ValueError as exc: - # Accumulated params (e.g. TOT_PREC1/6) are undefined at lead times - # shorter than their accumulation window -- normally excluded - # upstream (resolve_leadtimes), but handled here too as a - # defensive fallback rather than crashing the whole rule. - LOG.warning( - "%s Creating empty placeholder frames.", - exc, - ) + def _empty_frames(): for region_name in regions: outfn = outdir / f"frame_{lead_time}_{param}_{region_name}.png" outfn.touch() - return + + # Load state once — shared across all region plots + if is_fdb: + fdb_root = args.input[len("fdb:") :] + reftime = datetime.strptime(init_time, "%Y%m%d%H%M") + LOG.info("Loading forecast state from FDB %s", fdb_root) + try: + state = load_state_from_fdb( + fdb_root, reftime, lead_time, paramlist=paramlist + ) + except FileNotFoundError: + LOG.warning( + "No FDB data for lead_time=%s in %s (model may not write initial state)." + " Creating empty placeholder frames.", + lead_time, + fdb_root, + ) + _empty_frames() + return + except ValueError as exc: + # Accumulated params (e.g. TOT_PREC1/6) are undefined at lead times + # shorter than their accumulation window -- normally excluded + # upstream (resolve_leadtimes), but handled here too as a + # defensive fallback rather than crashing the whole rule. + LOG.warning( + "%s Creating empty placeholder frames.", + exc, + ) + _empty_frames() + return + else: + grib_dir = Path(args.input) + # TODO: fix file pattern & globbing + grib_file = list(grib_dir.glob(f"2*_{lead_time}.grib")) + if not grib_file: + grib_file = list(grib_dir.glob(f"2*_{lead_time:03}.grib")) + if not grib_file: + LOG.warning( + "No GRIB file found for lead_time=%s in %s (model may not write initial state)." + " Creating empty placeholder frames.", + lead_time, + grib_dir, + ) + _empty_frames() + return + grib_file = Path(grib_file[0]) + LOG.info("Loading grib file %s", grib_file) + try: + state = load_state_from_grib(grib_file, paramlist=paramlist) + except ValueError as exc: + # Accumulated params (e.g. TOT_PREC1/6) are undefined at lead times + # shorter than their accumulation window -- normally excluded + # upstream (resolve_leadtimes), but handled here too as a + # defensive fallback rather than crashing the whole rule. + LOG.warning( + "%s Creating empty placeholder frames.", + exc, + ) + _empty_frames() + return # tp is accumulated from start of forecast; de-accumulate to get period [lt-accu, lt] if base_param == "TOT_PREC": prev_lt = lead_time - accu if prev_lt > 0: - prev_grib_files = list(grib_dir.glob(f"2*_{prev_lt}.grib")) - if not prev_grib_files: - prev_grib_files = list(grib_dir.glob(f"2*_{prev_lt:03d}.grib")) - if not prev_grib_files: - raise FileNotFoundError( - f"No GRIB file for lead_time={prev_lt} in {grib_dir}, needed to" - f" de-accumulate {param} at lead_time={lead_time}." + if is_fdb: + try: + prev_state = load_state_from_fdb( + fdb_root, reftime, prev_lt, paramlist=paramlist + ) + except FileNotFoundError: + prev_state = None + else: + prev_grib_files = list(grib_dir.glob(f"2*_{prev_lt}.grib")) + if not prev_grib_files: + prev_grib_files = list(grib_dir.glob(f"2*_{prev_lt:03d}.grib")) + prev_state = ( + load_state_from_grib(Path(prev_grib_files[0]), paramlist=paramlist) + if prev_grib_files + else None + ) + if prev_state is not None: + LOG.info( + "De-accumulating TOT_PREC using previous lead_time=%s", prev_lt + ) + state["fields"]["TOT_PREC"] = ( + state["fields"]["TOT_PREC"] + - prev_state["fields"]["TOT_PREC"][ + : len(state["fields"]["TOT_PREC"]) + ] ) - prev_grib_file = Path(prev_grib_files[0]) - LOG.info( - "De-accumulating TOT_PREC: loading previous grib file %s", - prev_grib_file, - ) - prev_state = load_state_from_grib(prev_grib_file, paramlist=paramlist) - state["fields"]["TOT_PREC"] = ( - state["fields"]["TOT_PREC"] - - prev_state["fields"]["TOT_PREC"][: len(state["fields"]["TOT_PREC"])] - ) # Preprocess field once — shared across all region plots field, units_override = preprocess_field(base_param, state) diff --git a/workflow/scripts/verification_scoremaps.py b/workflow/scripts/verification_scoremaps.py index 2524e483..e664e876 100644 --- a/workflow/scripts/verification_scoremaps.py +++ b/workflow/scripts/verification_scoremaps.py @@ -77,12 +77,30 @@ def _season_of(dt: datetime) -> str: # --------------------------------------------------------------------------- -def iter_init_dirs(run_root: Path) -> list[tuple[datetime, Path]]: - """Return ``(reftime, grib_dir)`` pairs for every complete init time. +def iter_init_dirs( + run_root: Path, fdb_root: str | None = None +) -> list[tuple[datetime, Path | str]]: + """Return ``(reftime, source)`` pairs for every init-time directory found. Expects subdirectories named ``YYYYMMDDHHMI`` directly under *run_root*. - GRIB files may live either directly in the init-time directory or inside a - ``grib/`` subdirectory. + ``fdb_root`` being set or ``None`` is the caller's explicit declaration of + this run's backend (derived from its config's ``output: fdb``/``grib`` + field -- see ``workflow/rules/verification.smk``'s ``fdb_root`` param), + not a hint to probe with: every directory found is assigned a source the + same way, with no per-directory existence check of either backend. + - ``output: fdb`` (``fdb_root`` given): every reftime maps to an + ``"fdb:"`` marker string -- see ``data_input.load_forecast_data``'s + routing on that prefix, which every caller of this function's output + already goes through. + - ``output: grib`` (``fdb_root`` is ``None``): every reftime maps to its + GRIB directory (files may live directly in the init-time directory or + inside a ``grib/`` subdirectory -- a layout detail, not a backend guess). + + Completeness is the caller's job, not this function's: ``main()`` requires + every configured reftime to be discovered here, raising a hard error + otherwise (see its module docstring) -- consistent with the Snakemake + dependency on this reftime's inference_check_fdb/inference_execute okfile + already having established completeness before this script ever runs. """ result = [] for d in sorted(run_root.iterdir()): @@ -92,11 +110,11 @@ def iter_init_dirs(run_root: Path) -> list[tuple[datetime, Path]]: reftime = datetime.strptime(d.name, DATETIME_FMT) except ValueError: continue - grib_dir = d / "grib" if (d / "grib").is_dir() else d - if not any(grib_dir.glob("*.grib")): - LOG.debug("No GRIB files in %s, skipping", grib_dir) - continue - result.append((reftime, grib_dir)) + if fdb_root is not None: + result.append((reftime, f"fdb:{fdb_root}")) + else: + grib_dir = d / "grib" if (d / "grib").is_dir() else d + result.append((reftime, grib_dir)) return result @@ -147,7 +165,7 @@ def main(args: Namespace) -> None: ] LOG.info("Using %d baseline init times from --reftimes", len(init_items)) else: - init_items = iter_init_dirs(args.run_root) + init_items = iter_init_dirs(args.run_root, fdb_root=args.fdb_root) LOG.info("Found %d init time directories", len(init_items)) # Restrict to the experiment's configured init times, and require that @@ -466,6 +484,17 @@ def _strat_da(compute_fn) -> xr.DataArray: "--reftimes." ), ) + parser.add_argument( + "--fdb_root", + type=str, + default=None, + help=( + "FDB root holding this run's output, if any is configured for its " + "environment. Init times with no local GRIB output under --run_root " + "are additionally checked against it (see iter_init_dirs). Ignored " + "for --baseline_root." + ), + ) parser.add_argument( "--member", type=str, diff --git a/workflow/tools/config.schema.json b/workflow/tools/config.schema.json index 3eeb9726..edee5665 100644 --- a/workflow/tools/config.schema.json +++ b/workflow/tools/config.schema.json @@ -560,23 +560,34 @@ "title": "Extra Requirements", "type": "array" }, - "inference_resources": { + "disable_local_eccodes_definitions": { + "default": false, + "description": "If true, the ECCODES_DEFINITION_PATH environment variable will not be set to the COSMO local definitions.", + "title": "Disable Local Eccodes Definitions", + "type": "boolean" + }, + "output": { + "default": "grib", + "description": "Where this run's forecast output is stored, and therefore where downstream reads (verification, meteograms, frames, scoremaps) should look for it. Must match what the referenced inference `config`'s own output.tee actually writes to (an 'fdb' tee block requires output: fdb here, and vice versa) -- this is the single declared source of truth for backend routing; nothing probes for data presence or parses the inference config to guess it.", + "enum": [ + "grib", + "fdb" + ], + "title": "Output", + "type": "string" + }, + "root": { "anyOf": [ { - "$ref": "#/$defs/InferenceResources" + "type": "string" }, { "type": "null" } ], "default": null, - "description": "Resource requirements for inference jobs (optional; defaults handled externally)." - }, - "disable_local_eccodes_definitions": { - "default": false, - "description": "If true, the ECCODES_DEFINITION_PATH environment variable will not be set to the COSMO local definitions.", - "title": "Disable Local Eccodes Definitions", - "type": "boolean" + "description": "Only meaningful when output: fdb. Root directory of a shared/global FDB store this run reads from and writes into, used as-is (same pattern as baseline.root) -- the caller is responsible for pointing different checkpoints at different roots if they shouldn't share one. When unset (the default), this run's FDB output lives in a local per-checkpoint directory under the output root.", + "title": "Root" }, "config": { "anyOf": [ @@ -706,6 +717,19 @@ "default": null, "description": "Number of tasks per submission.", "title": "Tasks" + }, + "srun_prefix": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Optional command to prepend to the srun invocation, e.g. 'uenv run --view=realtime fdb/5.19:v2 --'.", + "title": "Srun Prefix" } }, "title": "InferenceResources", @@ -789,6 +813,18 @@ "title": "Executor", "type": "string" }, + "fdb": { + "anyOf": [ + { + "$ref": "#/$defs/InferenceResources" + }, + { + "type": "null" + } + ], + "default": null, + "description": "FDB and inference resource settings (optional; overrides defaults for all inference jobs)." + }, "global_resources": { "$ref": "#/$defs/GlobalResources" }, @@ -1172,23 +1208,34 @@ "title": "Extra Requirements", "type": "array" }, - "inference_resources": { + "disable_local_eccodes_definitions": { + "default": false, + "description": "If true, the ECCODES_DEFINITION_PATH environment variable will not be set to the COSMO local definitions.", + "title": "Disable Local Eccodes Definitions", + "type": "boolean" + }, + "output": { + "default": "grib", + "description": "Where this run's forecast output is stored, and therefore where downstream reads (verification, meteograms, frames, scoremaps) should look for it. Must match what the referenced inference `config`'s own output.tee actually writes to (an 'fdb' tee block requires output: fdb here, and vice versa) -- this is the single declared source of truth for backend routing; nothing probes for data presence or parses the inference config to guess it.", + "enum": [ + "grib", + "fdb" + ], + "title": "Output", + "type": "string" + }, + "root": { "anyOf": [ { - "$ref": "#/$defs/InferenceResources" + "type": "string" }, { "type": "null" } ], "default": null, - "description": "Resource requirements for inference jobs (optional; defaults handled externally)." - }, - "disable_local_eccodes_definitions": { - "default": false, - "description": "If true, the ECCODES_DEFINITION_PATH environment variable will not be set to the COSMO local definitions.", - "title": "Disable Local Eccodes Definitions", - "type": "boolean" + "description": "Only meaningful when output: fdb. Root directory of a shared/global FDB store this run reads from and writes into, used as-is (same pattern as baseline.root) -- the caller is responsible for pointing different checkpoints at different roots if they shouldn't share one. When unset (the default), this run's FDB output lives in a local per-checkpoint directory under the output root.", + "title": "Root" }, "config": { "anyOf": [ @@ -1213,7 +1260,7 @@ } ], "default": null, - "description": "Configuration for the forecaster run that this temporal downscaler is based on." + "description": "Configuration for the forecaster run that this temporal downscaler is based on. If the forecaster declares output: fdb, the downscaler reads its output directly from FDB (inference_prepare.py's _convert_cutout_to_fdb) instead of the local grib symlink." } }, "required": [ @@ -1317,7 +1364,8 @@ { "type": "null" } - ] + ], + "default": null }, "lapse_rate_correction": { "default": true, @@ -1376,7 +1424,6 @@ "description", "dates", "runs", - "truth", "locations", "profile" ],