#!/usr/bin/env python """Stage 20c — leakage-safe feature screening / selection over the candidate pool (stage 20b). For each lead block we rank the 28 candidate predictors by their genuine, OUT-OF-SAMPLE contribution and then select a stable subset — the expert screening step: 1. Nested LOYO permutation importance. For each held-out year y, an HGB classifier is trained on the other 19 years (train capped for speed; HGB handles NaN natively). Permutation importance is measured ON THE HELD-OUT YEAR with AUC scoring: the drop in held-out AUC when a predictor is shuffled. Because the model never saw y, this is an honest generalisation importance, not an in-sample gain. We aggregate mean +/- std across the 20 folds -> a consensus ranking and a stability estimate. 2. Univariate mutual information (filter) with the label, on a pooled subsample, as a model-free cross-check (NaNs median-imputed for MI only). 3. Selection rule (reported, not silent): keep every base forecast feature (the Stage-1 backbone) plus any candidate whose mean held-out permutation importance is positive and stable (mean > std). The selected set per block is written for the retrain. Outputs: results/tables/feature_importance_by_block.csv feature x block: perm mean/std, MI, rank, selected results/tables/feature_selection_by_block.csv per block: selected feature list + counts results/figures/feature_importance_heatmap.png perm-importance heatmap (feature x block) Descriptive ranking uses the stage-20b cache (full-baseline threshold). The honest skill of the selected set comes from the nested per-fold retrain, not from this stage. """ from __future__ import annotations import os, sys, glob, argparse, time import numpy as np import pandas as pd sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) from ehd.config import load_config, lead_blocks from sklearn.ensemble import HistGradientBoostingClassifier from sklearn.inspection import permutation_importance from sklearn.metrics import roc_auc_score from sklearn.feature_selection import mutual_info_classif HERE = os.path.dirname(os.path.abspath(__file__)) ROOT = os.path.dirname(HERE) TABLES = os.path.join(ROOT, "results", "tables") FIGS = os.path.join(ROOT, "results", "figures") BASE_BACKBONE = ["f_max", "f_mean", "f_min", "f_exceed", "f_maxanom", "lead_id", "sin_doy", "cos_doy", "lat", "lon"] IMP_FLOOR = 0.003 # noise floor on mean out-of-sample AUC-drop importance (see 20c docstring) def load_cache(cfg): paths = sorted(glob.glob(os.path.join(cfg["paths"]["outputs"], "candidates", "cand_*.npz"))) if not paths: raise SystemExit("no outputs/candidates/cand_*.npz — run stage 20b first") years, data, names = [], {}, None for p in paths: d = np.load(p, allow_pickle=True) yr = int(d["year"]); years.append(yr) data[yr] = dict(X=d["X"], y=d["y"], block_id=d["block_id"]) names = [str(s) for s in d["feat_names"]] return sorted(years), data, names def main(): ap = argparse.ArgumentParser() ap.add_argument("--cap", type=int, default=150000) ap.add_argument("--repeats", type=int, default=4) args = ap.parse_args() cfg = load_config(); blocks = lead_blocks(cfg) block_names = list(blocks); block_ids = {b: k + 1 for k, b in enumerate(block_names)} rng = np.random.default_rng(cfg["seed"]) years, data, names = load_cache(cfg) nfeat = len(names) print(f"[20c] {len(years)} folds, {nfeat} candidate features, blocks={block_names}") rows = [] for b in block_names: bid = block_ids[b] # per-fold out-of-sample permutation importance perm = np.full((len(years), nfeat), np.nan) for fi, yr in enumerate(years): tr_years = [t for t in years if t != yr] Xtr = np.concatenate([data[t]["X"][data[t]["block_id"] == bid] for t in tr_years]) ytr = np.concatenate([data[t]["y"][data[t]["block_id"] == bid] for t in tr_years]) m = data[yr]["block_id"] == bid Xte = data[yr]["X"][m]; yte = data[yr]["y"][m] ok_tr = np.isfinite(ytr); ok_te = np.isfinite(yte) Xtr, ytr = Xtr[ok_tr], ytr[ok_tr].astype(int) Xte, yte = Xte[ok_te], yte[ok_te].astype(int) if args.cap and Xtr.shape[0] > args.cap: idx = rng.choice(Xtr.shape[0], args.cap, replace=False) Xtr, ytr = Xtr[idx], ytr[idx] if len(np.unique(yte)) < 2: continue clf = HistGradientBoostingClassifier(max_iter=200, learning_rate=0.08, max_leaf_nodes=31, random_state=0) clf.fit(Xtr, ytr) # subsample held-out for permutation speed if Xte.shape[0] > 60000: s = rng.choice(Xte.shape[0], 60000, replace=False); Xpe, ype = Xte[s], yte[s] else: Xpe, ype = Xte, yte r = permutation_importance(clf, Xpe, ype, scoring="roc_auc", n_repeats=args.repeats, random_state=0, n_jobs=1) perm[fi] = r.importances_mean pmean = np.nanmean(perm, axis=0); pstd = np.nanstd(perm, axis=0) # univariate MI on a pooled subsample (median-impute NaN for MI only) Xall = np.concatenate([data[t]["X"][data[t]["block_id"] == bid] for t in years]) yall = np.concatenate([data[t]["y"][data[t]["block_id"] == bid] for t in years]) ok = np.isfinite(yall); Xall, yall = Xall[ok], yall[ok].astype(int) if Xall.shape[0] > 120000: s = rng.choice(Xall.shape[0], 120000, replace=False); Xall, yall = Xall[s], yall[s] col_med = np.nanmedian(Xall, axis=0) Xmi = np.where(np.isnan(Xall), col_med, Xall) mi = mutual_info_classif(Xmi, yall, discrete_features=False, random_state=0) order = np.argsort(-pmean) rank = np.empty(nfeat, int); rank[order] = np.arange(1, nfeat + 1) for j, nm in enumerate(names): # Selection: keep the Stage-1 forecast backbone, plus any candidate whose MEAN # out-of-sample importance clears the noise floor. We do NOT gate on stability # (mean>std): episodic teleconnections (ENSO/MJO) are the strongest long-lead # predictors yet legitimately have high year-to-year variance, so std is reported # as a diagnostic, not a filter. selected = (nm in BASE_BACKBONE) or (pmean[j] >= IMP_FLOOR) rows.append(dict(block=b, feature=nm, perm_imp_mean=round(float(pmean[j]), 5), perm_imp_std=round(float(pstd[j]), 5), mi=round(float(mi[j]), 5), rank=int(rank[j]), robust=bool(pmean[j] > pstd[j]), selected=bool(selected))) top = [names[k] for k in order[:8]] print(f" {b}: top-8 by perm-AUC importance -> {top}") df = pd.DataFrame(rows) os.makedirs(TABLES, exist_ok=True) df.to_csv(os.path.join(TABLES, "feature_importance_by_block.csv"), index=False) sel = (df[df.selected].groupby("block")["feature"].apply(list).reindex(block_names)) selrows = [dict(block=b, n_selected=len(sel[b]), selected=";".join(sel[b])) for b in block_names] pd.DataFrame(selrows).to_csv(os.path.join(TABLES, "feature_selection_by_block.csv"), index=False) print("[20c] wrote feature_importance_by_block.csv, feature_selection_by_block.csv") _heatmap(df, block_names, names) def _heatmap(df, block_names, names): import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt piv = df.pivot(index="feature", columns="block", values="perm_imp_mean").reindex( index=names, columns=block_names) fig, ax = plt.subplots(figsize=(7, 9)) im = ax.imshow(piv.values, aspect="auto", cmap="viridis") ax.set_xticks(range(len(block_names))); ax.set_xticklabels(block_names) ax.set_yticks(range(len(names))); ax.set_yticklabels(names, fontsize=8) ax.set_title("Out-of-sample permutation importance (AUC drop)\nby lead block") fig.colorbar(im, ax=ax, shrink=0.6, label="mean AUC drop when shuffled") os.makedirs(FIGS, exist_ok=True) fig.tight_layout(); fig.savefig(os.path.join(FIGS, "feature_importance_heatmap.png"), dpi=130) print("[20c] wrote results/figures/feature_importance_heatmap.png") if __name__ == "__main__": main()