"""True-spatial full-field ResU-Net (v4) — the fix for the qg1-hybrid-v3 HONEST LIMITATION. The v3 ResU-Net (gridmodels.TorchResUNet) only ever saw the VN 41x65 grid; the large-scale circulation entered as broadcast BOX-INDEX scalars, never real spatial patterns. This model consumes the stage-27 full-field cache (outputs/fullfield_cache/ff_.npz) as ACTUAL spatial channels on two grids and fuses them: VN branch — a residual U-Net over the native 41x65 grid whose input channels are the stage-27 `vn` fields (51: multi-level T/H/W/QUV + surface) STACKED with the Stage-1 per-cell forecast-Tx feature maps (f_max/f_mean/... rasterised). This is where the short-lead forecast signal lives. ext branch — a small CNN encoder over the extended 61x71 domain (Z500/T850/U/V200), global- average-pooled to a context vector and TILED onto the VN grid as extra channels, so every VN cell sees the synoptic/teleconnection state (the long-lead signal). Output is a per-cell P(yEHDE) map on the VN grid; land-cell values are read back in the SAME (init-major, cell-minor) order as the tabular zoo so this is a drop-in stack member. LEAKAGE RULE: channel standardisation stats, network weights and the isotonic calibration map are ALL fit on TRAIN arrays only; the held-out year is passed only to predict. """ from __future__ import annotations import numpy as np from sklearn.isotonic import IsotonicRegression _TORCH = None def _torch(): global _TORCH if _TORCH is None: import os as _os import torch torch.manual_seed(0) nt = _os.environ.get("TORCH_NUM_THREADS") if nt: torch.set_num_threads(int(nt)) _TORCH = torch return _TORCH class FullFieldResUNet: """Two-branch (VN ResU-Net + extended-domain context) fully-convolutional EHD classifier. fit/predict take image tensors, not tabular rows: vn_img : (ninit, C_vn, ny, nx) native-grid channels (already standardised? no — we do it) ext_img: (ninit, C_ext, ne, nw) extended-domain channels lab : (ninit, ny, nx) 0/1 EHD label (only where mask=1 is meaningful) mask : (ninit, ny, nx) 1 at valid land cells, 0 elsewhere predict returns land-cell probabilities flattened init-major, cell-minor (via iy,ix). """ def __init__(self, iy, ix, ny, nx, ch=32, ext_ch=16, ctx=8, lr=1e-3, max_epoch=40, batch=32, patience=6, seed=0): self.iy, self.ix, self.ny, self.nx = iy, ix, ny, nx self.cfg = dict(ch=ch, ext_ch=ext_ch, ctx=ctx, lr=lr, max_epoch=max_epoch, batch=batch, patience=patience, seed=seed) # ---- channel standardisation over finite entries (train stats), NaN -> channel mean (0) ---- def _fit_norm(self, vn_img, ext_img): self.vn_mu = np.nanmean(vn_img, axis=(0, 2, 3), keepdims=True) self.vn_sd = np.nanstd(vn_img, axis=(0, 2, 3), keepdims=True) + 1e-6 self.ex_mu = np.nanmean(ext_img, axis=(0, 2, 3), keepdims=True) self.ex_sd = np.nanstd(ext_img, axis=(0, 2, 3), keepdims=True) + 1e-6 def _norm(self, vn_img, ext_img): v = (np.where(np.isfinite(vn_img), vn_img, self.vn_mu) - self.vn_mu) / self.vn_sd e = (np.where(np.isfinite(ext_img), ext_img, self.ex_mu) - self.ex_mu) / self.ex_sd return v.astype("float32"), e.astype("float32") def _build(self, c_vn, c_ext): torch = _torch(); nn = torch.nn ch, ext_ch, ctx = self.cfg["ch"], self.cfg["ext_ch"], self.cfg["ctx"] ny, nx = self.ny, self.nx class ResBlock(nn.Module): def __init__(self, cin, cout): super().__init__() self.c1 = nn.Conv2d(cin, cout, 3, padding=1) self.c2 = nn.Conv2d(cout, cout, 3, padding=1) self.bn1 = nn.BatchNorm2d(cout); self.bn2 = nn.BatchNorm2d(cout) self.proj = nn.Conv2d(cin, cout, 1) if cin != cout else nn.Identity() self.act = nn.ReLU() def forward(self, x): h = self.act(self.bn1(self.c1(x))) h = self.bn2(self.c2(h)) return self.act(h + self.proj(x)) class Net(nn.Module): def __init__(self): super().__init__() self.ny, self.nx = ny, nx # extended-domain encoder -> context vector self.ext = nn.Sequential( nn.Conv2d(c_ext, ext_ch, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(ext_ch, ext_ch * 2, 3, padding=1), nn.ReLU(), nn.AdaptiveAvgPool2d(1)) self.ext_fc = nn.Linear(ext_ch * 2, ctx) # VN residual U-Net (input = vn channels + tiled context) nin = c_vn + ctx self.enc1 = ResBlock(nin, ch) self.enc2 = ResBlock(ch, ch * 2) self.pool = nn.MaxPool2d(2) self.bott = ResBlock(ch * 2, ch * 4) self.up2 = nn.ConvTranspose2d(ch * 4, ch * 2, 2, stride=2) self.dec2 = ResBlock(ch * 4, ch * 2) self.up1 = nn.ConvTranspose2d(ch * 2, ch, 2, stride=2) self.dec1 = ResBlock(ch * 2, ch) self.head = nn.Conv2d(ch, 1, 1) def forward(self, vn, ext): import torch.nn.functional as F ctxv = self.ext_fc(self.ext(ext).flatten(1)) # (B, ctx) ctx_map = ctxv[:, :, None, None].expand(-1, -1, vn.shape[-2], vn.shape[-1]) x = torch.cat([vn, ctx_map], dim=1) H, W = x.shape[-2], x.shape[-1] ph = (4 - H % 4) % 4; pw = (4 - W % 4) % 4 x = F.pad(x, (0, pw, 0, ph)) e1 = self.enc1(x) e2 = self.enc2(self.pool(e1)) b = self.bott(self.pool(e2)) d2 = self.dec2(torch.cat([self.up2(b), e2], 1)) d1 = self.dec1(torch.cat([self.up1(d2), e1], 1)) out = self.head(d1) return out[..., :H, :W] return Net() def fit(self, vn_img, ext_img, lab, mask): torch = _torch() self._fit_norm(vn_img, ext_img) vn, ext = self._norm(vn_img, ext_img) ninit = vn.shape[0] torch.manual_seed(self.cfg["seed"]); rng = np.random.default_rng(self.cfg["seed"]) nval = max(1, int(0.1 * ninit)); perm = rng.permutation(ninit) vi, ti = perm[:nval], perm[nval:] net = self._build(vn.shape[1], ext.shape[1]) opt = torch.optim.Adam(net.parameters(), lr=self.cfg["lr"]) pos = float((lab * mask).sum()); neg = float(mask.sum() - pos) pw = torch.tensor(neg / max(pos, 1.0), dtype=torch.float32) Vt = torch.from_numpy(vn); Et = torch.from_numpy(ext) Lt = torch.from_numpy(lab.astype("float32")); Mt = torch.from_numpy(mask.astype("float32")) def masked_loss(idx): out = net(Vt[idx], Et[idx]).squeeze(1) l = torch.nn.functional.binary_cross_entropy_with_logits( out, Lt[idx], reduction="none", pos_weight=pw) return (l * Mt[idx]).sum() / (Mt[idx].sum() + 1e-6) best, best_state, bad = np.inf, None, 0 for ep in range(self.cfg["max_epoch"]): net.train(); order = rng.permutation(ti) for s in range(0, len(order), self.cfg["batch"]): idx = order[s:s + self.cfg["batch"]] opt.zero_grad(); loss = masked_loss(idx); loss.backward(); opt.step() net.eval() with torch.no_grad(): vloss = masked_loss(vi).item() if vloss < best - 1e-4: best, best_state, bad = vloss, {k: v.clone() for k, v in net.state_dict().items()}, 0 else: bad += 1 if bad >= self.cfg["patience"]: break if best_state is not None: net.load_state_dict(best_state) self.net_ = net # isotonic calibration on TRAIN land-cell (pred, y) p_tr = self._raw(vn_img, ext_img) # (ninit, ncell) y_tr = lab[:, self.iy, self.ix] # (ninit, ncell) m_tr = mask[:, self.iy, self.ix] > 0 pf, yf = p_tr[m_tr], y_tr[m_tr] ok = np.isfinite(pf) & np.isfinite(yf) self.iso_ = IsotonicRegression(out_of_bounds="clip").fit(pf[ok], yf[ok].astype(int)) return self def _raw(self, vn_img, ext_img): torch = _torch() vn, ext = self._norm(vn_img, ext_img) Vt = torch.from_numpy(vn); Et = torch.from_numpy(ext) self.net_.eval() outs = [] with torch.no_grad(): for s in range(0, vn.shape[0], self.cfg["batch"]): logit = self.net_(Vt[s:s + self.cfg["batch"]], Et[s:s + self.cfg["batch"]]).squeeze(1) outs.append(torch.sigmoid(logit).numpy()) out = np.concatenate(outs, axis=0) # (ninit, ny, nx) return out[:, self.iy, self.ix] # (ninit, ncell) def predict_proba(self, vn_img, ext_img): """Return land-cell probs flattened init-major, cell-minor (matches tabular row order).""" p = self._raw(vn_img, ext_img) # (ninit, ncell) return self.iso_.transform(p.reshape(-1))