#!/usr/bin/env python3 """A convolutional network over the map, instead of one row per grid point. Every model up to here is point-wise: a row is one (initialisation, lead, point) triple and the model never sees what the neighbouring points are doing. A heat extreme is a synoptic-scale object, so this model keeps the map intact -- the input is the gridded forecast fields for one (initialisation, lead) and the output is a probability map, with the loss taken only over land points. It deliberately writes `oof_*.npz` and `pred_*.npz` in exactly the format the row-based pipeline uses, so `train_national.py finalize` scores it with the same thresholds, the same frames and the same netCDF writer. Nothing downstream has to know a CNN produced the numbers. Run with LD_LIBRARY_PATH=$CONDA_PREFIX/lib so torch finds a new enough libstdc++. """ from __future__ import annotations import argparse import os import sys import time 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 ai_data as D import train_national as N CNN_VARS = ["mx2t6", "t2m", "skt", "thk_1000_500", "thk_925_700", "msl", "tcc", "sm20", "cape", "tp", "dtd2m", "ssrd"] NSTATIC = 7 # land, lat, lon, thr_obs, doy_cos, doy_sin, lead NCHAN = len(CNN_VARS) + 1 + NSTATIC # + sv_mod class MapMaker: """Turns one initialisation into a (nlt, NCHAN, H, W) tensor of gridded fields.""" def __init__(self, store): self.s = store g = store.grid self.g = g self.H, self.W = g.nlat, g.nlon self.vi = [B.EC_VARS.index(v) for v in CNN_VARS] self.land = np.zeros((self.H, self.W), "f4") self.land[g.iy, g.ix] = 1.0 lat = np.zeros((self.H, self.W), "f4"); lat[g.iy, g.ix] = g.lat[g.iy] lon = np.zeros((self.H, self.W), "f4"); lon[g.iy, g.ix] = g.lon[g.ix] self.lat = (lat - 16.0) / 8.0 self.lon = (lon - 106.0) / 4.0 def scatter(self, v): """(nlt, nv, npt) land values -> (nlt, nv, H, W) with zeros off land.""" nk, nv, _ = v.shape out = np.zeros((nk, nv, self.H, self.W), "f4") out[:, :, self.g.iy, self.g.ix] = v return out def build(self, i): s, g = self.s, self.g nk = C.NLT z = np.asarray(s.ec_z[i][:, self.vi, :], dtype="f4") # (nk,nv,npt) sv = np.asarray(s.derived[i][:, 0:1, :], dtype="f4") / 5.0 dyn = self.scatter(np.concatenate([z, sv], axis=1)) doy = s.doy[i] # (nk,) thr = s.st["thr_obs_doy"][doy - 1, :] # (nk,npt) thr = self.scatter(((thr - 35.0) / 5.0)[:, None, :].astype("f4")) stat = np.empty((nk, NSTATIC, self.H, self.W), "f4") stat[:, 0] = self.land stat[:, 1] = self.lat stat[:, 2] = self.lon stat[:, 3] = thr[:, 0] stat[:, 4] = np.cos(2 * np.pi * doy / 365.0)[:, None, None] stat[:, 5] = np.sin(2 * np.pi * doy / 365.0)[:, None, None] stat[:, 6] = (np.arange(nk, dtype="f4") / 45.0)[:, None, None] x = np.concatenate([dyn, stat], axis=1) np.nan_to_num(x, copy=False) y = np.asarray(s.target[i][:, 2, :], dtype="f4") # evt at land points return x, y def make_net(nchan, nf=32): """Fully convolutional, no pooling, dilations widening the receptive field.""" import torch.nn as nn return nn.Sequential( nn.Conv2d(nchan, nf, 3, padding=1), nn.BatchNorm2d(nf), nn.ReLU(), nn.Conv2d(nf, nf, 3, padding=2, dilation=2), nn.BatchNorm2d(nf), nn.ReLU(), nn.Conv2d(nf, nf, 3, padding=4, dilation=4), nn.BatchNorm2d(nf), nn.ReLU(), nn.Conv2d(nf, nf, 3, padding=1), nn.BatchNorm2d(nf), nn.ReLU(), nn.Conv2d(nf, 1, 1)) def train(net, mm, inits, epochs, pos_weight, lr, seed, log): import torch torch.manual_seed(seed) opt = torch.optim.Adam(net.parameters(), lr=lr) lossf = torch.nn.BCEWithLogitsLoss(pos_weight=torch.tensor([float(pos_weight)]), reduction="none") iy = torch.from_numpy(mm.g.iy.astype("i8")) ix = torch.from_numpy(mm.g.ix.astype("i8")) rng = np.random.default_rng(seed) net.train() for ep in range(epochs): order = rng.permutation(inits) tot, nb, t0 = 0.0, 0, time.time() for i in order: x, y = mm.build(int(i)) xb = torch.from_numpy(x) yb = torch.from_numpy(y) good = torch.isfinite(yb) if not bool(good.any()): continue out = net(xb)[:, 0][:, iy, ix] # (nk, npt) land points only l = lossf(out, torch.nan_to_num(yb)) loss = (l * good).sum() / good.sum() opt.zero_grad(); loss.backward(); opt.step() tot += float(loss); nb += 1 print(" [cnn] epoch %d/%d loss %.5f %.1f min" % (ep + 1, epochs, tot / max(nb, 1), (time.time() - t0) / 60), flush=True) net.eval() return net def score(net, mm, inits, store): """Rows in the pipeline's (init, lead, point) order with all the info fields.""" import torch npt = mm.g.npt iy = torch.from_numpy(mm.g.iy.astype("i8")) ix = torch.from_numpy(mm.g.ix.astype("i8")) sc, ev, info = [], [], {k: [] for k in ("week", "lt", "point", "capable", "init", "region", "doy")} reg = store.st["region"] with torch.no_grad(): for i in inits: i = int(i) x, y = mm.build(i) out = net(torch.from_numpy(x))[:, 0][:, iy, ix] p = torch.sigmoid(out).numpy().astype("f4") # (nk,npt) sc.append(p.reshape(-1)); ev.append(y.reshape(-1)) doy = store.doy[i] lt = np.repeat(np.arange(C.NLT), npt) info["lt"].append(lt) info["week"].append(np.searchsorted(np.array([11, 18, 25, 32, 39]), lt, side="right")) info["point"].append(np.tile(np.arange(npt), C.NLT)) info["init"].append(np.full(C.NLT * npt, i, "i8")) d = np.repeat(doy, npt) info["doy"].append(d) info["region"].append(np.tile(reg, C.NLT)) thr = store.st["thr_obs_doy"][doy - 1, :] info["capable"].append(C.is_capable(thr).reshape(-1)) return (np.concatenate(sc), np.concatenate(ev), {k: np.concatenate(v) for k, v in info.items()}) def main(): ap = argparse.ArgumentParser() ap.add_argument("phase", choices=["fit", "predict"]) ap.add_argument("--fold", default="full") ap.add_argument("--tag", default="_cnn") ap.add_argument("--epochs", type=int, default=8) ap.add_argument("--nf", type=int, default=32) ap.add_argument("--lr", type=float, default=1e-3) ap.add_argument("--pos-weight", type=float, default=12.0) ap.add_argument("--threads", type=int, default=20) args = ap.parse_args() import torch torch.set_num_threads(args.threads) store = D.Store() mm = MapMaker(store) os.makedirs(N.shard_dir(args.tag), exist_ok=True) tr_all = N.train_inits(store) print("[cnn] %d channels, %dx%d grid, %d land points" % (NCHAN, mm.H, mm.W, mm.g.npt), flush=True) if args.phase == "fit": if args.fold == "full": net = train(make_net(NCHAN, args.nf), mm, tr_all, args.epochs, args.pos_weight, args.lr, 0, None) torch.save({"sd": net.state_dict(), "nf": args.nf}, os.path.join(N.shard_dir(args.tag), "cnn_full.pt")) print("[cnn full] saved", flush=True) return f = int(args.fold) hold = N.year_folds()[f] yr = store.init_year[tr_all] te = np.isin(yr, hold) net = train(make_net(NCHAN, args.nf), mm, tr_all[~te], args.epochs, args.pos_weight, args.lr, f, None) s, e, info = score(net, mm, tr_all[te], store) ok = np.isfinite(e) np.savez(os.path.join(N.shard_dir(args.tag), "oof_%d.npz" % f), score=s[ok].astype("f8"), evt=e[ok], week=info["week"][ok].astype("i4"), capable=info["capable"][ok].astype("i1"), region=info["region"][ok].astype("i4"), doy=info["doy"][ok].astype("i4")) print("[cnn fold %d] held out %s, %d rows" % (f, hold.tolist(), int(ok.sum())), flush=True) return # predict: the final model scores the held-out years ck = torch.load(os.path.join(N.shard_dir(args.tag), "cnn_full.pt"), weights_only=False) net = make_net(NCHAN, ck["nf"]); net.load_state_dict(ck["sd"]); net.eval() s, e, info = score(net, mm, N.valid_inits(store), store) np.savez(os.path.join(N.shard_dir(args.tag), "pred_0000.npz"), score=s, evt=e, **{k: v for k, v in info.items()}) print("[cnn predict] %d rows" % s.size, flush=True) if __name__ == "__main__": main()