#!/usr/bin/env python3 """EXP3: one model family per grid point, two-stage decision on p = M/N. Architecture, following the project note: stage 1 N models vote. p = M/N. p < P1 -> declare nEHDE and drop the sample. stage 2 a second ensemble is trained on the samples the gate KEPT (in the training years) and decides the survivors: p2 > P2 -> yEHDE. BOTH thresholds are fitted on the training years, never assumed. A fixed P1 = 0.3 was tried first and it destroyed the experiment: at W2 it kept 11% of the samples but threw away 82% of the events, because every member shrinks toward climatology so almost nothing scores above 0.3. A gate is only useful if it discards confident NON-events, so P1 is set to the largest threshold that still keeps RECALL1 of the training events, and the sample-closed / events-lost trade-off is reported. The votes those thresholds are fitted on must be OUT OF FOLD. Calibrating on in-sample votes was tried too and is worse than useless: trees reproduce their own training data, so the gate looked like it kept 95-100% of events while on the test years it threw away 66-80% of them. Folds are whole YEARS, since inits days apart share weather. Everything is scored over the whole test set, gated-out samples included as nEHDE, so the gate cannot flatter the result by hiding its own misses. Why per point at all: the EC post-mortem found the failure was SPATIAL DIFFERENTIATION, and a pooled national model cannot fix that by construction. This is the direct test of that diagnosis. Two ensembles are run side by side because the project note asks for Sev_obs as the target, and section 1 of the evaluation showed value regression collapses the event rate: reg N regressors on Sev = O - thr; a member votes yEHDE when its Sev_hat >= 0 clf N classifiers on the event itself; a member votes yEHDE when p_i >= 0.5 If the regression ensemble under-forecasts here the way D2 did, that is the same trap and it will show as FBI far below 1. Features are local AND large-scale, all z-scored against the model's own (week, doy, point) climatology from training years only: local 16 flux variables at the point (weekly_feats.npz) meso same 16 averaged over +-3 and +-8 grid cells (weekly_grid.f4) synoptic same 16 averaged over the whole window, plus N-S and E-W gradients of t2m and psfc across it upper air 27 derived pressure-level variables at the point (optional) calendar two annual harmonics """ from __future__ import annotations import argparse import json import os import sys import warnings import numpy as np warnings.filterwarnings("ignore") sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) import cfs_common as C import build_obs as B import build_daily as D from weekly_skill import auc, climz GRIDF = os.path.join(C.OUT, "Cube", "weekly_grid.f4") # BCSD downscaled Tx fields (d1) at the 541 points, per week: the cfs_0.25_sd # input. Row order matches weekly.npz (both written from the same cube). BCSDF = os.path.join(C.OUT, "AI_Results", "bcsd_fluxonly.json.w%d.fields.npz") RES = os.environ.get("EXP3_RES", os.path.join(C.OUT, "AI_Results", "exp3")) PGBF = os.environ.get("PGB_WEEKLY", os.path.join(C.OUT, "Cube", "pgb_weekly.npz")) RECALL1 = 0.95 # the gate must keep this share of TRAINING events NFOLD = 5 # year-block folds used to calibrate the thresholds NPT_REG = 3 # test points per climate region # ------------------------------------------------------------------ the members NMEMBER = int(os.environ.get("EXP3_NMEMBER", 60)) # target ensemble size def members(kind, seed, ncore): """~NMEMBER diverse models = the ensemble whose vote share is p = M/N. Same families for both ensembles so the only difference between `reg` and `clf` is the target, not the model zoo. Diversity comes from three sources stacked together, following the user's "pool đa dạng, N≈50-72" choice: * several model FAMILIES (gradient boosting, bagged trees, MLP, linear, kNN) * several HYPER-PARAMETER settings per family * several SEEDS per setting, each also perturbing the row subsample (bootstrap) and, for the tree learners, the column subsample (feature subset), so two members of the same setting still disagree. The list is built family-first then truncated to NMEMBER, so a smaller N still spans every family instead of front-loading one of them. """ import lightgbm as lgb from sklearn.ensemble import (ExtraTreesRegressor, ExtraTreesClassifier, RandomForestRegressor, RandomForestClassifier, HistGradientBoostingRegressor, HistGradientBoostingClassifier) from sklearn.linear_model import Ridge, LogisticRegression, ElasticNet from sklearn.neural_network import MLPRegressor, MLPClassifier from sklearn.neighbors import KNeighborsRegressor, KNeighborsClassifier R = kind == "reg" lgbm, etree, rforest, hgb, mlp, lin, knn = [], [], [], [], [], [], [] # ---- LightGBM: 4 settings x 6 seeds = 24, subsample+colsample jittered ---- for c, (lv, dp, lr) in enumerate([(31, -1, .05), (63, -1, .03), (15, 6, .08), (127, -1, .05)]): for s in range(6): ss = 0.7 + 0.07 * (s % 3) # 0.70/0.77/0.84 cs = 0.6 + 0.10 * ((s + 1) % 3) # 0.60/0.70/0.80 cl = lgb.LGBMRegressor if R else lgb.LGBMClassifier lgbm.append(("lgbm%d_%d" % (c, s), cl(n_estimators=400, num_leaves=lv, max_depth=dp, learning_rate=lr, subsample=ss, subsample_freq=1, colsample_bytree=cs, min_child_samples=40, verbose=-1, random_state=seed + c * 6 + s, n_jobs=ncore))) # ---- ExtraTrees & RandomForest: 2 settings x 4 seeds each = 8 + 8 ---- for c, (nt, msl) in enumerate([(300, 5), (500, 3)]): for s in range(4): cl = ExtraTreesRegressor if R else ExtraTreesClassifier etree.append(("etree%d_%d" % (c, s), cl(n_estimators=nt, min_samples_leaf=msl, max_features=0.6, bootstrap=True, random_state=seed + 100 + c * 4 + s, n_jobs=ncore))) cl = RandomForestRegressor if R else RandomForestClassifier rforest.append(("rforest%d_%d" % (c, s), cl(n_estimators=nt, min_samples_leaf=msl, max_features=0.6, random_state=seed + 200 + c * 4 + s, n_jobs=ncore))) # ---- HistGradientBoosting: 2 lr x 4 seeds = 8 ---- for c, lrr in enumerate((.05, .10)): for s in range(4): cl = HistGradientBoostingRegressor if R else HistGradientBoostingClassifier hgb.append(("hgb%d_%d" % (c, s), cl(max_iter=300, learning_rate=lrr, max_leaf_nodes=31 + 16 * s, random_state=seed + 300 + c * 4 + s))) # ---- MLP: 2 shapes x 3 seeds = 6 ---- for c, hl in enumerate([(64,), (128, 64)]): for s in range(3): cl = MLPRegressor if R else MLPClassifier mlp.append(("mlp%d_%d" % (c, s), cl(hidden_layer_sizes=hl, alpha=10.0 ** (-4 + s), max_iter=400, early_stopping=True, random_state=seed + 400 + c * 3 + s))) # ---- linear: 4 regularisations (deterministic, one each) ---- if R: lin = [("ridge0", Ridge(alpha=1.0)), ("ridge1", Ridge(alpha=5.0)), ("enet0", ElasticNet(alpha=.05, l1_ratio=.5)), ("enet1", ElasticNet(alpha=.01, l1_ratio=.3))] else: lin = [("logreg%d" % k, LogisticRegression(max_iter=2000, C=cc)) for k, cc in enumerate((1.0, .3, .1, 3.0))] # ---- kNN: 3 neighbourhood sizes = 3 ---- for c, nn in enumerate((15, 25, 40)): cl = KNeighborsRegressor if R else KNeighborsClassifier knn.append(("knn%d" % c, cl(n_neighbors=nn, weights="distance", n_jobs=ncore))) # round-robin interleave so truncating to NMEMBER still spans every family groups = [lgbm, etree, rforest, hgb, mlp, lin, knn] M = [] for i in range(max(len(g) for g in groups)): for g in groups: if i < len(g): M.append(g[i]) return M[:NMEMBER] def _fit_predict(m, Xtr, ytr_sev, ytr_ev, Xte, kind): if kind == "reg": m.fit(Xtr, ytr_sev) return (m.predict(Xte) >= 0).astype(np.float32) if len(np.unique(ytr_ev)) < 2: return None m.fit(Xtr, ytr_ev) return (m.predict_proba(Xte)[:, 1] >= .5).astype(np.float32) def vote(M, Xtr, ytr_sev, ytr_ev, Xte, kind): """Fit every member on all of Xtr and return the yEHDE vote share on Xte.""" v, used = [], [] for nm, m in M: try: r = _fit_predict(m, Xtr, ytr_sev, ytr_ev, Xte, kind) if r is None: continue v.append(r); used.append(nm) except Exception: continue if not v: return None, [] return np.mean(v, 0), used def vote_oof(X, sev, ev, kind, yrs, seed, ncore): """Out-of-fold vote share for every training sample, folds = year blocks. This is the only vote a threshold may be fitted on. An in-sample vote from a tree ensemble is close to the truth by construction and carries no information about what the gate will do on a year it has not seen. """ uy = np.unique(yrs) groups = np.array_split(uy, min(NFOLD, len(uy))) out = np.full(len(X), np.nan, dtype=np.float32) for g in groups: ho = np.isin(yrs, g) if ho.all() or not ho.any(): continue p, _ = vote(members(kind, seed, ncore), X[~ho], sev[~ho], ev[~ho], X[ho], kind) if p is not None: out[ho] = p return out # ------------------------------------------------------------------ features def pick_points(reg, lat, lon): """NPT_REG well-separated points inside each of the 7 regions.""" out = [] for r in np.unique(reg): idx = np.where(reg == r)[0] p = np.stack([lat[idx], lon[idx]], 1) c = [p[:, 0].argmin()] # southernmost, then farthest while len(c) < min(NPT_REG, len(idx)): d = np.min([np.hypot(p[:, 0] - p[j, 0], p[:, 1] - p[j, 1]) for j in c], 0) c.append(int(d.argmax())) out += [int(idx[j]) for j in c] return np.array(out) _CACHE = {} def _cached(): """weekly_feats is a 641 MB npz; loading it once per week wastes minutes.""" if not _CACHE: n, nlt, nv, ny, nx = D.shapes() _CACHE["grid"] = np.memmap(GRIDF, dtype="float32", mode="r", shape=(n, len(C.WEEKS), nv, ny, nx)) _CACHE["fx"] = np.load(os.path.join(C.OUT, "Cube", "weekly_feats.npz"))["X"] z = np.load(B.TX_NPZ) _CACHE["olat"] = z["lat"][z["jj"]] _CACHE["olon"] = z["lon"][z["ii"]] _CACHE["pg"] = np.load(PGBF) if os.path.exists(PGBF) else None _CACHE["bcsd"] = {} return _CACHE def _bcsd_week(w): """BCSD downscaled Tx (d1) and threshold at the 541 points for week w. Cached per week; the file also aligns row-for-row with weekly.npz.""" ca = _cached() if w not in ca["bcsd"]: f = np.load(BCSDF % w) ca["bcsd"][w] = (np.asarray(f["d1"]), np.asarray(f["thr"])) return ca["bcsd"][w] def build_features(pt, w, yrs, doy, train, track="ai"): """(ninit, nfeat) for one point and one week. Local + meso + synoptic. track="ai" (cfs_0.25_AI): the point tmax is the RAW nearest coarse cell, so the AI models learn the downscaling implicitly from the coarse state and its spatial context. track="sd" (cfs_0.25_sd): the point tmax is replaced by the BCSD downscaled Tx (d1), and its margin over the local threshold is added as a raw (non z-scored) feature so the zero-crossing stays meaningful. Every other feature is identical to track "ai", so the only thing that differs between the two experiments is how the point-scale Tx predictor was produced. """ n, nlt, nv, ny, nx = D.shapes() ca = _cached() grid, fx = ca["grid"], ca["fx"] olat, olon = ca["olat"], ca["olon"] g = C.flux_grid() pj = int(np.abs(olat[pt] - g["wlat"]).argmin()) pi = int(np.abs(olon[pt] - g["wlon"]).argmin()) blocks, names = [], [] def add(a, tag): """z-score one (ninit,) series against its own train climatology.""" zz = climz(a[:, None, None], doy[:, [w]], yrs, train)[:, 0, 0] blocks.append(zz[:, None]); names.append(tag) itmax = C.FLUX_NAMES.index("tmax") for v in range(nv): if v == itmax and track == "sd": d1, thrb = _bcsd_week(w) add(d1[:, pt], "loc_sd_tx") # BCSD Tx replaces raw tmax else: add(fx[:, w, v, pt], "loc_" + C.FLUX_NAMES[v]) if track == "sd": d1, thrb = _bcsd_week(w) blocks.append((d1[:, pt] - thrb[:, pt])[:, None].astype("float32")) names.append("sd_margin") # raw margin, 0 = threshold gw = np.asarray(grid[:, w]) # (ninit, nv, ny, nx) for half in (3, 8): j0, j1 = max(0, pj - half), min(ny, pj + half + 1) i0, i1 = max(0, pi - half), min(nx, pi + half + 1) for v in range(nv): add(np.nanmean(gw[:, v, j0:j1, i0:i1], (1, 2)), "box%d_%s" % (half, C.FLUX_NAMES[v])) for v in range(nv): add(np.nanmean(gw[:, v], (1, 2)), "dom_" + C.FLUX_NAMES[v]) for vn in ("t2m", "psfc"): v = C.FLUX_NAMES.index(vn) add(np.nanmean(gw[:, v, ny // 2:], (1, 2)) - np.nanmean(gw[:, v, :ny // 2], (1, 2)), "grad_ns_" + vn) add(np.nanmean(gw[:, v, :, nx // 2:], (1, 2)) - np.nanmean(gw[:, v, :, :nx // 2], (1, 2)), "grad_ew_" + vn) del gw pg = ca["pg"] if pg is not None: px, have = pg["X"], pg["have"] for v, nm in enumerate(pg["names"]): a = px[:, w, v, pt].copy() a[~have] = np.nan add(a, "ua_" + str(nm)) d = doy[:, w].astype(float) for k, f in ((1, np.cos), (1, np.sin), (2, np.cos), (2, np.sin)): blocks.append(f(2 * k * np.pi * d / 365)[:, None]) names.append("harm") return np.concatenate(blocks, 1), names # ------------------------------------------------------------------ thresholds def fit_gate(p1, ev): """Largest gate threshold that still keeps RECALL1 of the TRAINING events. Also returns the closed/lost trade-off, which is the number that decides whether a gate is worth having at all: closing 60% of the sample for 3% of the events is a good trade, closing 11% for 82% is not. """ cand = np.unique(np.round(p1, 4)) best, trade = 0.0, [] for t in cand: keep = p1 >= t rec = (ev & keep).sum() / max(ev.sum(), 1) trade.append({"P1": float(t), "closed": float(1 - keep.mean()), "recall": float(rec)}) if rec >= RECALL1: best = float(t) return best, trade def fit_p2(p2, ev2, keep_tr, ev_tr): """Stage-2 threshold: best training CSI SUBJECT TO FBI staying near 1. Unconstrained CSI maximisation is not usable here. When the signal is weak, CSI is maximised by forecasting the event everywhere -- that scores CSI equal to the base rate with POD 1 and POFD 1, which is not a forecast. The FBI band is what keeps the chosen point an actual decision. """ n = len(ev_tr) best_t, best_csi, fallback = 0.5, -1.0, 0.5 for t in np.unique(np.round(p2, 4)): fc = np.zeros(n, dtype=bool) fc[keep_tr] = p2 >= t s = scores(fc, ev_tr) if not np.isfinite(s["CSI"]): continue if 0.7 <= s["FBI"] <= 1.3 and s["CSI"] > best_csi: best_csi, best_t = s["CSI"], float(t) return best_t if best_csi > 0 else fallback # ------------------------------------------------------------------ scoring def scores(fc, ob): a = float(np.sum(fc & ob)); b = float(np.sum(fc & ~ob)) c = float(np.sum(~fc & ob)); d = float(np.sum(~fc & ~ob)) pod = a / (a + c) if a + c else np.nan pofd = b / (b + d) if b + d else np.nan far = b / (a + b) if a + b else np.nan o = dict(a=a, b=b, c=c, d=d, n=a + b + c + d, base=(a + c) / max(a + b + c + d, 1), POD=pod, POFD=pofd, FAR=far, FBI=(a + b) / (a + c) if a + c else np.nan, CSI=a / (a + b + c) if a + b + c else np.nan) H = min(max(pod if np.isfinite(pod) else 1e-6, 1e-6), 1 - 1e-6) F = min(max(pofd if np.isfinite(pofd) else 1e-6, 1e-6), 1 - 1e-6) lH, lF = np.log(H), np.log(F) o["SEDI"] = ((lF - lH - np.log(1 - F) + np.log(1 - H)) / (lF + lH + np.log(1 - F) + np.log(1 - H))) return o def run_point(pt, wk_list, yrs, doy, O, thr, ncore, track="ai"): train0 = np.isin(yrs, C.TRAIN_YEARS) valid0 = np.isin(yrs, C.VALID_YEARS) out = {} for w in wk_list: X, names = build_features(pt, w, yrs, doy, train0, track) sev = O[:, w, pt] - thr[:, w, pt] ev = sev >= 0 # Drop structurally undefined COLUMNS before dropping rows. Soil moisture # is NaN over water, so at a coastal point whose nearest coarse cell is at # sea one such column would otherwise wipe out every sample through # .all(1) -- which is exactly what happened to 4 of the 21 test points, # silently, while they still had 700+ events each. frac = np.isfinite(X).mean(0) keepc = frac > 0.95 X, names = X[:, keepc], [n for n, k in zip(names, keepc) if k] good = np.isfinite(X).all(1) & np.isfinite(sev) tr, te = train0 & good, valid0 & good if tr.sum() < 200 or te.sum() < 50 or ev[tr].sum() < 20: out[C.WEEK_NAMES[w]] = {"skip": True, "n_train": int(tr.sum()), "n_valid": int(te.sum()), "n_feat": int(X.shape[1]), "n_events_train": int(ev[tr].sum())} continue Xtr, Xte = X[tr], X[te] mu, sd = Xtr.mean(0), Xtr.std(0) + 1e-9 Xtr, Xte = (Xtr - mu) / sd, (Xte - mu) / sd row = {"n_train": int(tr.sum()), "n_valid": int(te.sum()), "n_feat": int(X.shape[1]), "n_feat_dropped": int((~keepc).sum()), "base": float(ev[te].mean())} for kind in ("reg", "clf"): p1tr = vote_oof(Xtr, sev[tr], ev[tr], kind, yrs[tr], 7, ncore) p1te, used1 = vote(members(kind, 7, ncore), Xtr, sev[tr], ev[tr], Xte, kind) if p1te is None or not np.isfinite(p1tr).any(): continue ok1 = np.isfinite(p1tr) P1, trade = fit_gate(p1tr[ok1], ev[tr][ok1]) keep_tr, keep_te = (p1tr >= P1) & ok1, p1te >= P1 gate = {"P1": float(P1), "trade": trade, "kept_frac_train": float(keep_tr.mean()), "kept_frac": float(keep_te.mean()), "events_lost": float((ev[te] & ~keep_te).sum()), "events_total": float(ev[te].sum()), "AUC_p1": float(auc(p1te, ev[te]))} fc = np.zeros(int(te.sum()), dtype=bool) stage2 = None if keep_tr.sum() >= 100 and len(np.unique(ev[tr][keep_tr])) > 1 \ and keep_te.sum() > 0: p2tr = vote_oof(Xtr[keep_tr], sev[tr][keep_tr], ev[tr][keep_tr], kind, yrs[tr][keep_tr], 99, ncore) p2te, used2 = vote(members(kind, 99, ncore), Xtr[keep_tr], sev[tr][keep_tr], ev[tr][keep_tr], Xte[keep_te], kind) if p2te is not None and np.isfinite(p2tr).any(): P2 = fit_p2(np.nan_to_num(p2tr, nan=0.0), ev[tr][keep_tr], keep_tr, ev[tr]) fc[keep_te] = p2te >= P2 stage2 = {"members": used2, "n_train2": int(keep_tr.sum()), "P2": float(P2)} else: # gate kept too little to fit a second stage fc[keep_te] = True row[kind] = {"gate": gate, "members1": used1, "stage2": stage2, "final": scores(fc, ev[te]), "one_stage": scores(p1te >= 0.5, ev[te]), "AUC_p1": float(auc(p1te, ev[te]))} out[C.WEEK_NAMES[w]] = row print(" pt %d %s done" % (pt, C.WEEK_NAMES[w]), flush=True) return out def main(): ap = argparse.ArgumentParser() ap.add_argument("--ipoint", type=int, default=-1, help="index into the 21 points") ap.add_argument("--weeks", default="0,1,2,3,4,5") ap.add_argument("--track", default="ai", choices=["ai", "sd"], help="ai=cfs_0.25_AI (raw coarse), sd=cfs_0.25_sd (BCSD input)") ap.add_argument("--ncore", type=int, default=int(os.environ.get("SLURM_CPUS_PER_TASK", 8))) a = ap.parse_args() wkz = np.load(os.path.join(C.OUT, "Cube", "weekly.npz")) yrs, O, doy = wkz["yrs"], wkz["O"], wkz["doy"] z = np.load(B.TX_NPZ) reg, lat, lon = z["reg"], z["lat"][z["jj"]], z["lon"][z["ii"]] thr_all = np.load(B.thr_path("train"))["thr"] thr = np.stack([thr_all[doy[:, w] - 1] for w in range(len(C.WEEKS))], 1) pts = pick_points(reg, lat, lon) if a.ipoint < 0: print("21 diem:", pts.tolist()); return pt = int(pts[a.ipoint]) wk_list = [int(x) for x in a.weeks.split(",")] print("diem %d (o luoi %d, vung %d, %.2fN %.2fE) track=%s N=%d" % (a.ipoint, pt, reg[pt], lat[pt], lon[pt], a.track, NMEMBER), flush=True) res = run_point(pt, wk_list, yrs, doy, O, thr, a.ncore, a.track) res["_meta"] = {"ipoint": a.ipoint, "point": pt, "region": int(reg[pt]), "lat": float(lat[pt]), "lon": float(lon[pt]), "track": a.track, "NMEMBER": NMEMBER, "RECALL1": RECALL1, "NFOLD": NFOLD} outdir = os.path.join(RES, a.track) os.makedirs(outdir, exist_ok=True) json.dump(res, open(os.path.join(outdir, "pt%02d.json" % a.ipoint), "w"), indent=1) print("saved", os.path.join(outdir, "pt%02d.json" % a.ipoint)) if __name__ == "__main__": main()