#!/usr/bin/env python3 """Build a full 1.5 deg predictor proxy from the existing 0.25 deg EC archive. Purpose: the EHD experiment is information-limited (see developing.txt A19), so the full predictor set is what matters, and the real 1.5 deg download needs ~30 h. The 0.25 deg archive already holds every variable for 2003-2022, so averaging it onto the true 1.5 deg MARS grid gives a faithful stand-in to develop the whole pipeline now. When the real data lands the same downstream code runs unchanged on it. Output: outputs/proxy/ec15_predictors.nc with each variable as (year, init, lead, lat, lon) = 20 x 104 x 46 x 13 x 10 Time conventions of the source archive: Srf_daily_ave 46 daily means, one per lead day -> used directly Srf_24hour 47 points at 0,24,..,1104 h. Accumulated fields (fluxes, radiation, precip) are cumulative from the start, so lead day d is value[d+1]-value[d]; instantaneous fields take value[d]. *_plev 47 instantaneous points; lead day d takes value[d]. """ import argparse import glob import os import sys from concurrent.futures import ProcessPoolExecutor import numpy as np import xarray as xr YEARS = list(range(2003, 2023)) NLEAD = 46 EC = "/home/tanpv/Data/QG.25.196/EC" OUT = "/home/tanpv/Progs/claude/HW_EC_10/outputs/proxy/ec15_predictors.nc" EC15_LON = 100.0 + 1.5 * np.arange(10) EC15_LAT = 24.0 - 1.5 * np.arange(13) ACCUM = {"sshf", "slhf", "ssrd", "strd", "ssr", "str", "ttr", "cp", "tp"} # Positive-definite members of ACCUM. De-accumulating by diff yields a small number # of negative daily values (~0.4% of steps for tp/cp, median magnitude ~1.1 mm, i.e. # not rounding noise) because the archived cumulative field is not exactly monotone # after interpolation. Negative precipitation is unphysical, so clip these at zero. NONNEG = {"cp", "tp", "ssrd", "ssr"} # group -> (subdir, filename pattern, {varname: level or None}) GROUPS = { "srf_daily": ("Srf_daily_ave", "Srf_daily_cf_{year}-{mmdd}.nc", {v: None for v in ["t2m", "d2m", "skt", "cape", "tcw", "tcc", "sm20", "sm100", "st20", "st100"]}), "srf_24h": ("Srf_24hour", "Srf_24hour_cf_{year}-{mmdd}.nc", {v: None for v in ["sshf", "slhf", "msl", "u10", "v10", "ssrd", "strd", "ssr", "str", "ttr", "cp", "tp"]}), "t_plev": ("T_plev", "T_plev_cf_{year}-{mmdd}.nc", {"t": [500, 700, 850, 925]}), "h_plev": ("H_plev", "H_plev_cf_{year}-{mmdd}.nc", {"gh": [500, 700, 850, 925, 1000]}), "quv": ("QUV_plev", "pquv_cf_{year}-{mmdd}.nc", {"u": [500, 700, 850, 925], "v": [500, 700, 850, 925], "q": [500, 700, 850, 925, 1000], "sp": None}), "w_plev": ("W_plev", "W_plev_cf_{year}-{mmdd}.nc", {"w": None}), } def ec15_index(lat, lon): jl = np.abs(lat[:, None] - EC15_LAT[None, :]).argmin(axis=1) il = np.abs(lon[:, None] - EC15_LON[None, :]).argmin(axis=1) flat = jl[:, None] * len(EC15_LON) + il[None, :] return flat.ravel(), np.bincount(flat.ravel(), minlength=len(EC15_LAT) * len(EC15_LON)) def to_ec15(a, flat, cnt): """(t, lat, lon) on the fine grid -> (t, 13, 10) cell means.""" nt = a.shape[0] out = np.full((nt, len(EC15_LAT) * len(EC15_LON)), np.nan, dtype="float32") for k in range(nt): s = np.bincount(flat, weights=np.nan_to_num(a[k]).ravel().astype("float64"), minlength=out.shape[1]) with np.errstate(invalid="ignore"): out[k] = s / np.maximum(cnt, 1) out[:, cnt == 0] = np.nan return out.reshape(nt, len(EC15_LAT), len(EC15_LON)) def load_year(a): gname, year, inits = a sub, pat, vars_ = GROUPS[gname] res = {} flat = cnt = None for k, mmdd in enumerate(inits): f = os.path.join(EC, sub, str(year), pat.format(year=year, mmdd=mmdd)) if not os.path.exists(f): hits = glob.glob(os.path.join(EC, sub, str(year), f"*{year}-{mmdd}*.nc")) if not hits: raise FileNotFoundError(f) f = hits[0] with xr.open_dataset(f) as d: if flat is None: flat, cnt = ec15_index(d.latitude.values, d.longitude.values) for v, levs in vars_.items(): if v not in d: continue da = d[v] sels = [(v, da)] if levs is None else \ [(f"{v}{L}", da.sel(level=L)) for L in levs if L in da.level.values] for name, arr in sels: x = arr.values.astype("float32") if gname == "srf_daily": x = x[:NLEAD] elif v in ACCUM: x = np.diff(x[:NLEAD + 1], axis=0) # de-accumulate to daily if v in NONNEG: np.clip(x, 0, None, out=x) else: x = x[:NLEAD] c = to_ec15(x, flat, cnt) if name not in res: res[name] = np.empty((len(inits),) + c.shape, dtype="float32") res[name][k] = c return year, res def main(): ap = argparse.ArgumentParser() ap.add_argument("--jobs", type=int, default=10) ap.add_argument("--groups", default=",".join(GROUPS)) ap.add_argument("--out", default=OUT) ap.add_argument("--hdates", default="/home/tanpv/Data/QG.25.196/hdate_2023.txt") args = ap.parse_args() with open(args.hdates) as fh: inits = sorted({t for t in fh.read().split() if len(t) == 5 and t[2] == "-"}) data = {} for gname in args.groups.split(","): print(f"=== {gname}", flush=True) tasks = [(gname, y, inits) for y in YEARS] with ProcessPoolExecutor(max_workers=args.jobs) as pool: for year, res in pool.map(load_year, tasks): for name, arr in res.items(): if name not in data: data[name] = np.empty((len(YEARS),) + arr.shape, dtype="float32") data[name][YEARS.index(year)] = arr print(f" {year}", flush=True) ds = xr.Dataset( {k: (("year", "init", "lead", "lat", "lon"), v) for k, v in data.items()}, coords={"year": YEARS, "init": inits, "lead": np.arange(NLEAD), "lat": EC15_LAT, "lon": EC15_LON}) ds.attrs["note"] = ("1.5 deg proxy built by averaging the 0.25 deg EC archive onto " "the true MARS 1.5 deg grid; for method development only.") os.makedirs(os.path.dirname(args.out), exist_ok=True) ds.to_netcdf(args.out, encoding={k: {"zlib": True, "complevel": 4} for k in data}) print(f"wrote {args.out} ({os.path.getsize(args.out)/1e9:.2f} GB), " f"{len(data)} variables: {sorted(data)}") return 0 if __name__ == "__main__": sys.exit(main())