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/cape.yaml b/config/cape.yaml new file mode 100644 index 00000000..d6c0d9a0 --- /dev/null +++ b/config/cape.yaml @@ -0,0 +1,57 @@ +# yaml-language-server: $schema=../workflow/tools/config.schema.json +description: | + Test run for the multi-output "secondary_decoder" anemoi architecture (ICON-CH1 + cutout forecaster with a diagnostic CAPE_MU stream on the ICON-CH1 1km grid). + +dates: + start: 2025-08-15T00:00 + end: 2025-08-16T00:00 + frequency: 24h + +runs: + - forecaster: + checkpoint: /scratch/mch/mbosch/devml/varda-forecaster-single-extra_decoders/data/output/checkpoint/a7e9ca1add4e4da8b1034f277e8d0681/inference-last.ckpt + label: cape_mu_bounded + steps: 0/24/6 + config: resources/inference/configs/sgm-multidataset-forecaster-cape-ich1.yaml + extra_requirements: + - "anemoi-inference @ git+https://github.com/martibosch/anemoi-inference@8c870ae7" + - anemoi-transform>=0.4.1 + + - baseline: + label: ICON-CH1-ctrl + root: /store_new/mch/msopr/osm/ICON-CH1-EPS + steps: 0/24/6 + +truth: + label: KENDA-CH1 + root: /store_new/mch/msopr/ml/datasets/mch-ich1-1km-2024-2025-1h-pl13-ifsnames-v1.0.zarr + +experiment: + params: + - CAPE_MU + stratification: + regions: + - jura + root: /scratch/mch/bhendj/regions/Prognoseregionen_LV95_20220517 + thresholds: + CAPE_MU: + gt: [100.0, 500.0, 1000.0] + dashboard: + stratification: + - season + +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 + jobs: 50 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-cape-ich1.yaml b/resources/inference/configs/sgm-multidataset-forecaster-cape-ich1.yaml new file mode 100644 index 00000000..ba21aef6 --- /dev/null +++ b/resources/inference/configs/sgm-multidataset-forecaster-cape-ich1.yaml @@ -0,0 +1,87 @@ +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" + modifiers: + - patches: + - variable: + U_10M: {"param": 165, "shortName": "10u"} + 10u: {"param": 165, "shortName": "10u"} + V_10M: {"param": 166, "shortName": "10v"} + 10v: {"param": 166, "shortName": "10v"} + TD_2M: {"param": 168, "shortName": "2d"} + 2d: {"param": 168, "shortName": "2d"} + T_2M: {"param": 167, "shortName": "2t"} + 2t: {"param": 167, "shortName": "2t"} + FR_LAND: {"param": 172, "shortName": "lsm"} + lsm: {"param": 172, "shortName": "lsm"} + PMSL: {"param": 151, "shortName": "msl"} + msl: {"param": 151, "shortName": "msl"} + PS: {"param": 134, "shortName": "sp"} + sp: {"param": 134, "shortName": "sp"} + SSO_SIGMA: {"param": 163, "shortName": "slor"} + slor: {"param": 163, "shortName": "slor"} + SSO_STDH: {"param": 160, "shortName": "sdor"} + sdor: {"param": 160, "shortName": "sdor"} + TOT_PREC: {"param": 228, "shortName": "tp"} + tp: {"param": 228, "shortName": "tp"} + z: {"param": 129, "shortName": "z"} + "^q_(\\d+)$": {"param": 133, "shortName": "q"} + "^t_(\\d+)$": {"param": 130, "shortName": "t"} + "^u_(\\d+)$": {"param": 131, "shortName": "u"} + "^v_(\\d+)$": {"param": 132, "shortName": "v"} + "^w_(\\d+)$": {"param": 135, "shortName": "w"} + "^z_(\\d+)$": {"param": 129, "shortName": "z"} + # Diagnostic CAPE_MU decoder on ICON grid only + secondary_decoder: + grib: + path: grib/secondary_decoder-{date}{time:04}_{step:03}.grib + # diagnostic-only branch has no initial state fields + 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 supplies CAPE_MU metadata +# for the diagnostic secondary_decoder stream. +patch_metadata: resources/sgm-cape-ich1-patch.yaml 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-cape-ich1-patch.yaml b/resources/inference/metadata/sgm-cape-ich1-patch.yaml new file mode 100644 index 00000000..ced8b90f --- /dev/null +++ b/resources/inference/metadata/sgm-cape-ich1-patch.yaml @@ -0,0 +1,765 @@ +# Metadata patch for the CAPE_MU multi-output forecaster. +# The `data:` block below is shared with sgm-windgust-ich1-patch.yaml — +# keep them in sync when changing primary-stream variable metadata. +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 CAPE_MU emitted by the multi-output "secondary_decoder" stream. + # CAPE_MU is instantaneous — no period/process keys needed. + secondary_decoder: + variables_metadata: + CAPE_MU: + mars: + date: 20050101 + levtype: sfc + param: CAPE_MU + step: 12 + time: 0 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..0e733a16 --- /dev/null +++ b/resources/inference/metadata/sgm-windgust-ich1-patch.yaml @@ -0,0 +1,770 @@ +# Metadata patch for the VMAX_10M multi-output forecaster. +# The `data:` block below is shared with sgm-cape-ich1-patch.yaml — +# keep them in sync when changing primary-stream variable metadata. +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=CAPE_MU.grib b/resources/inference/templates/icon-ch1-shortName=CAPE_MU.grib new file mode 100644 index 00000000..94493a2d Binary files /dev/null and b/resources/inference/templates/icon-ch1-shortName=CAPE_MU.grib differ 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..4390946c 100755 --- a/resources/inference/templates/icon-ch1_generate_templates.sh +++ b/resources/inference/templates/icon-ch1_generate_templates.sh @@ -19,3 +19,17 @@ 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 + +# template for CAPE_MU (instantaneous, atmMU level type) +# NOTE: system grib_set segfaults with COSMO definitions — generate via Python eccodes: +# python3 -c " +# import eccodes, os; from eccodes_cosmo_resources.path import get_definitions_path +# os.environ['ECCODES_DEFINITION_PATH'] = str(get_definitions_path()) +# with open('icon-ch1-typeOfLevel=surface.grib','rb') as f: h = eccodes.codes_grib_new_from_file(f) +# eccodes.codes_set(h,'shortName','CAPE_MU'); eccodes.codes_set(h,'centre','lssw') +# eccodes.codes_set_values(h,[0.0]*eccodes.codes_get_size(h,'values')) +# with open('icon-ch1-shortName=CAPE_MU.grib','wb') as f: eccodes.codes_write(h,f) +# eccodes.codes_release(h)" diff --git a/resources/inference/templates/templates_index_realch1.yaml b/resources/inference/templates/templates_index_realch1.yaml new file mode 100644 index 00000000..99470e3b --- /dev/null +++ b/resources/inference/templates/templates_index_realch1.yaml @@ -0,0 +1,6 @@ +# REAL-CH1 templates (ICON-CH1 1km grid) +# Used by diagnostic output streams of the multi-output anemoi architecture. +- - {levtype: sfc, param: [VMAX_10M]} + - resources/icon-ch1-shortName=VMAX_10M.grib +- - {levtype: sfc, param: [CAPE_MU]} + - resources/icon-ch1-shortName=CAPE_MU.grib diff --git a/src/data_input/__init__.py b/src/data_input/__init__.py index 2af07f7d..40eca6dc 100644 --- a/src/data_input/__init__.py +++ b/src/data_input/__init__.py @@ -411,16 +411,46 @@ def _load_analysis_data_from_zarr( return xr.Dataset(vars_out) -def _collect_ml_grib_files(root: Path, steps: list[int] | None = None) -> list[Path]: +# Known GRIB file prefixes that are NOT diagnostic decoder streams. +# The ``ifs-`` prefix is the global IFS output stream on the N320 grid — +# it must be excluded from auto-discovery of diagnostic streams. +_NON_DIAGNOSTIC_PREFIXES = frozenset({"ifs-"}) + + +def _discover_diagnostic_prefixes(root: Path) -> list[str]: + """Auto-discover diagnostic decoder GRIB stream prefixes in *root*. + + Diagnostic streams use a ``{prefix}20*.grib`` naming convention (e.g. + ``realv2-20250801…``, ``secondary_decoder-20250801…``). The main output + has no prefix (``20*.grib``) and the IFS stream uses ``ifs-``. This + function returns all prefixes that are neither main nor in + ``_NON_DIAGNOSTIC_PREFIXES``. + """ + prefixes: set[str] = set() + for f in root.glob("*-20*.grib"): + name = f.name + idx = name.index("20") + prefix = name[:idx] + if prefix not in _NON_DIAGNOSTIC_PREFIXES: + prefixes.add(prefix) + return sorted(prefixes) + + +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 a diagnostic + prefix like ``"realv2-"`` or ``"secondary_decoder-"`` matches the + corresponding 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 @@ -493,6 +523,7 @@ def variable_name_profile( "surface", "pressure", "entire_atmosphere", + "atmMU", ], ) -> dict[str, Any]: """Resolve variable name profile based on the level type.""" @@ -501,6 +532,7 @@ def variable_name_profile( "mean_sea", "surface", "entire_atmosphere", + "atmMU", ]: return {} elif level_type == "pressure": @@ -1370,10 +1402,32 @@ def load_forecast_data( load_steps = get_steps(steps, params) if any(root.glob("*.grib")): LOG.info("Loading forecasts from GRIB files...") - ds = _load_forecast_data_from_grib( + # First load from the main (un-prefixed) GRIB files, then auto-discover + # diagnostic decoder streams (prefixed GRIB files like realv2-*, secondary_decoder-*) + # and load any remaining params from there. + datasets = [] + main_ds = _load_forecast_data_from_grib( files=_collect_ml_grib_files(root, load_steps), params=load_params, ) + found_params = set(main_ds.data_vars) + if found_params: + datasets.append(main_ds) + + remaining_params = [p for p in load_params if p not in found_params] + if remaining_params: + for prefix in _discover_diagnostic_prefixes(root): + diag_files = _collect_ml_grib_files(root, load_steps, prefix=prefix) + if diag_files: + diag_ds = _load_forecast_data_from_grib(diag_files, remaining_params) + newly_found = set(diag_ds.data_vars) + if newly_found: + LOG.info("Found %s in diagnostic stream '%s'", newly_found, prefix) + datasets.append(diag_ds) + remaining_params = [p for p in remaining_params if p not in newly_found] + if not remaining_params: + break + ds = datasets[0] if len(datasets) == 1 else xr.merge(datasets) # Try to derive elevation from surface geopotential (FIS/z at step 0) # before falling back to ICON grid constants lookup. fis_files = _collect_ml_grib_files(root, [0]) diff --git a/src/evalml/config.py b/src/evalml/config.py index 3980046a..c3573830 100644 --- a/src/evalml/config.py +++ b/src/evalml/config.py @@ -509,6 +509,13 @@ 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." + ), + ) slurm_account: str | None = Field(None, description="SLURM account to charge.") gpus: int | None = Field( None, ge=0, description="Default GPU count per job (0 for non-GPU jobs)." @@ -516,11 +523,17 @@ class DefaultResources(BaseModel): def parsable(self) -> list[str]: """Convert the default resources to a string of key=value pairs.""" - return [ - f"{key}={value}" - for key, value in self.model_dump().items() - if value is not None - ] + 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 f411bbb9..4de8337e 100644 --- a/src/plotting/colormap_defaults.py +++ b/src/plotting/colormap_defaults.py @@ -60,6 +60,8 @@ def _precip_bias_map(accum_h: int) -> dict: | {"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, @@ -73,6 +75,12 @@ def _precip_bias_map(accum_h: int) -> dict: "extend": "both", }, "QV_925": load_ncl_colormap("RH_6lev.ct") | {"extend": "both"}, + "CAPE_MU": { + "cmap": plt.get_cmap("YlOrRd", 10), + "levels": [0, 100, 250, 500, 750, 1000, 1500, 2000, 3000, 4000, 5000], + "extend": "max", + "units": "J kg⁻¹", + }, "CLCT": { # extend="neither" relies on preprocess_field() clipping away from # exact 0/1 (see plot_forecast_frame.py) to avoid a tricontourf bug diff --git a/src/plotting/compat.py b/src/plotting/compat.py index 3f63313a..19ec4128 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, _discover_diagnostic_prefixes PARAMS_MAP = { @@ -32,6 +32,16 @@ def load_state_from_grib( | {PARAMS_MAP[p] for p in (paramlist or []) if p in PARAMS_MAP} ) ds = load_from_grib_file(file, {"parameter.variable": paramlist_extended}) + # Diagnostic params (e.g. VMAX_10M, CAPE_MU) may live in prefixed sibling + # GRIB files (realv2-*, secondary_decoder-*, …). If the main file doesn't + # contain the requested params, try each discovered diagnostic stream. + if paramlist and not any(p in ds for p in paramlist): + for prefix in _discover_diagnostic_prefixes(file.parent): + diag_file = file.with_name(f"{prefix}{file.name}") + if diag_file.exists(): + ds = load_from_grib_file(diag_file, {"parameter.variable": paramlist_extended}) + if any(p in ds for p in paramlist): + break # Rename any IFS shortnames back to COSMO names ifs_rename = { ifs: cosmo for ifs, cosmo in PARAMS_MAP_INV.items() if ifs in ds.data_vars @@ -79,8 +89,8 @@ def load_state_from_grib( ds["longitude"].values.size, np.nan, dtype=float ) global_file = str(file.parent / f"ifs-{file.stem}.grib") - if Path(global_file).exists(): - _paramlist_ecmwf = [PARAMS_MAP[p] for p in paramlist] + _paramlist_ecmwf = [PARAMS_MAP[p] for p in (paramlist or []) if p in PARAMS_MAP] + if Path(global_file).exists() and _paramlist_ecmwf: 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() diff --git a/workflow/rules/common.smk b/workflow/rules/common.smk index acf76746..37074af3 100644 --- a/workflow/rules/common.smk +++ b/workflow/rules/common.smk @@ -466,6 +466,10 @@ PARAMS_WITHOUT_STEP_ZERO_VALUE = { "lcc", "SSRD", "ssrd", + # period maximum: undefined at step 0 + "VMAX_10M", + # secondary decoder diagnostics: not output at step 0 + "CAPE_MU", } @@ -510,6 +514,8 @@ def resolve_leadtimes(steps_spec, requested="all", param=None): ) valid = wanted & supported - if param in PARAMS_WITHOUT_STEP_ZERO_VALUE: + # Strip trailing digits to match aggregated variants like TOT_PREC1 → TOT_PREC + 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: valid = {lt for lt in valid if lt >= step} return sorted(valid) diff --git a/workflow/rules/inference.smk b/workflow/rules/inference.smk index 13d7ccc0..ae9d5ca1 100644 --- a/workflow/rules/inference.smk +++ b/workflow/rules/inference.smk @@ -101,7 +101,7 @@ rule inference_prepare_env: trap "rm -rf $VENV_DIR" EXIT VENV=$VENV_DIR/.venv - PYTHON_VERSION=$(cat {input.metadata} | jq -r ".provenance_training.python") + PYTHON_VERSION=$(cat {input.metadata} | jq -r ".provenance_training.python" | cut -d. -f1,2) echo "[$(date)] Creating virtual environment with Python $PYTHON_VERSION in RAM (/dev/shm)..." uv venv --managed-python --python $PYTHON_VERSION --relocatable --link-mode=copy $VENV source $VENV/bin/activate @@ -289,6 +289,9 @@ rule inference_execute: if config["profile"]["default_resources"].get("slurm_account") else "" ), + expected_steps=lambda wc: ( + lambda s: len(range(int(s[0]), int(s[1]) + 1, int(s[2]))) + )(RUN_CONFIGS[wc.run_id]["steps"].split("/")), # fmt: off shell: """ @@ -312,21 +315,30 @@ rule inference_execute: CMD_ARGS+=(runner.parallel.cluster=slurm) fi - srun \ - --unbuffered \ - {params.account_flag} \ - --partition={resources.slurm_partition} \ - --cpus-per-task={resources.cpus_per_task} \ - --mem-per-cpu={resources.mem_mb_per_cpu} \ - --time={resources.runtime} \ - --gres={resources.gres} \ - --ntasks={resources.ntasks} \ - anemoi-inference run config.yaml "${{CMD_ARGS[@]}}" + anemoi-inference run config.yaml "${{CMD_ARGS[@]}}" }} export -f _run_inference - squashfs-mount {params.env_path}:/user-environment -- bash -c '_run_inference /user-environment' + srun \ + --unbuffered \ + {params.account_flag} \ + --partition={resources.slurm_partition} \ + --cpus-per-task={resources.cpus_per_task} \ + --mem-per-cpu={resources.mem_mb_per_cpu} \ + --time={resources.runtime} \ + --gres={resources.gres} \ + --ntasks={resources.ntasks} \ + squashfs-mount -s {params.env_path}:/user-environment -- bash -c '_run_inference /user-environment' \ + || echo "[$(date)] WARNING: anemoi-inference exited with non-zero status $?, checking if output was produced..." ) >{log} 2>&1 - touch {output.okfile} + + # Verify the expected number of main GRIB files was produced + ACTUAL=$(ls {params.workdir}/grib/20*.grib 2>/dev/null | wc -l) + if [ "$ACTUAL" -ge {params.expected_steps} ]; then + touch {output.okfile} + else + echo "ERROR: Expected at least {params.expected_steps} main GRIB files, found $ACTUAL in {params.workdir}/grib/" >> {log} + exit 1 + fi """ # fmt: on diff --git a/workflow/rules/plot.smk b/workflow/rules/plot.smk index 5afdefd8..96feec9b 100644 --- a/workflow/rules/plot.smk +++ b/workflow/rules/plot.smk @@ -187,7 +187,7 @@ rule plot_scoremaps: resources: slurm_partition="postproc", cpus_per_task=1, - runtime="10m", + runtime="30m", shell: """ export ECCODES_DEFINITION_PATH=$(realpath .venv/share/eccodes-cosmo-resources/definitions) 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/scripts/plot_forecast_frame.py b/workflow/scripts/plot_forecast_frame.py index ee8e564e..a6cf2678 100644 --- a/workflow/scripts/plot_forecast_frame.py +++ b/workflow/scripts/plot_forecast_frame.py @@ -10,6 +10,7 @@ from matplotlib.colors import Colormap import numpy as np +from data_input import parse_aggregated_param from plotting import DOMAINS from plotting import get_projection from plotting import StatePlotter @@ -87,7 +88,8 @@ def preprocess_field(param: str, state: dict): if param == "SP": return ekm_wind.speed(fields["U"], fields["V"]), "m/s" if param == "TOT_PREC": - return np.maximum(fields[param], 0), "mm" + # ML model outputs tp in meters (IFS convention); convert to mm. + return np.maximum(fields[param] * 1000, 0), "mm" if param in ("CLCT", "CLCL"): # Avoid exact 0/1 plateaus breaking tricontourf on orthographic # projections (tmp/reproduce_clct_bug.py). Pair with extend="neither". @@ -136,12 +138,18 @@ def main(): list(regions.keys()), ) - if param == "SP_10M": + # Decompose aggregated params (e.g. TOT_PREC6 → TOT_PREC, accu=6). + # Keep original param for output filenames; use load_param for GRIB loading. + load_param, agg_hours = parse_aggregated_param(param) + if agg_hours is not None: + accu = agg_hours + + if load_param == "SP_10M": paramlist = ["U_10M", "V_10M"] - elif param == "SP": + elif load_param == "SP": paramlist = ["U", "V"] else: - paramlist = [param] + paramlist = [load_param] # Load grib once — shared across all region plots # TODO: fix file pattern & globbing @@ -163,8 +171,11 @@ def main(): LOG.info("Loading grib file %s", grib_file) state = load_state_from_grib(grib_file, paramlist=paramlist) - # tp is accumulated from start of forecast; de-accumulate to get period [lt-accu, lt] - if param == "TOT_PREC": + # De-accumulate cumulative-from-start tp to get period [lt-accu, lt]. + # NOTE: anemoi-inference currently writes period accumulations (stepRange=0-6, + # 6-12, …), so this is a no-op for ML forecaster output but still needed for + # models that write cumulative tp. + if load_param == "TOT_PREC": prev_lt = lead_time - accu if prev_lt > 0: prev_grib_files = list(grib_dir.glob(f"2*_{prev_lt}.grib")) @@ -182,7 +193,7 @@ def main(): ) # Preprocess field once — shared across all region plots - field, units_override = preprocess_field(param, state) + field, units_override = preprocess_field(load_param, state) validtime = state["valid_time"].strftime("%Y%m%d%H%M") for region_name, region_cfg in regions.items(): @@ -220,7 +231,7 @@ def main(): field, title=f"{param}, time: {validtime}", gridline_labels=not region_cfg.get("rotate", False), - **get_style(param, units_override, accu=accu), + **get_style(load_param, units_override, accu=accu), ) if len(state["lam_envelope"]) > 0: subplot.ax.add_geometries( diff --git a/workflow/tools/config.schema.json b/workflow/tools/config.schema.json index ea2b84ba..ec146560 100644 --- a/workflow/tools/config.schema.json +++ b/workflow/tools/config.schema.json @@ -167,6 +167,19 @@ "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" + }, "slurm_account": { "anyOf": [ {