Support global models

README.md

@@ -81,6 +81,7 @@ truth:
 experiment:
   stratification:
     regions:
+      - icon  # first item will be used as default verification region
       - jura
       - mittelland
       - voralpen

config/aifs-single.yaml

@@ -18,16 +18,16 @@ runs:
       config: resources/inference/configs/aifs-single-forecaster.yaml
       extra_requirements:
         - torch-geometric==2.4.0
-        - anemoi-inference==0.6.3
+        - anemoi-inference==0.11.1
         - anemoi-models==0.5.0
         - https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp312-cp312-linux_x86_64.whl
       disable_local_eccodes_definitions: true
-      inference_resources:
-        slurm_partition: preemptible
-        gpus: 1
+
 truth:
-  label: KENDA-CH1
-  root: /store_new/mch/msopr/ml/datasets/mch-ich1-1km-2024-2025-1h-pl13-v1.0.zarr
+  label: ERA5-o96
+  root: /store_new/mch/msopr/ml/datasets/aifs-ea-an-oper-0001-mars-o96-1979-2024-1h-v3-with-era51.zarr
+
+lapse_rate_correction: false
 
 experiment:
   params:
@@ -41,7 +41,8 @@ experiment:
       # - init_hour
       - season
   stratification:
-    regions: []
+    regions:
+      - global
 
 locations:
   output_root: output/

config/forecasters-ich1-oper-fixed.yaml

@@ -46,6 +46,7 @@ experiment:
     - TOT_PREC6
   stratification:
     regions:
+      - icon
       - jura
       - mittelland
       - voralpen

config/forecasters-ich1-oper.yaml

@@ -44,6 +44,7 @@ experiment:
     - TOT_PREC6
   stratification:
     regions:
+      - icon
       - jura
       - mittelland
       - voralpen

config/forecasters-ich1.yaml

@@ -57,6 +57,7 @@ experiment:
     - PMSL
   stratification:
     regions:
+      - icon
       - jura
       - mittelland
       - voralpen

config/forecasters-ich1_mec_ffv2.yaml

@@ -44,6 +44,7 @@ experiment:
     - TOT_PREC
   stratification:
     regions:
+      - icon
       - jura
     root: /store_new/mch/msopr/ml/regions/Prognoseregionen_LV95_20220517
   thresholds:

config/varda-single-1.0.yaml

@@ -54,6 +54,7 @@ experiment:
     - PMSL
   stratification:
     regions:
+      - icon
       - mittelland
       - berge
       - alpennordseite

resources/inference/configs/aifs-single-forecaster.yaml

@@ -1,4 +1,6 @@
-input: test
+input:
+  test:
+    use_original_paths: true
 
 allow_nans: true
 
@@ -12,6 +14,11 @@ post_processors:
   - accumulate_from_start_of_forecast:
       accumulations:
         - tp
+  - forward_transform_filter:
+      rescale:
+        scale: 1000 # convert units from m to kg m-2
+        offset: 0
+        param: tp
 
 output:
   grib:

src/data_input/__init__.py

@@ -26,7 +26,16 @@
     "2d": "TD_2M",
     "sp": "PS",
     "lsm": "FR_LAND",
-    "z": "FSI",
+    "tcc": "CLCT",
+    "lcc": "CLCL",
+    # TODO: ssrd is treated as a plain per-step field (no de-accumulation),
+    # which only holds because it's not currently listed in any
+    # accumulate_from_start_of_forecast.accumulations in the inference
+    # configs (unlike tp, see _tot_prec_handling). If ssrd/strd are ever
+    # added there, this needs the same cumulative-since-start handling tp
+    # gets, or verification/plots will silently be wrong.
+    "ssrd": "SSRD",
+    "z": "FIS",
 }
 _ICON_TO_IFS = {v: k for k, v in _IFS_TO_ICON.items()}
 
@@ -324,6 +333,20 @@ def _open_analysis_zarr(root: Path, params: list[str]) -> xr.Dataset:
         )
         ds = ds.assign_coords(elevation=elevation).drop_vars(["FIS"])
 
+    # Drop grid points with undefined (NaN) coordinates. This can occur when
+    # xarray opens a zarr dataset whose lat/lon arrays have fill_value=0.0: any
+    # grid point sitting exactly on 0° longitude is masked to NaN by xarray even
+    # though it is a valid point in the raw zarr (e.g. aifs-ea-an-oper o96 ERA5
+    # dataset has 192 such points).
+    if "values" in ds.dims and "latitude" in ds.coords and "longitude" in ds.coords:
+        valid = np.isfinite(ds["latitude"].values) & np.isfinite(ds["longitude"].values)
+        if not valid.all():
+            LOG.warning(
+                "Dropping %d grid point(s) with undefined lat/lon from truth dataset.",
+                int((~valid).sum()),
+            )
+            ds = ds.isel(values=valid)
+
     return ds
 
 
