#!/usr/bin/env python """Stage 14 — EHD_ens: multi-model ensemble / stacking of the v2 zoo (leakage-safe). Combines the 7 zoo members (at cascade cutoff 1.0) into two ensemble forecasts: ens_mean unweighted mean of member probabilities ens_stack LOYO logistic stacker (meta-model fit on the other 19 years, applied to the held- out year) -- see src/ehd/ensemble.py for the leakage argument. Writes: outputs/loyo_pred_v2_ens/ens_pooled.npz per-sample block_id/cell/region/y + p_ensmean/p_ensstack + flag_ensmean/flag_ensstack results/tables/v2_scores_with_ensemble.csv zoo members (cutoff 1.0) + both ensembles results/figures/v2_ensemble_skill_by_lead.png SEDI/BSS by lead: members vs ensembles """ from __future__ import annotations import os, sys, glob import numpy as np import pandas as pd sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) from ehd.config import load_config, lead_blocks from ehd import verify as vf from ehd import ensemble as ens HERE = os.path.dirname(os.path.abspath(__file__)) ROOT = os.path.dirname(HERE) TABLES = os.path.join(ROOT, "results", "tables") FIGS = os.path.join(ROOT, "results", "figures") OUTDIR = os.path.join(ROOT, "outputs", "loyo_pred_v2_ens") def _scores(y, p, flag): m = np.isfinite(y) & np.isfinite(p) yv, pv, fv = y[m], p[m], flag[m].astype(bool) base = float(np.mean(yv)) brier = float(np.mean((pv - yv) ** 2)) bss = 1.0 - brier / (base * (1 - base)) if 0 < base < 1 else np.nan a = int(np.sum(fv & (yv == 1))); b = int(np.sum(fv & (yv == 0))) c = int(np.sum(~fv & (yv == 1))); d = int(np.sum(~fv & (yv == 0))) bs = vf.binary_scores_from_counts(a, b, c, d) return dict(n=int(m.sum()), base_rate=round(base, 3), brier=round(brier, 4), bss_vs_clim=round(bss, 3), auc=round(vf.auc(yv, pv), 3), POD=round(bs["POD"], 3), POFD=round(bs["POFD"], 3), FAR=round(bs["FAR"], 3), CSI=round(bs["CSI"], 3), ETS=round(bs["ETS"], 3), HSS=round(bs["HSS"], 3), Peirce=round(bs["Peirce"], 3), EDI=round(bs["EDI"], 3), SEDI=round(bs["SEDI"], 3), freq_bias=round(bs["freq_bias"], 3)) def main(): cfg = load_config() blocks = lead_blocks(cfg) block_names = list(blocks.keys()) block_ids = {b: k + 1 for k, b in enumerate(block_names)} region_ids = cfg["regions"]["ids"] years = list(cfg["cross_validation"]["years"]) paths = sorted(glob.glob(os.path.join(ROOT, "outputs", "loyo_pred_v2", "pred_v2_*.npz"))) files = {int(os.path.basename(p).split("_")[-1][:-4]): p for p in paths} cnn_files = {int(os.path.basename(p).split("_")[-1][:-4]): p for p in glob.glob(os.path.join(ROOT, "outputs", "loyo_pred_v2_cnn", "pred_v2cnn_*.npz"))} years = [y for y in years if y in files] data, used = ens.load_by_year(files, block_ids, extra_files=cnn_files, extra_members=ens.EXTRA_MEMBERS) print(f"[ens] {len(years)} folds, members={used}") result, decthr = ens.build_ensembles(data, block_names, region_ids, years) # flat per-sample arrays for downstream maps os.makedirs(OUTDIR, exist_ok=True) flat = {k: [] for k in ("block_id", "cell", "region", "y", "p_ensmean", "p_ensstack", "flag_ensmean", "flag_ensstack")} for b in block_names: r = result[b]; n = r["y"].size flat["block_id"].append(np.full(n, block_ids[b])) flat["cell"].append(r["cell"]); flat["region"].append(r["region"]); flat["y"].append(r["y"]) flat["p_ensmean"].append(r["p_mean"]); flat["p_ensstack"].append(r["p_stack"]) flat["flag_ensmean"].append(r["flag_mean"]); flat["flag_ensstack"].append(r["flag_stack"]) flat = {k: np.concatenate(v) for k, v in flat.items()} dt = np.array(decthr, dtype=[("tag", "U20"), ("block", "U4"), ("region", "i4"), ("thr", "f4")]) np.savez_compressed(os.path.join(OUTDIR, "ens_pooled.npz"), decthr=dt, **flat) print(f"[ens] wrote {os.path.join('outputs','loyo_pred_v2_ens','ens_pooled.npz')}") # ---- scores: zoo members (cutoff 1.0) + both ensembles ---- zoo = pd.read_csv(os.path.join(TABLES, "v2_cv_scores_by_model_block.csv")) zoo1 = zoo[zoo["cutoff"] == 1.0].copy() zoo1["model"] = zoo1["family"] keep = ["model", "block", "n", "base_rate", "brier", "bss_vs_clim", "auc", "POD", "POFD", "FAR", "CSI", "ETS", "HSS", "Peirce", "EDI", "SEDI", "freq_bias"] rows = [] for b in block_names: r = result[b] for name, p, fl in (("ens_mean", r["p_mean"], r["flag_mean"]), ("ens_stack", r["p_stack"], r["flag_stack"])): rows.append(dict(model=name, block=b, **_scores(r["y"], p, fl))) ens_df = pd.DataFrame(rows) allm = pd.concat([zoo1[keep], ens_df[keep]], ignore_index=True) allm["block_ord"] = allm["block"].map({b: i for i, b in enumerate(block_names)}) allm = allm.sort_values(["block_ord", "model"]).drop(columns="block_ord").reset_index(drop=True) allm.to_csv(os.path.join(TABLES, "v2_scores_with_ensemble.csv"), index=False) print(f"[ens] wrote v2_scores_with_ensemble.csv ({len(allm)} rows)") print("\n[ens] ensemble vs best member per block (SEDI / BSS_clim / AUC):") for b in block_names: sub = allm[allm.block == b] bestmem = sub[~sub.model.str.startswith("ens_")].sort_values("SEDI").iloc[-1] em = sub[sub.model == "ens_mean"].iloc[0]; es = sub[sub.model == "ens_stack"].iloc[0] print(f" {b}: best_member={bestmem.model:9s} SEDI={bestmem.SEDI:.3f} | " f"ens_mean SEDI={em.SEDI:.3f} BSS={em.bss_vs_clim:.3f} | " f"ens_stack SEDI={es.SEDI:.3f} BSS={es.bss_vs_clim:.3f} AUC={es.auc:.3f}") _figure(allm, block_names) def _figure(allm, block_names): import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt x = np.arange(len(block_names)) fig, ax = plt.subplots(1, 2, figsize=(12, 4.5)) members = [m for m in allm.model.unique() if not m.startswith("ens_")] for m in members: s = allm[allm.model == m].set_index("block").reindex(block_names) ax[0].plot(x, s["SEDI"], color="0.7", lw=1, marker=".", zorder=1) ax[1].plot(x, s["bss_vs_clim"], color="0.7", lw=1, marker=".", zorder=1) for m, col in (("ens_mean", "tab:blue"), ("ens_stack", "tab:red")): s = allm[allm.model == m].set_index("block").reindex(block_names) ax[0].plot(x, s["SEDI"], color=col, lw=2.2, marker="o", label=m, zorder=3) ax[1].plot(x, s["bss_vs_clim"], color=col, lw=2.2, marker="o", label=m, zorder=3) ax[0].plot([], [], color="0.7", label="zoo members") for a, t in zip(ax, ["SEDI", "BSS vs climatology"]): a.set_xticks(x); a.set_xticklabels(block_names); a.axhline(0, color="k", lw=.6) a.set_title(f"EHD_ens vs members — {t}"); a.set_xlabel("lead block"); a.grid(alpha=.3) a.legend(fontsize=8) os.makedirs(FIGS, exist_ok=True) fig.tight_layout(); fig.savefig(os.path.join(FIGS, "v2_ensemble_skill_by_lead.png"), dpi=130) print("[ens] wrote figure -> results/figures/v2_ensemble_skill_by_lead.png") if __name__ == "__main__": main()