"""Torch-CPU AI families for v2 Phase B: a deeper ANN and a fully-convolutional CNN. Both consume the SAME block-aggregated tabular features that features.assemble_year builds (so they slot into the LOYO engine next to the sklearn zoo) and both return an out-of-sample P(yEHDE) that is isotonic-recalibrated on TRAIN (matching the rest of the zoo). TorchANN a deeper multilayer perceptron ([128,64,32], ReLU, dropout) on standardized features -- the "ANN family" the user asked to add, distinct from the shallow sklearn `mlp`. TorchCNN a fully-convolutional network. The per-cell feature rows for one lead block are reshaped back onto the (init, channel, lat, lon) grid, so 3x3 convolutions let each cell borrow information from its spatial neighbours (the point of a CNN). Output is a per-cell probability map; values at the land cells are read back in the SAME sample order as the tabular members, so the CNN is a drop-in ensemble member. LEAKAGE RULE: everything (standardisation stats, weights, isotonic map, spatial padding) is fit on TRAIN arrays the caller passes; the held-out year is only ever passed 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) # under SLURM we pack several fold tasks per node; keep torch inside its cpu budget nt = _os.environ.get("TORCH_NUM_THREADS") if nt: torch.set_num_threads(int(nt)) _TORCH = torch return _TORCH def _isotonic(p_tr, y_tr, *extra): m = np.isfinite(p_tr) & np.isfinite(y_tr) iso = IsotonicRegression(out_of_bounds="clip").fit(p_tr[m], y_tr[m].astype(int)) return (iso.transform(p_tr), *(iso.transform(a) for a in extra)) # --------------------------------------------------------------------------- deep ANN (tabular) class TorchANN: is_sev = False def __init__(self, hidden=(128, 64, 32), dropout=0.2, lr=1e-3, max_epoch=60, batch=8192, patience=6, seed=0): self.cfg = dict(hidden=hidden, dropout=dropout, lr=lr, max_epoch=max_epoch, batch=batch, patience=patience, seed=seed) def _build(self, nin): torch = _torch(); nn = torch.nn layers, d = [], nin for h in self.cfg["hidden"]: layers += [nn.Linear(d, h), nn.ReLU(), nn.Dropout(self.cfg["dropout"])] d = h layers += [nn.Linear(d, 1)] return nn.Sequential(*layers) def fit(self, X, y): torch = _torch() m = np.isfinite(y) & np.all(np.isfinite(X), axis=1) X, y = X[m].astype("float32"), y[m].astype("float32") self.mu_ = X.mean(0); self.sd_ = X.std(0) + 1e-6 Xs = (X - self.mu_) / self.sd_ torch.manual_seed(self.cfg["seed"]) rng = np.random.default_rng(self.cfg["seed"]) nval = max(1, int(0.1 * len(Xs))) perm = rng.permutation(len(Xs)); vi, ti = perm[:nval], perm[nval:] net = self._build(Xs.shape[1]) opt = torch.optim.Adam(net.parameters(), lr=self.cfg["lr"]) # class imbalance -> pos_weight pos = float(y[ti].sum()); neg = float((y[ti] == 0).sum()) pw = torch.tensor([neg / max(pos, 1.0)], dtype=torch.float32) lossf = torch.nn.BCEWithLogitsLoss(pos_weight=pw) Xt = torch.from_numpy(Xs); yt = torch.from_numpy(y) 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() out = net(Xt[idx]).squeeze(1) loss = lossf(out, yt[idx]); loss.backward(); opt.step() net.eval() with torch.no_grad(): vloss = torch.nn.functional.binary_cross_entropy_with_logits( net(Xt[vi]).squeeze(1), yt[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 p_tr = self._raw(X) (self._p_tr_cal,) = _isotonic(p_tr, y) m2 = np.isfinite(p_tr) & np.isfinite(y) self.iso_ = IsotonicRegression(out_of_bounds="clip").fit(p_tr[m2], y[m2].astype(int)) return self def _raw(self, X): torch = _torch() Xs = ((X.astype("float32") - self.mu_) / self.sd_) self.net_.eval() # batch the forward pass: a single call on ~1e6 rows blows CPU memory for the # attention model (seq x seq x batch); 65536-row chunks keep it bounded. out = np.empty(Xs.shape[0], dtype="float32") with torch.no_grad(): for s in range(0, Xs.shape[0], 65536): xb = torch.from_numpy(Xs[s:s + 65536]) out[s:s + 65536] = torch.sigmoid(self.net_(xb).squeeze(1)).numpy() return out def predict_proba(self, X): r = np.full(X.shape[0], np.nan) ok = np.all(np.isfinite(X), axis=1) r[ok] = self.iso_.transform(self._raw(X[ok])) return r # ------------------------------------------------------------------------ fully-convolutional CNN class TorchCNN: """Fully-convolutional per-cell classifier. fit/predict take the SAME tabular (X, y) as the zoo PLUS the grid geometry needed to fold rows back into (init, channel, lat, lon).""" is_sev = False def __init__(self, iy, ix, ny, nx, ncell, channels=32, lr=1e-3, max_epoch=30, batch=64, patience=5, seed=0): self.iy, self.ix, self.ny, self.nx, self.ncell = iy, ix, ny, nx, ncell self.cfg = dict(channels=channels, lr=lr, max_epoch=max_epoch, batch=batch, patience=patience, seed=seed) # -- reshape tabular (ninit*ncell, nfeat) -> images (ninit, nfeat, ny, nx) + label/mask maps def _to_images(self, X, y=None): n, nf = X.shape ninit = n // self.ncell Xr = X.reshape(ninit, self.ncell, nf) img = np.zeros((ninit, nf, self.ny, self.nx), dtype="float32") for c in range(nf): img[:, c, self.iy, self.ix] = np.nan_to_num(Xr[:, :, c], nan=0.0).astype("float32") if y is None: return img, ninit yr = y.reshape(ninit, self.ncell) lab = np.zeros((ninit, self.ny, self.nx), dtype="float32") msk = np.zeros((ninit, self.ny, self.nx), dtype="float32") valid = np.isfinite(yr) lab[:, self.iy, self.ix] = np.where(valid, yr, 0.0).astype("float32") msk[:, self.iy, self.ix] = valid.astype("float32") return img, lab, msk, ninit def _build(self, nin, ch): torch = _torch(); nn = torch.nn return nn.Sequential( nn.Conv2d(nin, ch, 3, padding=1), nn.ReLU(), nn.Conv2d(ch, ch, 3, padding=1), nn.ReLU(), nn.Conv2d(ch, ch // 2, 3, padding=1), nn.ReLU(), nn.Conv2d(ch // 2, 1, 1)) def fit(self, X, y): torch = _torch() Xf = X.astype("float32").copy() # standardise channels over finite entries (train stats) self.mu_ = np.nanmean(np.where(np.isfinite(Xf), Xf, np.nan), axis=0) self.sd_ = np.nanstd(np.where(np.isfinite(Xf), Xf, np.nan), axis=0) + 1e-6 Xn = (np.where(np.isfinite(Xf), Xf, self.mu_) - self.mu_) / self.sd_ img, lab, msk, ninit = self._to_images(Xn, y) 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(img.shape[1], self.cfg["channels"]) opt = torch.optim.Adam(net.parameters(), lr=self.cfg["lr"]) pos = float((lab * msk).sum()); neg = float(msk.sum() - pos) pw = torch.tensor(neg / max(pos, 1.0), dtype=torch.float32) It = torch.from_numpy(img); Lt = torch.from_numpy(lab); Mt = torch.from_numpy(msk) def masked_loss(idx): out = net(It[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 p_tr = self._raw(X) m2 = np.isfinite(p_tr) & np.isfinite(y) self.iso_ = IsotonicRegression(out_of_bounds="clip").fit(p_tr[m2], y[m2].astype(int)) return self def _raw(self, X): torch = _torch() Xf = X.astype("float32") Xn = (np.where(np.isfinite(Xf), Xf, self.mu_) - self.mu_) / self.sd_ img, ninit = self._to_images(Xn) self.net_.eval() with torch.no_grad(): out = torch.sigmoid(self.net_(torch.from_numpy(img)).squeeze(1)).numpy() # (ninit,ny,nx) vals = out[:, self.iy, self.ix] # (ninit, ncell) return vals.reshape(-1) # init-major, cell-minor def predict_proba(self, X): return self.iso_.transform(self._raw(X)) # ------------------------------------------------------------------------ residual U-Net (spatial) class TorchResUNet(TorchCNN): """Residual U-Net over the block-aggregated channel grid — the hybrid-design note's high- priority spatial family for W3-W6 (multiscale spatial structure + teleconnection context). Reuses TorchCNN's tabular<->image reshape, masked-loss training and isotonic calibration; only the network is different: a 2-level encoder/decoder with skip connections and residual conv blocks, so the field is seen at three spatial scales. The input is zero-padded to a multiple of 4 (two 2x downsamples) and the output is cropped back to (ny, nx). Extra large-scale channels (e.g. gridded Z500/T850) can be stacked into X upstream and are handled transparently. """ is_sev = False def _build(self, nin, ch): torch = _torch(); nn = torch.nn 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 UNet(nn.Module): def __init__(self): super().__init__() self.ny, self.nx = ny, nx 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, x): import torch.nn.functional as F 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 UNet() # ---------------------------------------------------------------- FT-Transformer (tabular, per-cell) class TorchFTTransformer(TorchANN): """Feature-tokenizer Transformer on the per-cell lead-dependent predictor vector — an exploratory long-range-dependence challenger (doc §2, flagged high overfitting risk on 20 years, so kept compact: small d_model, few heads/layers, dropout, early stopping). Each numeric feature is projected to a d_model token (feature-specific weight+bias); a learnable [CLS] token is prepended; a shallow Transformer encoder mixes the feature tokens via self- attention; the CLS embedding maps to the logit. Reuses TorchANN's standardise/fit/calibrate. """ is_sev = False def __init__(self, d_model=32, heads=4, layers=2, dropout=0.1, lr=1e-3, max_epoch=60, batch=8192, patience=6, seed=0): super().__init__(hidden=(d_model,), dropout=dropout, lr=lr, max_epoch=max_epoch, batch=batch, patience=patience, seed=seed) self.tcfg = dict(d_model=d_model, heads=heads, layers=layers, dropout=dropout) def _build(self, nin): torch = _torch(); nn = torch.nn d = self.tcfg["d_model"]; heads = self.tcfg["heads"] layers = self.tcfg["layers"]; dropout = self.tcfg["dropout"] class FTT(nn.Module): def __init__(self): super().__init__() self.w = nn.Parameter(torch.randn(nin, d) * 0.02) # per-feature projection self.b = nn.Parameter(torch.zeros(nin, d)) self.cls = nn.Parameter(torch.randn(1, 1, d) * 0.02) enc = nn.TransformerEncoderLayer(d_model=d, nhead=heads, dim_feedforward=2 * d, dropout=dropout, batch_first=True) self.enc = nn.TransformerEncoder(enc, num_layers=layers) self.head = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, 1)) def forward(self, x): # x: (batch, nfeat) tok = x.unsqueeze(-1) * self.w + self.b # (batch, nfeat, d) cls = self.cls.expand(x.shape[0], -1, -1) h = torch.cat([cls, tok], dim=1) h = self.enc(h) return self.head(h[:, 0]) # (batch, 1) return FTT() def make_torch_model(name, geom=None): if name == "ann_torch": return TorchANN() if name == "cnn": assert geom is not None, "cnn needs grid geometry" return TorchCNN(**geom) if name == "resunet": assert geom is not None, "resunet needs grid geometry" return TorchResUNet(**geom) if name == "transformer": return TorchFTTransformer() raise KeyError(name)