diff --git a/config/showcase.yaml b/config/showcase.yaml index 29d02f2..e968aa5 100644 --- a/config/showcase.yaml +++ b/config/showcase.yaml @@ -26,10 +26,10 @@ runs: baselines: - baseline: - baseline_id: COSMO-E + baseline_id: COSMO-E-1h label: COSMO-E - root: /store_new/mch/msopr/ml/COSMO-E - steps: 0/120/6 + root: /store_new/mch/msopr/ml/COSMO-E_hourly + steps: 0/120/1 stratification: regions: @@ -43,7 +43,7 @@ stratification: analysis: label: COSMO KENDA - analysis_zarr: /scratch/mch/fzanetta/data/anemoi/datasets/mch-co2-an-archive-0p02-2015-2020-6h-v3-pl13.zarr + analysis_zarr: /scratch/mch/fzanetta/data/anemoi/datasets/mch-co2-an-archive-0p02-2015-2020-1h-v3-pl13.zarr locations: output_root: output/ diff --git a/pyproject.toml b/pyproject.toml index 8a46f80..376c556 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,6 +24,7 @@ dependencies = [ "earthkit-plots", "marimo>=0.16.5", "geopandas>=0.14.0", + "peakweather", ] [project.optional-dependencies] @@ -58,3 +59,6 @@ packages = [ "src/verification", "src/data_input" ] + +[tool.uv.sources] +peakweather = { git = "https://github.com/MeteoSwiss/PeakWeather.git" } diff --git a/uv.lock b/uv.lock index b00f135..e2df8d9 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 2 +revision = 3 requires-python = ">=3.11" resolution-markers = [ "python_full_version >= '3.12'", @@ -1006,6 +1006,7 @@ dependencies = [ { name = "meteodata-lab" }, { name = "mlflow" }, { name = "netcdf4" }, + { name = "peakweather" }, { name = "pydantic" }, { name = "pyproj" }, { name = "shapely" }, @@ -1041,6 +1042,7 @@ requires-dist = [ { name = "meteodata-lab", specifier = ">=0.4.0" }, { name = "mlflow", specifier = ">=3.1.1" }, { name = "netcdf4", specifier = ">=1.7.2" }, + { name = "peakweather", git = "https://github.com/MeteoSwiss/PeakWeather.git" }, { name = "pydantic", specifier = ">=2.11.7" }, { name = "pyproj", specifier = ">=3.7.2" }, { name = "shapely", specifier = ">=2.1.2" }, @@ -1351,6 +1353,8 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1f/8e/abdd3f14d735b2929290a018ecf133c901be4874b858dd1c604b9319f064/greenlet-3.2.4-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2523e5246274f54fdadbce8494458a2ebdcdbc7b802318466ac5606d3cded1f8", size = 587684, upload-time = "2025-08-07T13:18:25.164Z" }, { url = "https://files.pythonhosted.org/packages/5d/65/deb2a69c3e5996439b0176f6651e0052542bb6c8f8ec2e3fba97c9768805/greenlet-3.2.4-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:1987de92fec508535687fb807a5cea1560f6196285a4cde35c100b8cd632cc52", size = 1116647, upload-time = "2025-08-07T13:42:38.655Z" }, { url = "https://files.pythonhosted.org/packages/3f/cc/b07000438a29ac5cfb2194bfc128151d52f333cee74dd7dfe3fb733fc16c/greenlet-3.2.4-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:55e9c5affaa6775e2c6b67659f3a71684de4c549b3dd9afca3bc773533d284fa", size = 1142073, upload-time = "2025-08-07T13:18:21.737Z" }, + { url = "https://files.pythonhosted.org/packages/67/24/28a5b2fa42d12b3d7e5614145f0bd89714c34c08be6aabe39c14dd52db34/greenlet-3.2.4-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:c9c6de1940a7d828635fbd254d69db79e54619f165ee7ce32fda763a9cb6a58c", size = 1548385, upload-time = "2025-11-04T12:42:11.067Z" }, + { url = "https://files.pythonhosted.org/packages/6a/05/03f2f0bdd0b0ff9a4f7b99333d57b53a7709c27723ec8123056b084e69cd/greenlet-3.2.4-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:03c5136e7be905045160b1b9fdca93dd6727b180feeafda6818e6496434ed8c5", size = 1613329, upload-time = "2025-11-04T12:42:12.928Z" }, { url = "https://files.pythonhosted.org/packages/d8/0f/30aef242fcab550b0b3520b8e3561156857c94288f0332a79928c31a52cf/greenlet-3.2.4-cp311-cp311-win_amd64.whl", hash = "sha256:9c40adce87eaa9ddb593ccb0fa6a07caf34015a29bf8d344811665b573138db9", size = 299100, upload-time = "2025-08-07T13:44:12.287Z" }, { url = "https://files.pythonhosted.org/packages/44/69/9b804adb5fd0671f367781560eb5eb586c4d495277c93bde4307b9e28068/greenlet-3.2.4-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:3b67ca49f54cede0186854a008109d6ee71f66bd57bb36abd6d0a0267b540cdd", size = 274079, upload-time = "2025-08-07T13:15:45.033Z" }, { url = "https://files.pythonhosted.org/packages/46/e9/d2a80c99f19a153eff70bc451ab78615583b8dac0754cfb942223d2c1a0d/greenlet-3.2.4-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ddf9164e7a5b08e9d22511526865780a576f19ddd00d62f8a665949327fde8bb", size = 640997, upload-time = "2025-08-07T13:42:56.234Z" }, @@ -1360,6 +1364,8 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/19/0d/6660d55f7373b2ff8152401a83e02084956da23ae58cddbfb0b330978fe9/greenlet-3.2.4-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3b3812d8d0c9579967815af437d96623f45c0f2ae5f04e366de62a12d83a8fb0", size = 607586, upload-time = "2025-08-07T13:18:28.544Z" }, { url = "https://files.pythonhosted.org/packages/8e/1a/c953fdedd22d81ee4629afbb38d2f9d71e37d23caace44775a3a969147d4/greenlet-3.2.4-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:abbf57b5a870d30c4675928c37278493044d7c14378350b3aa5d484fa65575f0", size = 1123281, upload-time = "2025-08-07T13:42:39.858Z" }, { url = "https://files.pythonhosted.org/packages/3f/c7/12381b18e21aef2c6bd3a636da1088b888b97b7a0362fac2e4de92405f97/greenlet-3.2.4-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:20fb936b4652b6e307b8f347665e2c615540d4b42b3b4c8a321d8286da7e520f", size = 1151142, upload-time = "2025-08-07T13:18:22.981Z" }, + { url = "https://files.pythonhosted.org/packages/27/45/80935968b53cfd3f33cf99ea5f08227f2646e044568c9b1555b58ffd61c2/greenlet-3.2.4-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:ee7a6ec486883397d70eec05059353b8e83eca9168b9f3f9a361971e77e0bcd0", size = 1564846, upload-time = "2025-11-04T12:42:15.191Z" }, + { url = "https://files.pythonhosted.org/packages/69/02/b7c30e5e04752cb4db6202a3858b149c0710e5453b71a3b2aec5d78a1aab/greenlet-3.2.4-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:326d234cbf337c9c3def0676412eb7040a35a768efc92504b947b3e9cfc7543d", size = 1633814, upload-time = "2025-11-04T12:42:17.175Z" }, { url = "https://files.pythonhosted.org/packages/e9/08/b0814846b79399e585f974bbeebf5580fbe59e258ea7be64d9dfb253c84f/greenlet-3.2.4-cp312-cp312-win_amd64.whl", hash = "sha256:a7d4e128405eea3814a12cc2605e0e6aedb4035bf32697f72deca74de4105e02", size = 299899, upload-time = "2025-08-07T13:38:53.448Z" }, { url = "https://files.pythonhosted.org/packages/49/e8/58c7f85958bda41dafea50497cbd59738c5c43dbbea5ee83d651234398f4/greenlet-3.2.4-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:1a921e542453fe531144e91e1feedf12e07351b1cf6c9e8a3325ea600a715a31", size = 272814, upload-time = "2025-08-07T13:15:50.011Z" }, { url = "https://files.pythonhosted.org/packages/62/dd/b9f59862e9e257a16e4e610480cfffd29e3fae018a68c2332090b53aac3d/greenlet-3.2.4-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:cd3c8e693bff0fff6ba55f140bf390fa92c994083f838fece0f63be121334945", size = 641073, upload-time = "2025-08-07T13:42:57.23Z" }, @@ -1369,6 +1375,8 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ee/43/3cecdc0349359e1a527cbf2e3e28e5f8f06d3343aaf82ca13437a9aa290f/greenlet-3.2.4-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:23768528f2911bcd7e475210822ffb5254ed10d71f4028387e5a99b4c6699671", size = 610497, upload-time = "2025-08-07T13:18:31.636Z" }, { url = "https://files.pythonhosted.org/packages/b8/19/06b6cf5d604e2c382a6f31cafafd6f33d5dea706f4db7bdab184bad2b21d/greenlet-3.2.4-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:00fadb3fedccc447f517ee0d3fd8fe49eae949e1cd0f6a611818f4f6fb7dc83b", size = 1121662, upload-time = "2025-08-07T13:42:41.117Z" }, { url = "https://files.pythonhosted.org/packages/a2/15/0d5e4e1a66fab130d98168fe984c509249c833c1a3c16806b90f253ce7b9/greenlet-3.2.4-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:d25c5091190f2dc0eaa3f950252122edbbadbb682aa7b1ef2f8af0f8c0afefae", size = 1149210, upload-time = "2025-08-07T13:18:24.072Z" }, + { url = "https://files.pythonhosted.org/packages/1c/53/f9c440463b3057485b8594d7a638bed53ba531165ef0ca0e6c364b5cc807/greenlet-3.2.4-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:6e343822feb58ac4d0a1211bd9399de2b3a04963ddeec21530fc426cc121f19b", size = 1564759, upload-time = "2025-11-04T12:42:19.395Z" }, + { url = "https://files.pythonhosted.org/packages/47/e4/3bb4240abdd0a8d23f4f88adec746a3099f0d86bfedb623f063b2e3b4df0/greenlet-3.2.4-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:ca7f6f1f2649b89ce02f6f229d7c19f680a6238af656f61e0115b24857917929", size = 1634288, upload-time = "2025-11-04T12:42:21.174Z" }, { url = "https://files.pythonhosted.org/packages/0b/55/2321e43595e6801e105fcfdee02b34c0f996eb71e6ddffca6b10b7e1d771/greenlet-3.2.4-cp313-cp313-win_amd64.whl", hash = "sha256:554b03b6e73aaabec3745364d6239e9e012d64c68ccd0b8430c64ccc14939a8b", size = 299685, upload-time = "2025-08-07T13:24:38.824Z" }, { url = "https://files.pythonhosted.org/packages/22/5c/85273fd7cc388285632b0498dbbab97596e04b154933dfe0f3e68156c68c/greenlet-3.2.4-cp314-cp314-macosx_11_0_universal2.whl", hash = "sha256:49a30d5fda2507ae77be16479bdb62a660fa51b1eb4928b524975b3bde77b3c0", size = 273586, upload-time = "2025-08-07T13:16:08.004Z" }, { url = "https://files.pythonhosted.org/packages/d1/75/10aeeaa3da9332c2e761e4c50d4c3556c21113ee3f0afa2cf5769946f7a3/greenlet-3.2.4-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:299fd615cd8fc86267b47597123e3f43ad79c9d8a22bebdce535e53550763e2f", size = 686346, upload-time = "2025-08-07T13:42:59.944Z" }, @@ -1376,6 +1384,8 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/dc/8b/29aae55436521f1d6f8ff4e12fb676f3400de7fcf27fccd1d4d17fd8fecd/greenlet-3.2.4-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:b4a1870c51720687af7fa3e7cda6d08d801dae660f75a76f3845b642b4da6ee1", size = 694659, upload-time = "2025-08-07T13:53:17.759Z" }, { url = "https://files.pythonhosted.org/packages/92/2e/ea25914b1ebfde93b6fc4ff46d6864564fba59024e928bdc7de475affc25/greenlet-3.2.4-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:061dc4cf2c34852b052a8620d40f36324554bc192be474b9e9770e8c042fd735", size = 695355, upload-time = "2025-08-07T13:18:34.517Z" }, { url = "https://files.pythonhosted.org/packages/72/60/fc56c62046ec17f6b0d3060564562c64c862948c9d4bc8aa807cf5bd74f4/greenlet-3.2.4-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:44358b9bf66c8576a9f57a590d5f5d6e72fa4228b763d0e43fee6d3b06d3a337", size = 657512, upload-time = "2025-08-07T13:18:33.969Z" }, + { url = "https://files.pythonhosted.org/packages/23/6e/74407aed965a4ab6ddd93a7ded3180b730d281c77b765788419484cdfeef/greenlet-3.2.4-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:2917bdf657f5859fbf3386b12d68ede4cf1f04c90c3a6bc1f013dd68a22e2269", size = 1612508, upload-time = "2025-11-04T12:42:23.427Z" }, + { url = "https://files.pythonhosted.org/packages/0d/da/343cd760ab2f92bac1845ca07ee3faea9fe52bee65f7bcb19f16ad7de08b/greenlet-3.2.4-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:015d48959d4add5d6c9f6c5210ee3803a830dce46356e3bc326d6776bde54681", size = 1680760, upload-time = "2025-11-04T12:42:25.341Z" }, { url = "https://files.pythonhosted.org/packages/e3/a5/6ddab2b4c112be95601c13428db1d8b6608a8b6039816f2ba09c346c08fc/greenlet-3.2.4-cp314-cp314-win_amd64.whl", hash = "sha256:e37ab26028f12dbb0ff65f29a8d3d44a765c61e729647bf2ddfbbed621726f01", size = 303425, upload-time = "2025-08-07T13:32:27.59Z" }, ] @@ -2498,6 +2508,17 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c7/42/20119686047fd6caaa9fd8275bac45aae279866c51bee9f93850b3b89788/pdbufr-0.14.0-py3-none-any.whl", hash = "sha256:7b32ce6de5f1b67f6840e68c899aa4bb2853978294d90de05426a3453e823d38", size = 52625, upload-time = "2025-06-30T15:07:46.197Z" }, ] +[[package]] +name = "peakweather" +version = "0.2.0" +source = { git = "https://github.com/MeteoSwiss/PeakWeather.git#5edfeea1608dba3d08c5a93e3868aad079aa61ec" } +dependencies = [ + { name = "numpy" }, + { name = "pandas" }, + { name = "pyarrow" }, + { name = "tqdm" }, +] + [[package]] name = "pillow" version = "11.3.0" diff --git a/workflow/Snakefile b/workflow/Snakefile index f1d2fa5..eef2ea8 100644 --- a/workflow/Snakefile +++ b/workflow/Snakefile @@ -57,9 +57,16 @@ rule showcase_all: / "showcases/{run_id}/{init_time}/{init_time}_{param}_{region}.gif", init_time=[t.strftime("%Y%m%d%H%M") for t in REFTIMES], run_id=collect_all_candidates(), - param=["TOT_PREC", "T_2M", "SP_10M"], + param=["T_2M", "SP_10M"], region=["globe", "europe", "switzerland"], ), + expand( + OUT_ROOT / "showcases/{run_id}/{init_time}/{init_time}_{param}_{sta}.png", + init_time=[t.strftime("%Y%m%d%H%M") for t in REFTIMES], + run_id=collect_all_candidates(), + param=["T_2M", "SP_10M"], + sta=["GVE", "KLO", "LUG"], + ), expand( OUT_ROOT / "data/runs/{run_id}/summary.md", run_id=collect_all_candidates(), diff --git a/workflow/rules/data.smk b/workflow/rules/data.smk index bef818f..0696657 100644 --- a/workflow/rules/data.smk +++ b/workflow/rules/data.smk @@ -56,3 +56,14 @@ if "extract_cosmo1e" in config.get("include-optional-rules", []): --steps {params.steps} \ > {log} 2>&1 """ + + +rule download_obs_from_peakweather: + localrule: True + output: + peakweather=directory(OUT_ROOT / "data/observations/peakweather"), + run: + from peakweather.dataset import PeakWeatherDataset + + # Download the data from Huggingface + ds = PeakWeatherDataset(root=output.peakweather) diff --git a/workflow/rules/plot.smk b/workflow/rules/plot.smk index d4def1b..a372984 100644 --- a/workflow/rules/plot.smk +++ b/workflow/rules/plot.smk @@ -9,6 +9,52 @@ include: "common.smk" import pandas as pd +def _use_first_baseline_zarr(wc): + """Get the first available baseline zarr for the given init time.""" + for baseline_id in BASELINE_CONFIGS: + root = BASELINE_CONFIGS[baseline_id].get("root") + year = wc.init_time[2:4] + baseline_zarr = f"{root}/FCST{year}.zarr" + if Path(baseline_zarr).exists(): + return baseline_zarr + raise ValueError(f"No baseline zarr found for init time {wc.init_time}") + + +rule plot_meteogram: + input: + script="workflow/scripts/plot_meteogram.mo.py", + inference_okfile=rules.execute_inference.output.okfile, + analysis_zarr=config["analysis"].get("analysis_zarr"), + baseline_zarr=lambda wc: _use_first_baseline_zarr(wc), + peakweather_dir=rules.download_obs_from_peakweather.output.peakweather, + output: + OUT_ROOT / "showcases/{run_id}/{init_time}/{init_time}_{param}_{sta}.png", + # localrule: True + resources: + slurm_partition="postproc", + cpus_per_task=1, + runtime="5m", + params: + grib_out_dir=lambda wc: ( + Path(OUT_ROOT) / f"data/runs/{wc.run_id}/{wc.init_time}/grib" + ).resolve(), + shell: + """ + export ECCODES_DEFINITION_PATH=$(realpath .venv/share/eccodes-cosmo-resources/definitions) + python {input.script} \ + --forecast {params.grib_out_dir} --analysis {input.analysis_zarr} \ + --baseline {input.baseline_zarr} --peakweather {input.peakweather_dir} \ + --date {wildcards.init_time} --outfn {output[0]} \ + --param {wildcards.param} --station {wildcards.sta} + # interactive editing (needs to set localrule: True and use only one core) + # marimo edit {input.script} -- \ + # --forecast {params.grib_out_dir} --analysis {input.analysis_zarr} \ + # --baseline {input.baseline_zarr} --peakweather {input.peakweather_dir} \ + # --date {wildcards.init_time} --outfn {output[0]} \ + # --param {wildcards.param} --station {wildcards.sta} + """ + + rule plot_forecast_frame: input: script="workflow/scripts/plot_forecast_frame.mo.py", diff --git a/workflow/scripts/plot_meteogram.mo.py b/workflow/scripts/plot_meteogram.mo.py new file mode 100644 index 0000000..1738c0e --- /dev/null +++ b/workflow/scripts/plot_meteogram.mo.py @@ -0,0 +1,296 @@ +import marimo + +__generated_with = "0.16.5" +app = marimo.App(width="medium") + + +@app.cell +def _(): + from argparse import ArgumentParser + from pathlib import Path + + import matplotlib.pyplot as plt + import numpy as np + import pandas as pd + import xarray as xr + from meteodatalab import data_source, grib_decoder + from peakweather import PeakWeatherDataset + + from data_input import load_analysis_data_from_zarr, load_baseline_from_zarr + + return ( + ArgumentParser, + Path, + PeakWeatherDataset, + data_source, + grib_decoder, + load_analysis_data_from_zarr, + load_baseline_from_zarr, + np, + pd, + plt, + xr, + ) + + +@app.cell +def _(ArgumentParser, Path): + parser = ArgumentParser() + + parser.add_argument( + "--forecast", type=str, default=None, help="Directory to forecast grib data" + ) + parser.add_argument( + "--analysis", type=str, default=None, help="Path to analysis zarr data" + ) + parser.add_argument( + "--baseline", type=str, default=None, help="Path to baseline zarr data" + ) + parser.add_argument( + "--peakweather", type=str, default=None, help="Path to PeakWeather dataset" + ) + parser.add_argument("--date", type=str, default=None, help="reference datetime") + parser.add_argument("--outfn", type=str, help="output filename") + parser.add_argument("--param", type=str, help="parameter") + parser.add_argument("--station", type=str, help="station") + + args = parser.parse_args() + grib_dir = Path(args.forecast) + zarr_dir_ana = Path(args.analysis) + zarr_dir_base = Path(args.baseline) + peakweather_dir = Path(args.peakweather) + init_time = args.date + outfn = Path(args.outfn) + station = args.station + param = args.param + return ( + grib_dir, + init_time, + outfn, + param, + peakweather_dir, + station, + zarr_dir_ana, + zarr_dir_base, + ) + + +@app.cell +def _(np): + def preprocess_ds(ds, param: str): + ds = ds.copy() + # 10m wind speed + if param == "SP_10M": + ds[param] = np.sqrt(ds.U_10M**2 + ds.V_10M**2) + try: + units = ds["U_10M"].attrs["parameter"]["units"] + except KeyError: + units = None + ds[param].attrs["parameter"] = { + "shortName": "SP_10M", + "units": units, + "name": "10m wind speed", + } + ds = ds.drop_vars(["U_10M", "V_10M"]) + # wind speed from standard-level components + if param == "SP": + ds[param] = np.sqrt(ds.U**2 + ds.V**2) + units = ds.U.attrs["parameter"]["units"] + ds[param].attrs["parameter"] = { + "shortName": "SP", + "units": units, + "name": '"Wind speed', + } + ds = ds.drop_vars(["U", "V"]) + return ds + + return (preprocess_ds,) + + +@app.cell +def load_grib_data( + data_source, + grib_decoder, + grib_dir, + init_time, + load_analysis_data_from_zarr, + load_baseline_from_zarr, + param, + preprocess_ds, + xr, + zarr_dir_ana, + zarr_dir_base, +): + if param == "SP_10M": + paramlist = ["U_10M", "V_10M"] + elif param == "SP": + paramlist = ["U", "V"] + else: + paramlist = [param] + + grib_files = sorted(grib_dir.glob(f"{init_time}*.grib")) + fds = data_source.FileDataSource(datafiles=grib_files) + ds_fct = xr.Dataset(grib_decoder.load(fds, {"param": paramlist})) + ds_fct = preprocess_ds(ds_fct, param) + da_fct = ds_fct[param].squeeze() + + ds_ana = load_analysis_data_from_zarr(zarr_dir_ana, da_fct.valid_time, paramlist) + ds_ana = preprocess_ds(ds_ana, param) + da_ana = ds_ana[param].squeeze() + + steps = list( + range(da_fct.sizes["lead_time"]) + ) # FIX: this will fail if lead_time is not 0,1,2,... + ds_base = load_baseline_from_zarr(zarr_dir_base, da_fct.ref_time, steps, paramlist) + ds_base = preprocess_ds(ds_base, param) + da_base = ds_base[param].squeeze() + return da_ana, da_base, da_fct + + +@app.cell +def _(PeakWeatherDataset, da_fct, np, param, peakweather_dir, station): + if param == "T_2M": + parameter = "temperature" + offset = 273.15 # K to C + elif param == "SP_10M": + parameter = "wind_speed" + offset = 0 + elif param == "TOT_PREC": + parameter = "precipitation" + offset = 0 + else: + raise NotImplementedError( + f"The mapping for {param=} to PeakWeather is not implemented" + ) + + peakweather = PeakWeatherDataset(root=peakweather_dir, freq="1h") + obs, mask = peakweather.get_observations( + parameters=[parameter], + stations=station, + first_date=np.datetime_as_string(da_fct.valid_time.values[0]), + last_date=np.datetime_as_string(da_fct.valid_time.values[-1]), + return_mask=True, + ) + obs = obs.loc[:, mask.iloc[0]].droplevel("name", axis=1) + obs + return obs, offset, peakweather + + +@app.cell +def _(peakweather): + stations = peakweather.stations_table + stations.index.names = ["station"] + stations + return (stations,) + + +@app.cell +def _(da_ana, da_base, da_fct, np, pd, stations): + def nearest_yx_euclid(ds, lat_s, lon_s): + """ + Return (y_idx, x_idx) of the grid point nearest to (lat_s, lon_s) + using Euclidean distance in degrees. + """ + try: + lat2d = ds["lat"] # (y, x) + lon2d = ds["lon"] # (y, x) + except KeyError: + lat2d = ds["latitude"] # (y, x) + lon2d = ds["longitude"] # (y, x) + + dist2 = (lat2d - lat_s) ** 2 + (lon2d - lon_s) ** 2 + flat_idx = np.nanargmin(dist2.values) + y_idx, x_idx = np.unravel_index(flat_idx, dist2.shape) + return int(y_idx), int(x_idx) + + def get_fct_idx_row(row): + return nearest_yx_euclid(da_fct, row["latitude"], row["longitude"]) + + idxs_fct = stations.apply(get_fct_idx_row, axis=1, result_type="expand") + idxs_fct.columns = ["fct_y_idx", "fct_x_idx"] + + def get_ana_idx_row(row): + return nearest_yx_euclid(da_ana, row["latitude"], row["longitude"]) + + idxs_ana = stations.apply(get_ana_idx_row, axis=1, result_type="expand") + idxs_ana.columns = ["ana_y_idx", "ana_x_idx"] + + def get_base_idx_row(row): + return nearest_yx_euclid(da_base, row["latitude"], row["longitude"]) + + idxs_base = stations.apply(get_base_idx_row, axis=1, result_type="expand") + idxs_base.columns = ["base_y_idx", "base_x_idx"] + + sta_idxs = pd.concat([stations, idxs_fct, idxs_ana, idxs_base], axis=1) + sta_idxs + return (sta_idxs,) + + +@app.cell +def _(da_ana, da_base, da_fct, init_time, obs, offset, outfn, plt, sta_idxs, station): + # station indices + row = sta_idxs.loc[station] + fct_x_idx, fct_y_idx = row.fct_x_idx, row.fct_y_idx + ana_x_idx, ana_y_idx = row.ana_x_idx, row.ana_y_idx + base_x_idx, base_y_idx = row.base_x_idx, row.base_y_idx + + fig, ax = plt.subplots() + + # station + ax.plot( + obs.index.to_pydatetime(), + obs.to_numpy() + offset, + color="k", + ls="--", + label=station, + ) + + # analysis + ana2plot = da_ana.isel(x=ana_x_idx, y=ana_y_idx) + ax.plot( + ana2plot["time"].values, + ana2plot.values, + color="k", + ls="-", + label="analysis", + ) + + # baseline + base2plot = da_base.isel(x=base_x_idx, y=base_y_idx) + ax.plot( + base2plot["time"].values, + base2plot.values, + color="C1", + label="baseline", + ) + + # forecast + fct2plot = da_fct.isel(x=fct_x_idx, y=fct_y_idx) + ax.plot( + fct2plot["valid_time"].values, + fct2plot.values, + color="C0", + label="forecast", + ) + + ax.legend() + + param2plot = da_fct.attrs.get("parameter", {}) + short = param2plot.get("shortName", "") + units = param2plot.get("units", "") + name = param2plot.get("name", "") + + ax.set_ylabel(f"{short} ({units})" if short or units else "") + ax.set_title(f"{init_time} {name} at {station}") + + plt.savefig(outfn) + return + + +@app.cell +def _(): + return + + +if __name__ == "__main__": + app.run()