#!/usr/bin/env python """Stage 23 — diversity-driven, lead-dependent stacking of the v3 zoo (+ optional DL members). The hybrid-design note (§7-§8) makes two points this stage implements: * the final product should be a CALIBRATED PROBABILITY from a lead-dependent ensemble, with the stacking weights learned ONLY inside cross-validation; * member selection should prize ERROR DIVERSITY, judged by an error-correlation matrix and a leave-one-member-out (LOMO) contribution — three mildly-weak, decorrelated members can beat three near-identical strong ones. Per lead block we build three leakage-safe combinations (meta fit LEAVE-ONE-YEAR-OUT, so the held-out year never trains its own weights): ens_mean unweighted mean of member probabilities (no fitting). ens_stack L2 logistic meta on all members (dense blend). ens_l1 L1 (sparse) logistic meta — the weights zero out redundant members, so the stack performs data-driven, leakage-safe diversity selection per block automatically. Because members and weights differ by block, the ensemble composition is lead-dependent by construction (a strong short-lead member can get weight 0 at long lead and vice-versa). Diagnostics (pooled out-of-sample, for the paper / member-selection story): * per-member BSS and SEDI per block; * mean absolute error-correlation of each member with the others (lower = more diverse); * LOMO delta-BSS: change in the L2 stack's BSS when the member is dropped (contribution). Outputs: results/tables/v3_stack_scores_by_block.csv ens_mean/stack/l1 skill per block results/tables/v3_stack_diversity_by_block.csv member x block: bss, sedi, err-corr, LOMO results/figures/v3_stack_by_lead.png stack vs best member by lead outputs/loyo_pred_v3stack/v3stack_pooled.npz pooled y/region/block + p_mean/p_stack/p_l1 Consumes outputs/loyo_pred_v3zoo/pred_v3zoo_*.npz. Optional --dl-dir merges DL member npz (cnn/resunet/transformer) whose (block_id, cell) ordering matches stage 21. """ from __future__ import annotations import os, sys, glob, argparse, warnings import numpy as np import pandas as pd warnings.filterwarnings("ignore") from sklearn.linear_model import LogisticRegression from sklearn.preprocessing import StandardScaler from sklearn.pipeline import make_pipeline 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 ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) TABLES = os.path.join(ROOT, "results", "tables") FIGS = os.path.join(ROOT, "results", "figures") def _meta(penalty): # liblinear is far faster than saga for this small-feature meta problem on ~1e6 rows and # supports both L1 and L2 for binary logistic. (penalty= is deprecation-warned in sklearn # 1.9 but still functional; warning filtered at module load.) clf = LogisticRegression(solver="liblinear", penalty=penalty, C=1.0, max_iter=2000) return make_pipeline(StandardScaler(), clf) def load_years(indir, pattern, dl_dir=None, dl_pattern=None): files = sorted(glob.glob(os.path.join(indir, pattern))) if not files: raise SystemExit(f"no {pattern} in {indir}") data = {} fams = None for path in files: yr = int(os.path.splitext(os.path.basename(path))[0].split("_")[-1]) d = np.load(path, allow_pickle=True) fcols = sorted(c for c in d.files if c.startswith("p_")) fams = [c[2:] for c in fcols] data[yr] = dict(block_id=d["block_id"].astype(int), region=d["region"].astype(int), cell=d["cell"].astype(int), y=d["y"].astype(float), P={c[2:]: d[c].astype(float) for c in fcols}) if dl_dir: dp = os.path.join(dl_dir, dl_pattern.replace("*", str(yr))) if os.path.exists(dp): dd = np.load(dp, allow_pickle=True) assert np.array_equal(dd["cell"].astype(int), data[yr]["cell"]), \ f"DL cell ordering mismatch {yr}" for c in sorted(x for x in dd.files if x.startswith("p_")): data[yr]["P"][c[2:]] = dd[c].astype(float) members = sorted(next(iter(data.values()))["P"].keys()) return data, members def block_matrix(data, yr, bid, members): m = data[yr]["block_id"] == bid M = np.column_stack([data[yr]["P"][mm][m] for mm in members]) return M, data[yr]["y"][m], data[yr]["region"][m], data[yr]["cell"][m] def _fit_region_thr(y, p, region, region_ids): g, _ = vf.best_threshold(y, p, "Peirce") out = {int(r): g for r in region_ids} for r in region_ids: mm = region == r if mm.sum() > 200 and np.nansum(y[mm]) > 10: out[int(r)], _ = vf.best_threshold(y[mm], p[mm], "Peirce") return out def main(): ap = argparse.ArgumentParser() ap.add_argument("--indir", default=None) ap.add_argument("--pattern", default="pred_v3zoo_*.npz") ap.add_argument("--dl-dir", default=None, help="dir of DL member npz to merge") ap.add_argument("--dl-pattern", default="pred_*.npz") args = ap.parse_args() cfg = load_config(); blocks = lead_blocks(cfg); block_names = list(blocks) block_ids = {b: k + 1 for k, b in enumerate(block_names)} region_ids = cfg["regions"]["ids"] indir = args.indir or os.path.join(cfg["paths"]["outputs"], "loyo_pred_v3zoo") data, members = load_years(indir, args.pattern, args.dl_dir, args.dl_pattern) years = sorted(data.keys()) print(f"[23] {len(years)} folds | {len(members)} members: {members}") score_rows, div_rows = [], [] pooled = {k: [] for k in ("block_id", "region", "cell", "y", "p_mean", "p_stack", "p_l1")} fig_data = {} for b in block_names: bid = block_ids[b] if not any((data[yr]["block_id"] == bid).any() for yr in years): continue # block absent (e.g. smoke-test subset) # LOYO combinations agg = {t: {"y": [], "region": [], "cell": [], "mean": [], "stack": [], "l1": [], "members": []} for t in ["all"]} for yr in years: Mte, yte, rte, cte = block_matrix(data, yr, bid, members) tr_years = [t for t in years if t != yr] Mtr = np.concatenate([block_matrix(data, t, bid, members)[0] for t in tr_years]) ytr = np.concatenate([block_matrix(data, t, bid, members)[1] for t in tr_years]) ok = np.isfinite(ytr) & np.all(np.isfinite(Mtr), axis=1) m2 = _meta("l2").fit(Mtr[ok], ytr[ok].astype(int)) m1 = _meta("l1").fit(Mtr[ok], ytr[ok].astype(int)) Mte0 = np.nan_to_num(Mte, nan=0.0) agg["all"]["y"].append(yte); agg["all"]["region"].append(rte) agg["all"]["cell"].append(cte) agg["all"]["mean"].append(np.nanmean(Mte, axis=1)) agg["all"]["stack"].append(m2.predict_proba(Mte0)[:, 1]) agg["all"]["l1"].append(m1.predict_proba(Mte0)[:, 1]) agg["all"]["members"].append(Mte) y = np.concatenate(agg["all"]["y"]); reg = np.concatenate(agg["all"]["region"]) cell = np.concatenate(agg["all"]["cell"]) pmean = np.concatenate(agg["all"]["mean"]); pstack = np.concatenate(agg["all"]["stack"]) pl1 = np.concatenate(agg["all"]["l1"]); Mpool = np.concatenate(agg["all"]["members"]) base = float(np.nanmean(y)); clr = np.full(y.size, base) for tag, p in (("ens_mean", pmean), ("ens_stack", pstack), ("ens_l1", pl1)): thr, _ = vf.best_threshold(y, p, "Peirce") bs = vf.binary_scores(y, p, thr) score_rows.append(dict(block=b, combo=tag, base_rate=round(base, 4), bss=round(vf.brier_skill_score(y, p, clr), 4), auc=round(vf.auc(y, p), 4), SEDI=round(float(bs["SEDI"]), 4), Peirce=round(float(bs["Peirce"]), 4))) # per-member diagnostics + error-correlation + LOMO contribution to the L2 stack err = Mpool - y[:, None] # member errors ok = np.all(np.isfinite(err), axis=1) C = np.corrcoef(err[ok].T) if ok.sum() > 10 else np.full((len(members),) * 2, np.nan) stack_bss = vf.brier_skill_score(y, pstack, clr) for j, mm in enumerate(members): pj = Mpool[:, j] indiv_bss = vf.brier_skill_score(y, pj, clr) indiv_sedi = float(vf.binary_scores(y, pj, vf.best_threshold(y, pj, "Peirce")[0])["SEDI"]) offdiag = np.delete(C[j], j) if np.isfinite(C).any() else np.array([np.nan]) mean_abs_corr = float(np.nanmean(np.abs(offdiag))) # LOMO: refit the L2 stack LOYO without member j, measure BSS drop keep = [k for k in range(len(members)) if k != j] lomo_p = _lomo_stack(data, years, bid, members, keep) lomo_bss = vf.brier_skill_score(y, lomo_p, clr) div_rows.append(dict(block=b, member=mm, indiv_bss=round(indiv_bss, 4), indiv_sedi=round(indiv_sedi, 4), mean_abs_errcorr=round(mean_abs_corr, 3), lomo_delta_bss=round(stack_bss - lomo_bss, 4))) # pool for output pooled["block_id"].append(np.full(y.size, bid)); pooled["region"].append(reg) pooled["cell"].append(cell); pooled["y"].append(y) pooled["p_mean"].append(pmean); pooled["p_stack"].append(pstack); pooled["p_l1"].append(pl1) # best single member for the figure best_bss = max(vf.brier_skill_score(y, Mpool[:, j], clr) for j in range(len(members))) fig_data[b] = dict(mean=vf.brier_skill_score(y, pmean, clr), stack=stack_bss, l1=vf.brier_skill_score(y, pl1, clr), best_member=best_bss) print(f" {b}: best_member BSS={best_bss:.3f} stack={stack_bss:.3f} " f"l1={fig_data[b]['l1']:.3f} mean={fig_data[b]['mean']:.3f}") os.makedirs(TABLES, exist_ok=True) pd.DataFrame(score_rows).to_csv(os.path.join(TABLES, "v3_stack_scores_by_block.csv"), index=False) pd.DataFrame(div_rows).to_csv(os.path.join(TABLES, "v3_stack_diversity_by_block.csv"), index=False) outdir = os.path.join(cfg["paths"]["outputs"], "loyo_pred_v3stack"); os.makedirs(outdir, exist_ok=True) np.savez_compressed(os.path.join(outdir, "v3stack_pooled.npz"), **{k: np.concatenate(v) for k, v in pooled.items()}) print("[23] wrote v3_stack_scores_by_block.csv, v3_stack_diversity_by_block.csv, v3stack_pooled.npz") _plot(fig_data, block_names) def _lomo_stack(data, years, bid, members, keep): """LOYO L2 stack restricted to member indices `keep`; pooled OOS probs (year order).""" out = [] for yr in years: Mte = block_matrix(data, yr, bid, members)[0][:, keep] tr = [t for t in years if t != yr] Mtr = np.concatenate([block_matrix(data, t, bid, members)[0][:, keep] for t in tr]) ytr = np.concatenate([block_matrix(data, t, bid, members)[1] for t in tr]) ok = np.isfinite(ytr) & np.all(np.isfinite(Mtr), axis=1) m = _meta("l2").fit(Mtr[ok], ytr[ok].astype(int)) out.append(m.predict_proba(np.nan_to_num(Mte, nan=0.0))[:, 1]) return np.concatenate(out) def _plot(fig_data, block_names): import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt x = np.arange(len(block_names)) fig, ax = plt.subplots(figsize=(8, 4.5)) bn = [b for b in block_names if b in fig_data] x = np.arange(len(bn)) for key, lab in (("best_member", "best single member"), ("mean", "ens_mean"), ("stack", "ens_stack (L2)"), ("l1", "ens_l1 (sparse)")): ax.plot(x, [fig_data[b][key] for b in bn], marker="o", label=lab) ax.set_xticks(x); ax.set_xticklabels(bn); ax.axhline(0, color="k", lw=0.6) ax.set_ylabel("BSS vs climatology"); ax.grid(alpha=0.3); ax.legend() ax.set_title("v3 diversity stack vs best single member by lead") os.makedirs(FIGS, exist_ok=True) fig.tight_layout(); fig.savefig(os.path.join(FIGS, "v3_stack_by_lead.png"), dpi=130) print("[23] wrote results/figures/v3_stack_by_lead.png") if __name__ == "__main__": main()