#!/usr/bin/env python3 """EHD forecast model using the full 1.5 deg predictor set. Established earlier in this session (developing.txt A18/A19): thresholding a downscaled Tx reaches CSI ~0.14, and training a classifier on the SAME information does not help -- the limit is how much the inputs know, not how the problem is posed. This script therefore brings in every predictor: the 46 archived fields, the derived quantities of table 1 of HW_HPC_1_cla.docx, and the downscaled Tx. Predictors are z-scored against a (week, day-of-year) climatology built from the training years of each fold, so a test year influences neither the standardisation nor the model. Decision thresholds are picked per (region, week) on an inner split of the training years, at the operating point the user fixed: POD >= 0.6 and FAR <= 0.2, maximising CSI among the cuts that satisfy both. Usage: ehd_full.py [--subsample 6000000] [--folds 5] """ import argparse import glob import json import os import sys import numpy as np import xarray as xr import lightgbm as lgb ROOT = "/home/tanpv/Progs/claude/HW_EC_10" PRED = f"{ROOT}/outputs/proxy/ec15_predictors.nc" TXDS = f"{ROOT}/outputs/downscale/pred/pred_bcsd.npy" OBS = f"{ROOT}/outputs/OBS_4EC15/tx_obs_aligned.nc" STATIC = f"{ROOT}/outputs/Cli/static_fields_0.25.nc" CLI = f"{ROOT}/outputs/Cli" WEEKS = [("W1", 0, 11), ("W2", 11, 18), ("W3", 18, 25), ("W4", 25, 32), ("W5", 32, 39), ("W6", 39, 46)] FOLD_START = [2003, 2007, 2011, 2015, 2019] NDOY = 365 TARGET_POD, TARGET_FAR = 0.60, 0.20 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from downscale import bilinear_matrix, fill_coarse_nan # noqa: E402 from downscale_ai import window_climatology # noqa: E402 from evaluate_ehd import contingency # noqa: E402 def add_derived(d): """Derived predictors of table 1: thickness, stability, moisture transport.""" out = {} g = 9.80665 if {"t2m", "d2m"} <= d.keys(): out["dtd2m"] = d["t2m"] - d["d2m"] # dew-point depression if {"u10", "v10"} <= d.keys(): out["uv10"] = np.hypot(d["u10"], d["v10"]) if {"ssr", "str"} <= d.keys(): out["srn"] = d["ssr"] + d["str"] # net surface radiation if {"sshf", "slhf"} <= d.keys(): out["bowen"] = d["sshf"] / np.where(np.abs(d["slhf"]) < 1e4, np.nan, d["slhf"]) for a, b in [(1000, 500), (1000, 850), (850, 700), (925, 500), (925, 700)]: ka, kb = f"gh{a}", f"gh{b}" if ka in d and kb in d: out[f"thk{a}_{b}"] = d[kb] - d[ka] for a, b in [(850, 500), (925, 850), (700, 500)]: ka, kb = f"t{a}", f"t{b}" if ka in d and kb in d: out[f"dt{a}_{b}"] = d[ka] - d[kb] # lapse / stability # integrated vapour transport over the levels available levs = [L for L in (1000, 925, 850, 700, 500) if f"q{L}" in d and f"u{L}" in d and f"v{L}" in d] if len(levs) >= 2: iu = iv = 0.0 for i in range(len(levs) - 1): dp = (levs[i] - levs[i + 1]) * 100.0 iu = iu + 0.5 * (d[f"q{levs[i]}"] * d[f"u{levs[i]}"] + d[f"q{levs[i+1]}"] * d[f"u{levs[i+1]}"]) * dp / g iv = iv + 0.5 * (d[f"q{levs[i]}"] * d[f"v{levs[i]}"] + d[f"q{levs[i+1]}"] * d[f"v{levs[i+1]}"]) * dp / g out["ivt_u"], out["ivt_v"] = iu, iv out["ivt"] = np.hypot(iu, iv) return out PARTS = f"{ROOT}/outputs/downscale/fold_parts" def main(): ap = argparse.ArgumentParser() ap.add_argument("--folds", type=int, default=5) ap.add_argument("--subsample", type=int, default=6_000_000) ap.add_argument("--cover", type=float, default=0.20) ap.add_argument("--out", default=f"{ROOT}/outputs/downscale/ehd_full.json") ap.add_argument("--seed", type=int, default=0) # Multi-node fold parallelism (see memory hw-ec-10-multinode): # --fold K : train ONLY fold K on this node, dump its slice to PARTS/fold_K.npz # --reduce : load all PARTS/fold_*.npz, apply cuts, evaluate, write JSON # neither : run every fold in this process (original single-node behaviour) ap.add_argument("--fold", type=int, default=None) ap.add_argument("--reduce", action="store_true") ap.add_argument("--part-dir", default=PARTS) ap.add_argument("--pred", default=PRED, help="predictor cube; proxy by default, or the real 1.5 deg " "outputs/real/ec15_predictors.nc") ap.add_argument("--txds", default=TXDS, help="downscaled forecast Tx (sev_ds/tx_ds features); proxy BCSD " "by default, or the real-data BCSD to match --pred") ap.add_argument("--pcs", default=None, help="optional per-fold wide-domain EOF PC cube " "(outputs/real/wide_pcs.npz from build_wide_pcs.py)") args = ap.parse_args() o = xr.open_dataset(OBS) st = xr.open_dataset(STATIC) p = xr.open_dataset(args.pred) land = (np.isfinite(st.vn.values) & (st.vn.values > 0)).ravel() land_idx = np.where(land)[0] nland = len(land_idx) reg = st.reg.values.ravel()[land_idx] ny, ni, nl = o.sizes["year"], o.sizes["init"], o.sizes["lead"] obs = o.tx_obs.values.reshape(ny, ni, nl, -1)[..., land_idx] doy = o.doy.values txp = np.load(args.txds) # optional per-fold wide-domain EOF PCs (large-scale circulation features) pcs = None; pc_fields = []; pc_names = [] if args.pcs: z = np.load(args.pcs, allow_pickle=True) pc_names = [str(x) for x in z["feat_names"]] pc_fields = [k for k in z.files if k != "feat_names"] pcs = {k: z[k] for k in pc_fields} # each (folds, ny, ni, nl, nmode) print(f"loaded {len(pc_names)} PC features ({pc_fields}) from {args.pcs}", flush=True) thr = np.empty_like(obs) for k, y0 in enumerate(FOLD_START): f = np.arange(k * 4, k * 4 + 4) t = xr.open_dataset(f"{CLI}/Qp_tx_obs_fold_{y0}.nc" ).tx_thres_obs.values.reshape(NDOY, -1)[:, land_idx] thr[f] = t[doy[f] - 1] ehd_obs = obs >= thr # coarse predictor cube, plus derived fields, all on the 1.5 deg grid clat, clon = p.lat.values, p.lon.values ncell = len(clat) * len(clon) raw = {v: p[v].values.reshape(ny, ni, nl, ncell).astype("float32") for v in p.data_vars} raw.update(add_derived(raw)) for v in ["tp", "cp"]: if v in raw: raw[v] = np.log1p(np.maximum(raw[v], 0)) # precipitation -> ln(1+x) names = sorted(raw) print(f"{len(names)} coarse predictors, {nland} land points", flush=True) Wm = bilinear_matrix(clat, clon, o.lat.values, o.lon.values)[land_idx] # (nland, ncell) for v in names: raw[v] = fill_coarse_nan(raw[v].reshape(-1, ncell), clat, clon)[0 ].reshape(ny, ni, nl, ncell) stat = np.stack([st.dz.values.ravel(), st.orog.values.ravel(), st.orog_std.values.ravel(), st.lat2d.values.ravel(), st.lon2d.values.ravel()], axis=1)[land_idx].astype("float32") FEAT = (["sev_ds", "tx_ds", "thr"] + [f"z_{v}" for v in names] + ["dz", "orog", "orog_std", "lat", "lon", "lead", "doy_sin", "doy_cos", "reg"] + pc_names) print(f"{len(FEAT)} features", flush=True) wk_of_lead = np.zeros(nl, dtype=int) for w, (_, l0, l1) in enumerate(WEEKS): wk_of_lead[l0:l1] = w per = ny // args.folds prob = np.full((ny, ni, nl, nland), np.nan, dtype="float32") prob_val = np.full((ny, ni, nl, nland), np.nan, dtype="float32") cuts, gains = {}, np.zeros(len(FEAT)) os.makedirs(args.part_dir, exist_ok=True) if args.reduce: fold_list = [] # skip training; load parts below elif args.fold is not None: fold_list = [args.fold] # one fold on this node else: fold_list = list(range(args.folds)) # all folds in this process for fi in fold_list: rng = np.random.default_rng([args.seed, fi]) # per-fold, node-independent te = np.arange(fi * per, (fi + 1) * per) tr_all = np.array([i for i in range(ny) if i not in te]) inner_val = tr_all[rng.permutation(len(tr_all))[:max(1, len(tr_all) // 4)]] inner_fit = np.array([i for i in tr_all if i not in inner_val]) # climatology of every predictor on the COARSE grid, training years only cmean = {}; cstd = {} for v in names: mu = np.empty((len(WEEKS), NDOY, ncell), dtype="float32") sd = np.empty_like(mu) for w, (_, l0, l1) in enumerate(WEEKS): d_ = doy[inner_fit][:, :, l0:l1].ravel() x = raw[v][inner_fit][:, :, l0:l1].reshape(-1, ncell) mu[w] = window_climatology(x, d_, ncell) sd[w] = np.sqrt(np.maximum( window_climatology(x ** 2, d_, ncell) - mu[w] ** 2, 1e-6)) cmean[v], cstd[v] = mu, sd print(f"fold {fi+1}: predictor climatologies built", flush=True) def build(years, rows): R = len(rows) dd = doy[years].reshape(-1)[rows] ll = np.tile(np.arange(nl), len(years) * ni)[rows] wk = wk_of_lead[ll] tp_ = txp[years].reshape(-1, nland)[rows] th_ = thr[years].reshape(-1, nland)[rows] X = np.empty((R * nland, len(FEAT)), dtype="float32") X[:, 0] = (tp_ - th_).ravel(); X[:, 1] = tp_.ravel(); X[:, 2] = th_.ravel() for j, v in enumerate(names): zc = (raw[v][years].reshape(-1, ncell)[rows] - cmean[v][wk, dd - 1]) \ / cstd[v][wk, dd - 1] X[:, 3 + j] = (zc @ Wm.T).ravel() b = 3 + len(names) for k in range(5): X[:, b + k] = np.tile(stat[:, k], R) X[:, b + 5] = np.repeat(ll / 45.0, nland) X[:, b + 6] = np.repeat(np.sin(2 * np.pi * dd / NDOY), nland) X[:, b + 7] = np.repeat(np.cos(2 * np.pi * dd / NDOY), nland) X[:, b + 8] = np.tile(reg, R) if pcs is not None: # per-fold wide EOF PCs col = b + 9 for name in pc_fields: a_ = pcs[name][fi][years].reshape(-1, pcs[name].shape[-1])[rows] for k in range(a_.shape[1]): X[:, col] = np.repeat(a_[:, k], nland); col += 1 y = ehd_obs[years].reshape(-1, nland)[rows].ravel() return X, y ntime = len(inner_fit) * ni * nl want = max(1, args.subsample // nland) rows = np.sort(rng.choice(ntime, min(want, ntime), replace=False)) Xtr, ytr = build(inner_fit, rows) ok = np.isfinite(Xtr).all(axis=1) print(f"fold {fi+1}: fit on {ok.sum():,} rows, base {ytr[ok].mean():.3f}", flush=True) clf = lgb.LGBMClassifier(n_estimators=600, learning_rate=0.05, num_leaves=127, min_child_samples=200, subsample=0.8, subsample_freq=1, colsample_bytree=0.7, reg_lambda=1.0, n_jobs=16, verbose=-1) clf.fit(Xtr[ok], ytr[ok], feature_name=FEAT) g = clf.booster_.feature_importance(importance_type="gain") gains += g / max(g.sum(), 1) del Xtr, ytr for years, dest in ((inner_val, prob_val), (te, prob)): nt = len(years) * ni * nl out = np.empty((nt, nland), dtype="float32") for a in range(0, nt, 3000): r = np.arange(a, min(a + 3000, nt)) X, _ = build(years, r) m = np.isfinite(X).all(axis=1) pp = np.full(len(X), np.nan, dtype="float32") if m.any(): pp[m] = clf.predict_proba(X[m])[:, 1] out[a:a + len(r)] = pp.reshape(len(r), nland) dest[years] = out.reshape(len(years), ni, nl, nland) print(f"fold {fi+1}/{args.folds} scored", flush=True) # Tune the gridpoint cut on the REGIONAL EHDE metric that is actually # evaluated (coverage > args.cover), not on gridpoint POD/FAR. Priority: # 1) cuts satisfying POD>=0.6 & FAR<=0.2 -> max CSI # 2) else cuts with FBI<=1.0 -> max CSI (respect FBI~0.75 policy: # at long lead, under-forecast rather than blow FBI up to 3-4) # 3) else best CSI overall for wn, l0, l1 in WEEKS: for r in range(1, 8): m = reg == r if m.sum() == 0: continue pv = prob_val[inner_val][:, :, l0:l1][..., m] # (yv,ni,wl,npts) ov = ehd_obs[inner_val][:, :, l0:l1][..., m] fin = np.isfinite(pv) cnt = fin.sum(-1) cov_o = np.nanmean(ov, axis=-1).ravel() keep = np.isfinite(cov_o) & (cnt.ravel() > 0) cov_o = cov_o[keep] ehde_o = cov_o > args.cover best_c, best_csi = 0.5, -1.0 feas_c, feas_csi = None, -1.0 cap_c, cap_csi = None, -1.0 for cc in np.arange(0.02, 0.99, 0.01): cov_p = ((fin & (pv >= cc)).sum(-1) / np.maximum(cnt, 1)).ravel()[keep] ehde_p = cov_p > args.cover H = (ehde_p & ehde_o).sum(); M = (~ehde_p & ehde_o).sum() F = (ehde_p & ~ehde_o).sum() pod = H / max(H + M, 1); far = F / max(H + F, 1) fbi = (H + F) / max(H + M, 1); csi = H / max(H + M + F, 1) if csi > best_csi: best_csi, best_c = csi, cc if fbi <= 1.0 and csi > cap_csi: cap_csi, cap_c = csi, cc if pod >= TARGET_POD and far <= TARGET_FAR and csi > feas_csi: feas_csi, feas_c = csi, cc cuts[(fi, wn, r)] = (feas_c if feas_c is not None else cap_c if cap_c is not None else best_c) if args.fold is not None: # dump this node's fold fkeys = [(wn, r) for (ff, wn, r) in cuts if ff == fi] np.savez(os.path.join(args.part_dir, f"fold_{fi}.npz"), te=te, inner_val=inner_val, prob_te=prob[te], probval_iv=prob_val[inner_val], gains=gains, cut_keys=np.array([f"{wn}|{r}" for wn, r in fkeys]), cut_vals=np.array([cuts[(fi, wn, r)] for wn, r in fkeys], dtype="float32")) print(f"fold {fi}: wrote {args.part_dir}/fold_{fi}.npz", flush=True) if args.fold is not None: return 0 # part saved; the --reduce step finalises if args.reduce: # gather all node parts parts = sorted(glob.glob(os.path.join(args.part_dir, "fold_*.npz"))) if len(parts) < args.folds: print(f"WARNING: {len(parts)} parts for {args.folds} folds", flush=True) for pth in parts: z = np.load(pth) te = z["te"]; iv = z["inner_val"] fi = int(te[0] // per) prob[te] = z["prob_te"]; prob_val[iv] = z["probval_iv"] gains += z["gains"] for k, val in zip(z["cut_keys"], z["cut_vals"]): wn, r = str(k).split("|") cuts[(fi, wn, int(r))] = float(val) print(f"reduced {len(parts)} fold parts", flush=True) ehd_pred = np.zeros_like(ehd_obs) for fi in range(args.folds): te = np.arange(fi * per, (fi + 1) * per) for wn, l0, l1 in WEEKS: for r in range(1, 8): m = reg == r if m.sum() == 0: continue idx = np.where(m)[0] ehd_pred[np.ix_(te, np.arange(ni), np.arange(l0, l1), idx)] = \ prob[te][:, :, l0:l1][..., m] >= cuts[(fi, wn, r)] valid = np.isfinite(prob) & np.isfinite(obs) res = {"gridpoint": {}, "regional_EHDE": {}} for wn, l0, l1 in WEEKS + [("ALL", 0, nl)]: v = valid[:, :, l0:l1] res["gridpoint"][wn] = contingency(ehd_pred[:, :, l0:l1][v], ehd_obs[:, :, l0:l1][v]) po, pp = [], [] for r in range(1, 8): m = reg == r if m.sum() == 0: continue po.append(ehd_obs[:, :, l0:l1][..., m].mean(axis=-1).ravel()) pp.append(ehd_pred[:, :, l0:l1][..., m].mean(axis=-1).ravel()) res["regional_EHDE"][wn] = contingency(np.concatenate(pp) > args.cover, np.concatenate(po) > args.cover) for lvl in ("gridpoint", "regional_EHDE"): for wn in ["W1", "W3", "W6", "ALL"]: v = res[lvl][wn] print(f"{lvl:14s} {wn:4s} base={v['base']:.3f} POD={v['POD']:.3f} " f"POFD={v['POFD']:.3f} FAR={v['FAR']:.3f} FBI={v['FBI']:.3f} " f"CSI={v['CSI']:.3f}", flush=True) gains /= args.folds order = np.argsort(-gains) res["feature_gain"] = {FEAT[i]: float(gains[i]) for i in order[:40]} print("\ntop features:") for i in order[:20]: print(f" {FEAT[i]:16s} {gains[i]:.4f}") os.makedirs(os.path.dirname(args.out), exist_ok=True) with open(args.out, "w") as fh: json.dump({"full": res, "cuts": {f"{k[0]}|{k[1]}|S{k[2]}": float(v) for k, v in cuts.items()}}, fh, indent=2) np.save(f"{ROOT}/outputs/downscale/pred/prob_ehd_full.npy", prob) np.save(f"{ROOT}/outputs/downscale/pred/probval_ehd_full.npy", prob_val) np.save(f"{ROOT}/outputs/downscale/pred/ehdobs_full.npy", ehd_obs.astype("uint8")) np.save(f"{ROOT}/outputs/downscale/pred/reg_full.npy", reg.astype("int8")) print(f"wrote {args.out}") return 0 if __name__ == "__main__": sys.exit(main())