@@ -464,10 +487,21 @@ def load_from_grib_file(file: str | list[str], sel_kwargs):
 
 
 def variable_name_profile(
-    level_type: Literal["height_above_ground_level", "mean_sea", "surface", "pressure"],
+    level_type: Literal[
+        "height_above_ground_level",
+        "mean_sea",
+        "surface",
+        "pressure",
+        "entire_atmosphere",
+    ],
 ) -> dict[str, Any]:
     """Resolve variable name profile based on the level type."""
-    if level_type in ["height_above_ground_level", "mean_sea", "surface"]:
+    if level_type in [
+        "height_above_ground_level",
+        "mean_sea",
+        "surface",
+        "entire_atmosphere",
+    ]:
         return {}
     elif level_type == "pressure":
         return {

src/evalml/config.py

@@ -1,10 +1,15 @@
 from pathlib import Path
-from typing import Dict, List, Any, ClassVar, FrozenSet, Optional
+from typing import Dict, List, Any, ClassVar, FrozenSet, Optional, Union
 
 from pydantic import BaseModel, Field, RootModel, field_validator, model_validator
 
 PROJECT_ROOT = Path(__file__).parents[2]
 
+PREDEFINED_REGIONS: Dict[str, List[float]] = {
+    "global": [-180, 180, -90, 90],
+    "icon": [1.5, 16, 43, 49.5],
+}
+
 
 class Dates(BaseModel):
     """Start/stop of the hindcast period and the launch frequency."""
@@ -273,9 +278,29 @@ class DomainConfig(BaseModel):
         "orthographic",
         description="Projection name (must be a key in plotting._PROJECTIONS, e.g. 'orthographic').",
     )
+    rotate: bool = Field(
+        False,
+        description=(
+            "Rotate the viewpoint across animation frames as lead time advances. "
+            "Only valid for full-globe domains (extent: null)."
+        ),
+    )
+    hours_per_revolution: float = Field(
+        96.0,
+        gt=0,
+        description="Simulated lead-time hours for one full 360° rotation, when rotate is enabled.",
+    )
 
     model_config = {"extra": "forbid"}
 
+    @model_validator(mode="after")
+    def _rotate_requires_globe(self):
+        if self.rotate and self.extent is not None:
+            raise ValueError(
+                "rotate: true is only valid for full-globe domains (extent: null)."
+            )
+        return self
+
 
 class MeteogramConfig(BaseModel):
     """Configuration for meteogram generation."""
@@ -382,15 +407,45 @@ class Locations(BaseModel):
 class Stratification(BaseModel):
     """Stratification settings for the analysis."""
 
