aifs/initial_conditions.py

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)