#!/usr/bin/env python """Stage 20f — calibration recovery for the Stage-2 selected-predictor model. Stage 20d found the selected extended predictors improve DISCRIMINATION (AUC/SEDI) at W3-W6 but degrade calibrated BSS. Diagnosis: the isotonic map was fit on the SAME training rows as the HGB, so it inherited in-sample overconfidence. This stage fixes calibration properly, all nested in the outer LOYO fold: 1. Inner K-fold (by year) over the 19 training years -> OUT-OF-FOLD (OOF) probabilities for both the base (10-feature backbone) and sel (backbone+selected) HGBs on every training row. 2. Calibration/combination fit on those OOF probabilities (never in-sample): base_recal : isotonic(base_oof) -> y sel_recal : isotonic(sel_oof) -> y stack : logistic([base_oof, sel_oof]) -> y (calibrated blend; down-weights sel where it does not help, uses it where it does) 3. Base and sel HGBs refit on the FULL 19 training years; test-year probabilities pushed through the OOF-fit calibrators. Per-region x block decision thresholds fit on the OOF train probs. Output: outputs/loyo_pred_v3cal/pred_v3cal_.npz with y, region, cell, block_id and p_base/p_sel/p_stack (+ flags). --cap controls the train-row cap (default 200000; a --cap 0 run gives the no-cap sensitivity check). Scored by stage 20g. """ from __future__ import annotations import os, sys, logging, argparse, time import numpy as np import pandas as pd import xarray as xr from sklearn.ensemble import HistGradientBoostingClassifier from sklearn.isotonic import IsotonicRegression from sklearn.linear_model import LogisticRegression from sklearn.model_selection import KFold sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) from ehd.config import load_config, lead_blocks from ehd import climatology as clim, features as feat, features_ext as fext, verify as vf logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") log = logging.getLogger("stage20f") ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) FEAT_NAMES = list(feat.FEATURE_NAMES) + list(fext.EXT_FEATURE_NAMES) + list(fext.GRID_FEATURE_NAMES) NAME_IDX = {n: i for i, n in enumerate(FEAT_NAMES)} def preload(cfg): da = clim.load_gridded_obs(cfg["paths"]["gridded_obs"]) obs_vals, obs_lut = feat.build_obs_index(da) fc, ref = {}, None for y in cfg["forecast"]["init_years"]: fc[y] = xr.open_dataarray(os.path.join(cfg["paths"]["outputs"], "forecast_daily_tx", f"fc_daily_tx_{y}.nc")).load() ref = fc[y] 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) extgrid = {y: fext.load_extgrid(cfg["paths"]["outputs"], y) for y in cfg["forecast"]["init_years"]} return da, obs_vals, obs_lut, fc, iy, ix, region_of_cell, extgrid def assemble(years, thr, blocks, obs_vals, obs_lut, fc, iy, ix, nino34, omi, extgrid): """Per-block (X[28], y, cell, year) over `years`, with per-row year label for inner CV.""" ncell = iy.size; block_names = list(blocks) acc = {b: {"X": [], "y": [], "cell": [], "yr": []} for b in block_names} for y in years: thr_i = thr.reindex(lat=fc[y].lat, lon=fc[y].lon, method="nearest") out = feat.assemble_year(fc[y], obs_vals, obs_lut, thr_i, blocks, iy, ix) ext = fext.ext_features_year(fc[y], obs_vals, obs_lut, thr_i, iy, ix, nino34, omi)["X_ext"] eg = extgrid[y] for bi, b in enumerate(block_names): Xg = fext.ext_grid_block(eg, bi, ncell) X = np.column_stack([out[b]["X"], ext, Xg]).astype(np.float32) acc[b]["X"].append(X); acc[b]["y"].append(out[b]["y_primary"]) acc[b]["cell"].append(out[b]["cell_id"]); acc[b]["yr"].append(np.full(X.shape[0], y)) return {b: (np.concatenate(v["X"]), np.concatenate(v["y"]), np.concatenate(v["cell"]), np.concatenate(v["yr"])) for b, v in acc.items()} def region_thresholds(y, p, regions, region_ids, rng, cap=50000): def _bt(yy, pp): if cap and yy.size > cap: s = rng.choice(yy.size, cap, replace=False); yy, pp = yy[s], pp[s] return vf.best_threshold(yy, pp, "Peirce")[0] g = _bt(y, p); out = {} for r in region_ids: m = regions == r out[r] = _bt(y[m], p[m]) if (m.sum() > 200 and np.nansum(y[m]) > 10) else g return out def _hgb(): return HistGradientBoostingClassifier(max_iter=300, learning_rate=0.06, max_leaf_nodes=31, l2_regularization=1.0, random_state=0) def main(): ap = argparse.ArgumentParser() ap.add_argument("--folds", type=int, nargs="*", default=None) ap.add_argument("--cap", type=int, default=200000, help="train-row cap per block (0 = no cap)") ap.add_argument("--inner", type=int, default=4, help="inner K-fold (by year) for OOF calibration") 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"] outdir = os.path.join(cfg["paths"]["outputs"], "loyo_pred_v3cal"); os.makedirs(outdir, exist_ok=True) folds = args.folds or cfg["cross_validation"]["years"] rng = np.random.default_rng(cfg["seed"]); CAP = args.cap sel = pd.read_csv(os.path.join(ROOT, "results", "tables", "feature_selection_by_block.csv")) sel_cols = {r["block"]: [NAME_IDX[f] for f in str(r["selected"]).split(";")] for _, r in sel.iterrows()} base_cols = list(range(10)) nino34 = fext.load_nino34(os.environ.get("NINO34_TXT", os.path.join(ROOT, "..", "HW", "data", "external", "nino34.txt"))) omi = fext.load_omi(os.environ.get("OMI_TXT", os.path.join(ROOT, "..", "HW", "data", "external", "omi.txt"))) da, obs_vals, obs_lut, fc, iy, ix, region_of_cell, extgrid = preload(cfg) for yr in folds: t0 = time.time() base_years = [y for y in range(cfg["obs"]["gridded_years"][0], cfg["obs"]["gridded_years"][1] + 1) if y != yr] thr = clim.ehd_threshold(clim.tx90p_climatology(da, base_years, 90, 15), 35.0, "min") tr_years = [y for y in cfg["cross_validation"]["years"] if y != yr] tr = assemble(tr_years, thr, blocks, obs_vals, obs_lut, fc, iy, ix, nino34, omi, extgrid) te = assemble([yr], thr, blocks, obs_vals, obs_lut, fc, iy, ix, nino34, omi, extgrid) cols = {"block_id": [], "region": [], "cell": [], "y": []} pc = {k: [] for k in ("p_base", "p_sel", "p_stack", "flag_base", "flag_sel", "flag_stack")} for b in block_names: Xtr, ytr, ctr, yrtr = tr[b]; Xte, yte, cte = te[b][0], te[b][1], te[b][2] reg_tr = region_of_cell[ctr]; reg_te = region_of_cell[cte] m = np.isfinite(ytr) Xtr, ytr, ctr, yrtr, reg_tr = Xtr[m], ytr[m], ctr[m], yrtr[m], reg_tr[m] if CAP and ytr.size > CAP: si = rng.choice(ytr.size, CAP, replace=False) Xtr, ytr, yrtr, reg_tr = Xtr[si], ytr[si], yrtr[si], reg_tr[si] yi = ytr.astype(int) # --- inner K-fold by YEAR -> OOF base/sel probs on train --- uyears = np.array(sorted(set(yrtr.tolist()))) kf = KFold(n_splits=min(args.inner, len(uyears)), shuffle=True, random_state=0) oof_base = np.full(ytr.size, np.nan); oof_sel = np.full(ytr.size, np.nan) for tr_idx, va_idx in kf.split(uyears): va_years = set(uyears[va_idx].tolist()) va = np.isin(yrtr, list(va_years)); trn = ~va mb = _hgb().fit(Xtr[trn][:, base_cols], yi[trn]) ms = _hgb().fit(Xtr[trn][:, sel_cols[b]], yi[trn]) oof_base[va] = mb.predict_proba(Xtr[va][:, base_cols])[:, 1] oof_sel[va] = ms.predict_proba(Xtr[va][:, sel_cols[b]])[:, 1] ok = np.isfinite(oof_base) & np.isfinite(oof_sel) iso_b = IsotonicRegression(out_of_bounds="clip").fit(oof_base[ok], yi[ok]) iso_s = IsotonicRegression(out_of_bounds="clip").fit(oof_sel[ok], yi[ok]) stk = LogisticRegression(max_iter=1000).fit( np.column_stack([oof_base[ok], oof_sel[ok]]), yi[ok]) # OOF calibrated train probs (for decision thresholds, leakage-safe) ptr_base = iso_b.transform(oof_base[ok]); ptr_sel = iso_s.transform(oof_sel[ok]) ptr_stk = stk.predict_proba(np.column_stack([oof_base[ok], oof_sel[ok]]))[:, 1] thr_base = region_thresholds(yi[ok], ptr_base, reg_tr[ok], region_ids, rng) thr_sel = region_thresholds(yi[ok], ptr_sel, reg_tr[ok], region_ids, rng) thr_stk = region_thresholds(yi[ok], ptr_stk, reg_tr[ok], region_ids, rng) # --- refit on FULL train, predict test, apply OOF calibrators --- mb = _hgb().fit(Xtr[:, base_cols], yi); ms = _hgb().fit(Xtr[:, sel_cols[b]], yi) pb = iso_b.transform(mb.predict_proba(Xte[:, base_cols])[:, 1]) ps = iso_s.transform(ms.predict_proba(Xte[:, sel_cols[b]])[:, 1]) raw_b = mb.predict_proba(Xte[:, base_cols])[:, 1] raw_s = ms.predict_proba(Xte[:, sel_cols[b]])[:, 1] pst = stk.predict_proba(np.column_stack([raw_b, raw_s]))[:, 1] n = yte.size cols["block_id"].append(np.full(n, block_ids[b])); cols["region"].append(reg_te) cols["cell"].append(cte); cols["y"].append(yte) for tag, p, th in (("base", pb, thr_base), ("sel", ps, thr_sel), ("stack", pst, thr_stk)): tvec = np.array([th[int(r)] for r in reg_te]) pc[f"p_{tag}"].append(p.astype(np.float32)); pc[f"flag_{tag}"].append(p >= tvec) pred = {k: np.concatenate(v) for k, v in cols.items()} for k, v in pc.items(): pred[k] = np.concatenate(v) tag = "" if CAP else "_nocap" np.savez_compressed(os.path.join(outdir, f"pred_v3cal{tag}_{yr}.npz"), **pred) w = pred["block_id"] == 4 # W4 (a long-lead block where sel helped AUC) cl = np.full(int(w.sum()), float(np.nanmean(pred["y"][w]))) log.info("fold %d %.0fs | W4 BSS base=%.3f sel=%.3f stack=%.3f", yr, time.time() - t0, vf.brier_skill_score(pred["y"][w], pred["p_base"][w], cl), vf.brier_skill_score(pred["y"][w], pred["p_sel"][w], cl), vf.brier_skill_score(pred["y"][w], pred["p_stack"][w], cl)) if __name__ == "__main__": main()