Add plots for tp

Author: OpheliaMiralles

src/plotting/__init__.py

@@ -1,3 +1,4 @@
+from contextlib import contextmanager
 from functools import cached_property
 from pathlib import Path
 
@@ -19,7 +20,7 @@
 
 # Mapping of region names to their geographic extent and projection
 # extent [lon_min, lon_max, lat_min, lat_max] in PlateCarree coordinates
-REGIONS = {
+DOMAINS = {
     "globe": {
         "extent": None,  # full globe view
         "projection": _PROJECTIONS["orthographic"],
@@ -40,7 +41,7 @@
 
 
 class StatePlotter:
-    """A class to plot state fields on various REGIONS."""
+    """A class to plot state fields on various DOMAINS."""
 
     def __init__(
         self,
@@ -168,7 +169,6 @@ def plot_field(
         field = field[mask]
         field = field[-1] if field.ndim == 2 else field.squeeze()
         finite = np.isfinite(field)
-
         # TODO: clip data to domain would make plotting faster (especially tripcolor)
         # tried using Map.domain.extract() but too memory heavy (probably uses
         # meshgrid in the background), implement clipping with e.g.
@@ -183,24 +183,72 @@ def plot_field(
         # subplot.tripcolor(  # also works but is slower
         # have to overwrite _plot_kwargs to avoid earthkit-plots trying to pass transform
         # PlateCarree based on NumpySource
-        subplot._plot_kwargs = lambda source: {}
-        subplot.tricontourf(
-            x=x[finite],
-            y=y[finite],
-            z=field[finite],
-            style=style,
-            transform=proj,
-            **kwargs,  # for earthkit.plots to work properly cmap and norm are needed here
-        )
+
+        # Normalize style and color-related kwargs
+        style_to_use, plot_kwargs = self._prepare_plot_kwargs(style, kwargs)
+
+        # Temporarily suppress earthkit-plots internal source-based kwargs
+        with self._temporary_plot_kwargs_override(subplot):
+            subplot.tricontourf(
+                x=x[finite],
+                y=y[finite],
+                z=field[finite],
+                style=style_to_use,
+                transform=proj,
+                **plot_kwargs,
+            )  # for earthkit.plots to work properly cmap and norm are needed here
         # TODO: gridlines etc would be nicer to have in the init, but I didn't get
         # them to overlay the plot layer
+
         subplot.standard_layers()
 
         if colorbar:
             subplot.legend()
         if title:
             subplot.title(title)
 
+    def _prepare_plot_kwargs(
+        self,
+        style: ekp.styles.Style | None,
+        kwargs: dict,
+    ) -> tuple[ekp.styles.Style | None, dict]:
+        """Return a cleaned style and plot kwargs without mutating the input."""
+        plot_kwargs = dict(kwargs)
+
+        # Discrete colors mode: if explicit 'colors' provided, drop cmap
+        colors = plot_kwargs.get("colors", None)
+        if colors is not None:
+            plot_kwargs.pop("cmap", None)
+            plot_kwargs.setdefault(
+                "no_style", True
+            )  # 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)
+
+        return style, plot_kwargs
+
+    @contextmanager
+    def _temporary_plot_kwargs_override(self, subplot: ekp.Map):
+        """Temporarily override internal _plot_kwargs to avoid transform issues."""
+        has_attr = hasattr(subplot, "_plot_kwargs")
+        old = getattr(subplot, "_plot_kwargs", None)
+        subplot._plot_kwargs = lambda source: {}
+        try:
+            yield
+        finally:
+            if has_attr:
+                subplot._plot_kwargs = old
+            else:
+                try:
+                    delattr(subplot, "_plot_kwargs")
+                except Exception:
+                    pass
+
     @cached_property
     def _orthographic_tri(self) -> Triangulation:
         """Compute the triangulation for the orthographic projection."""

src/plotting/colormap_defaults.py

@@ -22,6 +22,28 @@ def _fallback():
     "t_850": {"cmap": plt.get_cmap("inferno", 11), "vmin": 220, "vmax": 310},
     "z_850": {"cmap": plt.get_cmap("coolwarm", 11), "vmin": 8000, "vmax": 17000},
     "q_925": load_ncl_colormap("RH_6lev.ct"),
+    "tp": {
+        "colors": [
+            "#ffffff",
+            "#04e9e7",
+            "#019ff4",
+            "#0300f4",
+            "#02fd02",
+            "#01c501",
+            "#008e00",
+            "#fdf802",
+            "#e5bc00",
+            "#fd9500",
+            "#fd0000",
+            "#d40000",
+            "#bc0000",
+            "#f800fd",
+        ],
+        "vmin": 0,
+        "vmax": 100,
+        "units": "mm",
+        "levels": [0, 0.05, 0.1, 0.25, 0.5, 1, 1.5, 2, 3, 4, 5, 6, 7, 100],
+    },
 }
 
 CMAP_DEFAULTS = defaultdict(_fallback, _CMAP_DEFAULTS)

src/plotting/compat.py

@@ -28,9 +28,12 @@ def load_state_from_grib(
     lam_hull = MultiPoint(list(zip(lons.tolist(), lats.tolist()))).convex_hull
     state["lam_envelope"] = gpd.GeoSeries([lam_hull], crs="EPSG:4326")
     state["fields"] = {}
-    for param in paramlist:
+    for param in paramlist or []:
         if param in ds:
             state["fields"][param] = ds[param].values.flatten()
+        else:
+            # initialize with NaNs to keep consistent length
+            state["fields"][param] = np.full(lats.size, np.nan, dtype=float)
     global_file = str(file.parent / f"ifs-{file.stem}.grib")
     if Path(global_file).exists():
         global_file = str(file.parent / f"ifs-{file.stem}.grib")
@@ -40,18 +43,30 @@ def load_state_from_grib(
             for u in fds_global
             if u.metadata("param") in paramlist
         }
-        global_lats = fds_global.metadata("latitudes")[0]
-        global_lons = fds_global.metadata("longitudes")[0]
-        if max(global_lons) > 180:
-            global_lons = ((global_lons + 180) % 360) - 180
-        mask = np.where(~np.isnan(ds_global[paramlist[0]]))[0]
-        state["longitudes"] = np.concatenate([state["longitudes"], global_lons[mask]])
-        state["latitudes"] = np.concatenate([state["latitudes"], global_lats[mask]])
-        for param in paramlist:
-            if param in ds and param in ds_global:
-                state["fields"][param] = np.concatenate(
-                    [state["fields"][param], ds_global[param][mask]]
+        # Use first key from ds_global instead of paramlist[0]
+        ref_key = next(iter(ds_global), None)
+        if ref_key is not None:
+            global_lats = fds_global.metadata("latitudes")[0]
+            global_lons = fds_global.metadata("longitudes")[0]
+            if max(global_lons) > 180:
+                global_lons = ((global_lons + 180) % 360) - 180
+            mask = np.where(~np.isnan(ds_global[ref_key]))[0]
+            n_add = int(mask.size)
+            state["longitudes"] = np.concatenate(
+                [state["longitudes"], global_lons[mask]]
+            )
+            state["latitudes"] = np.concatenate([state["latitudes"], global_lats[mask]])
+            for param in paramlist or state["fields"].keys():
+                add = (
+                    ds_global[param][mask]
+                    if param in ds_global
+                    else np.full(n_add, np.nan, dtype=float)
                 )
+                # ensure base array exists (in case param wasn't in local ds)
+                base = state["fields"].get(
+                    param, np.full(lats.size, np.nan, dtype=float)
+                )
+                state["fields"][param] = np.concatenate([base, add])
     return state
 
 

workflow/Snakefile

@@ -52,7 +52,7 @@ 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=["2t", "10sp"],
+            param=["tp", "2t", "10sp"],
             region=["globe", "europe", "switzerland"],
         ),
 

workflow/rules/plot.smk

@@ -42,6 +42,8 @@ rule plot_forecast_frame:
 def get_leadtimes(wc):
     """Get all lead times from the run config."""
     start, end, step = map(int, RUN_CONFIGS[wc.run_id]["steps"].split("/"))
+    if wc.param == "tp":
+        start += step  # skip lead time 0 for precipitation
     return [f"{i:03}" for i in range(start, end + 1, step)]
 
 

workflow/scripts/plot_forecast_frame.mo.py

@@ -14,7 +14,7 @@ def _():
     import earthkit.plots as ekp
     import numpy as np
 
-    from plotting import REGIONS
+    from plotting import DOMAINS
     from plotting import StatePlotter
     from plotting.colormap_defaults import CMAP_DEFAULTS
     from plotting.compat import load_state_from_grib
@@ -28,7 +28,7 @@ def _():
         load_state_from_grib,
         logging,
         np,
-        REGIONS,
+        DOMAINS,
         ccrs,
     )
 
@@ -97,12 +97,17 @@ def get_style(param, units_override=None):
         units = units_override if units_override is not None else cfg.get("units", "")
         return {
             "style": ekp.styles.Style(
-                levels=cfg.get("bounds", None),
+                levels=cfg.get("bounds", cfg.get("levels", None)),
                 extend="both",
                 units=units,
+                colors=cfg.get("colors", None),
             ),
-            "cmap": cfg["cmap"],
             "norm": cfg.get("norm", None),
+            "cmap": cfg.get("cmap", None),
+            "levels": cfg.get("levels", None),
+            "vmin": cfg.get("vmin", None),
+            "vmax": cfg.get("vmax", None),
+            "colors": cfg.get("colors", None),
         }
 
     return (get_style,)
@@ -132,6 +137,13 @@ def _ms_to_knots(arr):
             except Exception:
                 return arr * 1.943844
 
+        def _m_to_mm(arr):
+            # robust conversion with pint, fallback if dtype unsupported
+            try:
+                return (_ureg.Quantity(arr, _ureg.meter).to(_ureg.millimeter)).magnitude
+            except Exception:
+                return arr * 1000
+
     except Exception:
         LOG.warning("pint not available; falling back hardcoded conversions")
 
@@ -141,6 +153,9 @@ def _k_to_c(arr):
         def _ms_to_knots(arr):
             return arr * 1.943844
 
+        def _m_to_mm(arr):
+            return arr * 1000
+
     def preprocess_field(param: str, state: dict):
         """
         - Temperatures (2t, 2d, t, d): K -> °C
@@ -162,6 +177,8 @@ def preprocess_field(param: str, state: dict):
             u = fields["u"]
             v = fields["v"]
             return np.sqrt(u**2 + v**2), "m/s"
+        if param == "tp":
+            return _m_to_mm(fields[param]), "mm"  # convert m to mm
         # default: passthrough
         return fields[param], None
 
@@ -179,7 +196,7 @@ def _(
     preprocess_field,
     region,
     state,
-    REGIONS,
+    DOMAINS,
     ccrs,
 ):
     # plot individual fields
@@ -191,8 +208,8 @@ def _(
     fig = plotter.init_geoaxes(
         nrows=1,
         ncols=1,
-        projection=REGIONS[region]["projection"],
-        bbox=REGIONS[region]["extent"],
+        projection=DOMAINS[region]["projection"],
+        bbox=DOMAINS[region]["extent"],
         name=region,
         size=(6, 6),
     )