#!/usr/bin/env python """Stage 06 — full leave-one-year-out (LOYO) engine producing out-of-sample predictions. For each held-out reforecast year, EVERYTHING is fit without that year (leakage rule): 1. Tx90p threshold = min(Tx90p, 35), climatology fit on obs years excluding the holdout. 2. calibration model (per family) fit on the other 19 reforecast years. 3. isotonic recalibration (task 7) fit on TRAIN predictions, applied to the holdout. 4. decision probability threshold (task 7) chosen per REGION x BLOCK on TRAIN recalibrated probabilities (Peirce-maximising), applied to the holdout to make the 1/0 flag. Writes one NPZ of out-of-sample predictions per held-out year to outputs/loyo_pred/pred_.npz (arrays: block_id, region, cell, y, and probability columns per family/recalibration), plus the fitted decision thresholds. Scoring, bootstrap significance and figures are stage 08 (kept separate so predictions are computed once). Families: raw, climat (references) + logistic, hgb (calibrated; each also recalibrated). """ from __future__ import annotations import os, sys, logging, argparse, time import numpy as np 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 climatology as clim, features as feat, calibration as cal, verify as vf logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") log = logging.getLogger("stage06") CAL_FAMILIES = ["logistic", "hgb"] # trainable calibrators (each also recalibrated) REF_FAMILIES = ["raw", "climat"] # non-trainable references 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) return da, obs_vals, obs_lut, fc, iy, ix, region_of_cell def assemble(cfg, years, thr, blocks, obs_vals, obs_lut, fc, iy, ix): acc = {b: {"X": [], "y": [], "cell": []} for b in blocks} 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) for b in blocks: acc[b]["X"].append(out[b]["X"]) acc[b]["y"].append(out[b]["y_primary"]) acc[b]["cell"].append(out[b]["cell_id"]) return {b: (np.concatenate(v["X"]), np.concatenate(v["y"]), np.concatenate(v["cell"])) for b, v in acc.items()} def region_thresholds(y, p, regions, region_ids, metric="Peirce"): """Best decision threshold per region on (y,p). Falls back to global if a region is sparse.""" g_thr, _ = vf.best_threshold(y, p, metric) out = {} for r in region_ids: m = regions == r if m.sum() > 200 and np.nansum(y[m]) > 10: out[r], _ = vf.best_threshold(y[m], p[m], metric) else: out[r] = g_thr return out, g_thr def main(): ap = argparse.ArgumentParser() ap.add_argument("--folds", type=int, nargs="*", default=None) args = ap.parse_args() cfg = load_config(); blocks = lead_blocks(cfg) block_ids = {b: k + 1 for k, b in enumerate(blocks)} region_ids = cfg["regions"]["ids"] outdir = os.path.join(cfg["paths"]["outputs"], "loyo_pred") os.makedirs(outdir, exist_ok=True) folds = args.folds or cfg["cross_validation"]["years"] log.info("preloading obs + %d forecast years ...", len(cfg["forecast"]["init_years"])) da, obs_vals, obs_lut, fc, iy, ix, region_of_cell = preload(cfg) for yr in folds: t0 = time.time() base = [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, 90, 15), 35.0, "min") train_years = [y for y in cfg["cross_validation"]["years"] if y != yr] tr = assemble(cfg, train_years, thr, blocks, obs_vals, obs_lut, fc, iy, ix) te = assemble(cfg, [yr], thr, blocks, obs_vals, obs_lut, fc, iy, ix) cols = {"block_id": [], "region": [], "cell": [], "y": []} prob_cols = {} dec_thr_records = [] for b in blocks: Xtr, ytr, _ = tr[b] Xte, yte, cte = te[b] reg_te = region_of_cell[cte] reg_tr = region_of_cell[tr[b][2]] 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 fam in REF_FAMILIES: mdl = cal.make_model(fam).fit(Xtr, ytr) prob_cols.setdefault(fam, []).append(mdl.predict_proba(Xte)) for fam in CAL_FAMILIES: mdl = cal.make_model(fam).fit(Xtr, ytr) ptr = mdl.predict_proba(Xtr); pte = mdl.predict_proba(Xte) prob_cols.setdefault(fam, []).append(pte) # isotonic recalibration fit on TRAIN only mtr = np.isfinite(ytr) & np.isfinite(ptr) iso = IsotonicRegression(out_of_bounds="clip").fit(ptr[mtr], ytr[mtr]) pte_rc = iso.transform(pte); ptr_rc = iso.transform(ptr) prob_cols.setdefault(fam + "_recal", []).append(pte_rc) # per-region decision thresholds on TRAIN recalibrated probs rthr, gthr = region_thresholds(ytr[mtr], ptr_rc[mtr], reg_tr[mtr], region_ids) for r in region_ids: dec_thr_records.append((fam, b, r, rthr[r])) pred = {k: np.concatenate(v).astype(np.float32) for k, v in cols.items()} for fam, lst in prob_cols.items(): pred["p_" + fam] = np.concatenate(lst).astype(np.float32) dt = np.array(dec_thr_records, dtype=[("family", "U12"), ("block", "U4"), ("region", "i4"), ("thr", "f4")]) np.savez_compressed(os.path.join(outdir, f"pred_{yr}.npz"), decthr=dt, **pred) # quick fold diagnostics (pooled over blocks) for the log w1 = pred["block_id"] == 1 bss = vf.brier_skill_score(pred["y"][w1], pred["p_hgb_recal"][w1], np.full(w1.sum(), np.nanmean(tr["W1"][1]))) log.info("fold %d done in %.0fs | W1 hgb_recal BSS_clim=%.3f AUC=%.3f | -> pred_%d.npz", yr, time.time() - t0, bss, vf.auc(pred["y"][w1], pred["p_hgb_recal"][w1]), yr) if __name__ == "__main__": main()