{{ crps_note | e }}.
{% endif %}.CRPS =
.MAE``. + + Only when the dataset contains an ensemble source; otherwise returned unchanged. + Deterministic forecasts are the sources (candidate runs, baselines) with an MAE + and no ensemble scores; the truth source is never filled. The result is for + display only. Returns ``(dataset, filled)``, *filled* True if a source was filled. + """ + ens = ensemble_sources(ds) + if not ens: + return ds, False + with_mae = _sources_with_values(ds, _mae_vars(ds)) + deterministic = [ + s + for s in ds["source"].values.tolist() + if s in with_mae and s not in ens and not str(s).startswith(TRUTH_SOURCE_PREFIX) + ] + if not deterministic: + return ds, False + is_det = xr.DataArray( + np.isin(ds["source"].values, deterministic), + dims="source", + coords={"source": ds["source"]}, + ) + ds = ds.copy() + for mae in _mae_vars(ds): + if "source" not in ds[mae].dims: + continue + name = mae[: -len("MAE")] + "CRPS" + base = ds[name] if name in ds else xr.full_like(ds[mae], np.nan) + out = xr.where(is_det, ds[mae], base) + ds[name] = out.transpose(*ds[mae].dims) + ds[name].attrs = dict(ds[mae].attrs) + return ds, True + + +def fill_crps_from_mae(ds: xr.Dataset) -> tuple[xr.Dataset, bool]: + """Single-participant dataset: if it is deterministic, set ``
.CRPS =
.MAE``.""" + if is_ensemble(ds): + return ds, False + ds = ds.copy() + filled = False + for mae in _mae_vars(ds): + name = mae[: -len("MAE")] + "CRPS" + if name in ds and bool(ds[name].notnull().any()): + continue + ds[name] = ds[mae] + filled = True + return ds, filled diff --git a/tests/unit/data/aggregation_golden_deterministic.nc b/tests/unit/data/aggregation_golden_deterministic.nc new file mode 100644 index 00000000..1e1db271 Binary files /dev/null and b/tests/unit/data/aggregation_golden_deterministic.nc differ diff --git a/tests/unit/data/verify_golden_gridded.nc b/tests/unit/data/verify_golden_gridded.nc new file mode 100644 index 00000000..62fbdf51 Binary files /dev/null and b/tests/unit/data/verify_golden_gridded.nc differ diff --git a/tests/unit/data/verify_golden_stations.nc b/tests/unit/data/verify_golden_stations.nc new file mode 100644 index 00000000..f6817434 Binary files /dev/null and b/tests/unit/data/verify_golden_stations.nc differ diff --git a/tests/unit/data/verify_golden_stations_tolerant.nc b/tests/unit/data/verify_golden_stations_tolerant.nc new file mode 100644 index 00000000..8242a88d Binary files /dev/null and b/tests/unit/data/verify_golden_stations_tolerant.nc differ diff --git a/tests/unit/test_ensemble_consumers.py b/tests/unit/test_ensemble_consumers.py new file mode 100644 index 00000000..4c816e2a --- /dev/null +++ b/tests/unit/test_ensemble_consumers.py @@ -0,0 +1,236 @@ +"""Consumers of aggregated verification (dashboard, scorecard, plots) and ensemble scores.""" + +import sys +from argparse import Namespace +from pathlib import Path + +import numpy as np +import pytest +import xarray as xr + +ROOT = Path(__file__).parents[2] +sys.path.insert(0, str(ROOT / "workflow" / "scripts")) +import report_experiment_dashboard as dash # noqa: E402 +import report_scorecard as sc # noqa: E402 +import verification_plot_metrics as vpm # noqa: E402 +from verification.ensemble_report import ( # noqa: E402 + add_deterministic_crps, + drop_non_score_vars, + ensemble_sources, + fill_crps_from_mae, +) + +N = 3 +STEPS = np.array([6, 12], dtype="timedelta64[h]").astype("timedelta64[ns]") + + +def _coords(source): + return { + "source": [source], + "region": ["icon"], + "season": ["all"], + "init_hour": [-999], + "station_group": ["all"], + "step": STEPS, + } + + +def _var(val, extra=None): + dims = ["source", "region", "season", "init_hour", "station_group", "step"] + shape = [1, 1, 1, 1, 1, 2] + arr = np.full(shape, val, dtype=float) + if extra: + for d, n in extra.items(): + dims.append(d) + shape.append(n) + arr = np.random.default_rng(0).uniform(1, 5, shape) + return dims, arr + + +def make_file(source, ensemble=False, n=N): + ds = xr.Dataset(coords=_coords(source)) + for name, val in [("T_2M.RMSE", 1.0), ("T_2M.MAE", 0.8), ("T_2M.BIAS", 0.1)]: + ds[name] = _var(val) + ds["n_samples"] = (("season", "init_hour"), [[4]]) + if ensemble: + for name, val in [ + ("T_2M.CRPS", 0.5), + ("T_2M.SPREAD", 0.9), + ("T_2M.SSR", 1.1), + ("T_2M.MEMBER_RMSE", 1.2), + ("T_2M.BS_gt_0p1", 0.2), + ]: + ds[name] = _var(val) + ds["T_2M.rank_hist"] = _var(0, {"rank": n + 1}) + ds["T_2M.prob_table_gt_0p1"] = _var(0, {"n_members_yes": n + 1, "obs": 2}) + ds = ds.assign_coords(rank=np.arange(n + 1), n_members_yes=np.arange(n + 1)) + ds.attrs["n_members"] = n + return ds + + +def test_drop_non_score_vars(): + ds = drop_non_score_vars(make_file("ens", ensemble=True)) + assert "T_2M.rank_hist" not in ds and "T_2M.prob_table_gt_0p1" not in ds + assert {"T_2M.CRPS", "T_2M.BS_gt_0p1", "T_2M.SSR"} <= set(ds.data_vars) + + +def test_deterministic_crps_is_mae_for_display_only(): + ens, det = make_file("ens", True), make_file("det") + ds = vpm._concat_sources([ens, det]) + assert ensemble_sources(ds) == ["ens"] + out, filled = add_deterministic_crps(ds) + assert filled + assert out["T_2M.CRPS"].sel(source="det").isel(step=0).item() == 0.8 + assert out["T_2M.CRPS"].sel(source="ens").isel(step=0).item() == 0.5 + assert "T_2M.CRPS" not in det # input file untouched + + +def test_deterministic_only_is_unchanged(): + ds = vpm._concat_sources([make_file("a"), make_file("b")]) + out, filled = add_deterministic_crps(ds) + assert not filled and out is ds + single, filled = fill_crps_from_mae(make_file("a")) + assert filled and single["T_2M.CRPS"].equals(single["T_2M.MAE"]) + ens, filled = fill_crps_from_mae(make_file("e", True)) + assert not filled and float(ens["T_2M.CRPS"].max()) == 0.5 + + +def _run_dashboard(tmp_path, files): + paths = [] + for i, f in enumerate(files): + p = tmp_path / f"f{i}.nc" + f.to_netcdf(p) + paths.append(p) + out = tmp_path / "out" + dash.main( + Namespace( + verif_files=paths, + template=ROOT / "resources/report/dashboard/template.html.jinja2", + script=ROOT / "resources/report/dashboard/script.js", + header_text="", + stratification=[], + station_holdout=False, + configfile=tmp_path / "none.yaml", + output=out, + label_map={}, + ) + ) + return (out / "dashboard.html").read_text() + + +def test_dashboard(tmp_path): + html = _run_dashboard(tmp_path, [make_file("ens", True), make_file("det")]) + for m in ["CRPS", "SPREAD", "SSR", "MEMBER_RMSE", "BS > 0.1"]: + assert f'data-metric="{m}"' in html + assert "rank_hist" not in html and "prob_table" not in html + assert "CRPS = MAE for deterministic forecasts" in html + + +def test_dashboard_deterministic_only_has_no_ensemble_content(tmp_path): + html = _run_dashboard(tmp_path, [make_file("a"), make_file("b")]) + assert "CRPS" not in html + assert "deterministic forecasts" not in html + + +def test_scorecard_crps(tmp_path): + ens, det = make_file("ens", True), make_file("det") + pe, pd_ = tmp_path / "e.nc", tmp_path / "d.nc" + ens.to_netcdf(pe) + det.to_netcdf(pd_) + cfg = { + "model": {"path": str(pe), "source": "ens", "label": "ens"}, + "baseline": {"path": str(pd_), "source": "det", "label": "det"}, + "stratification": "region", + "lead_times": "6/12/6", + "all_metrics": sc.ALL_METRICS if hasattr(sc, "ALL_METRICS") else ["RMSE"], + "variables": {"T_2M": ["CRPS"]}, + } + diff = sc._load_relative_diff(cfg) + assert cfg["crps_filled"] + assert "T_2M.rank_hist" not in diff + # CRPS 0.5 vs MAE 0.8 -> -37.5 % + assert float(diff["T_2M.CRPS"].isel(step=0)) == pytest.approx(-37.5) + assert sc._is_model_better(-37.5, "CRPS", ["CRPS"], []) is True + + +def test_plot_ensemble_diagnostics(tmp_path): + ds = vpm._concat_sources([make_file("ens", True), make_file("det")]) + ds = ds.sel(station_group="all", drop=True) + args = Namespace(output_dir=tmp_path, label_map={}) + vpm.plot_ensemble_diagnostics(ds, args) + names = {p.name for p in tmp_path.iterdir()} + assert names == {"rank_hist_T_2M.png", "reliability_T_2M_gt_0p1.png"} + + +def test_plot_ensemble_diagnostics_per_source_members(tmp_path, monkeypatch): + """Two ensembles of 3 and 5 members: each keeps its own N + 1 ranks, 1/(N+1) + reference line and probabilities k/N (the outer join pads the smaller one).""" + ds = vpm._concat_sources( + [make_file("small", True, n=3), make_file("big", True, n=5)] + ) + ds = ds.sel(station_group="all", drop=True) + assert ds.sizes["rank"] == 6 + figures, close = [], vpm.plt.close + monkeypatch.setattr(vpm.plt, "close", figures.append) + vpm.plot_ensemble_diagnostics(ds, Namespace(output_dir=tmp_path, label_map={})) + rank_fig, rel_fig = figures + lines = rank_fig.axes[0].get_lines() + curves = {ln.get_label(): ln for ln in lines if not ln.get_label().startswith("_")} + assert set(curves) == {"small (N=3)", "big (N=5)"} + assert curves["small (N=3)"].get_xdata().tolist() == [0, 1, 2, 3] + assert curves["big (N=5)"].get_xdata().tolist() == [0, 1, 2, 3, 4, 5] + assert sum(curves["small (N=3)"].get_ydata()) == pytest.approx(1.0) + references = sorted( + ln.get_ydata()[0] for ln in lines if ln.get_label().startswith("_") + ) + assert references == pytest.approx([1 / 6, 1 / 4]) + rel = { + ln.get_label(): ln.get_xdata().tolist() + for ln in rel_fig.axes[0].get_lines() + if not ln.get_label().startswith("_") + } + assert rel["small (N=3)"] == pytest.approx([0, 1 / 3, 2 / 3, 1]) + assert rel["big (N=5)"] == pytest.approx([0, 0.2, 0.4, 0.6, 0.8, 1]) + for fig in figures: + close(fig) + + +def make_truth(source="truth-c305"): + """The truth source of an aggregated file: statistics only, no scores.""" + ds = xr.Dataset(coords=_coords(source)) + ds["T_2M.mean"] = _var(285.0) + return ds + + +def test_ensemble_detected_by_ensemble_variables_not_spread(): + """An ensemble without a finite SPREAD (e.g. only PTYPE, or step 0) is an ensemble.""" + ens = make_file("ens", True) + ens["T_2M.SPREAD"][:] = np.nan + ds = vpm._concat_sources([ens, make_file("det")]) + assert ensemble_sources(ds) == ["ens"] + single, filled = fill_crps_from_mae(ens) + assert not filled and float(single["T_2M.CRPS"].max()) == 0.5 + + +def test_truth_never_gets_crps(): + """Ensemble-only experiment: the truth source is not a deterministic forecast.""" + ds = vpm._concat_sources([make_file("ens", True), make_truth()]) + out, filled = add_deterministic_crps(ds) + assert not filled + assert np.isnan(out["T_2M.CRPS"].sel(source="truth-c305")).all() + # with a deterministic participant, only that one is filled + ds = vpm._concat_sources([make_file("ens", True), make_file("det"), make_truth()]) + out, filled = add_deterministic_crps(ds) + assert filled + assert out["T_2M.CRPS"].sel(source="det").isel(step=0).item() == 0.8 + assert np.isnan(out["T_2M.CRPS"].sel(source="truth-c305")).all() + + +def test_dashboard_ensemble_only_has_no_crps_note(tmp_path): + ens = make_file("ens", True) + truth = make_truth() + truth["n_samples"] = ens["n_samples"] + both = xr.concat([ens, truth], "source", join="outer", data_vars="minimal") + html = _run_dashboard(tmp_path, [both]) + assert 'data-metric="CRPS"' in html + assert "deterministic forecasts" not in html diff --git a/tests/unit/test_ensemble_wiring.py b/tests/unit/test_ensemble_wiring.py new file mode 100644 index 00000000..1b603027 --- /dev/null +++ b/tests/unit/test_ensemble_wiring.py @@ -0,0 +1,410 @@ +"""Wiring of ensemble verification: the common.smk helpers of the verification rules, +member loading and scoring in verification_metrics.py, and the valid mask of an +ensemble in verification_valid_mask.py.""" + +import ast +import sys +from argparse import Namespace +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest +import xarray as xr + +ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(ROOT / "workflow" / "scripts")) + +import verification_metrics # noqa: E402 +import verification_valid_mask # noqa: E402 +from verification import verify # noqa: E402 + +COMMON = ROOT / "workflow/rules/common.smk" +FUNCTIONS = { + "run_grib_root", + "run_grib_dir", + "ensemble_members", + "members_flag", + "ensemble_scaled", +} + + +def _common_namespace(run_configs, out_root): + tree = ast.parse(COMMON.read_text()) + body = [ + node + for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name in FUNCTIONS + ] + namespace = {"Path": Path, "RUN_CONFIGS": run_configs, "OUT_ROOT": out_root} + exec( + compile(ast.Module(body=body, type_ignores=[]), str(COMMON), "exec"), namespace + ) + return namespace + + +RUN_CONFIGS = { + "forecaster-aaaa-bbbb/cccc": {"steps": "0/24/6"}, + "forecaster-aaaa-bbbb/dddd": { + "steps": "0/24/6", + "ensemble": {"members": 4, "seed": 0}, + }, +} +DET, ENS = RUN_CONFIGS + + +def test_verification_reads_grib_root_deterministic_unchanged(tmp_path): + """Deterministic runs: the same grib/ directory as before, and no members flag.""" + ns = _common_namespace(RUN_CONFIGS, tmp_path) + grib = ns["run_grib_root"](DET, "202406150000") + assert grib == tmp_path / f"data/runs/{DET}/202406150000/grib" + assert grib == ns["run_grib_dir"](DET, "202406150000") + assert ns["members_flag"](Namespace(run_id=DET)) == "" + + +def test_verification_reads_all_members_of_an_ensemble(tmp_path): + ns = _common_namespace(RUN_CONFIGS, tmp_path) + assert ns["run_grib_root"](ENS, "202406150000") == ( + tmp_path / f"data/runs/{ENS}/202406150000/grib" + ) + assert ns["members_flag"](Namespace(run_id=ENS)) == "--members 4" + # plots, SAL, score maps and MEC keep member 001 + assert ns["run_grib_dir"](ENS, "202406150000") == ( + tmp_path / f"data/runs/{ENS}/202406150000/grib/001" + ) + + +def test_ensemble_scaled_resources(tmp_path): + ns = _common_namespace(RUN_CONFIGS, tmp_path) + mem = ns["ensemble_scaled"](80_000, 2_000) + runtime = ns["ensemble_scaled"](60, 6, unit="m") + assert mem(Namespace(run_id=DET)) == 80_000 + assert runtime(Namespace(run_id=DET)) == "60m" + assert mem(Namespace(run_id=ENS)) == 88_000 + assert runtime(Namespace(run_id=ENS)) == "84m" + + +# --- member loading ----------------------------------------------------------- + + +def _member_dataset(member, n_values=3): + values = np.full((1, 2, n_values), float(member)) + return xr.Dataset( + {"T_2M": (("forecast_reference_time", "step", "values"), values)}, + coords={ + "forecast_reference_time": [np.datetime64("2024-06-15T00:00", "ns")], + "step": pd.to_timedelta([0, 6], unit="h"), + "latitude": ("values", np.linspace(46.0, 47.0, n_values)), + "longitude": ("values", np.linspace(6.0, 9.0, n_values)), + "valid_time": ( + ("forecast_reference_time", "step"), + np.array( + [["2024-06-15T00:00", "2024-06-15T06:00"]], dtype="datetime64[ns]" + ), + ), + # differs per member on purpose: member 001's must be kept + "elevation": ("values", np.full(n_values, 100.0 * member)), + }, + ) + + +@pytest.fixture +def member_dirs(tmp_path, monkeypatch): + grib = tmp_path / "grib" + for m in (1, 2, 3): + (grib / f"{m:03d}").mkdir(parents=True) + (grib / f"{m:03d}" / "x_0.grib").touch() + calls = [] + + def fake_load(root, reftime, steps, params, member="000"): + calls.append((Path(root).name, params)) + return _member_dataset(int(Path(root).name)) + + monkeypatch.setattr(verification_metrics, "load_forecast_data", fake_load) + return grib, calls + + +def test_load_ensemble_forecast_concats_members(member_dirs): + grib, calls = member_dirs + fcst = verification_metrics.load_ensemble_forecast( + grib, None, [0, 6], ["T_2M"], members=3 + ) + assert [c[0] for c in calls] == ["001", "002", "003"] + assert fcst["member"].values.tolist() == [1, 2, 3] + assert fcst["T_2M"].sel(member=2).values.tolist() == [[[2.0] * 3] * 2] + assert "member" not in fcst["elevation"].dims + np.testing.assert_array_equal(fcst["elevation"].values, 100.0) + + +def test_load_ensemble_forecast_fails_on_missing_member(member_dirs): + grib, _ = member_dirs + with pytest.raises(FileNotFoundError, match="member 004 of 4"): + verification_metrics.load_ensemble_forecast( + grib, None, [0, 6], ["T_2M"], members=4 + ) + + +def test_load_aligned_data_routes_members(member_dirs, monkeypatch): + grib, _ = member_dirs + truth = xr.Dataset( + {"T_2M": (("time", "values"), np.zeros((2, 3)))}, + coords={ + "time": np.array( + ["2024-06-15T00:00", "2024-06-15T06:00"], dtype="datetime64[ns]" + ) + }, + ) + monkeypatch.setattr(verification_metrics, "load_truth_data", lambda *a: truth) + monkeypatch.setattr(verification_metrics, "map_forecast_to_truth", lambda f, t: f) + args = Namespace( + forecast=grib, reftime=None, steps=[0, 6], truth=None, members=3, member="000" + ) + fcst, _ = verification_metrics.load_aligned_data(args, ["T_2M"]) + assert fcst.sizes["member"] == 3 + + +# --- scoring -------------------------------------------------------------------- + +_REGION = [{"type": "bbox", "name": "all", "bbox": [1.5, 16, 43, 49.5]}] + + +def _stations(values): + """(step, values) station field inside _REGION; values (..., n).""" + values = np.asarray(values, dtype=float) + n = values.shape[-1] + dims = ("member", "step", "values") if values.ndim == 3 else ("step", "values") + coords = { + "step": pd.to_timedelta([6], unit="h"), + "values": np.arange(n), + "longitude": ("values", np.linspace(5.0, 10.0, n)), + "latitude": ("values", np.linspace(46.0, 47.0, n)), + "forecast_reference_time": np.datetime64("2024-06-15T00:00", "ns"), + } + if values.ndim == 3: + coords["member"] = np.arange(1, values.shape[0] + 1) + return xr.Dataset({"T_2M": (dims, values)}, coords=coords) + + +TRUTH = _stations([[1.0, 2.0, 3.0, 4.0]]) +MEMBERS = _stations( + [ + [[0.0, 2.0, 2.0, 6.0]], + [[2.0, 4.0, 3.0, 2.0]], + [[4.0, 0.0, 7.0, 4.0]], + ] +) + + +def _run(fcst, truth=TRUTH, **kwargs): + return verification_metrics.verify_ensemble_run( + fcst, + truth, + "ens", + "truth", + regions=_REGION, + threshold_dict={"T_2M": {"gt": [2.5]}}, + holdout_stations=None, + **kwargs, + ) + + +def test_verify_ensemble_run_scores_the_ensemble_mean(): + out = _run(MEMBERS) + ref = verify( + MEMBERS.mean("member"), + TRUTH, + "ens", + "truth", + regions=_REGION, + threshold_dict={"T_2M": {"gt": [2.5]}}, + num_workers=1, + ) + for name in ("T_2M.BIAS", "T_2M.MSE", "T_2M.MAE", "T_2M.contingency_table"): + xr.testing.assert_allclose(out[name], ref[name]) + # truth statistics are those of verify() + for name in ("T_2M.mean", "T_2M.var", "T_2M.min", "T_2M.max"): + xr.testing.assert_allclose( + out[name].sel(source="truth"), ref[name].sel(source="truth") + ) + for name in ("T_2M.CRPS", "T_2M.ENS_VAR", "T_2M.MEMBER_MSE", "T_2M.rank_hist"): + assert name in out + assert "T_2M.prob_table" in out + assert out.attrs["n_members"] == 3 + assert out["source"].values.tolist() == ["ens", "truth"] + + +def test_verify_ensemble_run_member_mean_statistics(): + out = _run(MEMBERS) + sel = dict(source="ens", region="all", station_group="all") + values = MEMBERS["T_2M"].values[:, 0] # member, values + assert out["T_2M.mean"].sel(**sel).item() == pytest.approx(values.mean()) + assert out["T_2M.var"].sel(**sel).item() == pytest.approx(values.var(axis=1).mean()) + assert out["T_2M.min"].sel(**sel).item() == pytest.approx(values.min(axis=1).mean()) + assert out["T_2M.max"].sel(**sel).item() == pytest.approx(values.max(axis=1).mean()) + + +def test_verify_ensemble_run_missing_member_value(): + """A point where one member is missing makes the scores NaN, as in verify().""" + members = MEMBERS.copy(deep=True) + members["T_2M"][1, 0, 2] = np.nan + out = _run(members) + sel = dict(source="ens", region="all", station_group="all") + for name in ("T_2M.MSE", "T_2M.CRPS", "T_2M.mean", "T_2M.var"): + assert np.isnan(out[name].sel(**sel).item()), name + assert np.isfinite(out["T_2M.mean"].sel(source="truth", region="all").item()) + + +def test_verify_ensemble_run_circular_ensemble_mean(): + """The ensemble mean of wind direction is circular: 350 and 10 deg give 0 deg.""" + members = _stations([[[350.0, 90.0, 180.0, 270.0]], [[10.0, 90.0, 180.0, 270.0]]]) + members = members.rename({"T_2M": "DD_10M"}) + truth = _stations([[0.0, 90.0, 180.0, 270.0]]).rename({"T_2M": "DD_10M"}) + out = verification_metrics.verify_ensemble_run( + members, + truth, + "ens", + "truth", + regions=_REGION, + threshold_dict=None, + holdout_stations=None, + ) + sel = dict(source="ens", region="all", station_group="all") + assert out["DD_10M.MAE"].sel(**sel).item() == pytest.approx(0.0, abs=1e-9) + + +# --- valid mask ----------------------------------------------------------------- + + +def test_valid_mask_of_an_ensemble_is_the_and_over_members(tmp_path, monkeypatch): + members = MEMBERS.copy(deep=True) + members["T_2M"][1, 0, 2] = np.nan # member 2 missing at point 2 + monkeypatch.setattr( + verification_valid_mask, "load_aligned_data", lambda args, p: (members, None) + ) + args = Namespace(params=["T_2M"], output=tmp_path / "mask.nc", members=3) + verification_valid_mask.main(args) + mask = xr.open_dataset(args.output) + assert "member" not in mask.dims + assert mask["T_2M"].values.ravel().tolist() == [True, True, False, True] + + +def test_verify_ensemble_run_categorical_has_no_ensemble_mean_scores(): + """PTYPE: the mean of class codes is no class (rain 1 + snow 5 -> 3, freezing + rain), so the run's deterministic scores are NaN; prob_table scores the ensemble.""" + members = _stations([[[1.0, 5.0, 0.0, 1.0]], [[5.0, 1.0, 0.0, 1.0]]]) + members = members.rename({"T_2M": "PTYPE"}) + truth = _stations([[3.0, 1.0, 0.0, 1.0]]).rename({"T_2M": "PTYPE"}) + out = verification_metrics.verify_ensemble_run( + members, + truth, + "ens", + "truth", + regions=_REGION, + threshold_dict={"PTYPE": {"eq": [1, 3, 5]}}, + holdout_stations=None, + ) + sel = dict(source="ens", region="all", station_group="all") + for name in ("PTYPE.BIAS", "PTYPE.MSE", "PTYPE.MAE", "PTYPE.CORR"): + assert np.isnan(out[name].sel(**sel)).all(), name + table = out["PTYPE.contingency_table"] + assert "source" in table.dims + assert np.isnan(table.sel(**sel)).all() + # the ensemble's categorical score: 2 members, eq_3 never forecast but observed once + prob = out["PTYPE.prob_table"].sel(**sel) + assert prob.sel({"PTYPE.threshold": "eq_3", "n_members_yes": 0, "obs": 1}) == 1 + assert int(prob.sum()) == 3 * 4 + # forecast statistics (averaged over members) and truth statistics are kept + assert out["PTYPE.mean"].sel(**sel).item() == pytest.approx(14 / 8) + assert np.isfinite(out["PTYPE.mean"].sel(source="truth", region="all").item()) + + +def test_verify_ensemble_run_keeps_continuous_contingency_tables(): + out = _run(MEMBERS) + sel = dict(source="ens", region="all", station_group="all") + assert np.isfinite(out["T_2M.contingency_table"].sel(**sel)).all() + + +def _reference_member_mean_statistics(fcst, truth, regions, holdout_stations, params): + """member_mean_statistics before the memory fix (all params in one compute).""" + from verification import ( + CIRCULAR_PARAMS, + _apply_masks, + _compute_statistics, + _region_and_station_masks, + _too_many_missing, + ) + + dim = ["values"] + fcst, truth = xr.align(fcst, truth, join="inner", copy=False) + masks, station_masks = _region_and_station_masks(truth, regions, holdout_stations) + stats = [] + for param in params: + members = fcst[param].where(fcst[param].notnull().all("member")) + too_many_missing = _too_many_missing( + members.isnull().any("member"), + _apply_masks(truth[param], masks, station_masks), + dim, + 0.0, + ) + member_stats = _compute_statistics( + _apply_masks(members, masks, station_masks), + prefix=param + ".", + dim=dim, + circular=param in CIRCULAR_PARAMS, + ) + stats.append( + member_stats.mean("member") + .squeeze("source", drop=True) + .where(~too_many_missing) + ) + return xr.merge(stats).compute() + + +def test_member_mean_statistics_unchanged_by_chunking(): + """Members chunked one by one, one param at a time: same statistics as before.""" + rng = np.random.default_rng(1) + n_members, n_steps, n = 4, 3, 50 + values = rng.normal(280.0, 5.0, (n_members, n_steps, n)) + values[2, 1, 7] = np.nan # one member missing at one point and step + coords = { + "member": np.arange(1, n_members + 1), + "step": pd.to_timedelta(np.arange(n_steps) * 6, unit="h"), + "values": [f"S{i:02d}" for i in range(n)], + "longitude": ("values", np.linspace(5.0, 10.0, n)), + "latitude": ("values", np.linspace(45.9, 47.5, n)), + } + fcst = xr.Dataset( + { + "T_2M": (("member", "step", "values"), values), + "TD_2M": (("member", "step", "values"), values - 5.0), + }, + coords=coords, + ).chunk({"step": 1}) + truth = fcst.isel(member=0, drop=True).compute() + 0.5 + regions = _REGION + [{"type": "bbox", "name": "west", "bbox": [5.0, 7.5, 43, 49.5]}] + holdout = ["S00", "S03", "S10"] + params = ["T_2M", "TD_2M"] + new = verification_metrics.member_mean_statistics( + fcst, truth, regions, holdout, params=params, num_workers=2 + ) + ref = _reference_member_mean_statistics(fcst, truth, regions, holdout, params) + xr.testing.assert_allclose(new, ref, rtol=1e-12) + # S07 is missing in one member at step 1: NaN where it counts, not in holdout + assert np.isnan(new["T_2M.mean"].isel(step=1).sel(station_group="all")).all() + assert np.isfinite(new["T_2M.mean"].isel(step=1).sel(station_group="holdout")).all() + assert np.isfinite(new["T_2M.mean"].isel(step=0)).all() + + +def test_verification_rules_pass_members_flag(): + """verification_metrics and verification_valid_mask read grib/ and pass the members + flag as a plain param (empty for deterministic runs).""" + text = (ROOT / "workflow/rules/verification.smk").read_text() + for rule in ("verification_valid_mask", "verification_metrics"): + body = text.split(f"rule {rule}:", 1)[1].split("\nrule ", 1)[0] + assert ( + "grib_out_dir=lambda wc: run_grib_root(wc.run_id, wc.init_time).resolve()" + in body + ) + assert "members_flag=members_flag" in body + assert "--forecast {params.grib_out_dir}" in body + assert "{params.members_flag}" in body diff --git a/tests/unit/test_verification_aggregation.py b/tests/unit/test_verification_aggregation.py new file mode 100644 index 00000000..5d6f0731 --- /dev/null +++ b/tests/unit/test_verification_aggregation.py @@ -0,0 +1,216 @@ +import sys +from pathlib import Path + +import numpy as np +import xarray as xr + +ROOT = Path(__file__).parents[2] +DATA = Path(__file__).parent / "data" +sys.path.insert(0, str(ROOT / "workflow" / "scripts")) +from verification_aggregation import aggregate_results # noqa: E402 + +N = 4 +INITS = np.array( + ["2024-01-01T00", "2024-01-02T00", "2024-07-01T12"], dtype="datetime64[h]" +) +DIMS = ["forecast_reference_time", "source", "region", "step"] +SHAPE = (3, 1, 2, 2) +COORDS = { + "forecast_reference_time": INITS, + "source": ["run"], + "region": ["a", "b"], + "step": [6, 12], +} + + +def _det_dataset(seed=0): + rng = np.random.default_rng(seed) + ds = xr.Dataset( + { + "T_2M.BIAS": (DIMS, rng.normal(size=SHAPE)), + "T_2M.MSE": (DIMS, rng.uniform(1, 2, size=SHAPE)), + "T_2M.CORR": (DIMS, rng.uniform(0.5, 1, size=SHAPE)), + "T_2M.var": (DIMS, rng.uniform(1, 2, size=SHAPE)), + "TOT_PREC.contingency_table": ( + DIMS + ["contingency", "TOT_PREC.threshold"], + rng.integers(0, 50, size=SHAPE + (5, 2)).astype(float), + ), + }, + coords={ + **COORDS, + "contingency": [ + "tp_count", + "tn_count", + "fp_count", + "fn_count", + "total_count", + ], + "TOT_PREC.threshold": ["gt_0p1", "gt_1p0"], + }, + ) + return ds + + +def _ens_dataset(): + rng = np.random.default_rng(1) + ds = _det_dataset() + # tables of counts per init: dims (..., TOT_PREC.threshold, n_members_yes, obs) + table = rng.integers(0, 20, size=SHAPE + (2, N + 1, 2)).astype(float) + ens_var = rng.uniform(0.5, 2, size=SHAPE) + ens_var[:, :, :, 0] = 0.0 # zero spread at the first step + ds["T_2M.ENS_VAR"] = (DIMS, ens_var) + ds["T_2M.CRPS"] = (DIMS, rng.uniform(0.1, 1, size=SHAPE)) + ds["T_2M.MEMBER_MSE"] = (DIMS, rng.uniform(1, 3, size=SHAPE)) + ds["T_2M.rank_hist"] = ( + DIMS + ["rank"], + rng.integers(0, 30, size=SHAPE + (N + 1,)).astype(float), + ) + ds["TOT_PREC.prob_table"] = ( + DIMS + ["TOT_PREC.threshold", "n_members_yes", "obs"], + table, + ) + ds = ds.assign_coords( + {"rank": np.arange(N + 1), "n_members_yes": np.arange(N + 1), "obs": [0, 1]} + ) + ds.attrs["n_members"] = N + return ds + + +def test_deterministic_output_unchanged(): + """Deterministic inputs give the output of the code before ensemble support. + + ``data/aggregation_golden_deterministic.nc`` was written by + ``aggregate_results(_det_dataset())`` before the ensemble scores were added. + """ + expected = xr.load_dataset(DATA / "aggregation_golden_deterministic.nc") + result = aggregate_results(_det_dataset()) + xr.testing.assert_identical(result, expected) + for var in expected.variables: + if expected[var].dtype.kind in "fiub": + np.testing.assert_array_equal(result[var].values, expected[var].values) + + +def test_ensemble_spread_ssr_and_means(): + ds = _ens_dataset() + out = aggregate_results(ds) + allsel = dict(season="all", init_hour=-999) + sel = out.sel(**allsel) + assert "T_2M.ENS_VAR" not in out + assert "T_2M.MEMBER_MSE" not in out and "T_2M.MEMBER_RMSE" in out + assert "T_2M.ENS_std" not in out + np.testing.assert_allclose( + sel["T_2M.CRPS"], ds["T_2M.CRPS"].mean("forecast_reference_time") + ) + exp_member = np.sqrt(ds["T_2M.MEMBER_MSE"].mean("forecast_reference_time")) + np.testing.assert_allclose(sel["T_2M.MEMBER_RMSE"], exp_member) + spread = np.sqrt(ds["T_2M.ENS_VAR"].mean("forecast_reference_time")) + rmse = np.sqrt(ds["T_2M.MSE"].mean("forecast_reference_time")) + np.testing.assert_allclose(sel["T_2M.SPREAD"], spread) + exp_ssr = (spread / rmse * np.sqrt((N + 1) / N)).where(spread != 0) + np.testing.assert_allclose(sel["T_2M.SSR"], exp_ssr) + # zero spread at first step -> NaN SSR, SPREAD is 0 + assert np.isnan(sel["T_2M.SSR"].isel(step=0)).all() + assert (sel["T_2M.SPREAD"].isel(step=0) == 0).all() + assert out.attrs["n_members"] == N + + +def test_ssr_nan_where_rmse_zero(): + ds = _ens_dataset() + ds["T_2M.MSE"] = ds["T_2M.MSE"] * 0 + ds["T_2M.BIAS"] = ds["T_2M.BIAS"] * 0 + out = aggregate_results(ds) + assert np.isnan(out["T_2M.SSR"]).all() + + +def test_ssr_n_inferred_from_rank(): + ds = _ens_dataset() + ds.attrs.pop("n_members") + out = aggregate_results(ds) + ref = aggregate_results(_ens_dataset()) + xr.testing.assert_allclose(out["T_2M.SSR"], ref["T_2M.SSR"]) + assert out.attrs["n_members"] == N + + +def test_counts_conversion_with_stratification(): + ds = _ens_dataset() + out = aggregate_results(ds) + # rank_hist: sum over inits of the per-init counts, in every stratum + total = ds["T_2M.rank_hist"].sum("forecast_reference_time") + np.testing.assert_allclose( + out.sel(season="all", init_hour=-999)["T_2M.rank_hist"], total + ) + djf = ( + ds["T_2M.rank_hist"] + .isel(forecast_reference_time=[0, 1]) + .sum("forecast_reference_time") + ) + np.testing.assert_allclose( + out.sel(season="DJF", init_hour=-999)["T_2M.rank_hist"], djf + ) + jja_12 = ds["T_2M.rank_hist"].isel(forecast_reference_time=2) + np.testing.assert_allclose( + out.sel(season="JJA", init_hour=12)["T_2M.rank_hist"], jja_12 + ) + # prob_table split per threshold, dims n_members_yes, obs + pt = out["TOT_PREC.prob_table_gt_0p1"] + assert set(pt.dims) >= {"n_members_yes", "obs"} + exp = ( + ds["TOT_PREC.prob_table"] + .isel({"TOT_PREC.threshold": 0}) + .sum("forecast_reference_time") + ) + np.testing.assert_allclose(pt.sel(season="all", init_hour=-999), exp) + assert "TOT_PREC.prob_table" not in out + + +def test_brier_and_murphy_decomposition(): + ds = _ens_dataset() + out = aggregate_results(ds).sel(season="all", init_hour=-999) + counts = ds["TOT_PREC.prob_table"].sum("forecast_reference_time") + for i, label in enumerate(["gt_0p1", "gt_1p0"]): + c = ( + counts.isel({"TOT_PREC.threshold": i}) + .transpose(..., "n_members_yes", "obs") + .values + ) + t = c.sum(axis=(-1, -2)) + f = np.arange(N + 1) / N + n1, n0 = c[..., 1], c[..., 0] + nk = n1 + n0 + bs = ((n1 * (f - 1) ** 2 + n0 * f**2).sum(-1)) / t + obar = n1.sum(-1) / t + ok = n1 / nk + rel = (nk * (f - ok) ** 2).sum(-1) / t + res = (nk * (ok - obar[..., None]) ** 2).sum(-1) / t + unc = obar * (1 - obar) + got = { + k: out[f"TOT_PREC.{k}_{label}"].transpose("source", "region", "step").values + for k in ["BS", "BS_REL", "BS_RES", "BS_UNC"] + } + for k, ref in zip(["BS", "BS_REL", "BS_RES", "BS_UNC"], [bs, rel, res, unc]): + np.testing.assert_allclose(got[k], ref[0] if ref.ndim == 4 else ref) + np.testing.assert_allclose( + got["BS"], got["BS_REL"] - got["BS_RES"] + got["BS_UNC"] + ) + + +def test_brier_hand_example(): + # one init pair, N=2; bins k=0,1,2 with (obs0, obs1) counts + ds = _ens_dataset().isel(forecast_reference_time=[0], region=[0], step=[0]) + tab = np.zeros((1, 1, 1, 1, 2, 3, 2)) + tab[..., 0, :, :] = [[6, 0], [1, 1], [0, 2]] # k=0: 6 no,0 yes; k=1: 1,1; k=2: 0,2 + tab[..., 1, :, :] = tab[..., 0, :, :] + ds = ds.isel(rank=[0, 1, 2], n_members_yes=[0, 1, 2]) + ds["TOT_PREC.prob_table"] = (ds["TOT_PREC.prob_table"].dims, tab) + ds.attrs["n_members"] = 2 + out = ( + aggregate_results(ds) + .sel(season="all", init_hour=-999) + .isel(source=0, region=0, step=0) + ) + # T=10, f=(0,.5,1): BS = (0.25*(1+1) + 0)/10 = 0.05 ; obar = 3/10 + assert np.isclose(out["TOT_PREC.BS_gt_0p1"], 0.05) + assert np.isclose(out["TOT_PREC.BS_UNC_gt_0p1"], 0.21) + # REL = (2*(0.5-0.5)^2 + 0)/10 = 0 ; RES = (6*.09 + 2*.04 + 2*.49)/10 = 0.16 + assert np.isclose(out["TOT_PREC.BS_REL_gt_0p1"], 0.0) + assert np.isclose(out["TOT_PREC.BS_RES_gt_0p1"], 0.16) diff --git a/tests/unit/test_verify_ensemble.py b/tests/unit/test_verify_ensemble.py new file mode 100644 index 00000000..bef01de4 --- /dev/null +++ b/tests/unit/test_verify_ensemble.py @@ -0,0 +1,428 @@ +"""verify_ensemble(): ensemble scores on synthetic data.""" + +import itertools + +import numpy as np +import pytest +import xarray as xr + +from verification import verify, verify_ensemble + +_FRT = np.datetime64("2024-01-01T00", "ns") +_ALL = {"type": "bbox", "name": "all", "bbox": [1.5, 16, 43, 49.5]} +_WEST = {"type": "bbox", "name": "west", "bbox": [1.5, 7.5, 43, 49.5]} + + +def _stations(n): + """Coordinates of n stations in the 'all' region; those with lon < 7.5 are 'west'.""" + return { + "values": [f"S{k:05d}" for k in range(n)], + "longitude": ("values", np.linspace(5.0, 10.0, n)), + "latitude": ("values", np.linspace(46.0, 47.0, n)), + "forecast_reference_time": _FRT, + } + + +def _datasets(members, truth, param="T_2M", coords=None, dims=("values",)): + """members: (N, *dims), truth: dims.""" + n_members = members.shape[0] + coords = coords if coords is not None else _stations(truth.shape[-1]) + fcst = xr.Dataset( + {param: (("member", *dims), members)}, + coords={**coords, "member": np.arange(1, n_members + 1)}, + ) + obs = xr.Dataset({param: (dims, truth)}, coords=coords) + return fcst, obs + + +def _run(fcst, obs, **kwargs): + kwargs.setdefault("regions", [_ALL]) + return verify_ensemble(fcst, obs, "ens", "truth", num_workers=1, **kwargs) + + +def _sel(result, var, **sel): + sel = {"source": "ens", "region": "all", "station_group": "all", **sel} + return result[var].sel(sel).squeeze("forecast_reference_time", drop=True) + + +def _fair_crps_brute(members, truth): + """Fair CRPS per point by the double sum (members on axis 0).""" + n = members.shape[0] + skill = np.abs(members - truth).mean(axis=0) + spread = np.abs(members[:, None] - members[None, :]).sum(axis=(0, 1)) + return skill - spread / (2 * n * (n - 1)) + + +# --------------------------------------------------------------------------- +# Scores against brute force +# --------------------------------------------------------------------------- + + +def test_scores_match_brute_force_per_region_and_station_group(): + rng = np.random.default_rng(1) + n, n_members = 40, 5 + members = rng.normal(size=(n_members, n)) + truth = rng.normal(size=n) + fcst, obs = _datasets(members, truth) + holdout = ["S00001", "S00004", "S00010", "S00033"] + + result = _run(fcst, obs, regions=[_ALL, _WEST], holdout_stations=holdout) + + lon = fcst["longitude"].values + is_holdout = np.isin(fcst["values"].values, holdout) + groups = {"all": np.ones(n, bool), "holdout": is_holdout, "holdin": ~is_holdout} + regions = {"all": np.ones(n, bool), "west": lon < 7.5} + crps = _fair_crps_brute(members, truth) + for (region, rmask), (group, gmask) in itertools.product( + regions.items(), groups.items() + ): + sel = rmask & gmask + at = dict(region=region, station_group=group) + assert _sel(result, "T_2M.CRPS", **at).item() == pytest.approx(crps[sel].mean()) + assert _sel(result, "T_2M.ENS_VAR", **at).item() == pytest.approx( + members[:, sel].var(axis=0, ddof=1).mean() + ) + assert _sel(result, "T_2M.MEMBER_MSE", **at).item() == pytest.approx( + ((members[:, sel] - truth[sel]) ** 2).mean() + ) + ranks = (truth[sel] > members[:, sel]).sum(axis=0) # no ties + np.testing.assert_allclose( + _sel(result, "T_2M.rank_hist", **at).values, + np.bincount(ranks, minlength=n_members + 1), + ) + + +def test_crps_matches_scores_library(): + import scores + + rng = np.random.default_rng(2) + members = rng.gamma(2.0, size=(4, 30)) + truth = rng.gamma(2.0, size=30) + fcst, obs = _datasets(members, truth) + + result = _run(fcst, obs) + + expected = scores.probability.crps_for_ensemble( + fcst["T_2M"], obs["T_2M"], "member", method="fair" + ) + assert _sel(result, "T_2M.CRPS").item() == pytest.approx(expected.item()) + + +def test_identical_members_crps_equals_mae(): + """With identical members the fair CRPS is the MAE and the spread is 0.""" + rng = np.random.default_rng(3) + single = rng.normal(size=25) + truth = rng.normal(size=25) + fcst, obs = _datasets(np.stack([single] * 3), truth) + + result = _run(fcst, obs) + det = verify(fcst.mean("member"), obs, "ens", "truth", [_ALL], num_workers=1) + + assert _sel(result, "T_2M.CRPS").item() == pytest.approx( + _sel(det, "T_2M.MAE").item() + ) + assert _sel(result, "T_2M.ENS_VAR").item() == pytest.approx(0.0, abs=1e-12) + assert _sel(result, "T_2M.MEMBER_MSE").item() == pytest.approx( + _sel(det, "T_2M.MSE").item() + ) + + +# --------------------------------------------------------------------------- +# Spread-skill ratio and rank histogram shape +# --------------------------------------------------------------------------- + + +def _calibration_case(spread_scale, n=20000, n_members=10, seed=4): + """Truth ~ N(mu, 1); members ~ N(mu, spread_scale**2), mu varies per point.""" + rng = np.random.default_rng(seed) + mu = rng.normal(scale=3.0, size=n) + truth = mu + rng.normal(size=n) + members = mu + spread_scale * rng.normal(size=(n_members, n)) + fcst, obs = _datasets(members, truth) + ens = _run(fcst, obs) + det = verify(fcst.mean("member"), obs, "ens", "truth", [_ALL], num_workers=1) + # The ensemble scores merge into verify()'s output of the ensemble mean + merged = xr.merge([det, ens], compat="no_conflicts", join="outer") + n_members = ens.attrs["n_members"] + ssr = np.sqrt( + _sel(merged, "T_2M.ENS_VAR").item() / _sel(merged, "T_2M.MSE").item() + ) * np.sqrt((n_members + 1) / n_members) + hist = _sel(merged, "T_2M.rank_hist").values + return ssr, hist / hist.sum() + + +def test_perfect_ensemble_ssr_one_flat_rank_histogram(): + ssr, hist = _calibration_case(1.0) + assert ssr == pytest.approx(1.0, abs=0.03) + np.testing.assert_allclose(hist, 1 / 11, atol=0.01) + + +def test_underdispersive_ensemble_ssr_below_one_u_shaped_histogram(): + ssr, hist = _calibration_case(0.5) + assert ssr < 0.7 + assert hist[0] > 2 / 11 and hist[-1] > 2 / 11 + assert hist[5] < 1 / 11 + + +def test_overdispersive_ensemble_ssr_above_one_dome_histogram(): + ssr, hist = _calibration_case(2.0) + assert ssr > 1.5 + assert hist[0] < 0.5 / 11 and hist[-1] < 0.5 / 11 + assert hist[5] > 1 / 11 + + +# --------------------------------------------------------------------------- +# Ties, missing values +# --------------------------------------------------------------------------- + + +def test_ties_are_split_evenly_between_ranks(): + """Truth 0 mm with members (0, 0, 1, 2): ranks 0, 1 and 2 share each point.""" + n = 9 + members = np.tile(np.array([0.0, 0.0, 1.0, 2.0])[:, None], (1, n)) + fcst, obs = _datasets(members, np.zeros(n), param="TOT_PREC") + hist = _sel(_run(fcst, obs), "TOT_PREC.rank_hist").values + np.testing.assert_allclose(hist, [3, 3, 3, 0, 0], atol=1e-12) + + +def test_all_tied_spreads_over_all_ranks(): + """Dry everywhere: the truth ties every member, so all N+1 ranks get 1/(N+1).""" + n = 10 + fcst, obs = _datasets(np.zeros((4, n)), np.zeros(n), param="TOT_PREC") + result = _run(fcst, obs) + np.testing.assert_allclose( + _sel(result, "TOT_PREC.rank_hist").values, np.full(5, n / 5), atol=1e-12 + ) + assert _sel(result, "TOT_PREC.CRPS").item() == 0.0 + + +def test_nan_member_makes_point_missing(): + rng = np.random.default_rng(5) + n = 20 + members = rng.normal(size=(4, n)) + truth = rng.normal(size=n) + members[2, 7] = np.nan # one member missing at one point + truth[3] = np.nan # truth missing elsewhere: not a forecast gap + fcst, obs = _datasets(members, truth) + thresholds = {"T_2M": {"gt": [0.0]}} + + strict = _run(fcst, obs, threshold_dict=thresholds) + for var in ("CRPS", "ENS_VAR", "MEMBER_MSE", "rank_hist", "prob_table"): + assert np.isnan(_sel(strict, f"T_2M.{var}").values).all(), var + + tolerant = _run(fcst, obs, threshold_dict=thresholds, max_missing_fraction=0.1) + keep = np.ones(n, bool) + keep[[3, 7]] = False + assert _sel(tolerant, "T_2M.CRPS").item() == pytest.approx( + _fair_crps_brute(members[:, keep], truth[keep]).mean() + ) + assert _sel(tolerant, "T_2M.rank_hist").sum().item() == keep.sum() + assert _sel(tolerant, "T_2M.prob_table").sum().item() == keep.sum() + + +def test_truth_gaps_do_not_mask_scores(): + rng = np.random.default_rng(6) + truth = rng.normal(size=10) + truth[:5] = np.nan + fcst, obs = _datasets(rng.normal(size=(3, 10)), truth) + result = _run(fcst, obs) + assert not np.isnan(_sel(result, "T_2M.CRPS").item()) + assert _sel(result, "T_2M.rank_hist").sum().item() == 5 + + +# --------------------------------------------------------------------------- +# Probability tables +# --------------------------------------------------------------------------- + + +def test_prob_table_counts_and_threshold_dim_match_verify(): + rng = np.random.default_rng(7) + n, n_members = 50, 4 + members = rng.gamma(0.8, size=(n_members, n)) + truth = rng.gamma(0.8, size=n) + fcst, obs = _datasets(members, truth, param="TOT_PREC") + thresholds = {"TOT_PREC": {"gt": [0.1, 1.0], "lt": [0.5]}} + holdout = ["S00002", "S00020", "S00041"] + + result = _run( + fcst, + obs, + regions=[_ALL, _WEST], + threshold_dict=thresholds, + holdout_stations=holdout, + ) + det = verify( + fcst.mean("member"), + obs, + "ens", + "truth", + [_ALL, _WEST], + threshold_dict=thresholds, + holdout_stations=holdout, + num_workers=1, + ) + + table = result["TOT_PREC.prob_table"] + assert table.dims == ( + "forecast_reference_time", + "source", + "region", + "station_group", + "TOT_PREC.threshold", + "n_members_yes", + "obs", + ) + assert list(table["TOT_PREC.threshold"].values) == list( + det["TOT_PREC.threshold"].values + ) + assert list(table["n_members_yes"].values) == list(range(n_members + 1)) + assert list(table["obs"].values) == [0, 1] + # Every valid point is counted once per threshold + totals = table.sum(["n_members_yes", "obs"]) + n_valid = det["TOT_PREC.contingency_table"].sel(contingency="total_count") + np.testing.assert_array_equal( + totals.squeeze("forecast_reference_time").sel(source="ens").values, + n_valid.squeeze("forecast_reference_time") + .sel(source="ens") + .transpose(*totals.dims[2:]) + .values, + ) + # Brute force for one threshold, region "west", all stations + west = fcst["longitude"].values < 7.5 + k = (members[:, west] > 1.0).sum(axis=0) + o = (truth[west] > 1.0).astype(int) + expected = np.zeros((n_members + 1, 2)) + np.add.at(expected, (k, o), 1) + np.testing.assert_array_equal( + _sel( + result, + "TOT_PREC.prob_table", + region="west", + **{"TOT_PREC.threshold": "gt_1p0"}, + ).values, + expected, + ) + + +# --------------------------------------------------------------------------- +# Grids, special parameters, dims +# --------------------------------------------------------------------------- + + +def test_gridded_truth(): + """INCA-like (y, x) truth: works and has the single 'all' station group.""" + ny, nx, n_members = 4, 5, 3 + rng = np.random.default_rng(8) + lat, lon = np.meshgrid( + np.linspace(46.0, 47.0, ny), np.linspace(5.0, 10.0, nx), indexing="ij" + ) + coords = { + "y": np.arange(ny), + "x": np.arange(nx), + "latitude": (("y", "x"), lat), + "longitude": (("y", "x"), lon), + "forecast_reference_time": _FRT, + "step": [0, 6], + } + members = rng.normal(size=(n_members, 2, ny, nx)) + truth = rng.normal(size=(2, ny, nx)) + fcst, obs = _datasets(members, truth, coords=coords, dims=("step", "y", "x")) + + result = _run( + fcst.chunk({"step": 1, "member": 1}), + obs, + regions=[_ALL, _WEST], + threshold_dict={"T_2M": {"gt": [0.0]}}, + ) + + assert list(result["station_group"].values) == ["all"] + west = lon < 7.5 + for step in range(2): + crps = _fair_crps_brute(members[:, step], truth[step]) + assert _sel(result, "T_2M.CRPS", step=step * 6).item() == pytest.approx( + crps.mean() + ) + assert _sel( + result, "T_2M.CRPS", step=step * 6, region="west" + ).item() == pytest.approx(crps[west].mean()) + assert (_sel(result, "T_2M.rank_hist").sum("rank") == ny * nx).all() + assert ( + _sel(result, "T_2M.prob_table", region="west").sum(["n_members_yes", "obs"]) + == west.sum() + ).all() + with pytest.raises(ValueError, match="Holdout stations"): + _run(fcst, obs, holdout_stations=["CHM"]) + + +def test_wind_direction_and_precipitation_type(): + """No CRPS/spread for angles and class codes; PTYPE still gets prob tables.""" + n = 12 + rng = np.random.default_rng(9) + codes = np.array([0, 1, 5, 7]) + pt_members = codes[rng.integers(0, 4, size=(4, n))].astype(float) + pt_truth = codes[rng.integers(0, 4, size=n)].astype(float) + dd_members = rng.uniform(0, 360, size=(4, n)) + coords = _stations(n) + fcst = xr.Dataset( + { + "PTYPE": (("member", "values"), pt_members), + "DD_10M": (("member", "values"), dd_members), + }, + coords={**coords, "member": np.arange(1, 5)}, + ) + obs = xr.Dataset( + { + "PTYPE": ("values", pt_truth), + "DD_10M": ("values", rng.uniform(0, 360, size=n)), + }, + coords=coords, + ) + + result = _run(fcst, obs, threshold_dict={"PTYPE": {"eq": [0, 1, 5, 7]}}) + + for param in ("PTYPE", "DD_10M"): + for var in ("CRPS", "ENS_VAR", "MEMBER_MSE", "rank_hist"): + assert np.isnan(_sel(result, f"{param}.{var}").values).all() + assert "DD_10M.prob_table" not in result + rain = _sel(result, "PTYPE.prob_table", **{"PTYPE.threshold": "eq_1"}) + expected = np.zeros((5, 2)) + np.add.at(expected, ((pt_members == 1).sum(0), (pt_truth == 1).astype(int)), 1) + np.testing.assert_array_equal(rain.values, expected) + + +def test_output_dims_match_verify_continuous_scores(): + rng = np.random.default_rng(10) + coords = {**_stations(8), "step": [0, 6, 12]} + fcst, obs = _datasets( + rng.normal(size=(3, 3, 8)), + rng.normal(size=(3, 8)), + coords=coords, + dims=("step", "values"), + ) + result = _run(fcst.chunk({"step": 1}), obs, holdout_stations=["S00001"]) + det = verify( + fcst.mean("member"), + obs, + "ens", + "truth", + [_ALL], + holdout_stations=["S00001"], + num_workers=1, + ) + + assert result.attrs["n_members"] == 3 + for var in ("CRPS", "ENS_VAR", "MEMBER_MSE"): + assert result[f"T_2M.{var}"].dims == det["T_2M.MAE"].dims + assert result["T_2M.rank_hist"].dims == (*det["T_2M.MAE"].dims, "rank") + assert list(result["rank"].values) == [0, 1, 2, 3] + assert "member" not in result.dims + merged = xr.merge([det, result], compat="no_conflicts", join="outer") + assert merged.sizes["source"] == 2 # ens and truth, no new source + + +def test_needs_a_member_dim_with_two_members(): + fcst, obs = _datasets(np.zeros((1, 5)), np.zeros(5)) + with pytest.raises(ValueError, match="2 members"): + _run(fcst, obs) + with pytest.raises(ValueError, match="member"): + _run(fcst.isel(member=0), obs) diff --git a/tests/unit/test_verify_regression.py b/tests/unit/test_verify_regression.py new file mode 100644 index 00000000..89259ec4 --- /dev/null +++ b/tests/unit/test_verify_regression.py @@ -0,0 +1,170 @@ +"""verify() output must not change when its internals are refactored (up to last-bit +floating-point differences between machines). + +The golden files under ``tests/unit/data/`` were written by verify() before the +region/station-group mask construction and the missing-forecast rule were factored out +into helpers shared with verify_ensemble() (commit be3dc5c). Golden files rather than +loading the old code from git: the test does not depend on the git history, and any +later change of verify()'s output shows up here too. + +The inputs are built from closed-form expressions (no random numbers), so they do not +depend on the numpy version. To regenerate the files after an intended change of +verify()'s output: ``uv run python tests/unit/test_verify_regression.py``. +""" + +from pathlib import Path + +import numpy as np +import pytest +import xarray as xr + +from verification import verify + +DATA = Path(__file__).parent / "data" +_FRT = np.datetime64("2024-01-01T00", "ns") +_REGIONS = [ + {"type": "bbox", "name": "all", "bbox": [1.5, 16, 43, 49.5]}, + {"type": "bbox", "name": "west", "bbox": [1.5, 7.5, 43, 49.5]}, +] + + +def station_inputs() -> tuple[xr.Dataset, xr.Dataset]: + """12 stations, 3 steps; gaps in forecast and truth; circular, categorical, + truth-only and forecast-only parameters.""" + n, steps = 12, np.array([0, 6, 12]) + i = np.arange(n)[None, :] + s = np.arange(steps.size)[:, None] + coords = { + "values": [f"S{k:02d}" for k in range(n)], + "longitude": ("values", np.linspace(5.0, 10.0, n)), + "latitude": ("values", np.linspace(46.0, 47.0, n)), + "forecast_reference_time": _FRT, + "step": steps, + } + dims = ("step", "values") + t_obs = 280.0 + 3.0 * np.sin(0.7 * i + 0.3 * s) + t_fcst = t_obs + 0.8 * np.cos(1.3 * i - 0.5 * s) + 0.2 + t_fcst[1, 4] = np.nan # a forecast gap at one step + t_obs[2, 7] = np.nan # a truth gap + tp_obs = np.maximum(0.0, 2.0 * np.sin(0.9 * i + s)) + tp_fcst = np.maximum(0.0, 2.0 * np.sin(0.9 * i + s + 0.4) + 0.3) + dd_obs = np.mod(40.0 * i + 25.0 * s, 360.0) + dd_fcst = np.mod(dd_obs + 30.0 * np.sin(i + s) + 5.0, 360.0) + pt_obs = np.array([0, 1, 5, 7])[(i + s) % 4] * np.ones_like(t_obs) + pt_fcst = np.array([0, 1, 5, 7])[(i + 2 * s + 1) % 4] * np.ones_like(t_obs) + fcst = xr.Dataset( + { + "T_2M": (dims, t_fcst), + "TOT_PREC": (dims, tp_fcst), + "DD_10M": (dims, dd_fcst), + "PTYPE": (dims, pt_fcst), + "FCST_ONLY": (dims, t_fcst), + }, + coords=coords, + ) + obs = xr.Dataset( + { + "T_2M": (dims, t_obs), + "TOT_PREC": (dims, tp_obs), + "DD_10M": (dims, dd_obs), + "PTYPE": (dims, pt_obs), + "CAPE_MU": (dims, 100.0 + 10.0 * i + s * np.ones_like(t_obs)), + }, + coords=coords, + ) + return fcst, obs + + +STATION_THRESHOLDS = { + "T_2M": {"lt": [280.0], "gt": [278.5, 282.0]}, + "TOT_PREC": {"gt": [0.0, 1.0]}, + "DD_10M": {"gt": [180.0]}, + "PTYPE": {"eq": [0, 1, 5, 7]}, +} + + +def gridded_inputs() -> tuple[xr.Dataset, xr.Dataset]: + """INCA-like truth on a regular (y, x) grid, with gaps.""" + ny, nx, steps = 5, 6, np.array([0, 6]) + y = np.arange(ny)[None, :, None] + x = np.arange(nx)[None, None, :] + s = np.arange(steps.size)[:, None, None] + lat, lon = np.meshgrid( + np.linspace(46.0, 47.0, ny), np.linspace(5.0, 10.0, nx), indexing="ij" + ) + coords = { + "y": np.arange(ny), + "x": np.arange(nx), + "latitude": (("y", "x"), lat), + "longitude": (("y", "x"), lon), + "forecast_reference_time": _FRT, + "step": steps, + } + dims = ("step", "y", "x") + t_obs = 275.0 + np.sin(0.5 * y + 0.8 * x + s) + t_fcst = t_obs + 0.5 * np.cos(x - y + s) + t_fcst[1, 0, 0] = np.nan + t_obs[0, 4, 5] = np.nan + tp_obs = np.maximum(0.0, np.sin(y * x + s)) + tp_fcst = np.maximum(0.0, np.sin(y * x + s + 0.2)) + fcst = xr.Dataset( + {"T_2M": (dims, t_fcst), "TOT_PREC": (dims, tp_fcst)}, coords=coords + ) + obs = xr.Dataset({"T_2M": (dims, t_obs), "TOT_PREC": (dims, tp_obs)}, coords=coords) + return fcst, obs + + +def cases() -> dict: + """name -> verify() keyword arguments.""" + st_fcst, st_obs = station_inputs() + gr_fcst, gr_obs = gridded_inputs() + common = dict(fcst_label="fcst", obs_label="obs", regions=_REGIONS, num_workers=1) + return { + "stations": dict( + fcst=st_fcst.chunk({"step": 1}), + obs=st_obs, + threshold_dict=STATION_THRESHOLDS, + holdout_stations=["S00", "S03", "S07", "S11"], + **common, + ), + "stations_tolerant": dict( + fcst=st_fcst, + obs=st_obs, + threshold_dict=STATION_THRESHOLDS, + max_missing_fraction=0.5, + **common, + ), + "gridded": dict( + fcst=gr_fcst.chunk({"step": 1}), + obs=gr_obs, + threshold_dict={"TOT_PREC": {"gt": [0.1, 0.5]}}, + **common, + ), + } + + +@pytest.mark.parametrize("name", list(cases())) +def test_verify_output_unchanged(name): + result = verify(**cases()[name]) + expected = xr.load_dataset(DATA / f"verify_golden_{name}.nc") + + # Same structure, attributes and NaN positions; values equal up to the last bits, + # which differ between machines for sin/cos (e.g. the circular variance of DD_10M). + xr.testing.assert_allclose(result, expected, rtol=1e-12, atol=0) + assert result.attrs == expected.attrs + assert set(result.variables) == set(expected.variables) + for var in expected.variables: + assert result[var].dims == expected[var].dims, var + assert result[var].attrs == expected[var].attrs, var + if expected[var].dtype.kind in "fiub": + assert result[var].dtype == expected[var].dtype, var + if expected[var].dtype.kind == "f": + np.testing.assert_array_equal( + np.isnan(result[var].values), np.isnan(expected[var].values), var + ) + + +if __name__ == "__main__": + DATA.mkdir(exist_ok=True) + for name, kwargs in cases().items(): + verify(**kwargs).to_netcdf(DATA / f"verify_golden_{name}.nc") diff --git a/workflow/rules/common.smk b/workflow/rules/common.smk index b979677a..0dd832b1 100644 --- a/workflow/rules/common.smk +++ b/workflow/rules/common.smk @@ -684,23 +684,58 @@ def scoremap_runtime(wc): return max(60, math.ceil(1.5 * len(REFTIMES) * seconds / 60)) +def run_grib_root(run_id, init_time): + """``grib/`` directory of a run's inference output (one sub-directory per member + for an ensemble). Verification reads it, with ``members_flag`` for an ensemble.""" + return Path(OUT_ROOT) / f"data/runs/{run_id}/{init_time}/grib" + + def run_grib_dir(run_id, init_time): """GRIB directory of a run's inference output, as read by downstream rules. - Deterministic runs write ``grib/``; ensemble runs write ``grib/{NNN}/`` per member - and, until ensemble verification lands, downstream rules use member 001 only. + Deterministic runs write ``grib/``; ensemble runs write ``grib/{NNN}/`` per member. + Plots, SAL, score maps and MEC read member 001 only; verification reads all + members (``run_grib_root`` and ``members_flag``). """ - base = Path(OUT_ROOT) / f"data/runs/{run_id}/{init_time}/grib" + base = run_grib_root(run_id, init_time) if RUN_CONFIGS[run_id].get("ensemble") is not None: return base / "001" return base +def ensemble_members(run_id) -> int | None: + """Number of members of an ensemble run, None for a deterministic run.""" + ensemble = RUN_CONFIGS[run_id].get("ensemble") + return None if ensemble is None else ensemble["members"] + + +def members_flag(wildcards): + """``--members N`` for an ensemble run (``{run_id}`` wildcard), else empty: the + verification scripts then read every ``grib/{NNN}/`` and fail on a missing member. + """ + members = ensemble_members(wildcards.run_id) + return "" if members is None else f"--members {members}" + + +def ensemble_scaled(base, per_member, unit=None): + """Resource callable: ``base`` for deterministic runs, ``base + N * per_member`` + for an ensemble of N members (``{run_id}`` wildcard); as a string with ``unit`` + appended (e.g. ``"m"`` for a runtime) when ``unit`` is given.""" + + def resource(wc): + members = ensemble_members(wc.run_id) + value = base if members is None else base + members * per_member + return value if unit is None else f"{value}{unit}" + + return resource + + for _run_id, _rc in RUN_CONFIGS.items(): if _rc.get("ensemble") is not None: logging.getLogger("snakemake").warning( - "Run %s is an ensemble of %d members: verification, plots and MEC use " - "member 001 only until ensemble verification is implemented.", + "Run %s is an ensemble of %d members: verification scores the ensemble " + "mean and the ensemble scores over all members; plots, SAL, score maps " + "and MEC use member 001 only.", _run_id, _rc["ensemble"]["members"], ) diff --git a/workflow/rules/report.smk b/workflow/rules/report.smk index 7dea3a93..5a4635dd 100644 --- a/workflow/rules/report.smk +++ b/workflow/rules/report.smk @@ -22,6 +22,7 @@ def make_header_text(): rule report_experiment_dashboard: input: "src/verification/__init__.py", + "src/verification/ensemble_report.py", script="workflow/scripts/report_experiment_dashboard.py", verif=EXPERIMENT_PARTICIPANTS.values(), template="resources/report/dashboard/template.html.jinja2", diff --git a/workflow/rules/verification.smk b/workflow/rules/verification.smk index 97014091..7fb5f21e 100644 --- a/workflow/rules/verification.smk +++ b/workflow/rules/verification.smk @@ -63,20 +63,26 @@ rule verification_valid_mask: OUT_ROOT / f"data/runs/{{run_id}}/{{init_time}}/valid_mask_{VALID_MASK_HASH}.nc", log: OUT_ROOT / "logs/verification_valid_mask/{run_id}-{init_time}.log", + # Ensembles: every member is loaded, so memory and runtime grow with N. resources: - cpus_per_task=balfrin_cpus(40_000, actual_cpus=8), - mem_mb=40_000, - runtime="30m", + cpus_per_task=lambda wc: balfrin_cpus( + ensemble_scaled(40_000, 1_000)(wc), actual_cpus=8 + ), + mem_mb=ensemble_scaled(40_000, 1_000), + runtime=ensemble_scaled(30, 3, unit="m"), params: fcst_steps=lambda wc: RUN_CONFIGS[wc.run_id]["steps"], truth=config["truth"]["root"], - grib_out_dir=lambda wc: run_grib_dir(wc.run_id, wc.init_time).resolve(), + grib_out_dir=lambda wc: run_grib_root(wc.run_id, wc.init_time).resolve(), + # ensembles: all members (grib/{NNN}/) are read + members_flag=members_flag, mask_params=",".join(COMMON_POINTS_PARAMS), shell: """ export ECCODES_DEFINITION_PATH=$(realpath .venv/share/eccodes-cosmo-resources/definitions) uv run {input.script} \ --forecast {params.grib_out_dir} \ + {params.members_flag} \ --truth "{params.truth}" \ --reftime {wildcards.init_time} \ --steps "{params.fcst_steps}" \ @@ -162,10 +168,13 @@ rule verification_metrics: OUT_ROOT / f"data/runs/{{run_id}}/{{init_time}}/verif_{VERIF_HASH}.nc", log: OUT_ROOT / "logs/verification_metrics/{run_id}-{init_time}.log", + # Ensembles: every member is loaded and scored, so memory and runtime grow with N. resources: - cpus_per_task=balfrin_cpus(80_000, actual_cpus=24), - mem_mb=80_000, - runtime="60m", + cpus_per_task=lambda wc: balfrin_cpus( + ensemble_scaled(80_000, 2_000)(wc), actual_cpus=24 + ), + mem_mb=ensemble_scaled(80_000, 2_000), + runtime=ensemble_scaled(60, 6, unit="m"), # wildcard_constraints: # run_id="^" # to avoid ambiguitiy with run_baseline_verif # TODO: implement logic to use experiment name instead of run_id as wildcard @@ -174,7 +183,9 @@ rule verification_metrics: truth=config["truth"]["root"], truth_source_id=f"truth-{TRUTH_HASH}", regions=REGIONS, - grib_out_dir=lambda wc: run_grib_dir(wc.run_id, wc.init_time).resolve(), + grib_out_dir=lambda wc: run_grib_root(wc.run_id, wc.init_time).resolve(), + # ensembles: all members (grib/{NNN}/) are read + members_flag=members_flag, experiment_params=",".join(EXPERIMENT_PARAMS), threshold_dict=EXPERIMENT_CONFIG.get("thresholds", {}), holdout_flag=lambda wc, input: ( @@ -191,6 +202,7 @@ rule verification_metrics: export ECCODES_DEFINITION_PATH=$(realpath .venv/share/eccodes-cosmo-resources/definitions) uv run {input.script} \ --forecast {params.grib_out_dir} \ + {params.members_flag} \ --truth "{params.truth}" \ --reftime {wildcards.init_time} \ --steps "{params.fcst_steps}" \ @@ -252,6 +264,7 @@ use rule verification_metrics_aggregation as verification_metrics_aggregation_ba rule verification_metrics_plot: input: "src/verification/__init__.py", + "src/verification/ensemble_report.py", script="workflow/scripts/verification_plot_metrics.py", verif=list(EXPERIMENT_PARTICIPANTS.values()), output: diff --git a/workflow/scripts/report_experiment_dashboard.py b/workflow/scripts/report_experiment_dashboard.py index 75a13efb..532028a4 100644 --- a/workflow/scripts/report_experiment_dashboard.py +++ b/workflow/scripts/report_experiment_dashboard.py @@ -14,6 +14,7 @@ _ensure_unique_lead_time, _select_best_sources, ) +from verification.ensemble_report import CRPS_NOTE, add_deterministic_crps, is_score_var from verification import decode_metric LOG = logging.getLogger(__name__) @@ -69,10 +70,16 @@ def main(args): dfs = [_ensure_unique_lead_time(d) for d in dfs] dfs = _select_best_sources(dfs) ds = _concat_sources(dfs) + ds, crps_filled = add_deterministic_crps(ds) LOG.info("Loaded verification netcdf: \n%s", ds) # extract only non-spatial variables to pd.DataFrame - nonspatial_vars = [d for d in ds.data_vars if "spatial" not in d and "." in d] + # (rank histograms and probability tables are counts for plots, not scores) + nonspatial_vars = [ + d + for d in ds.data_vars + if "spatial" not in d and "." in d and is_score_var(ds[d]) + ] df = ds[nonspatial_vars].to_array("stack").to_dataframe(name="value").reset_index() df[["param", "metric"]] = df["stack"].str.split(".", n=1, expand=True) df["metric"] = df.metric.apply(decode_metric) @@ -175,6 +182,7 @@ def _sanitize(v): station_holdout=station_holdout, station_groups=station_groups, header_text=args.header_text, + crps_note=CRPS_NOTE if crps_filled else "", configfile_content=open(args.configfile, "r").read() if args.configfile.is_file() else "", diff --git a/workflow/scripts/report_scorecard.py b/workflow/scripts/report_scorecard.py index 5d8b4662..a5637906 100644 --- a/workflow/scripts/report_scorecard.py +++ b/workflow/scripts/report_scorecard.py @@ -6,7 +6,7 @@ diff = (model − baseline) / |baseline| × 100 [%] Positive diff is better for higher-is-better metrics (R², ETS, POD) and worse -for lower-is-better metrics (RMSE, MAE, STDE, FAR). +for lower-is-better metrics (RMSE, MAE, STDE, FAR, CRPS). """ import logging @@ -20,6 +20,11 @@ import matplotlib.pyplot as plt from matplotlib.transforms import ScaledTranslation +from verification.ensemble_report import ( + CRPS_NOTE, + drop_non_score_vars, + fill_crps_from_mae, +) from verification import decode_metric LOG = logging.getLogger(__name__) @@ -202,7 +207,7 @@ def _build_config(args) -> dict: # All recognised metrics — every entry must also appear in metric_directions. "all_metrics": ["RMSE", "MAE", "STDE", "R2", "ETS", "POD", "FAR"], "metric_directions": { - "lower_is_better": ["RMSE", "MAE", "STDE", "FAR"], + "lower_is_better": ["RMSE", "MAE", "STDE", "FAR", "CRPS"], "higher_is_better": ["R2", "ETS", "POD"], }, # None value → use all_metrics for that variable; list → restrict to that subset. @@ -259,10 +264,14 @@ def _load_relative_diff(cfg: dict) -> xr.Dataset: f"Fix: delete '{fewer}' and rerun the pipeline." ) - model_ds = model_ds.sel(**sel_coords).squeeze(drop=True) - baseline_ds = baseline_ds.sel(**{**sel_coords, "source": baseline_source}).squeeze( - drop=True + model_ds = drop_non_score_vars(model_ds.sel(**sel_coords).squeeze(drop=True)) + baseline_ds = drop_non_score_vars( + baseline_ds.sel(**{**sel_coords, "source": baseline_source}).squeeze(drop=True) ) + # Deterministic participants have no CRPS: show their MAE (display only). + model_ds, model_filled = fill_crps_from_mae(model_ds) + baseline_ds, baseline_filled = fill_crps_from_mae(baseline_ds) + cfg["crps_filled"] = model_filled or baseline_filled common_vars = [v for v in model_ds.data_vars if v in baseline_ds.data_vars] common_leads = sorted(set(model_ds.step.values) & set(baseline_ds.step.values)) @@ -793,6 +802,9 @@ def _render_scorecard(diff: xr.Dataset, cfg: dict, outfn: Path): cfg, ) + if cfg.get("crps_filled") and any(m == "CRPS" for _, m in rows): + fig.text(0.01, 0.0, CRPS_NOTE, fontsize=small_fs, va="top", ha="left") + outfn.parent.mkdir(parents=True, exist_ok=True) plt.savefig(outfn, bbox_inches="tight") LOG.info("Saved %s", outfn) diff --git a/workflow/scripts/verification_aggregation.py b/workflow/scripts/verification_aggregation.py index 60c7c492..8d7acb5e 100644 --- a/workflow/scripts/verification_aggregation.py +++ b/workflow/scripts/verification_aggregation.py @@ -34,6 +34,40 @@ def get_season(time): return "SON" +def _brier_from_prob_table(table: xr.DataArray) -> xr.Dataset: + """Brier score and Murphy decomposition from a prob_table of counts. + + ``table`` has dims ``n_members_yes`` (0..N) and ``obs`` (0, 1) besides the others. + The forecast probability of bin k is k/N. BS = REL - RES + UNC holds exactly as + the forecast probability takes one value per bin. + """ + n_bins = table.sizes["n_members_yes"] + n_members = n_bins - 1 + f = xr.DataArray( + np.arange(n_bins) / n_members, + dims="n_members_yes", + coords={"n_members_yes": table["n_members_yes"]}, + ) + n_obs1 = table.isel(obs=1, drop=True) + n_obs0 = table.isel(obs=0, drop=True) + n_k = n_obs0 + n_obs1 + total = n_k.sum("n_members_yes", skipna=False) + total = total.where(total > 0) + o_k = (n_obs1 / n_k.where(n_k > 0)).fillna(0.0) + o_bar = n_obs1.sum("n_members_yes", skipna=False) / total + bs = (n_obs1 * (f - 1) ** 2 + n_obs0 * f**2).sum("n_members_yes", skipna=False) + rel = (n_k * (f - o_k) ** 2).sum("n_members_yes", skipna=False) + res = (n_k * (o_k - o_bar) ** 2).sum("n_members_yes", skipna=False) + return xr.Dataset( + { + "BS": bs / total, + "BS_REL": rel / total, + "BS_RES": res / total, + "BS_UNC": o_bar * (1 - o_bar), + } + ) + + def aggregate_results(ds: xr.Dataset) -> xr.Dataset: """Compute mean metric values aggregated by season""" LOG.info("Aggregation results") @@ -74,8 +108,23 @@ def aggregate_results(ds: xr.Dataset) -> xr.Dataset: ds_mean.append(ds_result) out = xr.merge(ds_mean, compat="no_conflicts", join="outer") - for var in out.data_vars: - if "contingency_table" in var: + n_members_attr = ds.attrs.get("n_members") + for var in list(out.data_vars): + if "rank_hist" in var: + # ensemble rank histogram: back to per-stratum counts + out[var] = out[var] * out["n_samples"] + if n_members_attr is None: + n_members_attr = out.sizes["rank"] - 1 + elif "prob_table" in var: + # forecast-probability/observation counts: back to per-stratum counts, + # then Brier score and its decomposition per threshold + out[var] = out[var] * out["n_samples"] + brier = _brier_from_prob_table(out[var]) + for name, da in brier.data_vars.items(): + out[var.replace("prob_table", name)] = da + if n_members_attr is None: + n_members_attr = out.sizes["n_members_yes"] - 1 + elif "contingency_table" in var: # convert contingency table means back to per-stratum counts out[var] = out[var] * out["n_samples"] contingency_manager = scores.categorical.BasicContingencyManager( @@ -128,6 +177,24 @@ def aggregate_results(ds: xr.Dataset) -> xr.Dataset: out[var] = np.sqrt(out[var]) out = out.rename(var_transform) + # Ensemble spread and spread-skill ratio (after the sqrt/rename above: needs RMSE) + for var in [v for v in out.data_vars if v.endswith(".ENS_VAR")]: + prefix = var[: -len("ENS_VAR")] + spread = np.sqrt(out[var]) + out[f"{prefix}SPREAD"] = spread + rmse_var = f"{prefix}RMSE" + if rmse_var in out.data_vars and n_members_attr is not None: + n = float(n_members_attr) + rmse = out[rmse_var] + ssr = spread / rmse.where(rmse != 0) * np.sqrt((n + 1) / n) + out[f"{prefix}SSR"] = ssr.where(spread != 0) + out = out.drop_vars(var) + if n_members_attr is not None and any( + v.endswith(".SPREAD") or "rank_hist" in v or "prob_table" in v + for v in out.data_vars + ): + out.attrs["n_members"] = int(n_members_attr) + LOG.info("Computed aggregation in %.2f seconds", time.time() - start) LOG.info("Aggregated results: \n %s", out) diff --git a/workflow/scripts/verification_metrics.py b/workflow/scripts/verification_metrics.py index 0b6335a1..29cb3480 100644 --- a/workflow/scripts/verification_metrics.py +++ b/workflow/scripts/verification_metrics.py @@ -5,10 +5,22 @@ from datetime import datetime from pathlib import Path +import numpy as np import pandas as pd import xarray as xr -from verification import verify, apply_lapse_rate_correction_inplace # noqa: E402 +from verification import ( # noqa: E402 + CATEGORICAL_PARAMS, + CIRCULAR_PARAMS, + _apply_masks, + _compute_statistics, + _region_and_station_masks, + _too_many_missing, + apply_lapse_rate_correction_inplace, + default_num_workers, + verify, + verify_ensemble, +) from verification.spatial import map_forecast_to_truth # noqa: E402 from verification.valid_mask import apply_valid_mask, combine_valid_masks # noqa: E402 from data_input import ( @@ -88,16 +100,58 @@ def program_summary_log(args): LOG.info("=" * 80) +def load_ensemble_forecast( + root: Path, reftime: datetime, steps: list[int], params: list[str], members: int +) -> xr.Dataset: + """Load the members ``root/001 .. root/{members:03d}`` of an ensemble run. + + Each member is loaded with ``load_forecast_data`` (same derived and period + parameters as a deterministic run) and the members are concatenated on a + ``member`` dimension (1..N). A missing member is an error. All members get + member 001's ``elevation`` (model orography, the same for every member), so + the lapse-rate correction uses one elevation for the whole ensemble. + """ + datasets = [] + for m in range(1, members + 1): + member_dir = Path(root) / f"{m:03d}" + if not any(member_dir.glob("*.grib")): + raise FileNotFoundError( + f"Ensemble member {m:03d} of {members}: no GRIB files in {member_dir}" + ) + datasets.append(load_forecast_data(member_dir, reftime, steps, params)) + elevation = datasets[0].coords.get("elevation") + datasets = [ds.drop_vars("elevation", errors="ignore") for ds in datasets] + fcst = xr.concat( + datasets, + dim=pd.Index(range(1, members + 1), name="member"), + coords="minimal", + compat="override", + join="exact", + ) + if elevation is not None: + fcst = fcst.assign_coords(elevation=elevation) + return fcst + + def load_aligned_data(args, params: list[str]) -> tuple[xr.Dataset, xr.Dataset]: """Load forecast and truth for ``params``, with the forecast mapped to the truth - points and the truth selected at the forecast valid times.""" + points and the truth selected at the forecast valid times. + + With ``args.members`` set, ``args.forecast`` is the ``grib/`` directory of an + ensemble run and the forecast has a ``member`` dimension (load_ensemble_forecast). + """ # get baseline forecast data now = datetime.now() - fcst = load_forecast_data( - args.forecast, args.reftime, args.steps, params, member=args.member - ) + if getattr(args, "members", None): + fcst = load_ensemble_forecast( + args.forecast, args.reftime, args.steps, params, args.members + ) + else: + fcst = load_forecast_data( + args.forecast, args.reftime, args.steps, params, member=args.member + ) LOG.info( "Loaded forecast data in %s seconds: \n%s", @@ -130,6 +184,146 @@ def load_aligned_data(args, params: list[str]) -> tuple[xr.Dataset, xr.Dataset]: return fcst, truth +def member_mean_statistics( + fcst: xr.Dataset, + truth: xr.Dataset, + regions: list[dict], + holdout_stations: list[str] | None = None, + params: list[str] | None = None, + num_workers: int | None = None, +) -> xr.Dataset: + """Forecast statistics ``{p}.{mean,var,min,max}`` of an ensemble, averaged over + the members (without a ``source`` dim). + + Each member is masked as ``verify`` masks the forecast (regions, station groups, + NaN where too many forecast values are missing; a point is missing if any member + is). The ensemble mean is smoother than any member, so its own variance and + extremes would not describe the forecast fields. + + Memory: the masked members have region x station_group x points values per member + and step. Each member is its own dask chunk and the parameters are computed one + at a time with *num_workers* threads, so a chunk is no larger than in ``verify``. + """ + num_workers = num_workers or default_num_workers() + dim = ["x", "y"] if "x" in fcst.dims and "y" in fcst.dims else ["values"] + fcst, truth = xr.align(fcst, truth, join="inner", copy=False) + masks, station_masks = _region_and_station_masks(truth, regions, holdout_stations) + stats = [] + for param in params or list(fcst.data_vars): + if param not in truth.data_vars: + continue + members = fcst[param] + members = members.where(members.notnull().all("member")) + too_many_missing = _too_many_missing( + members.isnull().any("member"), + _apply_masks(truth[param], masks, station_masks), + dim, + 0.0, + ) + if members.chunks: + members = members.chunk({"member": 1}) + member_stats = _compute_statistics( + _apply_masks(members, masks, station_masks), + prefix=param + ".", + dim=dim, + circular=param in CIRCULAR_PARAMS, + ) + member_stats = member_stats.assign(too_many_missing=too_many_missing).compute( + num_workers=num_workers, scheduler="threads" + ) + too_many_missing = member_stats["too_many_missing"] + stats.append( + member_stats.drop_vars("too_many_missing") + .mean("member") + .squeeze("source", drop=True) + .where(~too_many_missing) + ) + return xr.merge(stats) + + +def _nan_categorical_scores(results: xr.Dataset, params, source_id: str) -> xr.Dataset: + """NaN for the scores of the categorical params (class codes) under *source_id*. + + The ensemble mean of class codes is no class (rain 1 and snow 5 average to 3, + freezing rain), so the contingency tables of the ensemble mean are meaningless. + The forecast statistics (averaged over the members) and the truth are kept. + """ + statistics = ("mean", "var", "min", "max") + is_source = results["source"] == source_id + for param in CATEGORICAL_PARAMS.intersection(params): + for name in list(results.data_vars): + if not name.startswith(param + "."): + continue + if name.split(".", 1)[1] in statistics: + continue + current = results[name] + if "source" not in current.dims: + continue + results[name] = ( + current.where(~is_source) + .transpose(*current.dims) + .assign_attrs(current.attrs) + ) + return results + + +def verify_ensemble_run( + fcst: xr.Dataset, truth: xr.Dataset, source_id: str, truth_source_id: str, **kwargs +) -> xr.Dataset: + """Scores of an ensemble run (``fcst`` with a ``member`` dimension). + + The deterministic scores are those of the ensemble mean (``verify``), under the + run's own ``source``; the ensemble scores (``verify_ensemble``: CRPS, ENS_VAR, + MEMBER_MSE, rank_hist, prob_table) are merged next to them. A point where any + member is missing is missing in the ensemble mean; the ensemble mean of a circular + parameter (wind direction) is the direction of the mean unit vector. The run's forecast statistics + are averaged over the members (``member_mean_statistics``), except for circular + parameters (wind direction), which keep those of the ensemble mean. + Categorical parameters (precipitation type) have no ensemble mean: their + deterministic scores (contingency tables) are NaN; their ``prob_table`` is the + ensemble's categorical score. + """ + n_members = fcst.sizes["member"] + LOG.info("Verifying an ensemble of %d members", n_members) + ens_mean = fcst.mean("member", skipna=False) + for param in CIRCULAR_PARAMS.intersection(fcst.data_vars): + # wind direction: direction of the mean unit vector (350 and 10 deg -> 0 deg) + rad = np.deg2rad(fcst[param]) + ens_mean[param] = np.mod( + np.degrees( + np.arctan2( + np.sin(rad).mean("member", skipna=False), + np.cos(rad).mean("member", skipna=False), + ) + ), + 360.0, + ).assign_attrs(fcst[param].attrs) + results = verify(ens_mean, truth, source_id, truth_source_id, **kwargs) + results = _nan_categorical_scores(results, fcst.data_vars, source_id) + + stat_params = [p for p in fcst.data_vars if p not in CIRCULAR_PARAMS] + member_stats = member_mean_statistics( + fcst, + truth, + kwargs["regions"], + kwargs.get("holdout_stations"), + params=stat_params, + num_workers=kwargs.get("num_workers"), + ) + for name, value in member_stats.data_vars.items(): + current = results[name] + results[name] = ( + xr.where(current["source"] == source_id, value, current) + .transpose(*current.dims) + .assign_attrs(current.attrs) + ) + + ens_scores = verify_ensemble(fcst, truth, source_id, truth_source_id, **kwargs) + out = xr.merge([results, ens_scores], join="outer", compat="no_conflicts") + out.attrs = {**results.attrs, "n_members": n_members} + return out + + def main(args: ScriptConfig): """Main function to verify baseline forecast data.""" @@ -154,15 +348,19 @@ def main(args: ScriptConfig): # compute metrics and statistics now = datetime.now() - results = verify( - fcst, - truth, - args.source_id, - args.truth_source_id, + verify_kwargs = dict( regions=args.regions, threshold_dict=args.threshold_dict, holdout_stations=holdout_stations, ) + if "member" in fcst.dims: + results = verify_ensemble_run( + fcst, truth, args.source_id, args.truth_source_id, **verify_kwargs + ) + else: + results = verify( + fcst, truth, args.source_id, args.truth_source_id, **verify_kwargs + ) LOG.info( "Computed verification metrics in %s seconds", (datetime.now() - now).total_seconds(), @@ -268,6 +466,17 @@ def main(args: ScriptConfig): default="000", help="Ensemble member to load: '000' for control, 'median' for the pre-computed median, 'mean' to average all members, or any 3-digit member ID.", ) + parser.add_argument( + "--members", + type=int, + default=None, + help=( + "Number of members of an ensemble run. --forecast is then the run's grib/ " + "directory holding one sub-directory per member (001, 002, ...); the " + "ensemble mean is scored as the run's deterministic forecast and the " + "ensemble scores are added." + ), + ) parser.add_argument( "--valid_masks", type=Path, diff --git a/workflow/scripts/verification_plot_metrics.py b/workflow/scripts/verification_plot_metrics.py index dd0b6e54..de86a7ba 100644 --- a/workflow/scripts/verification_plot_metrics.py +++ b/workflow/scripts/verification_plot_metrics.py @@ -7,6 +7,12 @@ import matplotlib.pyplot as plt import pandas as pd import xarray as xr +from verification.ensemble_report import ( + add_deterministic_crps, + ensemble_sources, + is_score_var, + CRPS_NOTE, +) from verification import decode_metric LOG = logging.getLogger(__name__) @@ -111,9 +117,13 @@ def main(args: Namespace) -> None: dfs = _select_best_sources(dfs) # 3) Concatenate by source; outer join to keep the union of lead_times ds = _concat_sources(dfs) + ds, crps_filled = add_deterministic_crps(ds) # extract only non-spatial variables to pd.DataFrame - nonspatial_vars = [d for d in ds.data_vars if "spatial" not in d] + # (rank histograms and probability tables are not scores: drawn separately) + nonspatial_vars = [ + d for d in ds.data_vars if "spatial" not in d and is_score_var(ds[d]) + ] all_df = ( ds[nonspatial_vars].to_array("stack").to_dataframe(name="value").reset_index() ) @@ -156,6 +166,10 @@ def _subset_df(df): title = f"{metric} - {param} - {region}" title += f"- {season} - {init_hour}" if args.stratify else "" + if crps_filled and metric == "CRPS": + ax.annotate( + CRPS_NOTE, (0.0, -0.3), xycoords="axes fraction", fontsize=9, va="top" + ) for source, df in sub_df.groupby("source"): display_label = args.label_map.get(source, source) df.plot( @@ -181,6 +195,101 @@ def _subset_df(df): plt.savefig(args.output_dir / fn, bbox_inches="tight") plt.close(fig) + plot_ensemble_diagnostics(ds, args) + + +def _counts_for_plot( + da: xr.DataArray, source, region, keep: list[str], bins: str +) -> xr.DataArray: + """Counts of one source at the default stratum, summed over everything but *keep*. + + The participants are concatenated with an outer join, so an ensemble of N members + has NaN in the *bins* (``rank`` or ``n_members_yes``) above N when another ensemble + has more members: those bins are dropped, and the counts keep the source's N + 1 bins. + """ + da = da.sel(source=source) + for dim, val in (("region", region), ("season", "all"), ("init_hour", -999)): + if dim in da.dims: + da = da.sel({dim: val}, drop=True) + da = da.sum([d for d in da.dims if d not in keep], skipna=True, min_count=1) + return da.dropna(bins, how="all") + + +def plot_ensemble_diagnostics(ds: xr.Dataset, args: Namespace) -> None: + """Rank histograms and reliability diagrams, for ensemble participants only. + + Counts are summed over lead times, at the default region, all seasons and init hours. + Each ensemble has its own number of members N (N + 1 ranks, probabilities k/N). + """ + ens_sources = ensemble_sources(ds) + if not ens_sources: + return + region = ds["region"].values[0] + args.output_dir.mkdir(parents=True, exist_ok=True) + label = lambda s: args.label_map.get(s, s) # noqa: E731 + + for name in [v for v in ds.data_vars if v.endswith(".rank_hist")]: + param = name.rsplit(".", 1)[0] + fig, ax = plt.subplots(figsize=(10, 6)) + n_plotted = 0 + for source in ens_sources: + counts = _counts_for_plot(ds[name], source, region, ["rank"], "rank") + total = float(counts.sum()) + if not total > 0: + continue + n_ranks = counts.sizes["rank"] + (line,) = ax.plot( + counts["rank"].values, + counts.fillna(0).values / total, + marker="o", + label=f"{label(source)} (N={n_ranks - 1})", + ) + # flat histogram of this ensemble: 1 / (N + 1) + ax.axhline(1 / n_ranks, color=line.get_color(), ls="--", lw=1) + n_plotted += 1 + if n_plotted: + ax.set( + title=f"Rank histogram - {param} - {region}", + xlabel="Rank of observation among members", + ylabel="Relative frequency", + ) + ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.15), ncol=n_plotted) + fig.savefig(args.output_dir / f"rank_hist_{param}.png", bbox_inches="tight") + plt.close(fig) + + for name in [v for v in ds.data_vars if ".prob_table_" in v]: + param, rest = name.split(".prob_table_", 1) + fig, (ax, axc) = plt.subplots( + 2, 1, figsize=(10, 8), sharex=True, height_ratios=[3, 1] + ) + n_plotted = 0 + for source in ens_sources: + tab = _counts_for_plot( + ds[name], source, region, ["n_members_yes", "obs"], "n_members_yes" + ).fillna(0) + n_bin = tab.sum("obs") + if not bool(n_bin.sum() > 0): + continue + n_mem = tab.sizes["n_members_yes"] - 1 + freq = tab.sel(obs=1) / n_bin.where(n_bin > 0) + prob = tab["n_members_yes"].values / n_mem + ax.plot(prob, freq.values, marker="o", label=f"{label(source)} (N={n_mem})") + axc.plot(prob, n_bin.values, marker="o") + n_plotted += 1 + if n_plotted: + ax.plot([0, 1], [0, 1], color="grey", ls="--", lw=1) + ax.set( + title=f"Reliability - {param} {decode_metric(rest)} - {region}", + ylabel="Observed frequency", + ) + axc.set(xlabel="Forecast probability (k/N)", ylabel="Count", yscale="log") + ax.legend(loc="upper left") + fig.savefig( + args.output_dir / f"reliability_{param}_{rest}.png", + bbox_inches="tight", + ) + plt.close(fig) + if __name__ == "__main__": parser = ArgumentParser(description="Verify results from KENDA-1 data.") diff --git a/workflow/scripts/verification_valid_mask.py b/workflow/scripts/verification_valid_mask.py index ced13ece..ebd32c40 100644 --- a/workflow/scripts/verification_valid_mask.py +++ b/workflow/scripts/verification_valid_mask.py @@ -26,6 +26,10 @@ def main(args): # Loaded exactly as for verification, so the mask matches the scored points. fcst, _ = load_aligned_data(args, args.params) + if "member" in fcst.dims: + # An ensemble has a value where every member has one: the ensemble mean + # (scored by verification_metrics.py) is missing where any member is. + fcst = fcst.mean("member", skipna=False) mask = forecast_valid_mask(fcst, args.params) LOG.info("Valid mask for %s: \n%s", args.params, mask) args.output.parent.mkdir(parents=True, exist_ok=True) @@ -66,5 +70,14 @@ def main(args): default="000", help="Ensemble member to load (see verification_metrics.py).", ) + parser.add_argument( + "--members", + type=int, + default=None, + help=( + "Number of members of an ensemble run (--forecast is then its grib/ " + "directory): the mask is where every member has a value." + ), + ) parser.add_argument("--output", type=Path, required=True, help="Output NetCDF.") main(parser.parse_args())