diff --git a/config/forecasters-ich1-oper-fixed.yaml b/config/forecasters-ich1-oper-fixed.yaml index a182f49f..8a332955 100644 --- a/config/forecasters-ich1-oper-fixed.yaml +++ b/config/forecasters-ich1-oper-fixed.yaml @@ -85,5 +85,3 @@ profile: runtime: "1h" gpus: 0 jobs: 50 - batch_rules: - plot_forecast_frame: 8 diff --git a/config/forecasters-ich1-oper.yaml b/config/forecasters-ich1-oper.yaml index 28c02121..a2956738 100644 --- a/config/forecasters-ich1-oper.yaml +++ b/config/forecasters-ich1-oper.yaml @@ -82,5 +82,3 @@ profile: runtime: "1h" gpus: 0 jobs: 50 - batch_rules: - plot_forecast_frame: 8 diff --git a/config/forecasters-ich1.yaml b/config/forecasters-ich1.yaml index 05d2c6a1..1bf6f7be 100644 --- a/config/forecasters-ich1.yaml +++ b/config/forecasters-ich1.yaml @@ -94,5 +94,3 @@ profile: runtime: "1h" gpus: 0 jobs: 50 - batch_rules: - plot_forecast_frame: 8 diff --git a/config/interpolators-ich1.yaml b/config/interpolators-ich1.yaml index 38eb22fd..bd1268a6 100644 --- a/config/interpolators-ich1.yaml +++ b/config/interpolators-ich1.yaml @@ -116,11 +116,11 @@ showcase: animations: enabled: true domains: + # - globe - europe + # - alps + - icon-ch - switzerland - - name: alpine_arc - extent: [3.0, 17.0, 43.5, 48.5] - projection: orthographic locations: output_root: output/ @@ -136,5 +136,3 @@ profile: runtime: "1h" gpus: 0 jobs: 50 - batch_rules: - plot_forecast_frame: 8 diff --git a/src/plotting/__init__.py b/src/plotting/__init__.py index a60d2dc7..163bf74b 100644 --- a/src/plotting/__init__.py +++ b/src/plotting/__init__.py @@ -42,10 +42,22 @@ def get_projection(name: str) -> "ccrs.Projection": "extent": [-2.6, 19.5, 40.2, 52.3], "projection": _PROJECTIONS["orthographic"], }, - "switzerland": { + "icon-ch": { "extent": [0, 17.5, 40.5, 53.0], "projection": _PROJECTIONS["orthographic"], }, + "alps": { + "extent": [4.5, 16.0, 43.5, 48.5], + "projection": _PROJECTIONS["orthographic"], + }, + "radar": { + "extent": [3.2, 12.5, 43.6, 49.4], + "projection": _PROJECTIONS["orthographic"], + }, + "switzerland": { + "extent": [5.6, 10.8, 45.0, 48.6], + "projection": _PROJECTIONS["orthographic"], + }, } @@ -227,11 +239,12 @@ def _prepare_plot_kwargs( ) # avoid interpolation being performed by earthkit-plots resulting in an error return style, plot_kwargs - # Continuous mode: remove None entries to avoid matplotlib errors - if plot_kwargs.get("colors", None) is None: - plot_kwargs.pop("colors", None) - if plot_kwargs.get("levels", None) is None: - plot_kwargs.pop("levels", None) + # Continuous mode: remove None entries so configure_style doesn't call + # style.with_overrides(colors=None, vmin=None, ...) and wipe the style's + # pre-configured values. + for key in ("colors", "levels", "cmap", "norm", "vmin", "vmax"): + if plot_kwargs.get(key) is None: + plot_kwargs.pop(key, None) return style, plot_kwargs diff --git a/src/plotting/colormap_defaults.py b/src/plotting/colormap_defaults.py index 6794b82a..88c065e6 100644 --- a/src/plotting/colormap_defaults.py +++ b/src/plotting/colormap_defaults.py @@ -12,17 +12,35 @@ def _fallback(): _CMAP_DEFAULTS = { - "SP": {"cmap": plt.get_cmap("coolwarm", 11), "vmin": 800 * 100, "vmax": 1100 * 100}, - "TD_2M": load_ncl_colormap("t2m_29lev.ct"), - "T_2M": load_ncl_colormap("t2m_29lev.ct") | {"units": "degC"}, - "V_10M": load_ncl_colormap("modified_uv_17lev.ct") | {"units": "m/s"}, - "U_10M": load_ncl_colormap("modified_uv_17lev.ct") | {"units": "m/s"}, - "SP_10M": load_ncl_colormap("modified_uv_17lev.ct") | {"units": "m/s"}, - # "10si": {"cmap": plt.get_cmap("GnBu", 11), "vmin": 0, "vmax": 25}, - "T_850": {"cmap": plt.get_cmap("inferno", 11), "vmin": 220, "vmax": 310}, - "FI_850": {"cmap": plt.get_cmap("coolwarm", 11), "vmin": 8000, "vmax": 17000}, - "QV_925": load_ncl_colormap("RH_6lev.ct"), + "SP": { + "cmap": plt.get_cmap("coolwarm", 11), + "vmin": 800 * 100, + "vmax": 1100 * 100, + "extend": "both", + }, + "TD_2M": load_ncl_colormap("t2m_29lev.ct") | {"extend": "both"}, + "T_2M": load_ncl_colormap("t2m_29lev.ct") | {"units": "degC", "extend": "both"}, + "V_10M": load_ncl_colormap("modified_uv_17lev.ct") + | {"units": "m/s", "extend": "both"}, + "U_10M": load_ncl_colormap("modified_uv_17lev.ct") + | {"units": "m/s", "extend": "both"}, + "SP_10M": load_ncl_colormap("modified_uv_17lev.ct") + | {"units": "m/s", "extend": "max"}, + "T_850": { + "cmap": plt.get_cmap("inferno", 11), + "vmin": 220, + "vmax": 310, + "extend": "both", + }, + "FI_850": { + "cmap": plt.get_cmap("coolwarm", 11), + "vmin": 8000, + "vmax": 17000, + "extend": "both", + }, + "QV_925": load_ncl_colormap("RH_6lev.ct") | {"extend": "both"}, "TOT_PREC_1H": { + "extend": "max", "colors": [ "#ffffff", "#ebf6ff", @@ -68,6 +86,7 @@ def _fallback(): ], }, "TOT_PREC_6H": { + "extend": "max", "colors": [ "#ffffff", "#d6e2ff", diff --git a/src/plotting/colormap_loader.py b/src/plotting/colormap_loader.py index 11bd3a22..0b8f442c 100644 --- a/src/plotting/colormap_loader.py +++ b/src/plotting/colormap_loader.py @@ -1,6 +1,6 @@ import pathlib import numpy as np -from matplotlib.colors import BoundaryNorm +from matplotlib.colors import BoundaryNorm, ListedColormap # Base directory for colormap files BASE_DIR = ( @@ -50,6 +50,10 @@ def load_ncl_colormap(filename): if len(rgb) != n_levs + 1: raise ValueError(f"Expected {n_levs} RGB rows, got {len(rgb)}.") - norm = BoundaryNorm(boundaries=bounds, ncolors=len(rgb) - 1) + n_intervals = len(bounds) - 1 + cmap = ListedColormap(rgb[1:-1], name=filename) + cmap.set_under(rgb[0]) + cmap.set_over(rgb[-1]) + norm = BoundaryNorm(boundaries=bounds, ncolors=n_intervals) - return {"cmap": rgb[1:-1], "norm": norm, "bounds": bounds} + return {"cmap": cmap, "norm": norm, "bounds": bounds} diff --git a/tests/unit/test_colormaps.py b/tests/unit/test_colormaps.py index b815f6ab..9f8ab565 100644 --- a/tests/unit/test_colormaps.py +++ b/tests/unit/test_colormaps.py @@ -15,18 +15,20 @@ def test_load_valid_colormap(monkeypatch, tmp_path): monkeypatch.setattr(colormap_loader, "BASE_DIR", tmp_path) result = colormap_loader.load_ncl_colormap("test_colormap.ct") - # cmap is a list of RGB colors (earthkit-plots consumes it as `colors`) - assert isinstance(result["cmap"], np.ndarray) + assert isinstance(result["cmap"], ListedColormap) assert isinstance(result["norm"], BoundaryNorm) cmap = result["cmap"] norm = result["norm"] bounds = result["bounds"] - # cmap holds the n_levs-1 inner colors, under/over rows excluded - assert cmap.shape == (2, 3) - assert np.allclose(cmap[0], (40 / 255, 50 / 255, 60 / 255)) - assert np.allclose(cmap[1], (70 / 255, 80 / 255, 90 / 255)) + # cmap holds the n_levs-1 inner colors, under/over excluded + assert cmap.N == 2 + assert np.allclose(cmap(0), (*[x / 255 for x in (40, 50, 60)], 1.0)) + assert np.allclose(cmap(1), (*[x / 255 for x in (70, 80, 90)], 1.0)) + # under/over are set separately + assert np.allclose(cmap.get_under(), (*[x / 255 for x in (10, 20, 30)], 1.0)) + assert np.allclose(cmap.get_over(), (*[x / 255 for x in (100, 110, 120)], 1.0)) # bounds assert np.allclose(norm.boundaries, [0, 1, 2]) assert np.allclose(bounds, [0, 1, 2]) @@ -70,11 +72,6 @@ def test_cmap_defaults_smoke(field, var): vmin = var.get("vmin", None) vmax = var.get("vmax", None) - # NCL colormaps are stored as a list of RGB colors (earthkit-plots - # consumes them as `colors`); wrap into a ListedColormap to plot here. - if isinstance(cmap, np.ndarray): - cmap = ListedColormap(cmap) - # make some synthetic data data = np.linspace(0, 1, 100).reshape(10, 10) diff --git a/workflow/rules/plot.smk b/workflow/rules/plot.smk index f8b6f73a..91ff1080 100644 --- a/workflow/rules/plot.smk +++ b/workflow/rules/plot.smk @@ -143,7 +143,8 @@ def get_leadtimes(wc): rule make_forecast_animation: input: lambda wc: expand( - rules.plot_forecast_frame.output, + OUT_ROOT + / "data/runs/{run_id}/{init_time}/frames/frame_{leadtime}_{param}_{region}.png", run_id=wc.run_id, init_time=wc.init_time, param=wc.param, diff --git a/workflow/scripts/plot_forecast_frame.py b/workflow/scripts/plot_forecast_frame.py index 51a234b1..6b685ebb 100644 --- a/workflow/scripts/plot_forecast_frame.py +++ b/workflow/scripts/plot_forecast_frame.py @@ -7,6 +7,7 @@ from earthkit.meteo.utils.convert import kelvin_to_celsius import earthkit.meteo.wind as ekm_wind import earthkit.plots as ekp +from matplotlib.colors import Colormap import numpy as np from plotting import DOMAINS @@ -27,19 +28,47 @@ def get_style(param, units_override=None, accu=1): lookup = f"{param}_{accu}H" if param == "TOT_PREC" else param cfg = CMAP_DEFAULTS[lookup] units = units_override if units_override is not None else cfg.get("units", "") + + bounds = cfg.get("bounds", cfg.get("levels", None)) + prebuilt_cmap = cfg.get("cmap", None) + + # When the config provides a pre-built matplotlib Colormap (e.g. a + # ListedColormap), we must use the earthkit Style's vmin/vmax path, which + # handles isinstance(colors, Colormap) correctly. The levels path calls + # cmap_and_norm → len(Colormap) → TypeError. + # + # earthkit's configure_style intercepts _STYLE_KWARGS (cmap/colors/levels/vmin/vmax) + # from tricontourf kwargs before matplotlib sees them. To work around this: + # - embed the Colormap directly in the Style via colors= + # - inject bounds as 'levels' into style._kwargs so they survive to matplotlib + # - pass only norm= as a kwarg (not in _STYLE_KWARGS, so not intercepted) + extend = cfg.get("extend", "both") + + if isinstance(prebuilt_cmap, Colormap) and bounds is not None: + style = ekp.styles.Style( + colors=prebuilt_cmap, + vmin=bounds[0], + vmax=bounds[-1], + extend=extend, + units=units, + ) + style._kwargs["levels"] = list(bounds) + return { + "style": style, + "norm": cfg.get("norm", None), + } + return { "style": ekp.styles.Style( - levels=cfg.get("bounds", cfg.get("levels", None)), - extend="both", + levels=bounds, + extend=extend, units=units, colors=cfg.get("colors", None), ), "norm": cfg.get("norm", None), - "cmap": cfg.get("cmap", None), - "levels": cfg.get("levels", None), + "cmap": prebuilt_cmap, "vmin": cfg.get("vmin", None), "vmax": cfg.get("vmax", None), - "colors": cfg.get("colors", None), }