#!/usr/bin/env python3 """CNN residual downscaling of Tx, 1.5 deg -> 0.25 deg. Point-wise learners (LightGBM) see each grid point independently. A convolutional net sees the whole field, so it can use the shape of the coarse gradient and of the terrain around a point -- exactly the information that decides where a hot anomaly lands in a narrow country with a mountain spine. Input channels, all on the 0.25 deg grid: 0 bilinear(Tx_ec) the field to be corrected 1 bilinear anomaly vs its training climatology for that (week, doy) 2 clim_delta = obs clim - bil clim, the systematic part of the correction 3 dz orog - orog_coarse 4 orog / 1000 5 orog_std / 100 6 landfrac 7 lead / 45 (constant plane) 8 sin(2 pi doy / 365) (constant plane) 9 cos(2 pi doy / 365) (constant plane) Target: residual Tx_obs - bilinear, scored on land points only (masked loss). One fold per job (`--fold k`): torch here is CPU-only, so five folds in parallel on five nodes is far faster than five folds in sequence on one. Usage: downscale_cnn.py --fold 0..4 [--epochs 25] [--ntrain 30000] """ import argparse import os import sys import numpy as np import torch import torch.nn as nn import xarray as xr ROOT = "/home/tanpv/Progs/claude/HW_EC_10" COARSE = f"{ROOT}/outputs/proxy/tx_ec_1.5_proxy.nc" OBS = f"{ROOT}/outputs/OBS_4EC15/tx_obs_aligned.nc" STATIC = f"{ROOT}/outputs/Cli/static_fields_0.25.nc" WEEKS = [("W1", 0, 11), ("W2", 11, 18), ("W3", 18, 25), ("W4", 25, 32), ("W5", 32, 39), ("W6", 39, 46)] NDOY = 365 NCH = 10 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from downscale import bilinear_matrix, fill_coarse_nan # noqa: E402 from downscale_ai import window_climatology # noqa: E402 class ResidualCNN(nn.Module): """Small fully-convolutional net. No down/up-sampling: the input is already on the target grid, so the job is local correction, not resolution synthesis.""" def __init__(self, nch=NCH, width=64, depth=5): super().__init__() layers = [nn.Conv2d(nch, width, 3, padding=1), nn.ReLU(inplace=True)] for _ in range(depth - 2): layers += [nn.Conv2d(width, width, 3, padding=1), nn.BatchNorm2d(width), nn.ReLU(inplace=True)] layers += [nn.Conv2d(width, 1, 3, padding=1)] self.net = nn.Sequential(*layers) def forward(self, x): return self.net(x).squeeze(1) def main(): ap = argparse.ArgumentParser() ap.add_argument("--fold", type=int, required=True) ap.add_argument("--folds", type=int, default=5) ap.add_argument("--epochs", type=int, default=25) ap.add_argument("--ntrain", type=int, default=30000) ap.add_argument("--batch", type=int, default=48) ap.add_argument("--lr", type=float, default=1e-3) ap.add_argument("--threads", type=int, default=16) ap.add_argument("--coarse", default=COARSE) ap.add_argument("--outdir", default=f"{ROOT}/outputs/downscale/cnn") args = ap.parse_args() torch.set_num_threads(args.threads) rng = np.random.default_rng(args.fold) c = xr.open_dataset(args.coarse) o = xr.open_dataset(OBS) st = xr.open_dataset(STATIC) clat, clon = c.lat.values, c.lon.values flat, flon = o.lat.values, o.lon.values ny, ni, nl = c.sizes["year"], c.sizes["init"], c.sizes["lead"] H, Wd = len(flat), len(flon) npts = H * Wd land2d = (np.isfinite(st.vn.values) & (st.vn.values > 0)) landm = torch.tensor(land2d.astype("float32")) print(f"fold {args.fold}: grid {H}x{Wd}, {int(land2d.sum())} land points", flush=True) Wm = bilinear_matrix(clat, clon, flat, flon) cxf, ndead = fill_coarse_nan( c.tx_ec.values.reshape(-1, len(clat) * len(clon)), clat, clon) bil = (cxf @ Wm.T).reshape(ny, ni, nl, npts).astype("float32") obs = o.tx_obs.values.reshape(ny, ni, nl, npts) doy = o.doy.values del cxf per = ny // args.folds te = np.arange(args.fold * per, (args.fold + 1) * per) tr = np.array([i for i in range(ny) if i not in te]) wk_of_lead = np.zeros(nl, dtype=int) for w, (_, l0, l1) in enumerate(WEEKS): wk_of_lead[l0:l1] = w clim_obs = np.empty((len(WEEKS), NDOY, npts), dtype="float32") clim_bil = np.empty((len(WEEKS), NDOY, npts), dtype="float32") for w, (_, l0, l1) in enumerate(WEEKS): d = doy[tr][:, :, l0:l1].ravel() clim_obs[w] = window_climatology(obs[tr][:, :, l0:l1].reshape(-1, npts), d, npts) clim_bil[w] = window_climatology(bil[tr][:, :, l0:l1].reshape(-1, npts), d, npts) print("climatologies built (training years only)", flush=True) stat = np.stack([st.dz.values, st.orog.values / 1000.0, st.orog_std.values / 100.0, st.landfrac.values]).astype("float32") stat[0] /= 500.0 # dz to a comparable scale def make_batch(years, rows): b = bil[years].reshape(-1, npts)[rows] ob = obs[years].reshape(-1, npts)[rows] dd = doy[years].reshape(-1)[rows] ll = np.tile(np.arange(nl), len(years) * ni)[rows] wk = wk_of_lead[ll] cb = clim_bil[wk, dd - 1] co = clim_obs[wk, dd - 1] n = len(rows) x = np.empty((n, NCH, H, Wd), dtype="float32") x[:, 0] = (b / 30.0).reshape(n, H, Wd) x[:, 1] = (b - cb).reshape(n, H, Wd) x[:, 2] = np.nan_to_num(co - cb).reshape(n, H, Wd) for k in range(4): x[:, 3 + k] = stat[k] x[:, 7] = (ll / 45.0)[:, None, None] x[:, 8] = np.sin(2 * np.pi * dd / NDOY)[:, None, None] x[:, 9] = np.cos(2 * np.pi * dd / NDOY)[:, None, None] y = (ob - b).reshape(n, H, Wd) return (torch.from_numpy(np.nan_to_num(x)), torch.from_numpy(np.nan_to_num(y)), torch.from_numpy(np.isfinite(y).astype("float32")), b.reshape(n, H, Wd)) ntime_tr = len(tr) * ni * nl sel = rng.choice(ntime_tr, min(args.ntrain, ntime_tr), replace=False) nval = max(1, len(sel) // 10) val_rows, tr_rows = np.sort(sel[:nval]), np.sort(sel[nval:]) print(f"train fields {len(tr_rows)}, val fields {len(val_rows)}", flush=True) model = ResidualCNN() opt = torch.optim.Adam(model.parameters(), lr=args.lr) sched = torch.optim.lr_scheduler.ReduceLROnPlateau(opt, factor=0.5, patience=3) best, best_state = np.inf, None for ep in range(1, args.epochs + 1): model.train() perm = rng.permutation(len(tr_rows)) tot = nb = 0.0 for a in range(0, len(tr_rows), args.batch): rows = tr_rows[np.sort(perm[a:a + args.batch])] x, y, m, _ = make_batch(tr, rows) m = m * landm opt.zero_grad() pr = model(x) loss = ((pr - y) ** 2 * m).sum() / m.sum().clamp(min=1) loss.backward() opt.step() tot += float(loss) nb += 1 model.eval() with torch.no_grad(): vt = vn = 0.0 for a in range(0, len(val_rows), args.batch): rows = val_rows[a:a + args.batch] x, y, m, _ = make_batch(tr, rows) m = m * landm pr = model(x) vt += float(((pr - y) ** 2 * m).sum()) vn += float(m.sum()) vrmse = (vt / max(vn, 1)) ** 0.5 sched.step(vrmse) if vrmse < best: best, best_state = vrmse, {k: v.clone() for k, v in model.state_dict().items()} print(f" epoch {ep:3d} train {(tot/max(nb,1))**0.5:.3f} " f"val_rmse {vrmse:.3f}{' *' if vrmse == best else ''}", flush=True) model.load_state_dict(best_state) model.eval() ntime_te = len(te) * ni * nl out = np.empty((ntime_te, npts), dtype="float32") with torch.no_grad(): for a in range(0, ntime_te, 128): rows = np.arange(a, min(a + 128, ntime_te)) x, _, _, b = make_batch(te, rows) out[a:a + len(rows)] = (b + model(x).numpy()).reshape(len(rows), npts) os.makedirs(args.outdir, exist_ok=True) np.save(f"{args.outdir}/pred_fold{args.fold}.npy", out.reshape(len(te), ni, nl, npts)) print(f"fold {args.fold}: best val_rmse {best:.3f}, prediction saved", flush=True) return 0 if __name__ == "__main__": sys.exit(main())