#!/usr/bin/env python3 """Stage 1, section 4: verification of the raw EC S2S forecast of Tx / EHD / EHDE. Everything is accumulated per land point and per forecast week W1..W6 by streaming over the 2080 initialisations, so memory stays flat. Section 4.2.1 Tx errors ME / MAE / RMSE, three flavours: daily, weekly-mean, weekly-max Section 4.2.2 EHD day counts ME / MAE / RMSE of N1 (model) vs N2 (obs) Section 4.2.3 EHDE events POD, POFD, FAR, FBI, CSI, ETS and a ROC/AUC built by sweeping thresholds on Sev_mod (the EC control run is deterministic, so a probabilistic reliability diagram is not defined here; it is produced for the AI models in stage 2 instead). """ from __future__ import annotations import argparse import os import sys import time import numpy as np from joblib import Parallel, delayed from netCDF4 import Dataset sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) import ehd_common as C # Sev_mod bin edges used to trace the ROC curve. SEV_LO, SEV_HI, SEV_STEP = -25.0, 15.0, 0.25 SEV_EDGES = np.arange(SEV_LO, SEV_HI + SEV_STEP, SEV_STEP) NBIN = SEV_EDGES.size - 1 NW = len(C.WEEKS) # np.trapz was removed in numpy 2.0 in favour of np.trapezoid. _trapz = getattr(np, "trapezoid", None) or np.trapz def _new_acc(npt): z = lambda *s: np.zeros(s, dtype="f8") return { "n_daily": z(NW, npt), "s_daily": z(NW, npt), "a_daily": z(NW, npt), "q_daily": z(NW, npt), "n_wmean": z(NW, npt), "s_wmean": z(NW, npt), "a_wmean": z(NW, npt), "q_wmean": z(NW, npt), "n_wmax": z(NW, npt), "s_wmax": z(NW, npt), "a_wmax": z(NW, npt), "q_wmax": z(NW, npt), "n_cnt": z(NW, npt), "s_cnt": z(NW, npt), "a_cnt": z(NW, npt), "q_cnt": z(NW, npt), "hit": z(NW, npt), "fal": z(NW, npt), "mis": z(NW, npt), "cor": z(NW, npt), "roc_y": np.zeros((NW, npt, NBIN), dtype="i4"), "roc_n": np.zeros((NW, npt, NBIN), dtype="i4"), } def _add(acc, other): for k in acc: acc[k] += other[k] return acc def _process_year(year, grid, thr_obs, thr_mod): """Accumulators for one hindcast year.""" acc = _new_acc(grid.npt) pt_offset = (np.arange(grid.npt) * NBIN)[None, :] # flat index base for bincount for init in C.init_dates(year): tx_m = C.read_ec_points(init, "mx2t6", grid) # (46, npt) tx_o = C.read_obs_points(init, "tasmax", grid) days = C.valid_days(init) doy = np.array([C.doy_noleap(d) for d in days]) - 1 th_o = thr_obs[doy] # (46, npt) th_m = thr_mod[np.arange(C.NLT), doy] # (46, npt) sev_m = tx_m - th_m sev_o = tx_o - th_o ehd_m = sev_m >= 0.0 ehd_o = sev_o >= 0.0 valid = np.isfinite(tx_m) & np.isfinite(tx_o) & np.isfinite(th_o) & np.isfinite(th_m) for w, (_, idx) in enumerate(C.WEEKS): ok = valid[idx] # 4.2.1a daily Tx error d = np.where(ok, tx_m[idx] - tx_o[idx], 0.0) acc["n_daily"][w] += ok.sum(axis=0) acc["s_daily"][w] += d.sum(axis=0) acc["a_daily"][w] += np.abs(d).sum(axis=0) acc["q_daily"][w] += (d * d).sum(axis=0) full = ok.all(axis=0) # week complete at this point # 4.2.1b weekly-mean error dm = np.where(full, tx_m[idx].mean(axis=0) - tx_o[idx].mean(axis=0), 0.0) acc["n_wmean"][w] += full acc["s_wmean"][w] += dm acc["a_wmean"][w] += np.abs(dm) acc["q_wmean"][w] += dm * dm # 4.2.1c weekly-max error dx = np.where(full, tx_m[idx].max(axis=0) - tx_o[idx].max(axis=0), 0.0) acc["n_wmax"][w] += full acc["s_wmax"][w] += dx acc["a_wmax"][w] += np.abs(dx) acc["q_wmax"][w] += dx * dx # 4.2.2 EHD day-count error dc = np.where(full, (ehd_m[idx] & ok).sum(axis=0) - (ehd_o[idx] & ok).sum(axis=0), 0.0) acc["n_cnt"][w] += full acc["s_cnt"][w] += dc acc["a_cnt"][w] += np.abs(dc) acc["q_cnt"][w] += dc * dc # 4.2.3 daily contingency table m, o = ehd_m[idx] & ok, ehd_o[idx] & ok acc["hit"][w] += (m & o).sum(axis=0) acc["fal"][w] += (m & ~o & ok).sum(axis=0) acc["mis"][w] += (~m & o & ok).sum(axis=0) acc["cor"][w] += (~m & ~o & ok).sum(axis=0) # ROC support: distribution of Sev_mod split by observed outcome. # Flatten (day, point) into one bincount per week to keep this cheap. b = np.clip(np.digitize(sev_m[idx], SEV_EDGES) - 1, 0, NBIN - 1) flat = pt_offset + b # (nday, npt) nflat = grid.npt * NBIN acc["roc_y"][w] += np.bincount(flat[o], minlength=nflat).reshape(grid.npt, NBIN) acc["roc_n"][w] += np.bincount(flat[ok & ~o], minlength=nflat).reshape(grid.npt, NBIN) return acc def _rate(num, den): out = np.full(num.shape, np.nan, dtype="f4") good = den > 0 out[good] = num[good] / den[good] return out def finalize(acc, grid): """Turn accumulators into the map fields, each (NW, npt).""" m = {} for tag, label in [("daily", "tx_%s_daily"), ("wmean", "tx_%s_weekly_mean"), ("wmax", "tx_%s_weekly_max"), ("cnt", "ehd_count_%s")]: n = acc["n_" + tag] m[label % "me"] = _rate(acc["s_" + tag], n) m[label % "mae"] = _rate(acc["a_" + tag], n) m[label % "rmse"] = np.sqrt(_rate(acc["q_" + tag], n)) a, b, c, d = acc["hit"], acc["fal"], acc["mis"], acc["cor"] tot = a + b + c + d m["pod"] = _rate(a, a + c) m["pofd"] = _rate(b, b + d) m["far"] = _rate(b, a + b) m["fbi"] = _rate(a + b, a + c) m["csi"] = _rate(a, a + b + c) ar = _rate((a + b) * (a + c), tot) m["ets"] = _rate(a - ar, (a + b + c - ar)) m["base_rate"] = _rate(a + c, tot) # ROC/AUC by sweeping the Sev_mod threshold (trapezoidal rule on the ROC curve). ry = acc["roc_y"][:, :, ::-1].cumsum(axis=2) # events with Sev >= edge rn = acc["roc_n"][:, :, ::-1].cumsum(axis=2) ny = ry[:, :, -1][:, :, None] nn = rn[:, :, -1][:, :, None] with np.errstate(invalid="ignore", divide="ignore"): tpr = np.where(ny > 0, ry / np.maximum(ny, 1), np.nan) fpr = np.where(nn > 0, rn / np.maximum(nn, 1), np.nan) zero = np.zeros(tpr.shape[:2] + (1,)) tpr = np.concatenate([zero, tpr], axis=2) fpr = np.concatenate([zero, fpr], axis=2) auc = _trapz(tpr, fpr, axis=2) auc[~np.isfinite(auc)] = np.nan m["auc"] = auc.astype("f4") return m def main(): ap = argparse.ArgumentParser() ap.add_argument("--jobs", type=int, default=20) ap.add_argument("--years", type=int, nargs="*", default=C.YEARS) ap.add_argument("--clidir", default=os.path.join(C.OUT, "Cli")) ap.add_argument("--tag", default="") ap.add_argument("--out", default=os.path.join(C.OUT, "ehd_verification.nc")) args = ap.parse_args() grid = C.Grid() t0 = time.time() with Dataset(os.path.join(args.clidir, "Qp_tx_obs%s.nc" % args.tag)) as ds: thr_obs = grid.to_points(np.ma.filled(ds["Tx_thres_obs"][:], np.nan).astype("f4")) with Dataset(os.path.join(args.clidir, "Qp_tx_mod%s.nc" % args.tag)) as ds: thr_mod = grid.to_points(np.ma.filled(ds["Tx_thres_mod"][:], np.nan).astype("f4")) print("[thr] obs %s mod %s" % (thr_obs.shape, thr_mod.shape), flush=True) parts = Parallel(n_jobs=min(args.jobs, len(args.years)), verbose=1)( delayed(_process_year)(y, grid, thr_obs, thr_mod) for y in args.years) acc = parts[0] for p in parts[1:]: _add(acc, p) del parts print("[acc] %d years in %.1fs" % (len(args.years), time.time() - t0), flush=True) fields = finalize(acc, grid) varlist = [] for name, arr in fields.items(): varlist.append((name, ("week", "lat", "lon"), grid.to_grid(arr), {})) C.write_nc(args.out, [("week", np.arange(1, NW + 1, dtype="i4")), ("lat", grid.lat), ("lon", grid.lon)], varlist, {"weeks": " ".join(C.WEEK_NAMES), "years": str(args.years), "note": "EC S2S control-run verification, land points only"}) print("[done] %.1fs -> %s" % (time.time() - t0, args.out), flush=True) if __name__ == "__main__": main()