diff --git a/.gitignore b/.gitignore index 50b2ff7d..3995e831 100644 --- a/.gitignore +++ b/.gitignore @@ -55,3 +55,6 @@ uv.lock # evalml .evalml_snakemake_cmd.txt + +# Paper plotting scripts + their figure outputs (not part of the realv2 showcase) +paper_plots/ diff --git a/config/windgust-peakweather.yaml b/config/windgust-peakweather.yaml new file mode 100644 index 00000000..19ec72a3 --- /dev/null +++ b/config/windgust-peakweather.yaml @@ -0,0 +1,84 @@ +# yaml-language-server: $schema=../workflow/tools/config.schema.json +description: | + As windgust.yaml (realv2 VMAX_10M showcase), but verifies against PeakWeather + station observations instead of the REAL-CH1 gridded analysis. + +# One week of daily initialisations for verification. +dates: + start: 2024-02-01T00:00 + end: 2024-02-07T00:00 + frequency: 24h + +runs: + - forecaster: + checkpoint: /scratch/mch/rradev/output/checkpoint/9efa01f8c7464328897edb2c03a407c2/inference-last.ckpt + label: realv2_vmax10m + steps: 0/120/6 + config: resources/inference/configs/sgm-multidataset-forecaster-windgust-ich1.yaml + extra_requirements: + # checkpoint needs: (1) EmptyInput.create_input_state propagates the date + # (else KeyError 'date' in add_initial_forcings_to_input_state for the + - git+https://github.com/radiradev/anemoi-inference.git@fix/empty-input-propagate-date + + - eccodes==2.39.1 + - eccodes-cosmo-resources-python==2.38.3.1 + + - baseline: + label: ICON-CH1-ctrl + root: /store_new/mch/msopr/osm/ICON-CH1-EPS + steps: 0/33/6 + +truth: + label: PeakWeather + root: output/data/observations/peakweather + +experiment: + params: + - VMAX_10M + stratification: + regions: + - jura + - mittelland + - voralpen + - alpennordhang + - innerealpentaeler + - alpensuedseite + root: /scratch/mch/bhendj/regions/Prognoseregionen_LV95_20220517 + thresholds: + VMAX_10M: + gt: [10.0, 20.0, 30.0] + dashboard: + stratification: + - season + +showcase: + params: + - T_2M + - SP_10M + - VMAX_10M + meteograms: + enabled: false + stations: [JUN] + animations: + enabled: true + domains: + - icon-ch + - switzerland + +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_extra: "--exclusive" # whole nodes; avoid oversubscription + jobs: 50 + batch_rules: + plot_forecast_frame: 32 diff --git a/config/windgust.yaml b/config/windgust.yaml new file mode 100644 index 00000000..d0ecdb54 --- /dev/null +++ b/config/windgust.yaml @@ -0,0 +1,84 @@ +# yaml-language-server: $schema=../workflow/tools/config.schema.json +description: | + Showcase the multi-output "realv2" anemoi architecture (ICON-CH1 cutout forecaster + with a diagnostic VMAX_10M stream on the REAL-CH1 / ICON-CH1 1km grid). + +# One week of daily initialisations for verification. +dates: + start: 2024-02-01T00:00 + end: 2024-02-07T00:00 + frequency: 24h + +runs: + - forecaster: + checkpoint: /scratch/mch/rradev/output/checkpoint/9efa01f8c7464328897edb2c03a407c2/inference-last.ckpt + label: realv2_vmax10m + steps: 0/120/6 + config: resources/inference/configs/sgm-multidataset-forecaster-windgust-ich1.yaml + extra_requirements: + # checkpoint needs: (1) EmptyInput.create_input_state propagates the date + # (else KeyError 'date' in add_initial_forcings_to_input_state for the + - git+https://github.com/radiradev/anemoi-inference.git@fix/empty-input-propagate-date + + - eccodes==2.39.1 + - eccodes-cosmo-resources-python==2.38.3.1 + + - baseline: + label: ICON-CH1-ctrl + root: /store_new/mch/msopr/osm/ICON-CH1-EPS + steps: 0/33/6 + +truth: + label: REAL-CH1 + root: /store_new/mch/msopr/ml/datasets/mch-realch1-fdb-1km-2005-2025-1h-pl13-v2.0.zarr + +experiment: + params: + - VMAX_10M + stratification: + regions: + - jura + - mittelland + - voralpen + - alpennordhang + - innerealpentaeler + - alpensuedseite + root: /scratch/mch/bhendj/regions/Prognoseregionen_LV95_20220517 + thresholds: + VMAX_10M: + gt: [10.0, 20.0, 30.0] + dashboard: + stratification: + - season + +showcase: + params: + - T_2M + - SP_10M + - VMAX_10M + meteograms: + enabled: false + stations: [JUN] + animations: + enabled: true + domains: + - icon-ch + - switzerland + +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_extra: "--exclusive" # whole nodes; avoid oversubscription + jobs: 50 + batch_rules: + plot_forecast_frame: 32 diff --git a/resources/inference/configs/sgm-multidataset-forecaster-windgust-ich1.yaml b/resources/inference/configs/sgm-multidataset-forecaster-windgust-ich1.yaml new file mode 100644 index 00000000..03fcec5d --- /dev/null +++ b/resources/inference/configs/sgm-multidataset-forecaster-windgust-ich1.yaml @@ -0,0 +1,58 @@ +lead_time: 120h +write_initial_state: true +allow_nans: true + +env: + ANEMOI_INFERENCE_NUM_CHUNKS: 8 # OOM error if not set + +# inputs +input: + test: + use_original_paths: true + +output: + # Global cutout state (ICON-CH1 LAM + AIFS N320 global). + data: + tee: + - grib: + path: grib/{date}{time:04}_{step:03}.grib + encoding: + typeOfGeneratingProcess: 2 + centre: lssw + templates: + samples: resources/templates_index_icon.yaml + post_processors: + - extract_mask: # keep only LAM points + mask: "lam_0/cutout_mask" + as_slice: true + - grib: + path: grib/ifs-{date}{time:04}_{step:03}.grib + encoding: + typeOfGeneratingProcess: 2 + centre: ecmf + templates: + samples: resources/templates_index_ifs.yaml + post_processors: + - extract_mask: # removes LAM points + mask: "lam_0/cutout_mask" + as_slice: true + inverse: true + - assign_mask: # fill local/global overlapping points with nan + mask: "global/cutout_mask" + # Dianostic decoder on ICON grid only + realv2: + grib: + path: grib/realv2-{date}{time:04}_{step:03}.grib + # diagnostic is first valid at the model step (e.g. +6h) + write_initial_state: false + encoding: + typeOfGeneratingProcess: 2 + centre: lssw + templates: + samples: resources/templates_index_realch1.yaml + +# Remaps the `data` stream's IFS variable names (2t, 10u, tp, ...) to the COSMO +# shortNames expected by the ICON GRIB templates, and gives the realv2 VMAX_10M +# diagnostic a whole-hour max period (['6h', '12h']) so its GRIB time-processing +# is encoded correctly. +patch_metadata: resources/sgm-windgust-ich1-patch.yaml diff --git a/resources/inference/metadata/sgm-windgust-ich1-patch.yaml b/resources/inference/metadata/sgm-windgust-ich1-patch.yaml new file mode 100644 index 00000000..61a237b9 --- /dev/null +++ b/resources/inference/metadata/sgm-windgust-ich1-patch.yaml @@ -0,0 +1,767 @@ +config: + dataloader: + test: + datasets: + data: + dataset_config: + dataset: + cutout: + - dataset: /store_new/mch/msopr/ml/datasets/mch-ich1-1km-2024-2025-1h-pl13-ifsnames-v1.0.zarr + - dataset: /store_new/mch/msopr/ml/datasets/aifs-od-an-oper-0001-mars-n320-2016-2025-6h-v1-combined-land.zarr + start: null + end: null + +dataset: + data: + variables_metadata: + slor: + mars: + date: 20050101 + levtype: sfc + param: SSO_SIGMA + step: 12 + time: 0 + sdor: + mars: + date: 20050101 + levtype: sfc + param: SSO_STDH + step: 12 + time: 0 + 10u: + mars: + date: 20050101 + levtype: sfc + param: U_10M + step: 12 + time: 0 + 10v: + mars: + date: 20050101 + levtype: sfc + param: V_10M + step: 12 + time: 0 + 2d: + mars: + date: 20050101 + levtype: sfc + param: TD_2M + step: 12 + time: 0 + 2t: + mars: + date: 20050101 + levtype: sfc + param: T_2M + step: 12 + time: 0 + cos_julian_day: + computed_forcing: true + constant_in_time: false + cos_latitude: + computed_forcing: true + constant_in_time: true + cos_local_time: + computed_forcing: true + constant_in_time: false + cos_longitude: + computed_forcing: true + constant_in_time: true + insolation: + computed_forcing: true + constant_in_time: false + lsm: + constant_in_time: true + mars: + date: 20050101 + levtype: sfc + param: FR_LAND + step: 0 + time: 12 + msl: + mars: + date: 20050101 + levtype: sfc + param: PMSL + step: 12 + time: 0 + q_100: + mars: + date: 20050101 + levelist: 100 + levtype: pl + param: QV + step: 12 + time: 0 + q_1000: + mars: + date: 20050101 + levelist: 1000 + levtype: pl + param: QV + step: 12 + time: 0 + q_150: + mars: + date: 20050101 + levelist: 150 + levtype: pl + param: QV + step: 12 + time: 0 + q_200: + mars: + date: 20050101 + levelist: 200 + levtype: pl + param: QV + step: 12 + time: 0 + q_250: + mars: + date: 20050101 + levelist: 250 + levtype: pl + param: QV + step: 12 + time: 0 + q_300: + mars: + date: 20050101 + levelist: 300 + levtype: pl + param: QV + step: 12 + time: 0 + q_400: + mars: + date: 20050101 + levelist: 400 + levtype: pl + param: QV + step: 12 + time: 0 + q_50: + mars: + date: 20050101 + levelist: 50 + levtype: pl + param: QV + step: 12 + time: 0 + q_500: + mars: + date: 20050101 + levelist: 500 + levtype: pl + param: QV + step: 12 + time: 0 + q_600: + mars: + date: 20050101 + levelist: 600 + levtype: pl + param: QV + step: 12 + time: 0 + q_700: + mars: + date: 20050101 + levelist: 700 + levtype: pl + param: QV + step: 12 + time: 0 + q_850: + mars: + date: 20050101 + levelist: 850 + levtype: pl + param: QV + step: 12 + time: 0 + q_925: + mars: + date: 20050101 + levelist: 925 + levtype: pl + param: QV + step: 12 + time: 0 + sin_julian_day: + computed_forcing: true + constant_in_time: false + sin_latitude: + computed_forcing: true + constant_in_time: true + sin_local_time: + computed_forcing: true + constant_in_time: false + sin_longitude: + computed_forcing: true + constant_in_time: true + sp: + mars: + date: 20050101 + levtype: sfc + param: PS + step: 12 + time: 0 + t_100: + mars: + date: 20050101 + levelist: 100 + levtype: pl + param: T + step: 12 + time: 0 + t_1000: + mars: + date: 20050101 + levelist: 1000 + levtype: pl + param: T + step: 12 + time: 0 + t_150: + mars: + date: 20050101 + levelist: 150 + levtype: pl + param: T + step: 12 + time: 0 + t_200: + mars: + date: 20050101 + levelist: 200 + levtype: pl + param: T + step: 12 + time: 0 + t_250: + mars: + date: 20050101 + levelist: 250 + levtype: pl + param: T + step: 12 + time: 0 + t_300: + mars: + date: 20050101 + levelist: 300 + levtype: pl + param: T + step: 12 + time: 0 + t_400: + mars: + date: 20050101 + levelist: 400 + levtype: pl + param: T + step: 12 + time: 0 + t_50: + mars: + date: 20050101 + levelist: 50 + levtype: pl + param: T + step: 12 + time: 0 + t_500: + mars: + date: 20050101 + levelist: 500 + levtype: pl + param: T + step: 12 + time: 0 + t_600: + mars: + date: 20050101 + levelist: 600 + levtype: pl + param: T + step: 12 + time: 0 + t_700: + mars: + date: 20050101 + levelist: 700 + levtype: pl + param: T + step: 12 + time: 0 + t_850: + mars: + date: 20050101 + levelist: 850 + levtype: pl + param: T + step: 12 + time: 0 + t_925: + mars: + date: 20050101 + levelist: 925 + levtype: pl + param: T + step: 12 + time: 0 + tp: + mars: + date: 20050101 + levtype: sfc + param: TOT_PREC + step: 12 + time: 0 + period: + - 6h + - 12h + process: accumulation + u_100: + mars: + date: 20050101 + levelist: 100 + levtype: pl + param: U + step: 12 + time: 0 + u_1000: + mars: + date: 20050101 + levelist: 1000 + levtype: pl + param: U + step: 12 + time: 0 + u_150: + mars: + date: 20050101 + levelist: 150 + levtype: pl + param: U + step: 12 + time: 0 + u_200: + mars: + date: 20050101 + levelist: 200 + levtype: pl + param: U + step: 12 + time: 0 + u_250: + mars: + date: 20050101 + levelist: 250 + levtype: pl + param: U + step: 12 + time: 0 + u_300: + mars: + date: 20050101 + levelist: 300 + levtype: pl + param: U + step: 12 + time: 0 + u_400: + mars: + date: 20050101 + levelist: 400 + levtype: pl + param: U + step: 12 + time: 0 + u_50: + mars: + date: 20050101 + levelist: 50 + levtype: pl + param: U + step: 12 + time: 0 + u_500: + mars: + date: 20050101 + levelist: 500 + levtype: pl + param: U + step: 12 + time: 0 + u_600: + mars: + date: 20050101 + levelist: 600 + levtype: pl + param: U + step: 12 + time: 0 + u_700: + mars: + date: 20050101 + levelist: 700 + levtype: pl + param: U + step: 12 + time: 0 + u_850: + mars: + date: 20050101 + levelist: 850 + levtype: pl + param: U + step: 12 + time: 0 + u_925: + mars: + date: 20050101 + levelist: 925 + levtype: pl + param: U + step: 12 + time: 0 + v_100: + mars: + date: 20050101 + levelist: 100 + levtype: pl + param: V + step: 12 + time: 0 + v_1000: + mars: + date: 20050101 + levelist: 1000 + levtype: pl + param: V + step: 12 + time: 0 + v_150: + mars: + date: 20050101 + levelist: 150 + levtype: pl + param: V + step: 12 + time: 0 + v_200: + mars: + date: 20050101 + levelist: 200 + levtype: pl + param: V + step: 12 + time: 0 + v_250: + mars: + date: 20050101 + levelist: 250 + levtype: pl + param: V + step: 12 + time: 0 + v_300: + mars: + date: 20050101 + levelist: 300 + levtype: pl + param: V + step: 12 + time: 0 + v_400: + mars: + date: 20050101 + levelist: 400 + levtype: pl + param: V + step: 12 + time: 0 + v_50: + mars: + date: 20050101 + levelist: 50 + levtype: pl + param: V + step: 12 + time: 0 + v_500: + mars: + date: 20050101 + levelist: 500 + levtype: pl + param: V + step: 12 + time: 0 + v_600: + mars: + date: 20050101 + levelist: 600 + levtype: pl + param: V + step: 12 + time: 0 + v_700: + mars: + date: 20050101 + levelist: 700 + levtype: pl + param: V + step: 12 + time: 0 + v_850: + mars: + date: 20050101 + levelist: 850 + levtype: pl + param: V + step: 12 + time: 0 + v_925: + mars: + date: 20050101 + levelist: 925 + levtype: pl + param: V + step: 12 + time: 0 + w_100: + mars: + date: 20050101 + levelist: 100 + levtype: pl + param: OMEGA + step: 12 + time: 0 + w_1000: + mars: + date: 20050101 + levelist: 1000 + levtype: pl + param: OMEGA + step: 12 + time: 0 + w_150: + mars: + date: 20050101 + levelist: 150 + levtype: pl + param: OMEGA + step: 12 + time: 0 + w_200: + mars: + date: 20050101 + levelist: 200 + levtype: pl + param: OMEGA + step: 12 + time: 0 + w_250: + mars: + date: 20050101 + levelist: 250 + levtype: pl + param: OMEGA + step: 12 + time: 0 + w_300: + mars: + date: 20050101 + levelist: 300 + levtype: pl + param: OMEGA + step: 12 + time: 0 + w_400: + mars: + date: 20050101 + levelist: 400 + levtype: pl + param: OMEGA + step: 12 + time: 0 + w_50: + mars: + date: 20050101 + levelist: 50 + levtype: pl + param: OMEGA + step: 12 + time: 0 + w_500: + mars: + date: 20050101 + levelist: 500 + levtype: pl + param: OMEGA + step: 12 + time: 0 + w_600: + mars: + date: 20050101 + levelist: 600 + levtype: pl + param: OMEGA + step: 12 + time: 0 + w_700: + mars: + date: 20050101 + levelist: 700 + levtype: pl + param: OMEGA + step: 12 + time: 0 + w_850: + mars: + date: 20050101 + levelist: 850 + levtype: pl + param: OMEGA + step: 12 + time: 0 + w_925: + mars: + date: 20050101 + levelist: 925 + levtype: pl + param: OMEGA + step: 12 + time: 0 + z: + constant_in_time: true + mars: + date: 20050101 + levelist: null + levtype: sfc + param: FIS + step: 0 + time: 12 + z_100: + mars: + date: 20050101 + levelist: 100 + levtype: pl + param: FI + step: 12 + time: 0 + z_1000: + mars: + date: 20050101 + levelist: 1000 + levtype: pl + param: FI + step: 12 + time: 0 + z_150: + mars: + date: 20050101 + levelist: 150 + levtype: pl + param: FI + step: 12 + time: 0 + z_200: + mars: + date: 20050101 + levelist: 200 + levtype: pl + param: FI + step: 12 + time: 0 + z_250: + mars: + date: 20050101 + levelist: 250 + levtype: pl + param: FI + step: 12 + time: 0 + z_300: + mars: + date: 20050101 + levelist: 300 + levtype: pl + param: FI + step: 12 + time: 0 + z_400: + mars: + date: 20050101 + levelist: 400 + levtype: pl + param: FI + step: 12 + time: 0 + z_50: + mars: + date: 20050101 + levelist: 50 + levtype: pl + param: FI + step: 12 + time: 0 + z_500: + mars: + date: 20050101 + levelist: 500 + levtype: pl + param: FI + step: 12 + time: 0 + z_600: + mars: + date: 20050101 + levelist: 600 + levtype: pl + param: FI + step: 12 + time: 0 + z_700: + mars: + date: 20050101 + levelist: 700 + levtype: pl + param: FI + step: 12 + time: 0 + z_850: + mars: + date: 20050101 + levelist: 850 + levtype: pl + param: FI + step: 12 + time: 0 + z_925: + mars: + date: 20050101 + levelist: 925 + levtype: pl + param: FI + step: 12 + time: 0 + # Diagnostic VMAX_10M emitted by the multi-output "realv2" stream. Gives the + # diagnostic a whole-hour max period so its GRIB time-processing is encoded + # correctly. + realv2: + variables_metadata: + VMAX_10M: + mars: + date: 20050101 + levtype: sfc + param: VMAX_10M + step: 12 + time: 0 + period: + - 6h + - 12h + process: maximum diff --git a/resources/inference/templates/icon-ch1-shortName=VMAX_10M.grib b/resources/inference/templates/icon-ch1-shortName=VMAX_10M.grib new file mode 100644 index 00000000..5cc22b92 Binary files /dev/null and b/resources/inference/templates/icon-ch1-shortName=VMAX_10M.grib differ diff --git a/resources/inference/templates/icon-ch1_generate_templates.sh b/resources/inference/templates/icon-ch1_generate_templates.sh index ec1c47ef..f419cac6 100755 --- a/resources/inference/templates/icon-ch1_generate_templates.sh +++ b/resources/inference/templates/icon-ch1_generate_templates.sh @@ -19,3 +19,6 @@ grib_copy -w shortName=T,level=500 $PL_SAMPLE /dev/stdout | grib_set -d 0 - icon # template for typeOfLevel=meanSea grib_copy -w shortName=PMSL $SFC_SAMPLE /dev/stdout | grib_set -d 0 - icon-ch1-typeOfLevel=meanSea.grib + +#template for windgust +grib_set -s shortName=VMAX_10M,level=10 -d 0 icon-ch1-typeOfLevel=heightAboveGround.grib icon-ch1-shortName=VMAX_10M.grib diff --git a/resources/inference/templates/templates_index_realch1.yaml b/resources/inference/templates/templates_index_realch1.yaml new file mode 100644 index 00000000..16ec62cf --- /dev/null +++ b/resources/inference/templates/templates_index_realch1.yaml @@ -0,0 +1,5 @@ +# REAL-CH1 templates (ICON-CH1 1km grid) +# Used by the realv2 output stream of the multi-output anemoi architecture, which +# emits the diagnostic VMAX_10M (maximum 10 m wind speed) on the ICON-CH1 1km grid. +- - {levtype: sfc, param: [VMAX_10M]} + - resources/icon-ch1-shortName=VMAX_10M.grib diff --git a/src/data_input/__init__.py b/src/data_input/__init__.py index 6c14005a..6b6cb0db 100644 --- a/src/data_input/__init__.py +++ b/src/data_input/__init__.py @@ -64,6 +64,7 @@ def load_analysis_data_from_zarr( "PS": "sp", "PMSL": "msl", "TOT_PREC": "tp", + "VMAX_10M": "VMAX_10M", } tot_prec_string = "TOT_PREC_6H" if min(np.diff(steps)) == 6 else "TOT_PREC_1H" PARAMS_MAP_COSMO1 = { @@ -115,16 +116,26 @@ def load_analysis_data_from_zarr( return _select_valid_times(ds, times) -def _collect_ml_grib_files(root: Path, steps: list[int] | None = None) -> list[Path]: +# Diagnostic params emitted by the multi-output "realv2" stream. They are written to a +# sibling ``realv2-*.grib`` file (same LAM grid as the main output) rather than the main +# ``{date}{time}_{step}.grib``, so loaders must source them from there. +REALV2_PARAMS = frozenset({"VMAX_10M"}) + + +def _collect_ml_grib_files( + root: Path, steps: list[int] | None = None, prefix: str = "" +) -> list[Path]: """Return GRIB files for an ML inference run (flat directory layout). When `steps` is provided, the discovered files are filtered to those whose - name ends with ``_{step:03d}.grib``. + name ends with ``_{step:03d}.grib``. `prefix` selects an output stream: the + default ``""`` matches the main ``20*.grib`` outputs, while ``"realv2-"`` + matches the diagnostic ``realv2-20*.grib`` sibling files. """ # TODO: this glob pattern is a dirty fix for anemoi-inference writing outputs # with wrong formatting. Eventually we will either have to have a fix upstream # or write a single output file. - files = sorted(root.glob("20*.grib")) + files = sorted(root.glob(f"{prefix}20*.grib")) if steps is None: return files @@ -747,11 +758,27 @@ def load_forecast_data( root = Path(root) if any(root.glob("*.grib")): LOG.info("Loading forecasts from GRIB files...") - return load_forecast_data_from_grib( - # NOTE: root is already for a specific reftime - files=_collect_ml_grib_files(root, steps), - params=params, - ) + # Diagnostic "realv2" params (e.g. VMAX_10M) live in sibling realv2-*.grib + # files; load them separately and merge with the main output stream. + main_params = [p for p in params if p not in REALV2_PARAMS] + realv2_params = [p for p in params if p in REALV2_PARAMS] + datasets = [] + if main_params: + datasets.append( + load_forecast_data_from_grib( + # NOTE: root is already for a specific reftime + files=_collect_ml_grib_files(root, steps), + params=main_params, + ) + ) + if realv2_params: + datasets.append( + load_forecast_data_from_grib( + files=_collect_ml_grib_files(root, steps, prefix="realv2-"), + params=realv2_params, + ) + ) + return datasets[0] if len(datasets) == 1 else xr.merge(datasets) if "INCA" in root.parts: LOG.info("Loading INCA baseline from NetCDF files...") return load_INCA_baseline_from_netcdf(root, reftime, steps, params) diff --git a/src/evalml/config.py b/src/evalml/config.py index e101afc3..75bf501e 100644 --- a/src/evalml/config.py +++ b/src/evalml/config.py @@ -404,10 +404,27 @@ class DefaultResources(BaseModel): cpus_per_task: int = Field(..., ge=1, description="Number of CPUs per task.") mem_mb_per_cpu: int = Field(..., ge=1, description="Memory per CPU in MB.") runtime: str = Field(..., description="Maximum runtime, e.g. '1h'.") + slurm_extra: str | None = Field( + None, + description=( + "Extra sbatch flags applied to every job, e.g. '--exclusive' to " + "request whole nodes and avoid sharing (oversubscribing) them." + ), + ) def parsable(self) -> str: """Convert the default resources to a string of key=value pairs.""" - return [f"{key}={value}" for key, value in self.model_dump().items()] + items = [] + for key, value in self.model_dump().items(): + if value is None: + continue + if key == "slurm_extra": + # Snakemake evaluates resource values as Python expressions; wrap the + # flag string in quotes so e.g. --exclusive is taken as a literal string. + items.append(f'{key}="{value}"') + else: + items.append(f"{key}={value}") + return items class GlobalResources(BaseModel): diff --git a/src/plotting/colormap_defaults.py b/src/plotting/colormap_defaults.py index 88c065e6..4d4c26a4 100644 --- a/src/plotting/colormap_defaults.py +++ b/src/plotting/colormap_defaults.py @@ -26,6 +26,8 @@ def _fallback(): | {"units": "m/s", "extend": "both"}, "SP_10M": load_ncl_colormap("modified_uv_17lev.ct") | {"units": "m/s", "extend": "max"}, + "VMAX_10M": load_ncl_colormap("modified_uv_17lev.ct") + | {"units": "m/s", "extend": "max"}, "T_850": { "cmap": plt.get_cmap("inferno", 11), "vmin": 220, diff --git a/src/plotting/compat.py b/src/plotting/compat.py index ef14e473..a1211b2e 100644 --- a/src/plotting/compat.py +++ b/src/plotting/compat.py @@ -4,7 +4,7 @@ import geopandas as gpd import numpy as np from shapely.geometry import MultiPoint -from data_input import load_from_grib_file +from data_input import load_from_grib_file, REALV2_PARAMS PARAMS_MAP = { @@ -23,6 +23,10 @@ def load_state_from_grib( file: Path, paramlist: list[str] | None = None ) -> dict[str, np.ndarray | dict[str, np.ndarray] | gpd.GeoSeries]: + if paramlist and set(paramlist) <= REALV2_PARAMS: + realv2_file = file.with_name(f"realv2-{file.name}") + if realv2_file.exists(): + file = realv2_file ds = load_from_grib_file(file, {"parameter.variable": paramlist}) state = {} ref_param = next((p for p in (paramlist or []) if p in ds), None) diff --git a/workflow/rules/plot.smk b/workflow/rules/plot.smk index 91ff1080..b20f10fa 100644 --- a/workflow/rules/plot.smk +++ b/workflow/rules/plot.smk @@ -134,8 +134,8 @@ rule plot_forecast_frame: def get_leadtimes(wc): """Get all lead times from the run config.""" start, end, step = map(int, RUN_CONFIGS[wc.run_id]["steps"].split("/")) - # skip lead time 0 for diagnostic variables - if wc.param in ["tp", "TOT_PREC"] and start == 0: + # skip lead time 0 for diagnostic variables (accumulations and period maxima) + if wc.param in ["tp", "TOT_PREC", "VMAX_10M"] and start == 0: start += step return [f"{i}" for i in range(start, end + 1, step)] diff --git a/workflow/scripts/inference_extract_requirements.py b/workflow/scripts/inference_extract_requirements.py index dc3441ae..6ca043fa 100644 --- a/workflow/scripts/inference_extract_requirements.py +++ b/workflow/scripts/inference_extract_requirements.py @@ -9,6 +9,7 @@ import argparse import json +import re import sys import warnings from packaging.version import Version, InvalidVersion @@ -35,11 +36,17 @@ "torch-geometric", ] +def _requirement_name(token: str) -> str: + """Return the canonical package name from a requirement token. + + Strips any version specifier (``==``, ``>=``, ``<``, ``~=``, ``!=``, …) so that + e.g. ``eccodes>=2.44.0,<2.48.0`` and ``eccodes==2.39.1`` both key as ``eccodes``. + """ + return re.split(r"[<>=!~]", token, maxsplit=1)[0].strip() + + # Canonical names of BASE_DEPENDENCIES for membership tests (strips version pins). -_BASE_DEPENDENCY_NAMES: set[str] = set() -for _dep in BASE_DEPENDENCIES: - _base_name = _dep.split("==")[0].strip() if "==" in _dep else _dep.strip() - _BASE_DEPENDENCY_NAMES.add(_base_name) +_BASE_DEPENDENCY_NAMES: set[str] = {_requirement_name(_dep) for _dep in BASE_DEPENDENCIES} def load_provenance(metadata_path: str) -> dict: @@ -230,6 +237,12 @@ 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 re.search(r"[<>!~]", item): + # Non-`==` version specifier (e.g. ``eccodes>=2.44.0,<2.48.0``). Key by the + # canonical name so a later ``name==version`` override replaces it; keep the + # full specifier (incl. operator) as the value. + name = _requirement_name(item) + result[name] = item[len(name):].strip() else: result[item] = None @@ -318,7 +331,12 @@ def format_requirements( for name, version in sorted(pypi_requirements.items()): if name not in allowed: continue - line = f"{name}=={version}" if version else f"{name}" + if not version: + line = f"{name}" + elif version[0] in "<>=!~": # a PEP 508 specifier like ">=2.44.0,<2.48.0" + line = f"{name}{version}" + else: + line = f"{name}=={version}" line += " # Extra (not from checkpoint)" if name in overrides else "" lines.append(line) diff --git a/workflow/tools/config.schema.json b/workflow/tools/config.schema.json index d5760f85..4d35da24 100644 --- a/workflow/tools/config.schema.json +++ b/workflow/tools/config.schema.json @@ -159,6 +159,19 @@ "description": "Maximum runtime, e.g. '1h'.", "title": "Runtime", "type": "string" + }, + "slurm_extra": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Extra sbatch flags applied to every job, e.g. '--exclusive' to request whole nodes and avoid sharing (oversubscribing) them.", + "title": "Slurm Extra" } }, "required": [