#!/usr/bin/env python """Stage 24 — items 6 & 7: per-member regional event, tau(Rj,Wk), and regional probability. item 6 a(m,j,k) = AREA-weighted fraction of region Rj where member m votes EHD (weights = fractional cell-polygon overlap with Vietnam7.shp x cos(lat); ehd.area_weights). yEHDE(m,j,k) = 1 if a(m,j,k) > tau(Rj,Wk). The observed regional event O(j,k) = 1 if the AREA-weighted observed EHD coverage a_obs >= 0.5 (region-majority; stated, interpretable). tau fit INSIDE the LOYO folds by maximising PSS of the family-balanced mean coverage against O; the frequency-bias-calibrated tau (forecast event freq = observed freq) is reported as the alternative. Delivered as a 7x6 table with fold mean +/- std. item 7 P(yEHDE|Rj,Wk) = family-balanced fraction of members with yEHDE(m,j,k)=1, isotonic- recalibrated in-CV (route A, regional-first). Compared against route B (aggregate the grid-point probabilities over the region, then no threshold) — the two are not equivalent. Outputs: results/tables/v3_tau_by_region_week.csv tau_pss / tau_bias mean+/-std over folds, base rate results/tables/v3_regional_event_scores.csv 7x6 contingency scores for the regional event results/tables/v3_regional_prob_compare.csv route A vs route B Brier/BSS per region-week results/figures/v3_tau_heatmap.png , v3_regional_perf.png """ from __future__ import annotations import os, sys, logging import numpy as np import pandas as pd import xarray as xr from sklearn.isotonic import IsotonicRegression sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) from ehd.config import load_config, lead_blocks from ehd import members_v3 as mv, verify as vf, area_weights as aw logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") log = logging.getLogger("stage24") NCELL = 541 COVERAGE_EVENT = 0.5 # observed regional event = area-weighted obs EHD coverage >= this def coverage_arrays(data, members, b, w_fb, wcell, region_of_cell, region_ids, years): """Per year: reshape (ninit,ncell); return dict region-> dict(a_bar[list over yr of ninit], O[...], amemb[list of ninit x M], gridfrac[list of ninit x ncell]).""" out = {r: dict(a_bar=[], O=[], amemb=[], y=[]) for r in region_ids} grid = {"gridfrac": [], "yr": []} for y in years: r = data[y][b] n = r["y"].size; ninit = n // NCELL V = r["V"].reshape(ninit, NCELL, -1) # (init, cell, member) yv = r["y"].reshape(ninit, NCELL) cell = r["cell"].reshape(ninit, NCELL)[0] # cell ids 0..NCELL-1 (same each init) reg = region_of_cell[cell] wf = wcell[cell] gridfrac = V.astype(float) @ w_fb # (init, cell) family-balanced vote frac grid["gridfrac"].append((gridfrac, yv, reg, wf)) for rj in region_ids: cm = reg == rj wj = wf[cm]; sw = wj.sum() a_obs = (yv[:, cm] * wj).sum(1) / sw # (init,) a_memb = (V[:, cm, :] * wj[None, :, None]).sum(1) / sw # (init, member) a_bar = a_memb @ w_fb # family-balanced mean coverage O = (a_obs >= COVERAGE_EVENT).astype(float) out[rj]["a_bar"].append(a_bar); out[rj]["O"].append(O) out[rj]["amemb"].append(a_memb); out[rj]["y"].append(np.full(a_bar.size, y)) return out, grid def fit_tau(a_bar_tr, O_tr, kind, base_rate_obs): grid = np.linspace(0.02, 0.98, 49) if kind == "pss": best, bt = -np.inf, 0.5 for t in grid: s = vf.binary_scores(O_tr, (a_bar_tr >= t).astype(float), 0.5)["Peirce"] if np.isfinite(s) and s > best: best, bt = s, t return bt # freq-bias: tau so forecast event freq == observed freq diffs = [abs((a_bar_tr >= t).mean() - base_rate_obs) for t in grid] return float(grid[int(np.argmin(diffs))]) def main(): cfg = load_config() blocks = lead_blocks(cfg) region_ids = cfg["regions"]["ids"] tdir = os.path.join(cfg["paths"]["results"], "tables") fdir = os.path.join(cfg["paths"]["results"], "figures") log.info("loading v3 member pool ...") data, members = mv.load_pool(cfg, blocks, include_v3=True) mask = mv.active_mask(cfg, members, blocks) years = sorted(data) ref = xr.open_dataarray(os.path.join(cfg["paths"]["outputs"], "forecast_daily_tx", f"fc_daily_tx_{years[0]}.nc")) reg = xr.open_dataset(cfg["paths"]["mask"])["reg"].reindex(lat=ref.lat, lon=ref.lon, method="nearest") iy, ix = np.where(np.isfinite(reg.values)) region_of_cell = reg.values[iy, ix].astype(int) wcell, rule = aw.cell_area_weights(ref.lat.values, ref.lon.values, region_of_cell, iy, ix, os.path.join(cfg["paths"]["shapefile_dir"], "Vietnam7.shp")) log.info("area-weight rule: %s", rule) tau_rows, ev_rows, cmp_rows = [], [], [] pairs = {} # (region,week) -> dict(O, P, ev) pooled OOS, for stage 21 for b in blocks: w_fb, _ = mv.family_balanced_weights(members, mask[b]) cov, grid = coverage_arrays(data, members, b, w_fb, wcell, region_of_cell, region_ids, years) for rj in region_ids: a_bar_y = cov[rj]["a_bar"]; O_y = cov[rj]["O"] # in-fold tau (LOYO over years) + pooled OOS regional event forecast/prob taus_pss, taus_bias = [], [] P_all, O_all, evflag_all = [], [], [] for i, y in enumerate(years): tr = [j for j in range(len(years)) if j != i] a_tr = np.concatenate([a_bar_y[j] for j in tr]); O_tr = np.concatenate([O_y[j] for j in tr]) br = float(O_tr.mean()) tp = fit_tau(a_tr, O_tr, "pss", br); tb = fit_tau(a_tr, O_tr, "bias", br) taus_pss.append(tp); taus_bias.append(tb) # regional probability route A: fraction of members with coverage > tau_pss amemb_te = cov[rj]["amemb"][i] # (init, member) P_te = (amemb_te > tp).astype(float) @ w_fb # family-balanced member fraction # recalibrate P in-CV on training years' route-A P Ptr = np.concatenate([(cov[rj]["amemb"][j] > taus_pss[-1]).astype(float) @ w_fb for j in tr]) m = np.isfinite(Ptr) & np.isfinite(O_tr) iso = IsotonicRegression(out_of_bounds="clip").fit(Ptr[m], O_tr[m].astype(int)) P_all.append(iso.transform(P_te)); O_all.append(O_y[i]) evflag_all.append((a_bar_y[i] > tp).astype(int)) # deterministic regional event P = np.concatenate(P_all); O = np.concatenate(O_all); ev = np.concatenate(evflag_all) yr_all = np.concatenate([np.full(O_y[i].size, years[i]) for i in range(len(years))]) pairs[f"{rj}_{b}_O"] = O.astype(np.float32); pairs[f"{rj}_{b}_P"] = P.astype(np.float32) pairs[f"{rj}_{b}_ev"] = ev.astype(np.int8); pairs[f"{rj}_{b}_yr"] = yr_all.astype(np.int16) a, bb, c, d = (int(((ev == 1) & (O == 1)).sum()), int(((ev == 1) & (O == 0)).sum()), int(((ev == 0) & (O == 1)).sum()), int(((ev == 0) & (O == 0)).sum())) s = vf.binary_scores_from_counts(a, bb, c, d) n_pos = int(O.sum()) tau_rows.append(dict(region=rj, week=b, tau_pss_mean=round(np.mean(taus_pss), 3), tau_pss_std=round(np.std(taus_pss), 3), tau_bias_mean=round(np.mean(taus_bias), 3), tau_bias_std=round(np.std(taus_bias), 3), base_rate_O=round(float(O.mean()), 3), n_events=n_pos)) ev_rows.append(dict(region=rj, week=b, base_rate_O=round(float(O.mean()), 3), n_events=n_pos, POD=round(s["POD"], 3), POFD=round(s["POFD"], 3), FAR=round(s["FAR"], 3), CSI=round(s["CSI"], 3), PSS=round(s["Peirce"], 3), SEDI=round(s["SEDI"], 3), HSS=round(s["HSS"], 3), freq_bias=round(s["freq_bias"], 3), low_base_rate=(n_pos < 30))) # route B (grid-first): area-weighted mean of the grid vote fraction over region j, # per init, recalibrated in-CV against O. Compared to route A (regional-first) by Brier. regc = grid["gridfrac"][0][2]; wf = grid["gridfrac"][0][3] cm = regc == rj; wj = wf[cm]; sw = wj.sum() B_y = [((g[0][:, cm] * wj).sum(1) / sw) for g in grid["gridfrac"]] # per year (init,) B_all = [] for i in range(len(years)): tr = [j for j in range(len(years)) if j != i] Btr = np.concatenate([B_y[j] for j in tr]); Otr = np.concatenate([O_y[j] for j in tr]) m = np.isfinite(Btr) & np.isfinite(Otr) iso = IsotonicRegression(out_of_bounds="clip").fit(Btr[m], Otr[m].astype(int)) B_all.append(iso.transform(B_y[i])) PB = np.concatenate(B_all) clim = np.full_like(O, O.mean()) cmp_rows.append(dict(region=rj, week=b, base_rate_O=round(float(O.mean()), 3), routeA_brier=round(vf.brier(O, P), 4), routeA_bss=round(vf.brier_skill_score(O, P, clim), 4), routeB_brier=round(vf.brier(O, PB), 4), routeB_bss=round(vf.brier_skill_score(O, PB, clim), 4))) log.info("%s done", b) pd.DataFrame(tau_rows).to_csv(os.path.join(tdir, "v3_tau_by_region_week.csv"), index=False) pd.DataFrame(ev_rows).to_csv(os.path.join(tdir, "v3_regional_event_scores.csv"), index=False) pd.DataFrame(cmp_rows).to_csv(os.path.join(tdir, "v3_regional_prob_compare.csv"), index=False) import matplotlib; matplotlib.use("Agg"); import matplotlib.pyplot as plt tt = pd.DataFrame(tau_rows).pivot(index="region", columns="week", values="tau_pss_mean").reindex(region_ids) fig, ax = plt.subplots(figsize=(8, 5)) im = ax.imshow(tt.values, aspect="auto", cmap="viridis", vmin=0, vmax=1) ax.set_xticks(range(len(blocks))); ax.set_xticklabels(blocks) ax.set_yticks(range(len(region_ids))); ax.set_yticklabels([f"R{r}" for r in region_ids]) for i in range(len(region_ids)): for j in range(len(blocks)): ax.text(j, i, f"{tt.values[i,j]:.2f}", ha="center", va="center", color="w", fontsize=8) fig.colorbar(im, label="tau_pss (in-fold mean)"); ax.set_title("Stage-24 tau(Rj,Wk) PSS-optimal") fig.tight_layout(); fig.savefig(os.path.join(fdir, "v3_tau_heatmap.png"), dpi=130) ev = pd.DataFrame(ev_rows) print(ev.to_string(index=False)) log.info("regional event: mean PSS W1=%.2f W6=%.2f", ev[ev.week == "W1"].PSS.mean(), ev[ev.week == "W6"].PSS.mean()) if __name__ == "__main__": main()