-    regions: List[str] = Field(
+    regions: List[Union[str, Dict[str, List[float]]]] = Field(
         default_factory=list,
-        description="List of region names for stratification. Empty list means no spatial stratification.",
+        description=(
+            "List of region specs for spatial stratification. At least one region is required. "
+            f"String entries are either predefined region names ({list(PREDEFINED_REGIONS)}) or "
+            "shapefile names resolved against 'root'. Predefined names take precedence over shapefiles. "
+            "Dict entries map a custom region name to a bounding box [lon_min, lon_max, lat_min, lat_max]. "
+            "The first entry is the domain region used by the dashboard when region stratification is not active."
+        ),
     )
     root: Optional[str] = Field(
         None,
-        description="Root directory where the region shapefiles are stored. Required when regions is non-empty.",
+        description="Root directory where the region shapefiles are stored. Required when regions contains string entries.",
     )
 
+    @field_validator("regions")
+    @classmethod
+    def validate_regions(
+        cls, v: List[Union[str, Dict[str, List[float]]]]
+    ) -> List[Union[str, Dict[str, List[float]]]]:
+        if not v:
+            raise ValueError(
+                "At least one region must be specified. "
+                f"Add a domain region as the first entry, e.g. a predefined name "
+                f"({list(PREDEFINED_REGIONS)}), a custom bbox dict, or a shapefile name."
+            )
+        for entry in v:
+            if isinstance(entry, dict):
+                if len(entry) != 1:
+                    raise ValueError(
+                        f"Each bbox region dict must have exactly one key, got: {list(entry.keys())}"
+                    )
+                name, bbox = next(iter(entry.items()))
+                if len(bbox) != 4:
+                    raise ValueError(
+                        f"Bbox for region '{name}' must have exactly 4 values [lon_min, lon_max, lat_min, lat_max], got {len(bbox)}."
+                    )
+        return v
+
 
 class Dashboard(BaseModel):
     """Settings for the dashboard"""

src/plotting/__init__.py

@@ -1,5 +1,4 @@
 from contextlib import contextmanager
-from functools import cached_property
 from pathlib import Path
 
 import cartopy.crs as ccrs
@@ -148,6 +147,7 @@ def plot_field(
         style: ekp.styles.Style | None = None,
         colorbar: bool = True,
         title: str | None = None,
+        gridline_labels: bool = True,
         **kwargs,
     ):
         """Plot a field on a Map object.
@@ -164,6 +164,12 @@ def plot_field(
             Whether to plot a colorbar, by default True.
         title: str, optional
             Map subplot title.
+        gridline_labels : bool
+            Whether to draw lat/lon degree labels on the gridlines, by
+            default True. Set to False for views whose center rotates
+            between frames (e.g. a rotating globe animation), where the
+            labels' varying width would otherwise make the map's position
+            shift from frame to frame.
         kwargs : dict
             Additional keyword arguments to pass to ax.tripcolor, including cmap,
             vmin, vmax, etc.
@@ -174,8 +180,8 @@ def plot_field(
         # of the plotting function is a lot faster than letting tricontourf or
         # tripcolor handle it in general, but not sure if using earthkit
         # removed for now to simplify the workflow
-        if proj == _PROJECTIONS["orthographic"]:
-            triang, mask = self._orthographic_tri
+        if isinstance(proj, ccrs.Orthographic):
+            triang, mask = self._orthographic_tri(proj)
         else:
             triang, mask = self.tri, slice(None, None)
         x, y = triang.x, triang.y
@@ -215,7 +221,13 @@ def plot_field(
         # 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 gridline_labels:
+            subplot.standard_layers()
+        else:
+            subplot.land()
+            subplot.coastlines()
+            subplot.borders()
+            subplot.gridlines(draw_labels=False)
 
         if colorbar:
             subplot.legend()
@@ -265,13 +277,10 @@ def _temporary_plot_kwargs_override(self, subplot: ekp.Map):
                 except Exception:
                     pass
 
-    @cached_property
-    def _orthographic_tri(self) -> Triangulation:
-        """Compute the triangulation for the orthographic projection."""
-        x, y, _ = (
-            _PROJECTIONS["orthographic"]
-            .transform_points(ccrs.PlateCarree(), self.lon, self.lat)
-            .T
-        )
+    def _orthographic_tri(
+        self, proj: ccrs.Projection
+    ) -> tuple[Triangulation, np.ndarray]:
+        """Compute the triangulation for an orthographic-family projection."""
+        x, y, _ = proj.transform_points(ccrs.PlateCarree(), self.lon, self.lat).T
         mask = ~(np.isnan(x) | np.isnan(y))
         return Triangulation(x[mask], y[mask]), mask

src/plotting/colormap_defaults.py

@@ -73,6 +73,38 @@ def _precip_bias_map(accum_h: int) -> dict:
         "extend": "both",
     },
     "QV_925": load_ncl_colormap("RH_6lev.ct") | {"extend": "both"},
+    "CLCT": {
+        # extend="neither" relies on preprocess_field() clipping away from
+        # exact 0/1 (see plot_forecast_frame.py) to avoid a tricontourf bug
+        # on orthographic projections.
+        "cmap": plt.get_cmap("Blues_r"),
+        "vmin": 0,
+        "vmax": 1,
+        "extend": "neither",
+        "units": "",
+        "levels": list(np.linspace(0, 1, 21)),
+    },
+    "CLCL": {
+        "cmap": plt.get_cmap("Blues_r"),
+        "vmin": 0,
+        "vmax": 1,
+        "extend": "neither",
+        "units": "",
+        "levels": list(np.linspace(0, 1, 21)),
+    },
+    "SSRD": {
+        # tricontourf always bands regardless of "levels" being set (it falls
+        # back to an auto locator with ~7 bands otherwise) — use a fine level
+        # set here to approximate a smooth gradient instead. extend="max"
+        # only (not "both") since preprocess_field() already clips away from
+        # exact 0 — see CLCT.
+        "cmap": plt.get_cmap("YlOrRd"),
+        "vmin": 0,
+        "vmax": 4e6,
+        "extend": "max",
+        "units": "J m-2",
+        "levels": list(np.linspace(0, 4e6, 21)),
+    },
     "TOT_PREC_1H": {
         "extend": "max",
         "colors": [

src/plotting/compat.py

@@ -15,6 +15,9 @@
     "PS": "sp",
     "PMSL": "msl",
     "TOT_PREC": "tp",
+    "CLCT": "tcc",
+    "CLCL": "lcc",
+    "SSRD": "ssrd",
 }
 
 PARAMS_MAP_INV = {v: k for k, v in PARAMS_MAP.items()}
@@ -35,6 +38,11 @@ def load_state_from_grib(
     }
     if ifs_rename:
         ds = ds.rename(ifs_rename)
+        # TODO check if needed
+    # if "tp" in ifs_rename and "TOT_PREC" in ds:
+    #     # IFS/ECMWF convention: "tp" is accumulated precip in meters.
+    #     # Convert to kg m-2 (mm) to match the ICON-native convention.
+    #     ds["TOT_PREC"] = ds["TOT_PREC"] * 1000
     state = {}
     ref_param = next((p for p in (paramlist or []) if p in ds), None)
     if ref_param is None: