Unify regions stations

workflow/rules/inference.smk

@@ -334,4 +334,4 @@ rule inference_execute:
                 '
         ) >{log} 2>&1
         """
-    # fmt: on
+# fmt: on

workflow/rules/plot.smk

@@ -10,31 +10,33 @@ import pandas as pd
 
 
 def _get_available_baselines(wc) -> list[dict[str, str]]:
-    """Get all available baseline zarr datasets for the given init time."""
+    """Get all available baseline datasets for the given init time."""
     baselines = []
     for baseline_id in BASELINE_CONFIGS:
         root = BASELINE_CONFIGS[baseline_id].get("root")
         steps = BASELINE_CONFIGS[baseline_id].get("steps")
         label = BASELINE_CONFIGS[baseline_id].get("label", baseline_id)
-        year = wc.init_time[2:4]
-        baseline_zarr = f"{root}/FCST{year}.zarr"
-        if Path(baseline_zarr).exists():
-            baselines.append({"zarr": baseline_zarr, "steps": steps, "label": label})
+        if root and Path(root).exists():
+            baselines.append({"root": root, "steps": steps, "label": label})
     if not baselines:
-        raise ValueError(f"No baseline zarr found for init time {wc.init_time}")
+        raise ValueError(f"No baseline data found for init time {wc.init_time}")
     return baselines
 
 
 rule plot_meteogram:
     input:
-        script="workflow/scripts/plot_meteogram.mo.py",
+        script="workflow/scripts/plot_meteogram.py",
         inference_okfile=rules.inference_execute.output.okfile,
         truth=config["truth"]["root"],
         peakweather_dir=rules.data_download_obs_from_peakweather.output.root,
     output:
-        OUT_ROOT
-        / "results/{showcase}/{run_id}/{init_time}/{init_time}_{param}_{sta}.png",
-    # localrule: True
+        expand(
+            OUT_ROOT
+            / "results/{{showcase}}/{{run_id}}/{{init_time}}/{{init_time}}_{{param}}_{sta}.png",
+            sta=config["showcase"]["meteograms"]["stations"],
+        ),
+    log:
+        OUT_ROOT / "logs/{showcase}/{run_id}/{init_time}/plot_meteogram_{param}.log",
     resources:
         slurm_partition="postproc",
         cpus_per_task=1,
@@ -46,15 +48,21 @@ rule plot_meteogram:
         ).resolve(),
         fcst_steps=lambda wc: RUN_CONFIGS[wc.run_id]["steps"],
         fcst_label=lambda wc: RUN_CONFIGS[wc.run_id]["label"],
-        baseline_zarrs=lambda wc: [x["zarr"] for x in _get_available_baselines(wc)],
+        baseline_roots=lambda wc: [x["root"] for x in _get_available_baselines(wc)],
         baseline_steps=lambda wc: [x["steps"] for x in _get_available_baselines(wc)],
         baseline_labels=lambda wc: [x["label"] for x in _get_available_baselines(wc)],
+        outdir=lambda wc: str(
+            (
+                Path(OUT_ROOT) / f"results/{wc.showcase}/{wc.run_id}/{wc.init_time}"
+            ).resolve()
+        ),
+        stations=config["showcase"]["meteograms"]["stations"],
     shell:
         """
         set -euo pipefail
         export ECCODES_DEFINITION_PATH=$(realpath .venv/share/eccodes-cosmo-resources/definitions)
 
-        BASELINE_ZARRS=({params.baseline_zarrs:q})
+        BASELINE_ROOTS=({params.baseline_roots:q})
         BASELINE_STEPS=({params.baseline_steps:q})
         BASELINE_LABELS=({params.baseline_labels:q})
 
@@ -66,63 +74,58 @@ rule plot_meteogram:
             --analysis_label {params.ana_label:q}
             --peakweather {input.peakweather_dir:q}
             --date {wildcards.init_time:q}
-            --outfn {output[0]:q}
+            --outdir {params.outdir:q}
             --param {wildcards.param:q}
-            --station {wildcards.sta:q}
+            --stations {params.stations:q}
         )
 
-        for i in "${{!BASELINE_ZARRS[@]}}"; do
-            CMD_ARGS+=(--baseline "${{BASELINE_ZARRS[$i]}}")
+        for i in "${{!BASELINE_ROOTS[@]}}"; do
+            CMD_ARGS+=(--baseline "${{BASELINE_ROOTS[$i]}}")
             CMD_ARGS+=(--baseline_steps "${{BASELINE_STEPS[$i]}}")
             CMD_ARGS+=(--baseline_label "${{BASELINE_LABELS[$i]}}")
         done
 
-        python {input.script} "${{CMD_ARGS[@]}}"
-        # interactive editing (needs to set localrule: True and use only one core)
-        # marimo edit {input.script} -- "${{CMD_ARGS[@]}}"
+        python {input.script} "${{CMD_ARGS[@]}}" >{log} 2>&1
         """
 
 
 rule plot_forecast_frame:
     input:
-        script="workflow/scripts/plot_forecast_frame.mo.py",
+        script="workflow/scripts/plot_forecast_frame.py",
         inference_okfile=rules.inference_execute.output.okfile,
     output:
+        expand(
+            OUT_ROOT
+            / "data/runs/{{run_id}}/{{init_time}}/frames/frame_{{leadtime}}_{{param}}_{region}.png",
+            region=list(SHOWCASE_REGIONS.keys()),
+        ),
+    log:
         OUT_ROOT
-        / "data/runs/{run_id}/{init_time}/frames/frame_{leadtime}_{param}_{region}.png",
+        / "logs/{run_id}/{init_time}/plot_forecast_frame_{leadtime}_{param}.log",
     wildcard_constraints:
         leadtime=r"\d+",  # only digits
-        region="|".join(map(re.escape, SHOWCASE_REGIONS.keys())),
     resources:
         slurm_partition="postproc",
         cpus_per_task=1,
         runtime="10m",
     params:
-        grib_out_dir=lambda wc: (
-            Path(OUT_ROOT) / f"data/runs/{wc.run_id}/{wc.init_time}/grib"
-        ).resolve(),
-        region_extra=lambda wc: (
-            "--extent {} --projection {}".format(
-                " ".join(map(str, SHOWCASE_REGIONS[wc.region]["extent"])),
-                SHOWCASE_REGIONS[wc.region]["projection"],
-            )
-            if SHOWCASE_REGIONS.get(wc.region, {}).get("extent") is not None
-            else ""
+        grib_out_dir=lambda wc: str(
+            (Path(OUT_ROOT) / f"data/runs/{wc.run_id}/{wc.init_time}/grib").resolve()
+        ),
+        regions_json=json.dumps(SHOWCASE_REGIONS),
+        outdir=lambda wc: str(
+            (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]),
     shell:
         """
         export ECCODES_DEFINITION_PATH=$(realpath .venv/share/eccodes-cosmo-resources/definitions)
         python {input.script} \
-            --input {params.grib_out_dir} --date {wildcards.init_time} --outfn {output[0]} \
-            --param {wildcards.param} --leadtime {wildcards.leadtime} --region {wildcards.region} \
-            {params.region_extra} \
-            --accu {params.accu}
-        # interactive editing (needs to set localrule: True and use only one core)
-        # marimo edit {input.script} -- \
-        #     --input {params.grib_out_dir}  --date {wildcards.init_time} --outfn {output[0]}\
-        #     --param {wildcards.param} --leadtime {wildcards.leadtime} --region {wildcards.region}\
-        #     --accu {params.accu}\
+            --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
         """
 
 

workflow/scripts/inference_extract_requirements.py

@@ -20,7 +20,7 @@
 BASE_DEPENDENCIES = [
     "anemoi-inference",
     "eccodes==2.39.1",
-    "eccodes-cosmo-resources-python",
+    "eccodes-cosmo-resources-python==2.38.3.1",
 ]
 
 # Packages emitted in the output even when only found in provenance (not overrides).