"""Build the ML training table from the assembled reforecast cube + VnGAD obs. For every (model_version_date, hindcast_year, lead_day, valid 0.25deg cell) it assembles predictors and targets, mirroring the B1 chain but keeping ALL candidate predictors so the ablation in scripts/11 can select column subsets: context : lead, sin_doy, cos_doy, lat, lon, thr (thr = fixed P90 EHD threshold) raw : fc_tx (forecast daily-max Tx, coarse->bilinear to the cell) local : t2m, msl, u10, v10 (local surface, interp to cell) large : {t,gh,u,v,q}{850,700,500,200} (circulation interp to cell) + {..}_dm (large-domain spatial mean) targets : y_tx (obs Tx, degC) and y_ehd (obs Tx >= thr) Output: one compressed .npz per model-version date in data/processed/ml/, each holding only rows with a matched finite observation on a national cell. Columns and their group tags are saved so the trainer can build raw/+local/+large sets. Usage: python3 scripts/10_build_features.py --dates all (or N for a smoke run) """ import argparse import glob import os import sys import numpy as np import xarray as xr sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) from targets.climatology import threshold_for_dates, day_of_year_noleap # noqa: E402 from utils.timealign import align_by_date # noqa: E402 ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) DATA = os.path.join(ROOT, "data") PROC = os.path.join(DATA, "processed") CUBE = os.path.join(PROC, "cube") OUT = os.path.join(PROC, "ml") os.makedirs(OUT, exist_ok=True) PL_VARS = ["t", "gh", "u", "v", "q"] PL_LEVS = [850, 700, 500, 200] def log(*a): print(*a, flush=True) def _mmdd(path): return os.path.basename(path).split("_mv2023-")[1].replace(".nc", "") def _interp(da, vlat, vlon): """Bilinear interp of (lead,[level,]lat,lon) to the VnGAD 0.25deg grid.""" return da.interp(lat=vlat, lon=vlon, method="linear") def build_columns(): cols = ["lead", "sin_doy", "cos_doy", "lat", "lon", "thr", "fc_tx", "t2m", "msl", "u10", "v10"] grp = ["ctx"] * 6 + ["raw"] + ["loc"] * 4 for v in PL_VARS: for lv in PL_LEVS: cols.append(f"{v}{lv}"); grp.append("large") for v in PL_VARS: for lv in PL_LEVS: cols.append(f"{v}{lv}_dm"); grp.append("large") return cols, grp COLS, GRP = build_columns() def process_mvdate(mmdd, obs_tx, obs_days, thr_clim, valid, vlat, vlon): # load the four group cubes for this model-version date tx = xr.open_dataset(os.path.join(CUBE, "sfc_tx", f"sfc_tx_mv2023-{mmdd}.nc")) t2 = xr.open_dataset(os.path.join(CUBE, "sfc_t2m", f"sfc_t2m_mv2023-{mmdd}.nc")) dy = xr.open_dataset(os.path.join(CUBE, "sfc_dyn", f"sfc_dyn_mv2023-{mmdd}.nc")) pl = xr.open_dataset(os.path.join(CUBE, "pl_hi", f"pl_hi_mv2023-{mmdd}.nc")) years = tx["hindcast_year"].values leads = tx["lead"].values nL = len(leads) # number of lead days (LT0..LT45) vmask2d = valid # (61,31) national cells ii, jj = np.nonzero(vmask2d) # cell indices latcell, loncell = vlat[ii], vlon[jj] # large-domain spatial means (per hindcast_year, lead, level) -> broadcast pl_dm = {v: pl[v].mean(dim=("lat", "lon")).values for v in PL_VARS} # (yr,lead,lev) rows_X, rows_tx, rows_ehd, rows_meta = [], [], [], [] for yi, yr in enumerate(years): init = np.datetime64(f"{int(yr):04d}-{mmdd}") vdays = init + leads.astype("timedelta64[D]") # (lead,) # obs + threshold on valid days oidx, om = align_by_date(vdays, obs_days) otx = obs_tx[oidx] # (lead,61,31) thr = threshold_for_dates(thr_clim, vdays) # (lead,61,31) # interp forecast fields to VnGAD grid for this hindcast year f_tx = _interp(tx["mx2t6"].isel(hindcast_year=yi), vlat, vlon).values f_t2 = _interp(t2["2t"].isel(hindcast_year=yi), vlat, vlon).values f_ms = _interp(dy["msl"].isel(hindcast_year=yi), vlat, vlon).values f_u = _interp(dy["10u"].isel(hindcast_year=yi), vlat, vlon).values f_v = _interp(dy["10v"].isel(hindcast_year=yi), vlat, vlon).values f_pl = {v: _interp(pl[v].isel(hindcast_year=yi), vlat, vlon).values for v in PL_VARS} # (lead,lev,61,31) doy = day_of_year_noleap(vdays) sin_d = np.sin(2 * np.pi * doy / 365.0) cos_d = np.cos(2 * np.pi * doy / 365.0) for li in range(nL): m = om[li] & valid & np.isfinite(otx[li]) & np.isfinite(f_tx[li]) if not m.any(): continue ci, cj = np.nonzero(m) n = ci.size feat = [np.full(n, leads[li], "float32"), np.full(n, sin_d[li], "float32"), np.full(n, cos_d[li], "float32"), vlat[ci].astype("float32"), vlon[cj].astype("float32"), thr[li][ci, cj].astype("float32"), f_tx[li][ci, cj].astype("float32"), f_t2[li][ci, cj].astype("float32"), f_ms[li][ci, cj].astype("float32"), f_u[li][ci, cj].astype("float32"), f_v[li][ci, cj].astype("float32")] for vi, v in enumerate(PL_VARS): for lvi in range(len(PL_LEVS)): feat.append(f_pl[v][li, lvi][ci, cj].astype("float32")) for v in PL_VARS: for lvi in range(len(PL_LEVS)): feat.append(np.full(n, pl_dm[v][yi, li, lvi], "float32")) rows_X.append(np.stack(feat, 1)) rows_tx.append(otx[li][ci, cj].astype("float32")) rows_ehd.append((otx[li][ci, cj] >= thr[li][ci, cj]).astype("int8")) # meta: year, cell_id (flattened i*31+j), valid_date (days since epoch) cell_id = (ci * vmask2d.shape[1] + cj).astype("int32") rows_meta.append(np.stack([ np.full(n, int(yr), "int32"), cell_id, np.full(n, vdays[li].astype("datetime64[D]").astype("int32"), "int32")], 1)) for d in (tx, t2, dy, pl): d.close() if not rows_X: return None return (np.concatenate(rows_X), np.concatenate(rows_tx), np.concatenate(rows_ehd), np.concatenate(rows_meta)) def main(): ap = argparse.ArgumentParser() ap.add_argument("--dates", default="all", help="N (first N) or 'all'") ap.add_argument("--force", action="store_true") a = ap.parse_args() obs = xr.open_dataset(os.path.join(DATA, "VnGAD_daily-0.25_1961-2023_tasmax.nc")) obs_days = obs["time"].values.astype("datetime64[D]") obs_tx = obs["tasmax"].values vlat, vlon = obs["lat"].values, obs["lon"].values thr_clim = xr.open_dataset(os.path.join(PROC, "EHD_THRESHOLD_2003_2022.nc"))["EHD_THRESHOLD"].values valid = xr.open_dataset(os.path.join(PROC, "region_codes_vngad.nc"))["region"].values >= 0 obs.close() log(f"VnGAD {obs_tx.shape}, {int(valid.sum())} national cells; {len(COLS)} features") files = sorted(glob.glob(os.path.join(CUBE, "sfc_tx", "sfc_tx_mv2023-*.nc")), key=_mmdd) mmdds = [_mmdd(f) for f in files] if a.dates != "all": mmdds = mmdds[:int(a.dates)] np.save(os.path.join(OUT, "columns.npy"), np.array(COLS)) np.save(os.path.join(OUT, "column_groups.npy"), np.array(GRP)) for i, mmdd in enumerate(mmdds, 1): dst = os.path.join(OUT, f"feat_mv2023-{mmdd}.npz") if os.path.exists(dst) and not a.force: log(f"[{i}/{len(mmdds)}] skip {mmdd}"); continue out = process_mvdate(mmdd, obs_tx, obs_days, thr_clim, valid, vlat, vlon) if out is None: log(f"[{i}/{len(mmdds)}] {mmdd}: no rows"); continue X, ytx, yehd, meta = out np.savez_compressed(dst, X=X, y_tx=ytx, y_ehd=yehd, meta=meta) log(f"[{i}/{len(mmdds)}] {mmdd}: {X.shape[0]} rows -> {os.path.getsize(dst)/1e6:.1f} MB") log("DONE build_features") if __name__ == "__main__": main()