diff --git a/src/data_input/__init__.py b/src/data_input/__init__.py index 39df43bd..3c8c29ce 100644 --- a/src/data_input/__init__.py +++ b/src/data_input/__init__.py @@ -323,7 +323,7 @@ def _select_valid_times(ds, times: np.datetime64, strict: bool = False): ) return ds.sel(time=times_np[times_included]) else: - raise ValueError( + raise LookupError( "Valid times are not included in the dataset. " "Please check the valid times and the dataset." ) diff --git a/src/evalml/config.py b/src/evalml/config.py index 2a8a22f8..919d62cd 100644 --- a/src/evalml/config.py +++ b/src/evalml/config.py @@ -445,6 +445,55 @@ class AnimationsConfig(BaseModel): ) +class ComparisonsConfig(BaseModel): + """Configuration for side-by-side map comparisons of truth, runs and baselines.""" + + enabled: bool = Field( + default=False, + description=( + "Whether to generate side-by-side maps: one figure per init time, lead " + "time, parameter and domain, with a panel for the truth (if gridded), " + "each run and each selected baseline." + ), + ) + truth: bool = Field( + default=True, + description=( + "Whether to include a truth panel. Ignored for station truth " + "(jretrievedwh), which cannot be drawn as a map." + ), + ) + baselines: List[str] = Field( + default_factory=list, + description="Labels of the baselines to add as panels.", + ) + domains: List[str | DomainConfig] = Field( + default=["alps"], + description=( + "Domains to plot. Same format as `animations.domains`: a named domain " + "defined in plotting.DOMAINS or a custom domain dict." + ), + ) + lead_times: List[int] = Field( + default=[24], + description=( + "Lead times (hours) to plot. A run or baseline without data at a lead " + "time is shown as an empty panel." + ), + ) + + model_config = {"extra": "forbid"} + + @model_validator(mode="after") + def _enabled_requires_a_reference(self): + if self.enabled and not (self.truth or self.baselines): + raise ValueError( + "comparisons need something to compare the runs with: enable " + "`truth` or list `baselines`." + ) + return self + + class ScorecardConfig(BaseModel): """Configuration for a single named scorecard.""" @@ -558,6 +607,10 @@ class ShowcaseConfig(BaseModel): default_factory=AnimationsConfig, description="Configuration for animation generation.", ) + comparisons: ComparisonsConfig = Field( + default_factory=ComparisonsConfig, + description="Configuration for side-by-side map comparisons.", + ) class Locations(BaseModel): @@ -1028,6 +1081,19 @@ def validate_scorecard_references(self) -> "ConfigModel": ) return self + @model_validator(mode="after") + def validate_comparison_baselines(self) -> "ConfigModel": + labels = [ + item.baseline.label for item in self.runs if isinstance(item, BaselineItem) + ] + unknown = [b for b in self.showcase.comparisons.baselines if b not in labels] + if unknown: + raise ValueError( + f"showcase.comparisons.baselines {unknown} do not match the label of " + f"any baseline in `runs`: {labels}" + ) + return self + model_config = { "extra": "forbid", # fail on misspelled keys "populate_by_name": True, diff --git a/src/plotting/__init__.py b/src/plotting/__init__.py index b7336915..6c78cfdb 100644 --- a/src/plotting/__init__.py +++ b/src/plotting/__init__.py @@ -60,6 +60,32 @@ def get_projection(name: str) -> "ccrs.Projection": } +def resolve_region( + region_name: str, region_cfg: dict, lead_time: int = 0 +) -> tuple["ccrs.Projection", list[float] | None]: + """Return (projection, extent) for a showcase region. + + ``region_cfg`` is one entry of the dict built by ``parse_showcase_regions``. + A custom ``extent`` takes precedence; otherwise the named domain in + ``DOMAINS`` is used, rotated by ``lead_time`` when ``rotate`` is set. + """ + if region_cfg.get("extent") is not None: + projection = get_projection(region_cfg.get("projection") or "orthographic") + return projection, region_cfg["extent"] + domain = DOMAINS[region_name] + if region_cfg.get("rotate"): + base = domain["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"] + ) + return projection, domain["extent"] + return domain["projection"], domain["extent"] + + class StatePlotter: """A class to plot state fields on various DOMAINS.""" @@ -208,16 +234,19 @@ def plot_field( # 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 + # An all-NaN field (e.g. accumulations at lead 0) leaves nothing to + # contour: only the map layers below are drawn. + if finite.any(): + # 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 diff --git a/src/plotting/styles.py b/src/plotting/styles.py new file mode 100644 index 00000000..68965bcb --- /dev/null +++ b/src/plotting/styles.py @@ -0,0 +1,116 @@ +"""Shared plotting styles and field preprocessing for map plots.""" + +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.colormap_defaults import CMAP_DEFAULTS + + +def _precip_lookup(accu): + """Key of the TOT_PREC palette for an accumulation window of ``accu`` hours. + + Palettes exist only for some windows (1h, 6h); others use the longest + defined window not exceeding ``accu`` (e.g. 3h -> 1h, 24h -> 6h). + """ + windows = sorted( + int(key[len("TOT_PREC_") : -1]) + for key in CMAP_DEFAULTS + if key.startswith("TOT_PREC_") and key.endswith("H") + ) + window = max((w for w in windows if w <= accu), default=windows[0]) + return f"TOT_PREC_{window}H" + + +def get_style(param, units_override=None, accu=1): + """Get style and colormap settings for the plot.""" + lookup = _precip_lookup(accu) 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=bounds, + extend=extend, + units=units, + colors=cfg.get("colors", None), + ), + "norm": cfg.get("norm", None), + "cmap": prebuilt_cmap, + "vmin": cfg.get("vmin", None), + "vmax": cfg.get("vmax", None), + } + + +def preprocess_field(param: str, state: dict): + """ + - Wind speed: sqrt(u^2 + v^2) + - Everything else: see :func:`prepare_values` + Returns: (field_array, units_override or None) + """ + fields = state["fields"] + if param == "SP_10M": + values = ekm_wind.speed(fields["U_10M"], fields["V_10M"]) + elif param == "SP": + values = ekm_wind.speed(fields["U"], fields["V"]) + else: + values = fields[param] + return prepare_values(param, values) + + +def prepare_values(param: str, values): + """Unit conversion and clipping of an already derived ``param`` field. + + - Temperatures: K -> °C + - Precipitation: m -> mm + Used directly by callers that load derived params (e.g. ``SP_10M``). + Returns: (field_array, units_override or None) + """ + if param in ("T_2M", "TD_2M", "TW_2M", "SKT", "T", "TD"): + return kelvin_to_celsius(values), "°C" + if param in ("SP_10M", "SP"): + return values, "m/s" + if param == "TOT_PREC": + return np.maximum(values, 0), "mm" + if param in ("CLCT", "CLCL", "CLCM", "CLCH"): + # Cloud cover is in percent. Avoid exact 0/100 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(values, 1e-4, 100 - 1e-4), None + if param == "SSRD": + # Same issue, bottom boundary only (night-side plateau). + return np.maximum(values, 1e-6), None + return values, None diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index 00befe23..11dd16f5 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -90,6 +90,46 @@ def test_experiment_optional(example_config): assert cfg.experiment is None +def test_comparisons_truth_only(example_config): + example_config["showcase"]["comparisons"] = {"enabled": True, "baselines": []} + cfg = ConfigModel.model_validate(example_config) + assert cfg.showcase.comparisons.truth is True + + +def test_comparisons_need_a_reference(example_config): + example_config["showcase"]["comparisons"] = { + "enabled": True, + "truth": False, + "baselines": [], + } + with pytest.raises(ValueError, match="enable `truth` or list `baselines`"): + ConfigModel.model_validate(example_config) + + +def test_comparison_baselines_must_be_baselines(example_config): + baseline_label = next( + item["baseline"]["label"] + for item in example_config["runs"] + if "baseline" in item + ) + forecaster_label = next( + item[k]["label"] + for item in example_config["runs"] + for k in item + if k in ("forecaster", "temporal_downscaler") + ) + example_config["showcase"]["comparisons"] = { + "enabled": True, + "baselines": [baseline_label], + } + ConfigModel.model_validate(example_config) + + for label in ("no-such-label", forecaster_label): + example_config["showcase"]["comparisons"]["baselines"] = [label] + with pytest.raises(ValueError, match="do not match the label of any baseline"): + ConfigModel.model_validate(example_config) + + def test_score_on_common_points(example_config): """score_on_common_points accepts a subset of experiment.params only.""" diff --git a/tests/unit/test_plot_forecast_comparison.py b/tests/unit/test_plot_forecast_comparison.py new file mode 100644 index 00000000..0ee06968 --- /dev/null +++ b/tests/unit/test_plot_forecast_comparison.py @@ -0,0 +1,68 @@ +import sys +from pathlib import Path + +import numpy as np +import pytest + +sys.path.insert(0, str(Path(__file__).parents[2] / "workflow" / "scripts")) +import plot_forecast_comparison # noqa: E402 +from plot_forecast_comparison import parse_source # noqa: E402 + + +def test_parse_source(): + assert parse_source("ICON-CH1=/store/x=y") == ("ICON-CH1", "/store/x=y") + + +@pytest.mark.parametrize("text", ["no-separator", "=root", "label="]) +def test_parse_source_rejects(text): + with pytest.raises(ValueError): + parse_source(text) + + +@pytest.fixture +def run_main(monkeypatch, tmp_path): + """Run main() without plotting; loaders raise for roots listed in `errors`.""" + loaded = [] + + def run(args, errors): + def loader(root, *_): + if root in errors: + raise errors[root] + loaded.append(root) + return np.zeros(1), np.zeros(1), np.zeros(1) + + monkeypatch.setattr(plot_forecast_comparison, "load_source", loader) + monkeypatch.setattr(plot_forecast_comparison, "load_truth", loader) + monkeypatch.setattr(plot_forecast_comparison, "StatePlotter", lambda *_: None) + monkeypatch.setattr( + sys, + "argv", + ["plot_forecast_comparison.py", *args] + + ["--date", "202401010000", "--leadtime", "6", "--param", "T_2M"] + + ["--regions_json", "{}", "--outdir", str(tmp_path)], + ) + plot_forecast_comparison.main() + return loaded + + return run + + +def test_missing_baseline_and_truth_are_skipped(run_main): + args = ["--source", "ML=run", "--baseline", "ICON=icon", "--truth", "truth"] + errors = {"icon": FileNotFoundError("icon"), "truth": LookupError("truth")} + assert run_main(args, errors) == ["run"] + + +def test_candidate_failure_raises(run_main): + with pytest.raises(FileNotFoundError): + run_main(["--source", "ML=run"], {"run": FileNotFoundError("run")}) + + +def test_unexpected_baseline_failure_raises(run_main): + with pytest.raises(ValueError): + run_main(["--baseline", "ICON=icon"], {"icon": ValueError("corrupt")}) + + +def test_nothing_loaded_raises(run_main): + with pytest.raises(RuntimeError): + run_main(["--baseline", "ICON=icon"], {"icon": FileNotFoundError("icon")}) diff --git a/workflow/Snakefile b/workflow/Snakefile index 6895cb9a..7a3809bc 100644 --- a/workflow/Snakefile +++ b/workflow/Snakefile @@ -283,6 +283,17 @@ rule showcase_all: if config["showcase"]["meteograms"]["enabled"] else [] ), + ( + expand( + rules.plot_forecast_comparison.output, + init_time=[t.strftime("%Y%m%d%H%M") for t in REFTIMES], + param=SHOWCASE_CONFIG["params"], + leadtime=COMPARISONS_CONFIG["lead_times"], + showcase=EXPERIMENT_NAME, + ) + if COMPARISONS_CONFIG["enabled"] + else [] + ), rule sandbox_all: diff --git a/workflow/rules/common.smk b/workflow/rules/common.smk index dc8a4324..2ad240ec 100644 --- a/workflow/rules/common.smk +++ b/workflow/rules/common.smk @@ -1,6 +1,5 @@ import logging import copy -import re from datetime import datetime, timedelta import yaml import hashlib @@ -9,6 +8,7 @@ from urllib.parse import urlparse from evalml.config import MODEL_REGISTRY_ROOT from data_input import inca, truth_type +from evalml.leadtimes import resolve_leadtimes CONFIG_ROOT = Path("config").resolve() OUT_ROOT = Path(config["locations"]["output_root"]) @@ -124,9 +124,11 @@ def parse_regions(): return json.dumps(result) -def parse_showcase_regions(): +def parse_showcase_regions(domains=None): """Parse showcase domains from config. + ``domains`` defaults to ``showcase.animations.domains``. + 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. @@ -134,12 +136,10 @@ def parse_showcase_regions(): rotate/hours_per_revolution only take effect when extent is None (full-globe domains). """ + if domains is None: + domains = config["showcase"]["animations"]["domains"] result = {} - for r in ( - config.get("showcase", {}) - .get("animations", {}) - .get("domains", ["globe", "europe", "switzerland"]) - ): + for r in domains: if isinstance(r, str): result[r] = { "extent": None, @@ -149,10 +149,10 @@ def parse_showcase_regions(): } 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), + "extent": r["extent"], + "projection": r["projection"], + "rotate": r["rotate"], + "hours_per_revolution": r["hours_per_revolution"], } return result @@ -522,11 +522,20 @@ HOLDOUT_FILE = ( ) HOLDOUT_WORKDIR_NAME = "holdout_stations.csv" REGIONS = parse_regions() -_showcase = config.get("showcase", {}) +_showcase = config["showcase"] SHOWCASE_CONFIG = { "regions": parse_showcase_regions(), - "params": _showcase.get("params", ["T_2M", "SP_10M"]), - "fps": _showcase.get("animations", {}).get("frames_per_second", 2.0), + "params": _showcase["params"], + "fps": _showcase["animations"]["frames_per_second"], +} +_comparisons = _showcase["comparisons"] +COMPARISONS_CONFIG = { + "enabled": _comparisons["enabled"], + "regions": parse_showcase_regions(_comparisons["domains"]), + "baselines": _comparisons["baselines"], + "lead_times": _comparisons["lead_times"], + # station truth (jretrievedwh) cannot be drawn as a map + "truth": _comparisons["truth"] and TRUTH_TYPE != "jretrieve", } EXPERIMENT_PARAMS = EXPERIMENT_CONFIG.get( "params", ["T_2M", "TD_2M", "SP_10M", "PS", "PMSL", "TOT_PREC6"] @@ -557,31 +566,6 @@ SCORECARD_CONFIGS = ( ) -# 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 in AIFS) aren't part of the model's input state, -# so they're simply absent from the initial-state GRIB file written at step 0. -# Cloud cover is prognostic in some models (Varda-single RC); listing it here -# only drops their 0h animation frame. -# Short and canonical names both appear across the workflow (showcases vs maps). -PARAMS_WITHOUT_STEP_ZERO_VALUE = { - "TOT_PREC", - "tp", - "CLCT", - "tcc", - "CLCL", - "lcc", - "CLCM", - "mcc", - "CLCH", - "hcc", - "SSRD", - "ssrd", -} - - def balfrin_cpus(mem_mb, actual_cpus=1, node_mem_mb=512_000, node_cpus=256): """Minimum cores to request on Balfrin based on memory usage. @@ -598,49 +582,6 @@ def balfrin_cpus(mem_mb, actual_cpus=1, node_mem_mb=512_000, node_cpus=256): return max(actual_cpus, mem_cpus) -def resolve_leadtimes(steps_spec, requested="all", param=None): - """Lead times to compute for a single participant. - - A run or baseline produces only the lead times in its own ``steps`` spec - (``start/stop/step``, hours). This returns those of the ``requested`` - selection that the participant actually produces — the literal ``"all"`` - (every produced lead time) or an explicit list of ints — so a 36h lead is - never requested of an ICON-CH1 baseline (steps ``0/33/6``), nor a >120h - lead of ICON-CH2. Explicitly requested lead times the participant cannot - produce are skipped with a warning. For accumulated ``param``s, lead times - shorter than one step spacing are dropped (no accumulation window). - """ - start, end, step = map(int, steps_spec.split("/")) - supported = set(range(start, end + 1, step)) - wanted = supported if requested == "all" else set(requested) - - unsupported = sorted(wanted - supported) - if unsupported: - logging.getLogger("snakemake").warning( - "Skipping lead time(s) %sh: not produced by forecast steps '%s'.", - unsupported, - steps_spec, - ) - - valid = wanted & supported - # Showcase-facing names encode an accumulation window as a numeric suffix - # (e.g. "TOT_PREC1", "TOT_PREC6"); strip it before checking membership so - # those aggregated variants are recognized the same as their base param. - base_param = re.sub(r"\d+$", "", param) if param else param - if ( - param in PARAMS_WITHOUT_STEP_ZERO_VALUE - or base_param in PARAMS_WITHOUT_STEP_ZERO_VALUE - ): - # For aggregated params (e.g. TOT_PREC6), the accumulation window encoded - # in the suffix is the minimum valid lead time — not the run's step spacing. - # At shorter leads the de-accumulation window doesn't fit (prev_lt < 0), - # so the frame would fall back to showing a raw cumulative total. - agg_match = re.search(r"\d+$", param) if param and base_param != param else None - min_lt = int(agg_match.group()) if agg_match else step - valid = {lt for lt in valid if lt >= min_lt} - return sorted(valid) - - def scoremap_runtime(wc): """SLURM runtime (minutes) for a scoremap job, which loops over all reference times. Accumulated parameters (TOT_PREC*) re-accumulate hourly analysis diff --git a/workflow/rules/plot.smk b/workflow/rules/plot.smk index b28b16a2..7b30124e 100644 --- a/workflow/rules/plot.smk +++ b/workflow/rules/plot.smk @@ -168,6 +168,8 @@ rule make_forecast_animation: for f in input if Path(f).stat().st_size > 0 ] + if not frames: + raise ValueError(f"No non-empty frames to animate for {output[0]}") frames[0].save( output[0], save_all=True, @@ -222,3 +224,86 @@ use rule plot_scoremaps as plot_scoremaps_baseline with: log: OUT_ROOT / "logs/plot_scoremaps/{experiment}/{baseline_id}-{param}-{score}-{region}-{season}-{init_hour}-{leadtime}.log", + + +def _comparison_sources(wc): + """'label=root' for each candidate run panel, after the truth.""" + return [ + f"{RUN_CONFIGS[run_id]['label']}=" + f"{(Path(OUT_ROOT) / f'data/runs/{run_id}/{wc.init_time}/grib').resolve()}" + for run_id in CANDIDATES + ] + + +def _comparison_baselines(wc): + """'label=root' for each selected baseline panel, after the runs.""" + # baselines only: config validation checks the labels against them + baselines_by_label = {cfg["label"]: cfg for cfg in BASELINE_CONFIGS.values()} + return [ + f"{label}={baselines_by_label[label]['root']}" + for label in COMPARISONS_CONFIG["baselines"] + ] + + +rule plot_forecast_comparison: + input: + script="workflow/scripts/plot_forecast_comparison.py", + inference_okfiles=lambda wc: expand( + rules.inference_execute.output.okfile, + run_id=CANDIDATES, + init_time=wc.init_time, + ), + truth=lambda wc: truth_file_dep(wc) if COMPARISONS_CONFIG["truth"] else [], + output: + expand( + OUT_ROOT + / "results/{{showcase}}/{{init_time}}/compare/{{init_time}}_{{leadtime}}_{{param}}_{region}.png", + region=list(COMPARISONS_CONFIG["regions"].keys()), + ), + log: + OUT_ROOT + / "logs/plot_forecast_comparison/{showcase}-{init_time}-{leadtime}-{param}.log", + wildcard_constraints: + showcase=r"[^/]+", + init_time=r"\d{12}", + leadtime=r"\d+", + resources: + slurm_partition="postproc", + cpus_per_task=balfrin_cpus(32_000, actual_cpus=4), + mem_mb=32_000, + runtime="30m", + params: + sources=_comparison_sources, + baselines=_comparison_baselines, + truth=( + [ + "--truth", + config["truth"]["root"], + "--truth_label", + config["truth"]["label"], + ] + if COMPARISONS_CONFIG["truth"] + else [] + ), + regions_json=json.dumps(COMPARISONS_CONFIG["regions"]), + outdir=lambda wc: str( + ( + Path(OUT_ROOT) / f"results/{wc.showcase}/{wc.init_time}/compare" + ).resolve() + ), + shell: + """ + export ECCODES_DEFINITION_PATH=$(realpath .venv/share/eccodes-cosmo-resources/definitions) + SOURCE_ARGS=() + for s in {params.sources:q}; do + SOURCE_ARGS+=(--source "$s") + done + for s in {params.baselines:q}; do + SOURCE_ARGS+=(--baseline "$s") + done + uv run python {input.script} \ + "${{SOURCE_ARGS[@]}}" {params.truth:q} \ + --date {wildcards.init_time:q} --leadtime {wildcards.leadtime:q} \ + --param {wildcards.param:q} --regions_json {params.regions_json:q} \ + --outdir {params.outdir:q} >{log} 2>&1 + """ diff --git a/workflow/scripts/plot_forecast_comparison.py b/workflow/scripts/plot_forecast_comparison.py new file mode 100644 index 00000000..a1cd6396 --- /dev/null +++ b/workflow/scripts/plot_forecast_comparison.py @@ -0,0 +1,185 @@ +"""Plot side-by-side maps of the truth, ML runs and baselines at a common valid time.""" + +import json +import logging +import math +from argparse import ArgumentParser +from datetime import datetime +from datetime import timedelta +from pathlib import Path + +import cartopy.crs as ccrs +import earthkit.plots as ekp +import numpy as np + +from data_input import load_forecast_data +from data_input import load_truth_data +from data_input import parse_aggregated_param +from plotting import resolve_region +from plotting import StatePlotter +from plotting.styles import get_style +from plotting.styles import prepare_values + +LOG = logging.getLogger(__name__) +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", +) + + +# What the loaders raise when a baseline or the truth does not cover the lead +# time: missing ICON archive file, INCA step outside the file, valid time +# outside the truth zarr. Anything else is a real failure. +NOT_AVAILABLE = (FileNotFoundError, LookupError) + + +def _flat(ds, param): + base_param, _ = parse_aggregated_param(param) + values, _ = prepare_values(base_param, np.asarray(ds[param].values).ravel()) + lon = np.asarray(ds["longitude"].values).ravel() + lat = np.asarray(ds["latitude"].values).ravel() + return lon, lat, values + + +def load_source(root, init_time, lead_time, param): + """Return (lon, lat, values) as flat arrays.""" + ds = load_forecast_data(root, init_time, [lead_time], [param]).squeeze() + return _flat(ds, param) + + +def load_truth(root, init_time, lead_time, param): + """Return (lon, lat, values) of the truth at the valid time as flat arrays.""" + ds = load_truth_data(Path(root), init_time, [lead_time], [param]) + valid_time = np.datetime64(init_time + timedelta(hours=lead_time)) + return _flat(ds.sel(time=valid_time).squeeze(), param) + + +def max_in_view(subplot, lon, lat, values): + """Max of the values at points inside the subplot's visible map extent.""" + xyz = subplot._crs.transform_points(ccrs.PlateCarree(), lon, lat) + x, y = xyz[:, 0], xyz[:, 1] + x0, x1, y0, y1 = subplot.ax.get_extent() + # Points on the far side of an orthographic globe project to non-finite x/y. + inside = (x >= x0) & (x <= x1) & (y >= y0) & (y <= y1) & np.isfinite(values) + return np.max(values[inside]) if inside.any() else np.nan + + +def parse_source(text): + """'label=root' -> (label, root); the label may not contain '='.""" + label, sep, root = text.partition("=") + if not sep or not label or not root: + raise ValueError(f"Expected 'label=root', got {text!r}") + return label, root + + +def main(): + parser = ArgumentParser() + parser.add_argument( + "--source", + action="append", + default=[], + type=parse_source, + help="ML run panel as 'label=grib dir', repeatable. Must load.", + ) + parser.add_argument( + "--baseline", + action="append", + default=[], + type=parse_source, + help="Baseline panel as 'label=root', repeatable. " + "Left empty if it has no data at the lead time.", + ) + parser.add_argument("--truth", type=str, default=None, help="Truth root (zarr)") + parser.add_argument("--truth_label", type=str, default="Truth") + parser.add_argument("--date", type=str, help="reference datetime") + parser.add_argument("--leadtime", type=int, help="leadtime") + parser.add_argument("--param", type=str, help="parameter") + parser.add_argument( + "--regions_json", + type=str, + help="JSON dict mapping region name -> {extent, projection}", + ) + parser.add_argument("--outdir", type=str, help="output directory") + args = parser.parse_args() + + init_time = datetime.strptime(args.date, "%Y%m%d%H%M") + lead_time = args.leadtime + param = args.param + regions = json.loads(args.regions_json) + outdir = Path(args.outdir) + outdir.mkdir(parents=True, exist_ok=True) + valid_time = init_time + timedelta(hours=lead_time) + + panels = [(label, root, load_source, False) for label, root in args.source] + panels += [(label, root, load_source, True) for label, root in args.baseline] + if args.truth: + panels.insert(0, (args.truth_label, args.truth, load_truth, True)) + if not panels: + parser.error("no panels to plot") + data = {} + for label, root, loader, optional in panels: + LOG.info("Loading '%s' from %s", label, root) + try: + lon, lat, values = loader(root, init_time, lead_time, param) + except NOT_AVAILABLE: + if not optional: + raise + LOG.warning( + "'%s' has no data at +%sh, plotting an empty panel", + label, + lead_time, + exc_info=True, + ) + continue + # Triangulate each grid once, shared by all regions. + data[label] = (StatePlotter(lon, lat, outdir), values) + if not data: + raise RuntimeError(f"No source has data at +{lead_time}h") + + base_param, accu = parse_aggregated_param(param) + _, units = prepare_values(base_param, np.empty(0)) + style = get_style(base_param, units, accu=accu or 1) + + for region_name, region_cfg in regions.items(): + LOG.info("Plotting region %s", region_name) + projection, extent = resolve_region(region_name, region_cfg, lead_time) + columns = min(len(panels), 4) + rows = math.ceil(len(panels) / columns) + fig = ekp.Figure( + crs=projection, + domain=extent, + rows=rows, + columns=columns, + figsize=(5 * columns, 5 * rows), + ) + for i, (label, *_) in enumerate(panels): + subplot = fig.add_map(row=i // columns, column=i % columns) + if label not in data: + subplot.standard_layers() + subplot.title(f"{label} (not available)") + continue + plotter, values = data[label] + vmax = max_in_view(subplot, plotter.lon, plotter.lat, values) + plotter.plot_field( + subplot, + values, + title=f"{label} (max {vmax:.1f} {units})", + colorbar=False, + **style, + ) + if data: + # One colorbar shared by all panels, which use identical levels. + fig.legend(location="bottom") + # Figure.title() passes an explicit y, which constrained layout ignores + # when reserving space, so the suptitle can overlap the panel titles. + fig.fig.suptitle( + f"{param} valid: {valid_time:%Y-%m-%d %H:%M} UTC " + f"(init: {init_time:%Y-%m-%d %H:%M} UTC, +{lead_time}h)" + ) + outfn = outdir / f"{args.date}_{lead_time}_{param}_{region_name}.png" + fig.save(outfn, dpi=200) + LOG.info("saved: %s", outfn) + + +if __name__ == "__main__": + main() diff --git a/workflow/scripts/plot_forecast_frame.py b/workflow/scripts/plot_forecast_frame.py index 9b4b4dac..5f895a1b 100644 --- a/workflow/scripts/plot_forecast_frame.py +++ b/workflow/scripts/plot_forecast_frame.py @@ -4,18 +4,13 @@ from pathlib import Path import cartopy.crs as ccrs -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 data_input import parse_aggregated_param -from plotting import DOMAINS -from plotting import get_projection +from plotting import resolve_region from plotting import StatePlotter -from plotting.colormap_defaults import CMAP_DEFAULTS from plotting.compat import load_state_from_grib +from plotting.styles import get_style +from plotting.styles import preprocess_field LOG = logging.getLogger(__name__) logging.basicConfig( @@ -24,83 +19,6 @@ ) -def get_style(param, units_override=None, accu=1): - """Get style and colormap settings for the plot.""" - 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=bounds, - extend=extend, - units=units, - colors=cfg.get("colors", None), - ), - "norm": cfg.get("norm", None), - "cmap": prebuilt_cmap, - "vmin": cfg.get("vmin", None), - "vmax": cfg.get("vmax", None), - } - - -def preprocess_field(param: str, state: dict): - """ - - Temperatures: K -> °C - - Wind speed: sqrt(u^2 + v^2) - - Precipitation: m -> mm - Returns: (field_array, units_override or None) - """ - fields = state["fields"] - if param in ("T_2M", "TD_2M", "TW_2M", "SKT", "T", "TD"): - return kelvin_to_celsius(fields[param]), "°C" - if param == "SP_10M": - return ekm_wind.speed(fields["U_10M"], fields["V_10M"]), "m/s" - if param == "SP": - 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", "CLCM", "CLCH"): - # Avoid exact 0/100 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-4, 100 - 1e-4), None - if param == "SSRD": - # Same issue, bottom boundary only (night-side plateau). - return np.maximum(fields[param], 1e-6), None - return fields[param], None - - def main(): parser = ArgumentParser() parser.add_argument( @@ -190,6 +108,11 @@ def main(): prev_grib_files = list(grib_dir.glob(f"2*_{prev_lt}.grib")) if not prev_grib_files: prev_grib_files = list(grib_dir.glob(f"2*_{prev_lt:03d}.grib")) + if not prev_grib_files: + raise FileNotFoundError( + f"No GRIB file for lead_time={prev_lt} in {grib_dir}, needed to" + f" de-accumulate {param} at lead_time={lead_time}." + ) prev_grib_file = Path(prev_grib_files[0]) LOG.info( "De-accumulating TOT_PREC: loading previous grib file %s", @@ -209,22 +132,7 @@ def main(): 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"] + projection, extent = resolve_region(region_name, region_cfg, lead_time) fig = plotter.init_geoaxes( nrows=1, ncols=1, diff --git a/workflow/tools/config.schema.json b/workflow/tools/config.schema.json index e6e33d60..3eeb9726 100644 --- a/workflow/tools/config.schema.json +++ b/workflow/tools/config.schema.json @@ -88,6 +88,63 @@ "title": "BaselineItem", "type": "object" }, + "ComparisonsConfig": { + "additionalProperties": false, + "description": "Configuration for side-by-side map comparisons of truth, runs and baselines.", + "properties": { + "enabled": { + "default": false, + "description": "Whether to generate side-by-side maps: one figure per init time, lead time, parameter and domain, with a panel for the truth (if gridded), each run and each selected baseline.", + "title": "Enabled", + "type": "boolean" + }, + "truth": { + "default": true, + "description": "Whether to include a truth panel. Ignored for station truth (jretrievedwh), which cannot be drawn as a map.", + "title": "Truth", + "type": "boolean" + }, + "baselines": { + "description": "Labels of the baselines to add as panels.", + "items": { + "type": "string" + }, + "title": "Baselines", + "type": "array" + }, + "domains": { + "default": [ + "alps" + ], + "description": "Domains to plot. Same format as `animations.domains`: a named domain defined in plotting.DOMAINS or a custom domain dict.", + "items": { + "anyOf": [ + { + "type": "string" + }, + { + "$ref": "#/$defs/DomainConfig" + } + ] + }, + "title": "Domains", + "type": "array" + }, + "lead_times": { + "default": [ + 24 + ], + "description": "Lead times (hours) to plot. A run or baseline without data at a lead time is shown as an empty panel.", + "items": { + "type": "integer" + }, + "title": "Lead Times", + "type": "array" + } + }, + "title": "ComparisonsConfig", + "type": "object" + }, "Dashboard": { "description": "Settings for the dashboard", "properties": { @@ -968,6 +1025,10 @@ }, "animations": { "$ref": "#/$defs/AnimationsConfig" + }, + "comparisons": { + "$ref": "#/$defs/ComparisonsConfig", + "description": "Configuration for side-by-side map comparisons." } }, "title": "ShowcaseConfig",