diff --git a/README.md b/README.md index 1111586a..9644065a 100644 --- a/README.md +++ b/README.md @@ -81,6 +81,7 @@ truth: experiment: stratification: regions: + - icon # first item will be used as default verification region - jura - mittelland - voralpen diff --git a/config/aifs-single.yaml b/config/aifs-single.yaml index 3a806361..043c5d13 100644 --- a/config/aifs-single.yaml +++ b/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/ diff --git a/config/forecasters-ich1-oper-fixed.yaml b/config/forecasters-ich1-oper-fixed.yaml index d7d4d4ef..9409b55a 100644 --- a/config/forecasters-ich1-oper-fixed.yaml +++ b/config/forecasters-ich1-oper-fixed.yaml @@ -46,6 +46,7 @@ experiment: - TOT_PREC6 stratification: regions: + - icon - jura - mittelland - voralpen diff --git a/config/forecasters-ich1-oper.yaml b/config/forecasters-ich1-oper.yaml index 984d3778..d723e7f8 100644 --- a/config/forecasters-ich1-oper.yaml +++ b/config/forecasters-ich1-oper.yaml @@ -44,6 +44,7 @@ experiment: - TOT_PREC6 stratification: regions: + - icon - jura - mittelland - voralpen diff --git a/config/forecasters-ich1.yaml b/config/forecasters-ich1.yaml index a3a8debf..9492939b 100644 --- a/config/forecasters-ich1.yaml +++ b/config/forecasters-ich1.yaml @@ -57,6 +57,7 @@ experiment: - PMSL stratification: regions: + - icon - jura - mittelland - voralpen diff --git a/config/forecasters-ich1_mec_ffv2.yaml b/config/forecasters-ich1_mec_ffv2.yaml index d8993364..8d3b4ce5 100644 --- a/config/forecasters-ich1_mec_ffv2.yaml +++ b/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: diff --git a/config/varda-single-1.0.yaml b/config/varda-single-1.0.yaml index 21a6c794..4ee63d57 100644 --- a/config/varda-single-1.0.yaml +++ b/config/varda-single-1.0.yaml @@ -54,6 +54,7 @@ experiment: - PMSL stratification: regions: + - icon - mittelland - berge - alpennordseite diff --git a/resources/inference/configs/aifs-single-forecaster.yaml b/resources/inference/configs/aifs-single-forecaster.yaml index f5003c13..ed40d2f2 100644 --- a/resources/inference/configs/aifs-single-forecaster.yaml +++ b/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: diff --git a/src/data_input/__init__.py b/src/data_input/__init__.py index 6909c0ed..074921e0 100644 --- a/src/data_input/__init__.py +++ b/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 { diff --git a/src/evalml/config.py b/src/evalml/config.py index c8dc538b..981a67e7 100644 --- a/src/evalml/config.py +++ b/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""" diff --git a/src/plotting/__init__.py b/src/plotting/__init__.py index 524834c6..63b6be12 100644 --- a/src/plotting/__init__.py +++ b/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 diff --git a/src/plotting/colormap_defaults.py b/src/plotting/colormap_defaults.py index fd9c7223..f411bbb9 100644 --- a/src/plotting/colormap_defaults.py +++ b/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": [ diff --git a/src/plotting/compat.py b/src/plotting/compat.py index f88dd51a..caf33be9 100644 --- a/src/plotting/compat.py +++ b/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: diff --git a/src/verification/__init__.py b/src/verification/__init__.py index f2c68773..9ab3777d 100644 --- a/src/verification/__init__.py +++ b/src/verification/__init__.py @@ -3,8 +3,6 @@ import re import time -from pathlib import Path - import cartopy.crs as ccrs from cartopy.io.shapereader import Reader @@ -112,27 +110,26 @@ class ShapefileSpatialAggregationMasks(SpatialAggregationMasks): regions: dict[str, list[Polygon]] def __init__( - self, shp: str | list[str], src_crs=ccrs.epsg(2056), dst_crs=ccrs.PlateCarree() + self, + regions: list[dict], + src_crs=ccrs.epsg(2056), + dst_crs=ccrs.PlateCarree(), ): proj = pyproj.Transformer.from_crs( src_crs.proj4_init, dst_crs.proj4_init, always_xy=True ).transform - regions = {} - # add inner region for ML evaluation - # this is the extent of the largest lat/lon box that is fully within the radar/INCA domain - regions["all"] = [ - Polygon(list(zip([1.5, 16, 16, 1.5, 1.5], [43, 43, 49.5, 49.5, 43]))) - ] - if shp and shp != [""]: - shp = [shp] if isinstance(shp, str) else shp - for shapefile in shp: - region_name = Path(shapefile).stem - reader = Reader(shapefile) - regions[region_name] = [ + self.regions = {} + for spec in regions: + name = spec["name"] + if spec["type"] == "bbox": + lon_min, lon_max, lat_min, lat_max = spec["bbox"] + self.regions[name] = [_bbox_polygon(lon_min, lon_max, lat_min, lat_max)] + elif spec["type"] == "shp": + reader = Reader(spec["path"]) + self.regions[name] = [ transform(proj, record.geometry) for record in reader.records() ] - self.regions = regions def get_masks(self, lat: xr.DataArray, lon: xr.DataArray) -> xr.DataArray: masks = [] @@ -301,12 +298,24 @@ def _merge_metrics(ds: xr.Dataset, num_workers: int = 4) -> xr.Dataset: return out +def _bbox_polygon(lon_min, lon_max, lat_min, lat_max) -> Polygon: + return Polygon( + [ + (lon_min, lat_min), + (lon_max, lat_min), + (lon_max, lat_max), + (lon_min, lat_max), + (lon_min, lat_min), + ] + ) + + def verify( fcst: xr.Dataset, obs: xr.Dataset, fcst_label: str, obs_label: str, - regions: list[str] | None = None, + regions: list[dict] | None = None, dim: list[str] | None = None, threshold_dict: dict[str, dict[str, list[float]]] | None = None, num_workers: int | None = None, @@ -328,8 +337,12 @@ def verify( Label for the forecast source (used in output dataset). obs_label : str Label for the observation source (used in output dataset). - regions : list[str] or None, optional - List of shapefile paths or region names to use for spatial aggregation. If None, uses default region ('all'). + regions : list[dict] + Ordered list of region specs. Each entry is either + ``{"type": "bbox", "name": ..., "bbox": [lon_min, lon_max, lat_min, lat_max]}`` or + ``{"type": "shp", "name": ..., "path": ...}``. The list order is preserved in the + output NetCDF region coordinate; the first entry is the domain region used by + dashboards and scorecards when region stratification is not active. dim : list[str] or None, optional List of dimension names to reduce over when computing metrics/statistics. If None, tries to infer from fcst. threshold_dict : dict[str, dict[str, list[float]]] or None, optional @@ -361,8 +374,14 @@ def verify( else: dim = ["values"] + if not regions: + raise ValueError( + "At least one region must be specified. " + "Provide an ordered list of region specs via the 'regions' argument." + ) + fcst_aligned, obs_aligned = xr.align(fcst, obs, join="inner", copy=False) - region_polygons = ShapefileSpatialAggregationMasks(shp=regions) + region_polygons = ShapefileSpatialAggregationMasks(regions=regions) masks = region_polygons.get_masks( lon=obs_aligned["longitude"], lat=obs_aligned["latitude"] ) diff --git a/src/verification/spatial.py b/src/verification/spatial.py index 3ada9ea7..b58c5c61 100644 --- a/src/verification/spatial.py +++ b/src/verification/spatial.py @@ -141,7 +141,7 @@ def map_forecast_to_truth(fcst: xr.Dataset, truth: xr.Dataset) -> xr.Dataset: "latitude": (fcst["latitude"].dims, truth["latitude"].data), "longitude": (fcst["longitude"].dims, truth["longitude"].data), } - if "values" in fcst.dims and "values" in truth.dims: + if "values" in fcst.dims and "values" in truth.coords: coords["values"] = truth["values"].data return fcst.assign_coords(coords) @@ -163,6 +163,7 @@ def map_forecast_to_truth(fcst: xr.Dataset, truth: xr.Dataset) -> xr.Dataset: fcst = fcst.drop_vars(["x", "y", "values"], errors="ignore") fcst = fcst.assign_coords(longitude=("values", truth["longitude"].data)) fcst = fcst.assign_coords(latitude=("values", truth["latitude"].data)) + # Restore the multi-index on values (needed for unstack) without pulling in # truth's other coordinates (e.g. elevation), which would overwrite fcst's. if truth_is_grid: diff --git a/tests/unit/test_verification.py b/tests/unit/test_verification.py index 751f69da..96399004 100644 --- a/tests/unit/test_verification.py +++ b/tests/unit/test_verification.py @@ -44,7 +44,7 @@ def _make_verif_dataset(forecast_reference_times): }, coords={ "forecast_reference_time": forecast_reference_times, - "region": ["all"], + "region": ["global"], "source": ["fcst"], }, ) diff --git a/workflow/Snakefile b/workflow/Snakefile index 7657d27f..67838e97 100644 --- a/workflow/Snakefile +++ b/workflow/Snakefile @@ -25,14 +25,18 @@ include: "rules/verif_obs.smk" # about workflow # ----------------------------------------------------- +CANDIDATES = collect_all_candidates() +BASELINES = collect_all_baselines() + RUN_CMD = Path(".evalml_snakemake_cmd.txt").read_text().strip() CONFIG_HASH = master_hash() WHEN = datetime.now().strftime("%Y%m%d") CONFIG_FILE = workflow.config_settings.configfiles[0] CONFIG_LABEL = config["config_label"] or CONFIG_FILE.stem EXPERIMENT_NAME = f"{WHEN}_{CONFIG_LABEL}_{CONFIG_HASH}" -CANDIDATES = collect_all_candidates() -BASELINES = collect_all_baselines() +CANDIDATE_LABELS = ", ".join(v.get("label", k) for k, v in CANDIDATES.items()) +BASELINE_LABELS = ", ".join(v.get("label", k) for k, v in BASELINES.items()) +GROUNDTRUTH_LABEL = config["truth"].get("label", config["truth"]["root"]) DATA_DIR = OUT_ROOT / "data" LOGS_DIR = OUT_ROOT / "logs" @@ -83,26 +87,19 @@ onstart: print() _hr() print(_c("🚀 EvalML workflow started", "1")) - print(_c(f" Time: {when_iso}", "90")) - print(_c(f" Snakemake: {RUN_CMD}", "90")) - print(_c(f" Workdir: {Path.cwd()}", "90")) - print(_c(f" Config: {CONFIG_FILE.name}", "90")) - print(_c(f" Experiment: {EXPERIMENT_NAME}", "90")) - print( - _c( - f" Candidates: {", ".join(v.get("label", k) for k, v in CANDIDATES.items())}", - "90", - ) - ) - print( - _c( - f" Baselines: {", ".join(v.get("label", k) for k, v in BASELINES.items())}", - "90", - ) - ) - print(_c(f" Data dir: {DATA_DIR}", "90")) - print(_c(f" Logs dir: {LOGS_DIR}", "90")) - print(_c(f" Results dir: {RESULTS_DIR}", "90")) + print(_c(f" Time: {when_iso}", "90")) + print(_c(f" Snakemake: {RUN_CMD}", "90")) + print(_c(f" Workdir: {Path.cwd()}", "90")) + print(_c(f" Config: {CONFIG_FILE.name}", "90")) + print(_c(f" Experiment: {EXPERIMENT_NAME}", "90")) + if CANDIDATES: + print(_c(f" Candidates: {CANDIDATE_LABELS}", "90")) + if BASELINES: + print(_c(f" Baselines: {BASELINE_LABELS}", "90")) + print(_c(f" Ground truth: {GROUNDTRUTH_LABEL}", "90")) + print(_c(f" Data dir: {DATA_DIR}", "90")) + print(_c(f" Logs dir: {LOGS_DIR}", "90")) + print(_c(f" Results dir: {RESULTS_DIR}", "90")) _hr() print() diff --git a/workflow/rules/common.smk b/workflow/rules/common.smk index 23b31d15..1e34f407 100644 --- a/workflow/rules/common.smk +++ b/workflow/rules/common.smk @@ -80,22 +80,33 @@ def parse_reference_times(): def parse_regions(): - """Parse regions from the configuration.""" + """Return a JSON list of region specs in config order. + + Each entry is either ``{"type": "bbox", "name": ..., "bbox": [...]}`` + or ``{"type": "shp", "name": ..., "path": ...}``, preserving the original + order so the NetCDF region coordinate matches the config. + """ cfg = config["experiment"]["stratification"] - region_names = cfg.get("regions", []) - if not region_names: - return "" - regions = [f"{cfg['root']}/{region}.shp" for region in region_names] - return ",".join(regions) + root = cfg.get("root", "") + result = [] + for entry in cfg.get("regions", []): + if isinstance(entry, str): + result.append({"type": "shp", "name": entry, "path": f"{root}/{entry}.shp"}) + elif isinstance(entry, dict): + name, bbox = next(iter(entry.items())) + result.append({"type": "bbox", "name": name, "bbox": bbox}) + return json.dumps(result) def parse_showcase_regions(): """Parse showcase domains from config. - Returns a dict mapping domain name -> {extent, projection}. + Returns a dict mapping domain name -> {extent, projection, rotate, hours_per_revolution}. Named domains (strings) have extent=None and projection=None, meaning the plot script will fall back to the DOMAINS lookup. Custom domains carry their explicit extent and projection. + rotate/hours_per_revolution only take effect when extent is None + (full-globe domains). """ result = {} for r in ( @@ -104,11 +115,18 @@ def parse_showcase_regions(): .get("domains", ["globe", "europe", "switzerland"]) ): if isinstance(r, str): - result[r] = {"extent": None, "projection": None} + result[r] = { + "extent": None, + "projection": None, + "rotate": False, + "hours_per_revolution": 96.0, + } else: result[r["name"]] = { "extent": r.get("extent"), "projection": r.get("projection", "orthographic"), + "rotate": r.get("rotate", False), + "hours_per_revolution": r.get("hours_per_revolution", 96.0), } return result @@ -392,8 +410,8 @@ if "jretrieve" in str(config["truth"]["root"]): TRUTH_HASH = truth_hash(config["truth"]) -VERIF_HASH = verif_hash(config) REGIONS = parse_regions() +VERIF_HASH = verif_hash(config) _showcase = config.get("showcase", {}) SHOWCASE_CONFIG = { "regions": parse_showcase_regions(), @@ -414,10 +432,23 @@ SCORECARD_CONFIGS = ( ) -# Period-accumulated params verify a [lead - period, lead] window, so they have -# no value at lead times shorter than one step spacing (e.g. no 0h precip map). +# Params with no value at lead time 0. Two distinct reasons land here: +# - period-accumulated params (TOT_PREC/tp) verify a [lead - period, lead] +# window, so they have no value at lead times shorter than one step spacing +# (e.g. no 0h precip map). +# - diagnostic params (e.g. CLCT/tcc) aren't part of the model's input state, +# so they're simply absent from the initial-state GRIB file written at step 0. # Short and canonical names both appear across the workflow (showcases vs maps). -ACCUMULATED_PARAMS = {"TOT_PREC", "tp"} +PARAMS_WITHOUT_STEP_ZERO_VALUE = { + "TOT_PREC", + "tp", + "CLCT", + "tcc", + "CLCL", + "lcc", + "SSRD", + "ssrd", +} def resolve_leadtimes(steps_spec, requested="all", param=None): @@ -445,6 +476,6 @@ def resolve_leadtimes(steps_spec, requested="all", param=None): ) valid = wanted & supported - if param in ACCUMULATED_PARAMS: + if param in PARAMS_WITHOUT_STEP_ZERO_VALUE: valid = {lt for lt in valid if lt >= step} return sorted(valid) diff --git a/workflow/rules/plot.smk b/workflow/rules/plot.smk index 34cd7c8a..ed008601 100644 --- a/workflow/rules/plot.smk +++ b/workflow/rules/plot.smk @@ -12,6 +12,8 @@ import pandas as pd def _get_available_baselines(wc) -> list[dict[str, str]]: """Get all available baseline datasets for the given init time.""" baselines = [] + if not BASELINE_CONFIGS: + return baselines for baseline_id in BASELINE_CONFIGS: root = BASELINE_CONFIGS[baseline_id].get("root") steps = BASELINE_CONFIGS[baseline_id].get("steps") diff --git a/workflow/rules/verification.smk b/workflow/rules/verification.smk index 9b503d00..4886a5f1 100644 --- a/workflow/rules/verification.smk +++ b/workflow/rules/verification.smk @@ -47,7 +47,7 @@ rule verification_metrics_baseline: --steps "{params.baseline_steps}" \ --source_id "{wildcards.baseline_id}" \ --truth_source_id "{params.truth_source_id}" \ - --regions "{params.regions}" \ + --regions '{params.regions}' \ --params "{params.experiment_params}" \ --threshold_dict "{params.threshold_dict}" \ --member "{params.member}" \ @@ -106,7 +106,7 @@ rule verification_metrics: --steps "{params.fcst_steps}" \ --source_id "{wildcards.run_id}" \ --truth_source_id "{params.truth_source_id}" \ - --regions "{params.regions}" \ + --regions '{params.regions}' \ --params "{params.experiment_params}" \ --threshold_dict "{params.threshold_dict}" \ {params.lapse_rate_flag} \ diff --git a/workflow/scripts/plot_forecast_frame.py b/workflow/scripts/plot_forecast_frame.py index 981f397e..ee8e564e 100644 --- a/workflow/scripts/plot_forecast_frame.py +++ b/workflow/scripts/plot_forecast_frame.py @@ -88,6 +88,15 @@ def preprocess_field(param: str, state: dict): return ekm_wind.speed(fields["U"], fields["V"]), "m/s" if param == "TOT_PREC": return np.maximum(fields[param], 0), "mm" + if param in ("CLCT", "CLCL"): + # Avoid exact 0/1 plateaus breaking tricontourf on orthographic + # projections (tmp/reproduce_clct_bug.py). Pair with extend="neither". + # Any new bounded field with silent-blank or GeometryCollection-crash + # globe frames likely needs the same clip-away-from-boundary fix. + return np.clip(fields[param], 1e-6, 1 - 1e-6), None + if param == "SSRD": + # Same issue, bottom boundary only (night-side plateau). + return np.maximum(fields[param], 1e-6), None return fields[param], None @@ -178,10 +187,21 @@ def main(): for region_name, region_cfg in regions.items(): LOG.info("Plotting region %s", region_name) + outfn = outdir / f"frame_{lead_time}_{param}_{region_name}.png" plotter = StatePlotter(state["longitudes"], state["latitudes"], outdir) if region_cfg.get("extent") is not None: projection = get_projection(region_cfg.get("projection") or "orthographic") extent = region_cfg["extent"] + elif region_cfg.get("rotate"): + base = DOMAINS[region_name]["projection"].proj4_params + central_longitude = ( + base["lon_0"] + 360.0 * lead_time / region_cfg["hours_per_revolution"] + ) % 360 + # central_longitude=central_longitude, central_latitude=0.0 for zero-centered + projection = ccrs.Orthographic( + central_longitude=central_longitude, central_latitude=base["lat_0"] + ) + extent = DOMAINS[region_name]["extent"] else: projection = DOMAINS[region_name]["projection"] extent = DOMAINS[region_name]["extent"] @@ -196,7 +216,11 @@ def main(): subplot = fig.add_map(row=0, column=0) plotter.plot_field( - subplot, field, **get_style(param, units_override, accu=accu) + subplot, + field, + title=f"{param}, time: {validtime}", + gridline_labels=not region_cfg.get("rotate", False), + **get_style(param, units_override, accu=accu), ) if len(state["lam_envelope"]) > 0: subplot.ax.add_geometries( @@ -205,10 +229,13 @@ def main(): facecolor="none", crs=ccrs.PlateCarree(), ) - fig.title(f"{param}, time: {validtime}") - outfn = outdir / f"frame_{lead_time}_{param}_{region_name}.png" - fig.save(outfn, bbox_inches="tight", dpi=200) + # earthkit.plots' Figure.save() defaults bbox_inches to "tight", which + # crops to each frame's own content extent — that extent varies with the + # rotating globe's gridline labels, making frames jump around when + # stitched into a GIF. Pass bbox_inches=None explicitly to override that + # default and always save the fixed full canvas. + fig.save(outfn, dpi=200, bbox_inches=None) LOG.info("saved: %s", outfn) diff --git a/workflow/scripts/report_experiment_dashboard.py b/workflow/scripts/report_experiment_dashboard.py index 24fd0953..b8d3abf8 100644 --- a/workflow/scripts/report_experiment_dashboard.py +++ b/workflow/scripts/report_experiment_dashboard.py @@ -84,7 +84,7 @@ def main(args): # retain only rows relevant for the active stratifications stratification = args.stratification if "region" not in stratification: - df = df[df["region"] == "all"] + df = df[df["region"] == df["region"].unique()[0]] if "season" not in stratification: df = df[df["season"] == "all"] if "init_hour" not in stratification: diff --git a/workflow/scripts/report_scorecard.py b/workflow/scripts/report_scorecard.py index 68096ed9..93a4a8d6 100644 --- a/workflow/scripts/report_scorecard.py +++ b/workflow/scripts/report_scorecard.py @@ -25,7 +25,7 @@ # Sentinel values that select the "aggregate over all" slice for each # stratification dimension that is not the active stratification axis. -_STRAT_ALL_VALUES = {"region": "all", "season": "all", "init_hour": -999} +_STRAT_ALL_VALUES = {"season": "all", "init_hour": -999} DEFAULT_PLOT_CFG = { "rcparams": { @@ -228,6 +228,9 @@ def _load_relative_diff(cfg: dict) -> xr.Dataset: model_ds = xr.open_dataset(cfg["model"]["path"]) baseline_ds = xr.open_dataset(cfg["baseline"]["path"]) + if strat_dim != "region": + sel_coords["region"] = model_ds["region"].values[0] + for label, ds in [("model", model_ds), ("baseline", baseline_ds)]: if "n_samples" not in ds.data_vars: raise ValueError( diff --git a/workflow/scripts/verification_metrics.py b/workflow/scripts/verification_metrics.py index fa8ac417..85a79879 100644 --- a/workflow/scripts/verification_metrics.py +++ b/workflow/scripts/verification_metrics.py @@ -1,3 +1,4 @@ +import json import logging from argparse import ArgumentParser from argparse import Namespace @@ -94,7 +95,7 @@ def main(args: ScriptConfig): truth, args.source_id, args.truth_source_id, - args.regions, + regions=args.regions, threshold_dict=args.threshold_dict, ) LOG.info( @@ -172,9 +173,13 @@ def main(args: ScriptConfig): ) parser.add_argument( "--regions", - type=lambda x: [r for r in x.split(",") if r], - help="Comma-separated list of shapefile paths defining regions for stratification.", - default="", + type=json.loads, + help=( + "JSON list of region specs in config order. " + 'Each entry is {"type": "bbox", "name": ..., "bbox": [...]} ' + 'or {"type": "shp", "name": ..., "path": ...}.' + ), + default="[]", ) parser.add_argument( "--threshold_dict", diff --git a/workflow/scripts/verification_plot_metrics.py b/workflow/scripts/verification_plot_metrics.py index e2e3f02f..64d316e1 100644 --- a/workflow/scripts/verification_plot_metrics.py +++ b/workflow/scripts/verification_plot_metrics.py @@ -112,7 +112,9 @@ def main(args: Namespace) -> None: metrics = all_df["metric"].unique() params = all_df["param"].unique() seasons = all_df["season"].unique() if args.stratify else ["all"] - regions = all_df["region"].unique() if args.stratify else ["all"] + regions = ( + all_df["region"].unique() if args.stratify else [all_df["region"].unique()[0]] + ) init_hours = ( all_df["init_hour"].unique() if args.stratify else [-999] ) # numeric code to indicate all init hours diff --git a/workflow/tools/config.schema.json b/workflow/tools/config.schema.json index 7204e44a..ea2b84ba 100644 --- a/workflow/tools/config.schema.json +++ b/workflow/tools/config.schema.json @@ -234,6 +234,19 @@ "description": "Projection name (must be a key in plotting._PROJECTIONS, e.g. 'orthographic').", "title": "Projection", "type": "string" + }, + "rotate": { + "default": false, + "description": "Rotate the viewpoint across animation frames as lead time advances. Only valid for full-globe domains (extent: null).", + "title": "Rotate", + "type": "boolean" + }, + "hours_per_revolution": { + "default": 96.0, + "description": "Simulated lead-time hours for one full 360\u00b0 rotation, when rotate is enabled.", + "exclusiveMinimum": 0, + "title": "Hours Per Revolution", + "type": "number" } }, "required": [ @@ -856,9 +869,22 @@ "description": "Stratification settings for the analysis.", "properties": { "regions": { - "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. String entries are either predefined region names (['global', 'icon']) 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.", "items": { - "type": "string" + "anyOf": [ + { + "type": "string" + }, + { + "additionalProperties": { + "items": { + "type": "number" + }, + "type": "array" + }, + "type": "object" + } + ] }, "title": "Regions", "type": "array" @@ -873,7 +899,7 @@ } ], "default": null, - "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.", "title": "Root" } },