#!/usr/bin/env python3 """Combine several finished national runs into one ensemble forecast. The three national models fail in different ways -- logistic regression is the best calibrated, the boosted trees discriminate best, the net sits between them -- which is the situation where an ensemble usually beats all of its members. Their scores are on incompatible scales, so each model's score is first mapped through its own out-of-fold empirical distribution (a probability integral transform). The mapping is built ONLY from out-of-fold rows inside the training years and the very same mapping is applied to the held-out years, so nothing about the validation years enters the combination. The ensemble score is then the mean of the transformed scores. The result is written as another set of shards, so `train_national.py finalize` scores the ensemble with the identical thresholds, frames and netCDF writer. """ from __future__ import annotations import argparse import glob import os import sys import numpy as np sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from train_national import shard_dir CARRY = ("evt", "week", "capable", "region", "doy", "lt", "point", "init") def pit_reference(tags): """Sorted pooled out-of-fold scores per tag: the reference distribution.""" ref = {} for t in tags: files = sorted(glob.glob(os.path.join(shard_dir(t), "oof_*.npz"))) if not files: raise SystemExit("no out-of-fold shards for tag '%s'" % t) s = np.concatenate([np.load(f)["score"] for f in files]) ref[t] = np.sort(s[np.isfinite(s)]) return ref def pit(x, sorted_ref): """Map scores onto [0,1] through the reference empirical distribution.""" return np.searchsorted(sorted_ref, x).astype("f8") / max(sorted_ref.size, 1) def simplex_grid(k, step=0.1): """All weight vectors on the k-simplex with the given step.""" n = int(round(1.0 / step)) out = [] def rec(prefix, left, depth): if depth == k - 1: out.append(prefix + [left / n]) return for i in range(left + 1): rec(prefix + [i / n], left - i, depth + 1) rec([], n, 0) return np.array(out) def tune_weights(tags, ref, step=0.1): """Per week, the member weights maximising out-of-fold AUC. Weights are chosen on AUC rather than on CSI because AUC needs no threshold: discrimination is settled first, then the FBI~1 rule sets the threshold on the combined score. Only training-year out-of-fold rows are used. """ import ehd_metrics as M import ehd_common as C u, evt, week = [], None, None for t in tags: files = sorted(glob.glob(os.path.join(shard_dir(t), "oof_*.npz"))) parts = [np.load(f) for f in files] s = np.concatenate([p["score"] for p in parts]).astype("f8") u.append(pit(np.nan_to_num(s, nan=-1e18), ref[t])) if evt is None: evt = np.concatenate([p["evt"] for p in parts]) week = np.concatenate([p["week"] for p in parts]).astype("i4") u = np.vstack(u) grid = simplex_grid(len(tags), step) best = {} for k in range(len(C.WEEKS)): m = week == k if m.sum() == 0: best[k] = np.full(len(tags), 1.0 / len(tags)) continue uu, ob = u[:, m], evt[m] > 0.5 sc = [M.auc_score(w @ uu, ob) for w in grid] j = int(np.argmax(sc)) best[k] = grid[j] print("[ens] %s weights %s out-of-fold AUC %.4f (equal %.4f)" % (C.WEEK_NAMES[k], np.round(grid[j], 2).tolist(), sc[j], M.auc_score(np.full(len(tags), 1.0 / len(tags)) @ uu, ob)), flush=True) return best def combine(tags, ref, name, out_tag, weights=None): """Average the transformed scores of one shard across the tags.""" paths = {t: sorted(glob.glob(os.path.join(shard_dir(t), name))) for t in tags} n = len(paths[tags[0]]) if any(len(v) != n for v in paths.values()): raise SystemExit("tags disagree on the number of %s shards" % name) for i in range(n): parts = [np.load(paths[t][i]) for t in tags] base = parts[0] us, bads = [], np.zeros(base["score"].size, bool) for t, p in zip(tags, parts): s = p["score"].astype("f8") bads |= ~np.isfinite(s) us.append(pit(np.nan_to_num(s, nan=-1e18), ref[t])) us = np.vstack(us) if weights is None: acc = us.mean(axis=0) else: wk = base["week"].astype("i4") acc = np.zeros(us.shape[1]) for k, w in weights.items(): m = wk == k if m.any(): acc[m] = w @ us[:, m] acc[bads] = np.nan # Members were produced at different times and the older runs predate the # region/doy fields, so each carry-field is taken from the first member # that actually has it rather than from parts[0] alone. keep = {} for k in CARRY: for p_ in parts: if k in p_.files: keep[k] = p_[k] break np.savez(os.path.join(shard_dir(out_tag), os.path.basename(paths[tags[0]][i])), score=acc, **keep) print("[ens] %d x %s -> %s" % (n, name, shard_dir(out_tag)), flush=True) def main(): ap = argparse.ArgumentParser() ap.add_argument("--tags", nargs="+", required=True, help="shard tags to combine; '-' means the untagged run") ap.add_argument("--out-tag", default="_ens") ap.add_argument("--tune", action="store_true", help="pick per-week member weights on out-of-fold AUC") args = ap.parse_args() tags = ["" if t == "-" else t for t in args.tags] os.makedirs(shard_dir(args.out_tag), exist_ok=True) for f in glob.glob(os.path.join(shard_dir(args.out_tag), "*.npz")): os.remove(f) ref = pit_reference(tags) for t in tags: print("[ens] '%s': %d out-of-fold rows in the reference distribution" % (t or "(untagged)", ref[t].size), flush=True) w = tune_weights(tags, ref) if args.tune else None combine(tags, ref, "oof_*.npz", args.out_tag, w) combine(tags, ref, "pred_*.npz", args.out_tag, w) if all(glob.glob(os.path.join(shard_dir(t), "oofpred_*.npz")) for t in tags): combine(tags, ref, "oofpred_*.npz", args.out_tag, w) if w is not None: np.savez(os.path.join(shard_dir(args.out_tag), "weights.npz"), tags=np.array(tags), **{str(k): v for k, v in w.items()}) if __name__ == "__main__": main()