#!/usr/bin/env python3 """Assemble model-ready design matrices from the AI input store. One row = one (initialisation, lead time, grid point) triple, laid out in the sample order of section 5.3 so a row index maps back to the docx sample number. Observation variables never enter the feature matrix; they are targets only. """ from __future__ import annotations import os import sys import numpy as np sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) import ehd_common as C import build_cube as B import build_ai_inputs as A from calc_eof import EOF_DIR, NMODE from external_feats import EXT_FILE STATIC_NAMES = ["lead", "doy_cos", "doy_sin"] GEO_NAMES = ["lat", "lon", "region", "thr_obs", "thr_mod"] FEATURE_SETS = { "F0": "sv_mod, ev_mod + the three mandatory static variables", "F1": "F0 + the 40 z-scored EC variables", "F2": "F1 + the 20 EOF principal components", "F3": "F2 + grid-point geography and the local climatological thresholds", "F4": "F3 + running means of the heat-relevant fields along the lead axis", "F5": "F3 + a wavelet decomposition of those fields along the lead axis", "F6": "F3 + observed history before init and the MJO/ENSO state (PA1)", "F7": "F6 minus ENSO, which acts as a year label with only 16 training years", "F8": "F3 + the MJO state alone, the only external block with no year signal", } # Which external columns each set takes. Nino3.4 carries R^2 = 0.46 against year # identity, so with sixteen independent training years a tree can use it to # recognise WHICH YEAR a row came from and memorise that year's climate: adding it # cost lgbm_clf 0.025 AUC at W4 while it ranked as the third most important feature # of all. MJO sits at R^2 = 0.002 against year and is safe; the observed-history # block is in between (h_ehd30 at 0.30), so F8 isolates MJO on its own. ENSO_NAMES = ("x_nino34", "x_nino34_3m") MJO_NAMES = ("x_omi1", "x_omi2", "x_omi_amp", "x_omi1_m15", "x_omi2_m15") EXT_SETS = {"F6": None, "F7": "no_enso", "F8": "mjo_only"} # F6 is the first feature set carrying information from OUTSIDE the ECMWF run: # what the observed surface was doing before the forecast was issued, and where the # two slow modes that modulate a Vietnamese heat spell -- MJO and ENSO -- stood at # initialisation. Everything in F0-F5 comes from one EC run, so the model has # never known whether yesterday was hot. Built by src/external_feats.py; the # observed windows stop one day BEFORE the initialisation date, so nothing at or # after the forecast valid time is visible. # Every F0-F3 feature is the value on the target day alone. Three to six weeks # out the day-to-day detail of a forecast is mostly noise while the surrounding # regime still carries signal, so F4 adds running means along the lead axis: a # one-week and a three-week window centred on the target day, over the fields # that drive a heat extreme, plus the model severity itself. The deviation of # the day from its own window is added for the two most direct fields, because a # tree can split on it but cannot form the difference by itself. WIN_LEN = (7, 21) WIN_VARS = ["mx2t6", "t2m", "skt", "thk_1000_500", "thk_925_700", "msl", "tcc", "sm20", "cape", "tp", "dtd2m", "ssrd"] DEV_VARS = ["mx2t6", "sv_mod"] # F5 replaces F4's two fixed windows with a proper multi-resolution view of the # same lead axis: a stationary (undecimated) wavelet transform, which keeps one # coefficient per lead at every scale instead of decimating. Each row then sees # its own day (the finest detail), the synoptic scale and the slow background # separately, rather than one pre-chosen smoothing width. # # OPERATIONALLY VALID: the transform runs along the LEAD axis of a single # forecast, every value of which exists at initialisation time. A wavelet taken # along real time instead -- across successive initialisations, or over the # observed record -- would need data from after the forecast was issued near the # right-hand edge, and could not be reproduced in real time. SWT_WAVELET = "db2" SWT_LEVEL = 3 SWT_PAD = 48 # 46 leads padded up to a multiple of 2**SWT_LEVEL SWT_VARS = ["mx2t6", "t2m", "thk_1000_500", "msl", "sm20", "tcc"] def swt_features(a): """Stationary wavelet coefficients along the lead axis of (ni, nlt, nv, npt). Returns (ni, nlt, nv * (SWT_LEVEL + 1), npt): for every variable the coarse approximation followed by the detail at each scale, all still one value per lead so they line up with the rows of the design matrix. """ import pywt ni, nk, nv, npt = a.shape pad = SWT_PAD - nk lo = pad // 2 ap = np.pad(a, ((0, 0), (lo, pad - lo), (0, 0), (0, 0)), mode="edge") coeffs = pywt.swt(ap, SWT_WAVELET, level=SWT_LEVEL, axis=1, trim_approx=True, norm=True) out = np.concatenate([c[:, lo:lo + nk] for c in coeffs], axis=2) return out.astype("f4") def running_mean(a, win): """Centred running mean along the lead axis of an (ninit, nlt, ...) array. The window shrinks rather than wraps at the ends of the forecast, so lead 0 averages only the leads that exist. Done with a cumulative sum because the lead axis is short but the point axis is not. """ nk = a.shape[1] cs = np.cumsum(a, axis=1, dtype="f8") cs = np.concatenate([np.zeros_like(cs[:, :1]), cs], axis=1) h = win // 2 lo = np.clip(np.arange(nk) - h, 0, nk) hi = np.clip(np.arange(nk) + h + 1, 0, nk) n = (hi - lo).reshape([1, nk] + [1] * (a.ndim - 2)) return ((cs[:, hi] - cs[:, lo]) / n).astype("f4") class Store: """Lazy handles on the memmaps written by build_ai_inputs.py.""" def __init__(self): self.grid = C.Grid() self.st = np.load(A.STATIC, allow_pickle=True) self.ninit = int(self.st["ninit"][0]) self.npt = int(self.st["npt"][0]) n, k, p = self.ninit, C.NLT, self.npt self.ec_z = np.memmap(A.EC_Z, "f4", "r", shape=(n, k, len(B.EC_VARS), p)) self.derived = np.memmap(A.DERIVED, "f4", "r", shape=(n, k, 2, p)) self.target = np.memmap(A.TARGET, "f4", "r", shape=(n, k, 3, p)) self.pcs = {} for name in ("H500", "T850"): z = np.load(os.path.join(EOF_DIR, "eof_%s.npz" % name)) self.pcs[name] = (z["pcs"], [str(s) for s in z["names"]]) self.ext = np.load(EXT_FILE, allow_pickle=True) if os.path.exists(EXT_FILE) else None self.doy = self.st["doy"] self.init_year = self.st["init_year"] # ------------------------------------------------------------------ names def feature_names(self, fset): names = list(A.DERIVED_NAMES) + STATIC_NAMES if fset in ("F1", "F2", "F3", "F4", "F5", "F6", "F7", "F8"): names += ["z_" + v for v in B.EC_VARS] if fset in ("F2", "F3", "F4", "F5", "F6", "F7", "F8"): for k in ("H500", "T850"): names += self.pcs[k][1] if fset in ("F3", "F4", "F5", "F6", "F7", "F8"): names += GEO_NAMES if fset in EXT_SETS: names += self._ext_names(fset)[0] if fset == "F4": for w in WIN_LEN: names += ["w%d_sv_mod" % w] + ["w%d_%s" % (w, v) for v in WIN_VARS] names += ["dev_%s" % v for v in DEV_VARS] if fset == "F5": lvl = ["a%d" % SWT_LEVEL] + ["d%d" % j for j in range(SWT_LEVEL, 0, -1)] for v in ["sv_mod"] + SWT_VARS: names += ["%s_%s" % (L, v) for L in lvl] return names def _ext_names(self, fset): """External column names for `fset`, plus the index arrays that select them.""" hn = [str(x) for x in self.ext["hist_names"]] gn = [str(x) for x in self.ext["idx_names"]] mode = EXT_SETS[fset] if mode == "mjo_only": hkeep = [] gkeep = [j for j, n in enumerate(gn) if n in MJO_NAMES] elif mode == "no_enso": hkeep = list(range(len(hn))) gkeep = [j for j, n in enumerate(gn) if n not in ENSO_NAMES] else: hkeep = list(range(len(hn))) gkeep = list(range(len(gn))) return ([hn[j] for j in hkeep] + [gn[j] for j in gkeep], np.array(hkeep, dtype=int), np.array(gkeep, dtype=int)) # ------------------------------------------------------------------- rows @staticmethod def _gather(cube, isel, points, nvar): """Read (len(isel), 46, nvar, len(points)) one initialisation at a time. Slicing the memmap as cube[isel][..., points] would first materialise the whole 8 GB selection; taking one initialisation slab at a time keeps the footprint at the size of the result. """ out = np.empty((isel.size, C.NLT, nvar, points.size), dtype="f4") for a, i in enumerate(isel): out[a] = cube[i][:, :, points] return out def build(self, points, years, fset="F2"): """Design matrix and targets for the given points and initialisation years.""" isel = np.nonzero(np.isin(self.init_year, years))[0] return self.build_isel(points, isel, fset) def build_isel(self, points, isel, fset="F2"): """Same as build() but for an explicit list of initialisation indices. Callers that cannot hold the whole national matrix at once walk the initialisations in chunks through this entry point. """ points = np.asarray(points) isel = np.asarray(isel) ni, nk, npt = isel.size, C.NLT, points.size nrow = ni * nk * npt cols = [] der = self._gather(self.derived, isel, points, 2) # (ni,nk,2,npt) cols.append(der.transpose(0, 1, 3, 2).reshape(nrow, 2)) if fset not in ("F4", "F5"): del der lead = np.broadcast_to((np.arange(nk, dtype="f4") / 45.0)[None, :, None], (ni, nk, npt)).reshape(nrow, 1) doy = self.doy[isel] # (ni,nk) dcos = np.broadcast_to(np.cos(2 * np.pi * doy / 365.0).astype("f4")[:, :, None], (ni, nk, npt)).reshape(nrow, 1) dsin = np.broadcast_to(np.sin(2 * np.pi * doy / 365.0).astype("f4")[:, :, None], (ni, nk, npt)).reshape(nrow, 1) cols += [lead, dcos, dsin] if fset in ("F1", "F2", "F3", "F4", "F5", "F6", "F7", "F8"): z = self._gather(self.ec_z, isel, points, len(B.EC_VARS)) cols.append(z.transpose(0, 1, 3, 2).reshape(nrow, len(B.EC_VARS))) if fset not in ("F4", "F5"): del z if fset in ("F2", "F3", "F4", "F5", "F6", "F7", "F8"): for k in ("H500", "T850"): pc = self.pcs[k][0][isel] # (ni,nk,NMODE) cols.append(np.broadcast_to(pc[:, :, None, :], (ni, nk, npt, NMODE)).reshape(nrow, NMODE)) if fset in ("F3", "F4", "F5", "F6", "F7", "F8"): lat = np.broadcast_to(self.st["lat"][points][None, None, :], (ni, nk, npt)).reshape(nrow, 1) lon = np.broadcast_to(self.st["lon"][points][None, None, :], (ni, nk, npt)).reshape(nrow, 1) reg = np.broadcast_to(self.st["region"][points].astype("f4")[None, None, :], (ni, nk, npt)).reshape(nrow, 1) tho = self.st["thr_obs_doy"][:, points][doy - 1] # (ni,nk,npt) thm = np.stack([self.st["thr_mod_lt_doy"][k][:, points][doy[:, k] - 1] for k in range(nk)], axis=1) cols += [lat, lon, reg, tho.reshape(nrow, 1), thm.reshape(nrow, 1)] if fset in EXT_SETS: # hist varies with (init, point) and idx with init alone; both describe # the state AT initialisation, so both are constant along the lead axis _, hkeep, gkeep = self._ext_names(fset) if hkeep.size: h = self.ext["hist"][np.ix_(isel, points)][:, :, hkeep] cols.append(np.broadcast_to(h[:, None, :, :], (ni, nk, npt, h.shape[2]) ).reshape(nrow, h.shape[2])) if gkeep.size: g = self.ext["idx"][isel][:, gkeep] cols.append(np.broadcast_to(g[:, None, None, :], (ni, nk, npt, g.shape[1]) ).reshape(nrow, g.shape[1])) if fset == "F5": vi5 = [B.EC_VARS.index(v) for v in SWT_VARS] sw = swt_features(np.concatenate([der[:, :, 0:1, :], z[:, :, vi5, :]], axis=2)) cols.append(sw.transpose(0, 1, 3, 2).reshape(nrow, sw.shape[2])) del der, z, sw if fset == "F4": vi = [B.EC_VARS.index(v) for v in WIN_VARS] sv = der[:, :, 0:1, :] # model severity for w in WIN_LEN: rm = running_mean(np.concatenate([sv, z[:, :, vi, :]], axis=2), w) cols.append(rm.transpose(0, 1, 3, 2).reshape(nrow, 1 + len(vi))) # how far the target day departs from its own three-week window base = running_mean(np.concatenate( [z[:, :, B.EC_VARS.index("mx2t6"):B.EC_VARS.index("mx2t6") + 1, :], sv], axis=2), WIN_LEN[-1]) dev = np.concatenate( [z[:, :, B.EC_VARS.index("mx2t6"):B.EC_VARS.index("mx2t6") + 1, :], sv], axis=2) - base cols.append(dev.transpose(0, 1, 3, 2).reshape(nrow, len(DEV_VARS))) del der, z X = np.concatenate([c.astype("f4", copy=False) for c in cols], axis=1) tgt = self._gather(self.target, isel, points, 3) y_tx = tgt[:, :, 0, :].reshape(nrow) y_sev = tgt[:, :, 1, :].reshape(nrow) y_evt = tgt[:, :, 2, :].reshape(nrow) info = { "year": np.broadcast_to(self.init_year[isel][:, None, None], (ni, nk, npt)).reshape(nrow), "init": np.broadcast_to(isel[:, None, None], (ni, nk, npt)).reshape(nrow), "lt": np.broadcast_to(np.arange(nk)[None, :, None], (ni, nk, npt)).reshape(nrow), "point": np.broadcast_to(points[None, None, :], (ni, nk, npt)).reshape(nrow), "doy": np.broadcast_to(doy[:, :, None], (ni, nk, npt)).reshape(nrow), "region": self.st["region"][np.broadcast_to( points[None, None, :], (ni, nk, npt)).reshape(nrow)], } info["week"] = np.searchsorted( np.array([11, 18, 25, 32, 39]), info["lt"], side="right") # a cell is "heat-capable" when the local 31-day Tx90p already reaches 35 C, # i.e. the 35 C cap is what defines the threshold there thr_o = self.st["thr_obs_doy"][info["doy"] - 1, info["point"]] info["thr_obs"] = thr_o info["capable"] = C.is_capable(thr_o) return X, {"tx": y_tx, "sev": y_sev, "evt": y_evt}, info def sample_points(grid, per_region=3, seed=0): """A few land points from each of the seven climatic sub-regions.""" rng = np.random.default_rng(seed) out = [] for r in range(1, 8): idx = np.nonzero(grid.reg == r)[0] out.append(rng.choice(idx, size=min(per_region, idx.size), replace=False)) return np.sort(np.concatenate(out))