"""Train & evaluate Tx forecast-correction / EHD models on the assembled cube. Design (locked with the user, 2026-10-17): * dual heads : Tx regression + EHD (extreme-heat-day) classification; * validation : leave-one-year-out over the 20 hindcast years (out-of-sample predictions pooled across folds, scored by lead day); * ablation : predictor scope raw -> +local -> +large (which is best?); * models : Ridge (linear MOS) and HistGradientBoosting (trees) for Tx; HistGradientBoosting classifier for EHD probability. Baseline = raw ECMWF forecast (fc_tx) thresholded, i.e. the B1 chain, computed on the identical rows for an apples-to-apples comparison. Outputs: outputs/metrics/model_ablation_by_lead.csv, outputs/metrics/model_summary.json, outputs/figures/model_leadcurves.png, outputs/MODEL_RESULTS.html. Usage: python3 scripts/11_train_eval.py [--train-cap 5000000] [--max-iter 250] """ import argparse import glob import json import os import sys import numpy as np import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt from sklearn.ensemble import HistGradientBoostingRegressor, HistGradientBoostingClassifier from sklearn.linear_model import Ridge from sklearn.preprocessing import StandardScaler ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) ML = os.path.join(ROOT, "data", "processed", "ml") OUTM = os.path.join(ROOT, "outputs", "metrics") OUTF = os.path.join(ROOT, "outputs", "figures") for d in (OUTM, OUTF): os.makedirs(d, exist_ok=True) MAXLT = 45 ABLATIONS = {"raw": ("ctx", "raw"), "local": ("ctx", "raw", "loc"), "large": ("ctx", "raw", "loc", "large")} def log(*a): print(*a, flush=True) def load_all(): cols = list(np.load(os.path.join(ML, "columns.npy"))) grp = list(np.load(os.path.join(ML, "column_groups.npy"))) files = sorted(glob.glob(os.path.join(ML, "feat_mv2023-*.npz"))) Xs, tx, ehd, meta = [], [], [], [] for f in files: d = np.load(f) Xs.append(d["X"]); tx.append(d["y_tx"]); ehd.append(d["y_ehd"]); meta.append(d["meta"]) X = np.concatenate(Xs); y_tx = np.concatenate(tx) y_ehd = np.concatenate(ehd).astype("int8"); meta = np.concatenate(meta) log(f"loaded {X.shape[0]:,} rows x {X.shape[1]} cols from {len(files)} files") return cols, grp, X, y_tx, y_ehd, meta def contingency(fcst_ev, obs_ev): H = int(((fcst_ev == 1) & (obs_ev == 1)).sum()); M = int(((fcst_ev == 0) & (obs_ev == 1)).sum()) FA = int(((fcst_ev == 1) & (obs_ev == 0)).sum()); CN = int(((fcst_ev == 0) & (obs_ev == 0)).sum()) pod = H / (H + M) if H + M else np.nan far = FA / (H + FA) if H + FA else np.nan csi = H / (H + M + FA) if H + M + FA else np.nan fbi = (H + FA) / (H + M) if H + M else np.nan return pod, far, csi, fbi def tx_metrics(pred, obs): e = pred - obs bias = float(e.mean()); mae = float(np.abs(e).mean()); rmse = float(np.sqrt((e * e).mean())) corr = float(np.corrcoef(pred, obs)[0, 1]) if pred.size > 2 else np.nan return bias, mae, rmse, corr def main(): ap = argparse.ArgumentParser() ap.add_argument("--train-cap", type=int, default=5_000_000) ap.add_argument("--max-iter", type=int, default=250) a = ap.parse_args() rng = np.random.default_rng(0) cols, grp, X, y_tx, y_ehd, meta = load_all() lead = X[:, cols.index("lead")].astype(int) thr = X[:, cols.index("thr")] fc_tx = X[:, cols.index("fc_tx")] year = meta[:, 0] years = np.unique(year) col_idx = {ab: [i for i, g in enumerate(grp) if g in gs] for ab, gs in ABLATIONS.items()} # out-of-sample prediction holders N = X.shape[0] pred_tx = {ab: {"ridge": np.full(N, np.nan, "float32"), "hgb": np.full(N, np.nan, "float32")} for ab in ABLATIONS} pred_ehd_p = {ab: np.full(N, np.nan, "float32") for ab in ABLATIONS} for yi, yr in enumerate(years): te = year == yr tr = ~te tr_idx = np.nonzero(tr)[0] if tr_idx.size > a.train_cap: # subsample training rows tr_idx = rng.choice(tr_idx, a.train_cap, replace=False) ytx_tr = y_tx[tr_idx]; yehd_tr = y_ehd[tr_idx] log(f"[fold {yi+1}/20] test year {yr}: train {tr_idx.size:,} test {int(te.sum()):,}") for ab in ABLATIONS: ci = col_idx[ab] Xtr = X[np.ix_(tr_idx, ci)]; Xte = X[te][:, ci] # Ridge (standardized) Tx sc = StandardScaler().fit(Xtr) rid = Ridge(alpha=10.0).fit(sc.transform(Xtr), ytx_tr) pred_tx[ab]["ridge"][te] = rid.predict(sc.transform(Xte)).astype("float32") # HGB Tx hgr = HistGradientBoostingRegressor(max_iter=a.max_iter, learning_rate=0.07, max_leaf_nodes=63, l2_regularization=1.0, random_state=0).fit(Xtr, ytx_tr) pred_tx[ab]["hgb"][te] = hgr.predict(Xte).astype("float32") # HGB EHD probability hgc = HistGradientBoostingClassifier(max_iter=a.max_iter, learning_rate=0.07, max_leaf_nodes=63, l2_regularization=1.0, random_state=0).fit(Xtr, yehd_tr) pred_ehd_p[ab][te] = hgc.predict_proba(Xte)[:, 1].astype("float32") # ----- evaluate by lead day (pooled out-of-sample) ----- log("evaluating by lead day ...") # EHD probability decision threshold: pick the global value maximizing CSI # for each config on the pooled out-of-sample predictions. def best_ehd_thresh(p, obs): best_t, best_csi = 0.5, -1 for t in np.linspace(0.05, 0.6, 12): _, _, csi, _ = contingency((p >= t).astype(int), obs) if csi is not np.nan and csi > best_csi: best_csi, best_t = csi, t return best_t ehd_t = {ab: best_ehd_thresh(pred_ehd_p[ab], y_ehd) for ab in ABLATIONS} rows = ["config,model,lead,n,tx_bias,tx_mae,tx_rmse,tx_corr,ehd_POD,ehd_FAR,ehd_CSI,ehd_FBI,ehd_brier"] def emit(config, model, mask_tx_pred, ehd_fcst, ehd_prob): for l in range(MAXLT + 1): m = lead == l if not m.any(): continue if mask_tx_pred is not None: b, mae, rmse, cr = tx_metrics(mask_tx_pred[m], y_tx[m]) else: b = mae = rmse = cr = np.nan pod, far, csi, fbi = contingency(ehd_fcst[m], y_ehd[m]) brier = float(((ehd_prob[m] - y_ehd[m]) ** 2).mean()) if ehd_prob is not None else np.nan rows.append(f"{config},{model},{l},{int(m.sum())},{b:.3f},{mae:.3f},{rmse:.3f},{cr:.3f}," f"{pod:.3f},{far:.3f},{csi:.3f},{fbi:.3f},{brier:.4f}") # baseline: raw forecast (B1) emit("baseline", "raw_ecmwf", fc_tx, (fc_tx >= thr).astype(int), None) for ab in ABLATIONS: # Tx heads; EHD derived from corrected Tx (operational, == B1 mechanism) emit(ab, "ridge_tx", pred_tx[ab]["ridge"], (pred_tx[ab]["ridge"] >= thr).astype(int), None) emit(ab, "hgb_tx", pred_tx[ab]["hgb"], (pred_tx[ab]["hgb"] >= thr).astype(int), None) # direct EHD classifier (probability head) emit(ab, "hgb_ehd", None, (pred_ehd_p[ab] >= ehd_t[ab]).astype(int), pred_ehd_p[ab]) with open(os.path.join(OUTM, "model_ablation_by_lead.csv"), "w") as fo: fo.write("\n".join(rows) + "\n") # ----- headline summary ----- arr = np.genfromtxt(os.path.join(OUTM, "model_ablation_by_lead.csv"), delimiter=",", names=True, dtype=None, encoding=None) def pick(config, model, lead_val, field): s = arr[(arr["config"] == config) & (arr["model"] == model) & (arr["lead"] == lead_val)] return float(s[field][0]) if s.size else float("nan") def meanfield(config, model, field, lo=0, hi=45): s = arr[(arr["config"] == config) & (arr["model"] == model) & (arr["lead"] >= lo) & (arr["lead"] <= hi)] return float(np.nanmean(s[field])) if s.size else float("nan") summary = {"n_rows": int(N), "train_cap": a.train_cap, "ehd_thresholds": ehd_t, "baseline": {"tx_mae_LT0": pick("baseline", "raw_ecmwf", 0, "tx_mae"), "tx_mae_LT45": pick("baseline", "raw_ecmwf", 45, "tx_mae"), "ehd_CSI_mean": meanfield("baseline", "raw_ecmwf", "ehd_CSI")}, "configs": {}} for ab in ABLATIONS: for model in ("ridge_tx", "hgb_tx", "hgb_ehd"): summary["configs"][f"{ab}/{model}"] = { "tx_mae_LT0": pick(ab, model, 0, "tx_mae"), "tx_mae_LT7": pick(ab, model, 7, "tx_mae"), "tx_mae_LT28": pick(ab, model, 28, "tx_mae"), "tx_mae_LT45": pick(ab, model, 45, "tx_mae"), "ehd_CSI_mean": meanfield(ab, model, "ehd_CSI"), "ehd_brier_mean": meanfield(ab, model, "ehd_brier")} with open(os.path.join(OUTM, "model_summary.json"), "w") as fo: json.dump(summary, fo, indent=2, default=lambda o: None if o != o else o) # ----- figure: Tx MAE (hgb_tx) and EHD CSI (hgb_tx-derived) vs lead ----- leads = np.arange(MAXLT + 1) fig, ax = plt.subplots(1, 2, figsize=(12, 4.5)) def series(config, model, field): return np.array([pick(config, model, l, field) for l in leads]) ax[0].plot(leads, series("baseline", "raw_ecmwf", "tx_mae"), "k--", label="raw ECMWF (B1)") for ab in ABLATIONS: ax[0].plot(leads, series(ab, "hgb_tx", "tx_mae"), "-", label=f"HGB Tx [{ab}]") ax[0].set_xlabel("lead day"); ax[0].set_ylabel("Tx MAE (degC)") ax[0].set_title("Tx error vs lead"); ax[0].legend(fontsize=8) ax[1].plot(leads, series("baseline", "raw_ecmwf", "ehd_CSI"), "k--", label="raw ECMWF (B1)") for ab in ABLATIONS: ax[1].plot(leads, series(ab, "hgb_ehd", "ehd_CSI"), "-", label=f"HGB EHD [{ab}]") ax[1].set_xlabel("lead day"); ax[1].set_ylabel("EHD CSI"); ax[1].set_ylim(0, 1) ax[1].set_title("EHD skill vs lead"); ax[1].legend(fontsize=8) fig.tight_layout(); fig.savefig(os.path.join(OUTF, "model_leadcurves.png"), dpi=120) # ----- HTML report ----- def row_html(name, d): return (f"{name}{d['tx_mae_LT0']:.2f}{d.get('tx_mae_LT7', float('nan')):.2f}" f"{d.get('tx_mae_LT28', float('nan')):.2f}{d['tx_mae_LT45']:.2f}" f"{d['ehd_CSI_mean']:.3f}{d.get('ehd_brier_mean', float('nan')):.4f}") b = summary["baseline"] body = [f"raw ECMWF (B1){b['tx_mae_LT0']:.2f}––" f"{b['tx_mae_LT45']:.2f}{b['ehd_CSI_mean']:.3f}–"] body += [row_html(k, v) for k, v in summary["configs"].items()] html = f"""Tx/EHD models

Tx forecast-correction & EHD models — leave-one-year-out

Out-of-sample over 20 hindcast years ({N:,} cell-day samples), scored by lead day. Ablation: raw forecast → +local surface → +large-scale circulation. EHD = daily-max Tx ≥ fixed P90 threshold.

{''.join(body)}
config / modelTx MAE LT0LT7LT28LT45 EHD CSI (mean)EHD Brier (mean)

Tx MAE in °C (lower better); CSI higher better; Brier lower better. hgb_* = HistGradientBoosting, ridge_tx = linear MOS. EHD from *_tx rows is derived by thresholding corrected Tx; hgb_ehd is the direct probability classifier.

Full by-lead table: outputs/metrics/model_ablation_by_lead.csv

""" with open(os.path.join(ROOT, "outputs", "MODEL_RESULTS.html"), "w") as fo: fo.write(html) log("DONE train_eval"); log(json.dumps(summary, indent=2)) if __name__ == "__main__": main()