import datetime import random import time import warnings from collections import defaultdict from pathlib import Path import numpy as np warnings.filterwarnings("ignore", category=DeprecationWarning) warnings.filterwarnings("ignore", category=UserWarning) # ── Meteorological variable lists ───────────────────────────────────────────── #: Surface parameters (levtype=sfc) PARAM_SFC = [ "10u", "10v", "2d", "2t", "msl", "skt", "sp", "tcw", "lsm", "z", "slor", "sdor", "sd", ] #: Soil parameters (levtype=sfc, levelist=[1,2]) PARAM_SOIL = ["vsw", "sot"] SOIL_LEVELS = [1, 2] #: Ocean-wave parameters (stream=wave) PARAM_WAVE = [ "wmb", "h1012", "h1214", "h1417", "h1721", "h2125", "h2530", "mwd", "cdww", "mwp", "swh", ] #: Pressure-level parameters PARAM_PL = ["gh", "t", "u", "v", "q"] LEVELS = [1000, 925, 850, 700, 600, 500, 400, 300, 250, 200, 150, 100, 50, 10] SOURCE = "ecmwf" #: ECMWF Open Data's primary endpoint ("ecmwf") only keeps a rolling ~4-day #: window. The "aws" named source (an S3 mirror, still via the same #: ecmwf-opendata/earthkit-data client) retains a much deeper archive #: (observed back to 2023-01-18) — used for historical initial conditions. HISTORICAL_SOURCE = "aws" #: Valid ECMWF synoptic run hours (UTC) — the archive only has data at these. VALID_RUN_HOURS = (0, 6, 12, 18) #: Earliest date the "aws" historical archive has been observed to serve. EARLIEST_HISTORICAL_DATE = datetime.date(2023, 1, 18) #: Before this date, ECMWF's 06/18 UTC runs used a reduced product ("scda" for #: pressure levels, "scwv" for waves) instead of the full "oper"/"wave" stream #: — confirmed missing: all pressure-level vars at 10 hPa, and everything in #: PARAM_WAVE except mwd/mwp/swh. 00/12 UTC always used the full stream, even #: before this cutover. This was unified across all run hours starting on this #: date, but that's *after* EarthMover's free ERA5 archive's coverage ends (see #: aifs.era5_verify) — so for this app's actual use case (historical init + #: ERA5 verification), 06/18 UTC is never a usable choice anyway. REDUCED_PRODUCT_CUTOVER = datetime.date(2026, 5, 12) REDUCED_PRODUCT_HOURS = (6, 18) #: Run hours guaranteed to carry the full field set, at any historical date. FULL_FIELD_RUN_HOURS = (0, 12) # ── Cache helpers ───────────────────────────────────────────────────────────── DEFAULT_CACHE_DIR = Path("ic_cache") def _cache_path(date: datetime.datetime, cache_dir: Path) -> Path: return cache_dir / f"ic_{date.strftime('%Y%m%dT%H%M%S')}.npz" def _save(date: datetime.datetime, fields: dict, cache_dir: Path) -> Path: cache_dir.mkdir(parents=True, exist_ok=True) path = _cache_path(date, cache_dir) np.savez_compressed(str(path), **fields) return path def _try_load(date: datetime.datetime, cache_dir: Path): """Return ``(fields_dict, path)`` if cached, else ``(None, None)``.""" path = _cache_path(date, cache_dir) if path.exists(): try: return dict(np.load(str(path))), path except Exception: path.unlink() return None, None def list_cached(cache_dir: Path = DEFAULT_CACHE_DIR) -> list[Path]: """Return all cached .npz files, newest first.""" if not cache_dir.exists(): return [] return sorted(cache_dir.glob("ic_*.npz"), reverse=True) # ── Download helpers ────────────────────────────────────────────────────────── _RETRIABLE_KEYWORDS = ("429", "rate limit", "too many requests", "timeout", "connection reset", "503", "service unavailable") def _fetch_with_retry(ekd, ekr, date, param, max_retries: int = 6, **kwargs): """ Generator wrapper around _fetch_fields with exponential backoff. Yields ("log", str) on each retry, then ("result", dict) on success. Raises the last exception if all retries are exhausted. """ for attempt in range(max_retries): try: result = _fetch_fields(ekd, ekr, date, param, **kwargs) yield "result", result return except Exception as exc: msg = str(exc).lower() retriable = any(k in msg for k in _RETRIABLE_KEYWORDS) if retriable and attempt < max_retries - 1: wait = min(5 * (2 ** attempt) + random.uniform(0, 3), 120) yield "log", f"⚠️ Server busy — retrying in {wait:.0f}s (attempt {attempt + 2}/{max_retries})…" time.sleep(wait) else: raise def _fetch_fields(ekd, ekr, date, param, levelist=None, source=SOURCE, **kwargs) -> dict: """ Download ``param`` for two time-steps (t-6h, t) and return a dict ``{variable_name: np.ndarray shape (2, N320_nodes)}``. """ levelist = levelist or [] raw: dict[str, list] = defaultdict(list) for t in [date - datetime.timedelta(hours=6), date]: dataset = ekd.from_source( "ecmwf-open-data", date=t, param=param, levelist=levelist, source=source, **kwargs, ) for field in dataset: assert field.to_numpy().shape == (721, 1440), ( f"Unexpected grid shape for {field.metadata('param')}: " f"{field.to_numpy().shape}" ) # Shift lon from [0,360) to [-180,180) then regrid to N320 Gaussian values = np.roll(field.to_numpy(), -field.shape[1] // 2, axis=1) values = ekr.interpolate(values, {"grid": (0.25, 0.25)}, {"grid": "N320"}) if levelist: name = f"{field.metadata('param')}_{field.metadata('levelist')}" else: name = field.metadata("param") raw[name].append(values) return {k: np.stack(v) for k, v in raw.items()} def _build_fields(ekd, ekr, date: datetime.datetime, source: str = SOURCE): """Download and transform all required fields for ``date``.""" fields: dict = {} log_lines: list[str] = [] def log(msg: str): """Append a line and yield the whole history so far.""" log_lines.append(msg) return "log", "\n".join(log_lines) def fetch(label, *args, **kwargs): """Yield log messages then return the result dict.""" yield log(label) result = None kwargs.setdefault("source", source) for kind, payload in _fetch_with_retry(ekd, ekr, *args, **kwargs): if kind == "log": yield log(f" {payload}") else: result = payload return result # Surface fields sfc = yield from fetch("⬇ Surface fields …", date, PARAM_SFC, levtype="sfc") fields.update(sfc) # Ocean-wave fields wave = yield from fetch("⬇ Wave fields …", date, PARAM_WAVE, stream="wave") fields.update(wave) # Soil fields (kept separate for renaming below) soil = yield from fetch("⬇ Soil fields …", date, PARAM_SOIL, levelist=SOIL_LEVELS) # Pressure-level fields pl = yield from fetch("⬇ Pressure-level fields …", date, PARAM_PL, levelist=LEVELS) fields.update(pl) yield log("✓ All fields fetched.") # ── Transformations ─────────────────────────────────────────────────────── # Wave direction: decompose scalar angle into sin/cos components mwd = fields.pop("mwd") mwd_rad = np.deg2rad(mwd) fields["cos_mwd"] = np.cos(mwd_rad) fields["sin_mwd"] = np.sin(mwd_rad) # Rename soil fields to ECMWF short-names expected by AIFS _soil_rename = { "sot_1": "stl1", "sot_2": "stl2", "vsw_1": "swvl1", "vsw_2": "swvl2", } for src, dst in _soil_rename.items(): fields[dst] = soil[src] # Remove q levels that AIFS does not use fields.pop("q_10", None) fields.pop("q_50", None) # Apply land-sea mask to snow depth and soil moisture (ocean → NaN) try: lsm = ekd.from_source("file", "lsm.grib")[0].to_numpy(flatten=True) ocean_mask = np.equal(lsm, 0) for var in ("sd", "swvl1", "swvl2"): if var in fields: fields[var][:, ocean_mask] = np.nan except Exception: pass # lsm.grib not found; skip masking # Convert geopotential height → geopotential (Z = gh × g) G = 9.80665 for level in LEVELS: gh = fields.pop(f"gh_{level}", None) if gh is not None: fields[f"z_{level}"] = gh * G yield "result" , fields # ── Public API ──────────────────────────────────────────────────────────────── def load_ics( cache_dir: Path | str = DEFAULT_CACHE_DIR, force: bool = False, date: datetime.datetime | None = None, ): """ Generator version of load_ics. Parameters ---------- date: If ``None`` (default), fetch the latest available run from ECMWF Open Data's primary endpoint — unchanged live behaviour. If given, fetch that specific historical run instead, from the "aws" S3 mirror (a much deeper archive than the ~4-day rolling window of the primary endpoint — observed back to 2023-01-18). Must fall on a synoptic run hour (00/06/12/18 UTC). Yields ------ ("log", str) -- progress messages ("result", (fields, date)) -- final payload (always the last item yielded) """ import earthkit.data as ekd import earthkit.regrid as ekr ekd.config.set({"cache-policy": "user"}) cache_dir = Path(cache_dir) if date is None: from ecmwf.opendata import Client as OpendataClient date = OpendataClient(SOURCE).latest() yield "log", f"📅 Latest ECMWF run: {date}" source = SOURCE else: if date.hour not in VALID_RUN_HOURS or date.minute or date.second or date.microsecond: raise ValueError( f"date must fall on a synoptic run hour {VALID_RUN_HOURS} UTC, got {date}" ) if date.hour in REDUCED_PRODUCT_HOURS and date.date() < REDUCED_PRODUCT_CUTOVER: raise ValueError( f"{date} is a 06/18 UTC run before {REDUCED_PRODUCT_CUTOVER.isoformat()} — " "ECMWF served a reduced product at those hours back then (missing 10 hPa " "pressure levels and most wave fields AIFS needs). Use a 00 or 12 UTC run instead." ) yield "log", f"📅 Historical ECMWF run: {date} (via '{HISTORICAL_SOURCE}' archive)" source = HISTORICAL_SOURCE if not force: cached, path = _try_load(date, cache_dir) if cached is not None: sz_mb = path.stat().st_size / 1e6 yield "log", f"✅ Loaded from cache ({sz_mb:.0f} MB) → {path}" yield "result", (cached, date) return yield "log", "⬇️ Downloading initial conditions …" fields = None for kind, payload in _build_fields(ekd, ekr, date, source=source): if kind == "log": yield "log", payload else: # "result" fields = payload path = _save(date, fields, cache_dir) sz_mb = path.stat().st_size / 1e6 yield "log", f"💾 Saved to {path} ({sz_mb:.0f} MB)" yield "result", (fields, date)