#!/usr/bin/env python3 """Per-fold EOF principal components of the wide-domain (80-150E, 15S-45N) fields. The native 1.5 deg VN forecast lacks the sub-1.5 deg structure that let the proxy cluster EHD into regional coverage (A31); the large-scale circulation over the surrounding domain is the remaining lever. We reproduce the proven F3 approach of the ../HW session (calc_eof.py / eof_freeze.py): 1. read each wide field aligned to (year, init, lead=46), flip lat ascending; 2. remove the per-(lead, day-of-year) climatology fit on TRAINING years only; 3. weight by sqrt(cos lat) for equal area; 4. build the space-space covariance over TRAINING samples only -> top-10 EOFs; 5. project every sample onto the 10 modes, standardise each PC by its TRAIN std. Because ehd_full.py uses 5-fold-by-year cross-validation, the basis is refit for every fold with that fold's test years excluded -- so the reported test metric never sees an EOF basis fit on its own years (no leakage). Output: outputs/real/wide_pcs.npz with one (5, ny, ni, nl, 10) array per variable (key = pc field name) plus `feat_names`. ehd_full.py --pcs loads it and appends the columns, indexing pcs[fold] inside the fold loop. """ import argparse import glob import os import sys import numpy as np import xarray as xr from concurrent.futures import ProcessPoolExecutor sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from downscale_ai import window_climatology # 31-day circular window, (365,npts) ROOT = "/home/tanpv/Progs/claude/HW_EC_10" NC = f"{ROOT}/data/nc" OBS = f"{ROOT}/outputs/OBS_4EC15/tx_obs_aligned.nc" OUT = f"{ROOT}/outputs/real/wide_pcs.npz" YEARS = list(range(2003, 2023)) NLEAD = 46 NMODE = 10 FOLDS = 5 FOLD_START = [2003, 2007, 2011, 2015, 2019] # matches ehd_full.py # pc field name -> (group, nc variable, pressure level or None for single-level) FIELDS = { "gh500": ("wide_plev", "gh", 500), "t850": ("wide_plev", "t", 850), } def _read_one(a): group, ncvar, level, year, mmdd = a f = os.path.join(NC, group, str(year), f"{group}_cf_{year}-{mmdd}.nc") if not os.path.exists(f): hits = glob.glob(os.path.join(NC, group, str(year), f"*{year}-{mmdd}*.nc")) if not hits: raise FileNotFoundError(f) f = hits[0] with xr.open_dataset(f) as d: da = d[ncvar] if level is not None: da = da.sel(level=level) a = da.values[:NLEAD].astype("float32") # (lead, lat, lon) return a[:, ::-1, :] # flip lat -> ascending def load_field(group, ncvar, level, inits, jobs): """(ny, ni, nl, nlat, nlon) plus ascending lat.""" tasks = [(group, ncvar, level, y, m) for y in YEARS for m in inits] with ProcessPoolExecutor(max_workers=jobs) as pool: flat = list(pool.map(_read_one, tasks)) ni = len(inits) nlat, nlon = flat[0].shape[1:] field = np.stack(flat).reshape(len(YEARS), ni, NLEAD, nlat, nlon) # lat coordinate (ascending) from the first file g0 = _sample_path(group, YEARS[0], inits[0]) with xr.open_dataset(g0) as d: lat = d.latitude.values[::-1].astype("float32") return field, lat def _sample_path(group, year, mmdd): f = os.path.join(NC, group, str(year), f"{group}_cf_{year}-{mmdd}.nc") if os.path.exists(f): return f return glob.glob(os.path.join(NC, group, str(year), f"*{year}-{mmdd}*.nc"))[0] def remove_clim(field, doy, train_mask): """Subtract per-(lead, doy) climatology fit on TRAIN years. Returns (ny,ni,nl,ng).""" ny, ni, nl, nlat, nlon = field.shape ng = nlat * nlon flat = field.reshape(ny, ni, nl, ng).copy() for k in range(nl): vals = flat[:, :, k, :].reshape(ny * ni, ng) d = doy[:, :, k].reshape(ny * ni) tr = np.repeat(train_mask, ni) clim = window_climatology(vals[tr], d[tr], ng) # (365, ng) flat[:, :, k, :] = (vals - clim[d - 1]).reshape(ny, ni, ng) return flat def eof_pcs(anom, lat, nlon, train_mask, ni): """anom (ny,ni,nl,ng) -> standardised PCs (ny,ni,nl,NMODE), train-fit basis.""" ny, niL, nl, ng = anom.shape w = np.repeat(np.sqrt(np.cos(np.deg2rad(lat))).astype("float32"), nlon) # (ng,) X = np.nan_to_num(anom.reshape(-1, ng)) * w[None, :] # (ns, ng) ns = X.shape[0] samp_year = np.repeat(np.arange(ny), ni * nl) tr = np.isin(samp_year, np.where(train_mask)[0]) cov = (X[tr].astype("f8").T @ X[tr].astype("f8")) / max(tr.sum() - 1, 1) vals, vecs = np.linalg.eigh(cov) order = np.argsort(vals)[::-1][:NMODE] basis = vecs[:, order].astype("float32") # (ng, NMODE) pcs = X @ basis # (ns, NMODE) sd = pcs[tr].std(axis=0) sd[sd == 0] = 1.0 pcs = (pcs / sd).astype("float32") evr = vals[order] / max(np.trace(cov), 1e-9) return pcs.reshape(ny, ni, nl, NMODE), evr def main(): ap = argparse.ArgumentParser() ap.add_argument("--jobs", type=int, default=10) ap.add_argument("--out", default=OUT) ap.add_argument("--fields", nargs="*", default=list(FIELDS)) args = ap.parse_args() o = xr.open_dataset(OBS) inits = [str(x) for x in o.init.values] doy = o.doy.values # (ny, ni, nl) ni = len(inits) per = len(YEARS) // FOLDS out = {} feat_names = [] for name in args.fields: group, ncvar, level = FIELDS[name] print(f"[{name}] loading {group}/{ncvar}@{level} ...", flush=True) field, lat = load_field(group, ncvar, level, inits, args.jobs) nlon = field.shape[-1] pc_cube = np.empty((FOLDS, len(YEARS), ni, NLEAD, NMODE), dtype="float32") for fi in range(FOLDS): te = np.arange(fi * per, fi * per + per) train_mask = ~np.isin(np.arange(len(YEARS)), te) anom = remove_clim(field, doy, train_mask) pcs, evr = eof_pcs(anom, lat, nlon, train_mask, ni) pc_cube[fi] = pcs print(f" [{name}] fold {fi+1}/{FOLDS} evr(top10)={evr.sum():.3f} " f"mode1={evr[0]:.3f}", flush=True) out[name] = pc_cube feat_names += [f"pc_{name}_{k}" for k in range(NMODE)] out["feat_names"] = np.array(feat_names) os.makedirs(os.path.dirname(args.out), exist_ok=True) np.savez_compressed(args.out, **out) sz = os.path.getsize(args.out) / 1e6 print(f"wrote {args.out} ({sz:.0f} MB); {len(feat_names)} PC features") if __name__ == "__main__": sys.exit(main())