#!/usr/bin/env python """Stage 22b — v3 deep-learning members: CNN, residual U-Net, FT-Transformer, deep ANN. Generates the spatial / attention members that the hybrid-design note ranks highest for the extended range (doc §2-§3), on the SAME leakage-safe LOYO footing and the SAME lead-dependent predictor slices as the sklearn zoo (stage 21): cnn fully-convolutional per-cell classifier (v2 Phase B, retained for continuity). resunet residual U-Net (encoder/decoder + skips) — multiscale spatial structure, W3-W6. transformer feature-tokenizer Transformer on the per-cell predictor vector — long-range dependence challenger (compact; overfitting-controlled). ann_torch deeper torch MLP (distinct from the shallow sklearn `mlp`). Each family trains per fold x block on the config `predictors` slice, self-calibrates with an isotonic map on TRAIN (the meta-stack in stage 23 recalibrates at the ensemble level), and gets a per-region x block Peirce decision threshold on TRAIN. Writes outputs/loyo_pred_v3dl/ pred_v3dl_.npz with block_id/region/cell/y matching stage 21's row order, plus p_ columns and decthr. Merge into the stack with: scripts/23_v3_diversity_stack.py --dl-dir outputs/loyo_pred_v3dl --dl-pattern pred_v3dl_*.npz """ from __future__ import annotations import os, sys, logging, argparse, time import numpy as np import xarray as xr sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) from ehd.config import load_config, lead_blocks, block_predictor_index, block_predictors from ehd import climatology as clim, features as feat, features_ext as fext, verify as vf from ehd import gridmodels as gm logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") log = logging.getLogger("stage22b") 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) DL_FAMILIES = ["cnn", "resunet", "transformer", "ann_torch"] GRID_FAMILIES = {"cnn", "resunet"} # need grid geometry 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) ny, nx = reg.values.shape 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, ny, nx, extgrid def assemble(years, thr, blocks, obs_vals, obs_lut, fc, iy, ix, nino34, omi, extgrid): ncell = iy.size; block_names = list(blocks) acc = {b: {"X": [], "y": [], "cell": []} 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"]) 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, rng, cap=50000): def _bt(yy, pp): ok = np.isfinite(yy) & np.isfinite(pp); yy, pp = yy[ok], pp[ok] 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 main(): ap = argparse.ArgumentParser() ap.add_argument("--folds", type=int, nargs="*", default=None) ap.add_argument("--families", nargs="*", default=DL_FAMILIES) ap.add_argument("--blocks", nargs="*", default=None) ap.add_argument("--cap", type=int, default=0, help="train-row cap for tabular families (0=none)") args = ap.parse_args() cfg = load_config(); blocks = lead_blocks(cfg); block_names = args.blocks or list(blocks) block_ids = {b: k + 1 for k, b in enumerate(list(blocks))} region_ids = cfg["regions"]["ids"] outdir = os.path.join(cfg["paths"]["outputs"], "loyo_pred_v3dl"); os.makedirs(outdir, exist_ok=True) folds = args.folds or cfg["cross_validation"]["years"] rng = np.random.default_rng(cfg["seed"]) blk_idx = {b: block_predictor_index(cfg, b, FEAT_NAMES) for b in block_names} 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, ny, nx, extgrid = preload(cfg) ncell = iy.size log.info("families=%s blocks=%s grid=%dx%d ncell=%d", args.families, block_names, ny, nx, ncell) 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": []} prob_cols = {}; dec_records = [] for b in block_names: Xtr, ytr, ctr = tr[b]; Xte, yte, cte = te[b] reg_tr = region_of_cell[ctr]; reg_te = region_of_cell[cte] idx = blk_idx[b] 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 args.families: if fam in GRID_FAMILIES: geom = dict(iy=iy, ix=ix, ny=ny, nx=nx, ncell=ncell) model = gm.make_torch_model(fam, geom) # grid families keep ALL cells per init (image reshape); slice channels only model.fit(Xtr[:, idx], ytr) p_te = model.predict_proba(Xte[:, idx]) p_tr = model.predict_proba(Xtr[:, idx]) else: Xtr_b, ytr_b = Xtr[:, idx], ytr if args.cap and ytr_b.shape[0] > args.cap: m = np.isfinite(ytr_b) si = rng.choice(np.where(m)[0], args.cap, replace=False) Xtr_b, ytr_b = Xtr_b[si], ytr_b[si] model = gm.make_torch_model(fam) model.fit(Xtr_b, ytr_b) p_te = model.predict_proba(Xte[:, idx]) p_tr = model.predict_proba(Xtr[:, idx]) prob_cols.setdefault("p_" + fam, []).append(p_te.astype(np.float32)) rthr = region_thresholds(ytr, p_tr, reg_tr, region_ids, rng) for r in region_ids: dec_records.append((fam, b, r, rthr[r])) pred = {k: np.concatenate(v).astype(np.float32) for k, v in cols.items()} for k, lst in prob_cols.items(): pred[k] = np.concatenate(lst) dt = np.array(dec_records, dtype=[("tag", "U20"), ("block", "U4"), ("region", "i4"), ("thr", "f4")]) np.savez_compressed(os.path.join(outdir, f"pred_v3dl_{yr}.npz"), decthr=dt, **pred) wsel = "W4" if "W4" in block_names else block_names[0] w = pred["block_id"] == block_ids[wsel] cl = np.full(int(w.sum()), float(np.nanmean(pred["y"][w]))) msg = " ".join(f"{f}={vf.brier_skill_score(pred['y'][w], pred['p_'+f][w], cl):.3f}" for f in args.families) log.info("fold %d %.0fs | %s BSS: %s -> pred_v3dl_%d.npz", yr, time.time() - t0, wsel, msg, yr) if __name__ == "__main__": main()