aifs/weathernext2.py

"""
aifs.weathernext2
==================
Standalone pipeline for Google DeepMind's WeatherNext 2, via the (currently
unmerged) Hugging Face ``transformers`` fork:
https://huggingface.co/kashif/weathernext2

Deliberately independent from the AIFS pipeline (aifs.initial_conditions /
aifs.forecast) for now, even though both ultimately pull from ECMWF Open
Data — WeatherNext2 uses the raw regular 0.25° 721x1440 grid directly (no
regridding to AIFS's N320 grid), a different set of variables/levels, and a
completely different inference stack (``transformers``, not ``anemoi``).

Supports multi-step autoregressive rollout (see run_forecast), via the
feature extractor's own advance_state. WeatherNext2 is architecturally
generative (a "Functional Generative Network"): one run with a fixed seed
gives one reproducible sample, not necessarily the same thing as AIFS's
single deterministic forecast.

Field/level requirements below are read directly from the model's own HF
config (``WeatherNext2Config``) and the feature extractor's documented
input contract — not guessed from the paper.
"""

from __future__ import annotations

import datetime
from pathlib import Path

import numpy as np

CHECKPOINT = "kashif/weathernext2"
STEP_HOURS = 6  # WeatherNext2 advances in fixed 6-hour steps, same cadence as AIFS

DEFAULT_STATIC_CACHE = Path("weathernext2_static_cache.npz")

#: WeatherNext2's own required pressure levels (from its HF config) — notably
#: does NOT include 10 hPa, unlike AIFS's 14-level list.
PRESSURE_LEVELS = [50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000]

#: ECMWF Open Data pressure-level params needed -> WeatherNext2 variable name.
#: "gh" (geopotential height) is converted to true geopotential (Z = gh * g)
#: before use, same conversion aifs.initial_conditions already does for AIFS.
PRESSURE_PARAM_MAP = {
    "gh": "geopotential",
    "t":  "temperature",
    "u":  "u_component_of_wind",
    "v":  "v_component_of_wind",
    "w":  "vertical_velocity",
    "q":  "specific_humidity",
}

#: ECMWF Open Data surface params needed (time-varying) -> WeatherNext2 name.
SURFACE_PARAM_MAP = {
    "2t":   "2m_temperature",
    "msl":  "mean_sea_level_pressure",
    "10v":  "10m_v_component_of_wind",
    "10u":  "10m_u_component_of_wind",
    # ECMWF Open Data has no raw "sst" param at all (confirmed against the
    # live API — it suggests "skt" instead). Skin temperature is the closest
    # available proxy over ocean points; flagged here as an approximation,
    # not a real substitute for a proper SST analysis.
    "skt":  "sea_surface_temperature",
    "100u": "100m_u_component_of_wind",
    "100v": "100m_v_component_of_wind",
}

#: Static (time-invariant) surface params -> WeatherNext2 name. Fetched once,
#: cached to disk, and reused for any date — orography/land-sea mask don't change.
STATIC_PARAM_MAP = {
    "z":   "geopotential_at_surface",
    "lsm": "land_sea_mask",
}

G = 9.80665  # geopotential height (m) -> geopotential (m^2/s^2)

#: Calendar-forcing variables. The model's config lists these as ordinary
#: `input_variables` (not just target-step forcings) — meaning the feature
#: extractor expects them supplied for BOTH conditioning frames (t-6h, t) as
#: regular time-varying inputs, computed the same way for any timestamp.
#: `seconds_since_epoch` at call time only auto-computes them for the
#: *target* step being predicted, not these past frames — confirmed by
#: reading the extractor's source directly (a first attempt trusted an
#: AI-summarized description of it, which was wrong: it read as "computed
#: internally" without this past/future distinction).
_SECONDS_PER_DAY = 24 * 3600
_AVG_DAYS_PER_YEAR = 365.24219

#: ECMWF's native regular grid — same convention as aifs.era5_verify's ERA5
#: grid (lat descends 90 -> -90, lon ascends 0 -> 359.75, both 0.25° steps).
GRID_LATITUDES = np.linspace(90.0, -90.0, 721)
GRID_LONGITUDES = np.linspace(0.0, 359.75, 1440)

#: Meteorological (non-cyclone-diagnostic) output fields worth plotting.
PLOTTABLE_ATMOSPHERIC = ["temperature", "geopotential", "u_component_of_wind", "v_component_of_wind", "vertical_velocity", "specific_humidity"]
PLOTTABLE_SURFACE = ["2m_temperature", "mean_sea_level_pressure", "10m_u_component_of_wind", "10m_v_component_of_wind", "sea_surface_temperature", "100m_u_component_of_wind", "100m_v_component_of_wind", "total_precipitation_6hr"]


def _forcings_for_frame(dt: datetime.datetime) -> dict:
    """
    Calendar-forcing values for one conditioning frame — same formula as the
    feature extractor's own ``compute_forcings``, replicated here so IC
    fetching doesn't need to load the model/processor first.

    Returns ``year_progress_sin/cos`` shaped ``(1,)`` (global — no lat/lon)
    and ``day_progress_sin/cos`` shaped ``(1, 1440)`` (varies by longitude,
    since it encodes local solar time).
    """
    seconds = np.array([dt.replace(tzinfo=datetime.timezone.utc).timestamp()], dtype=np.int64)

    year_progress = np.mod(seconds / _SECONDS_PER_DAY / _AVG_DAYS_PER_YEAR, 1.0).astype(np.float32)
    year = year_progress * (2 * np.pi)

    greenwich = np.mod(seconds, _SECONDS_PER_DAY) / _SECONDS_PER_DAY
    lon_offsets = np.deg2rad(GRID_LONGITUDES) / (2 * np.pi)
    day_progress = np.mod(greenwich[..., None] + lon_offsets, 1.0).astype(np.float32)
    day = day_progress * (2 * np.pi)

    return {
        "year_progress_sin": np.sin(year),
        "year_progress_cos": np.cos(year),
        "day_progress_sin": np.sin(day),
        "day_progress_cos": np.cos(day),
    }


# ── Fetching (native 0.25° grid — no regridding, unlike AIFS) ──────────────────

def _fetch_params(ekd, date: datetime.datetime, param: list[str], **kwargs) -> dict:
    """
    Download `param` at a single time `date` on ECMWF's native regular grid.
    Returns {ecmwf_param_name: (721, 1440) array} for single-level params, or
    {f"{ecmwf_param_name}_{level}": (721, 1440) array} when `levelist` is given.
    """
    dataset = ekd.from_source("ecmwf-open-data", date=date, param=param, source="ecmwf", **kwargs)
    fields = dataset.to_fieldlist() if hasattr(dataset, "to_fieldlist") else dataset

    out = {}
    for field in fields:
        values = field.to_numpy()
        assert values.shape == (721, 1440), (
            f"Unexpected grid shape for {field.metadata('param')}: {values.shape} "
            "(expected ECMWF's native 721x1440 regular grid)"
        )
        # ECMWF's raw grib longitude axis starts at 180° (confirmed via
        # longitudeOfFirstGridPointInDegrees), not 0° — column 0 is lon=180,
        # wrapping through 359.75/0 back to 179.75 at the last column. Rolling
        # by half the width re-indexes it to ascending 0->359.75 (column 0 =
        # lon 0), matching GRID_LONGITUDES and what WeatherNext2 expects —
        # same fix aifs.initial_conditions already applies for AIFS, for the
        # same reason. Without this, every field fed to the model is
        # geographically offset by 180°, not just the plot.
        values = np.roll(values, -values.shape[1] // 2, axis=1)
        name = field.metadata("param")
        if kwargs.get("levelist"):
            name = f"{name}_{field.metadata('levelist')}"
        out[name] = values
    return out


def _fetch_pressure_frame(ekd, date: datetime.datetime) -> dict:
    """One time frame of all pressure-level vars, shape (levels, 721, 1440) each."""
    raw = _fetch_params(ekd, date, list(PRESSURE_PARAM_MAP), levtype="pl", levelist=PRESSURE_LEVELS)

    frame = {}
    for ecmwf_param, wn2_name in PRESSURE_PARAM_MAP.items():
        stacked = np.stack([raw[f"{ecmwf_param}_{level}"] for level in PRESSURE_LEVELS])
        frame[wn2_name] = stacked * G if ecmwf_param == "gh" else stacked
    return frame


def _fetch_surface_frame(ekd, date: datetime.datetime) -> dict:
    """One time frame of all time-varying surface vars, shape (721, 1440) each."""
    raw = _fetch_params(ekd, date, list(SURFACE_PARAM_MAP), levtype="sfc")
    return {wn2_name: raw[ecmwf_param] for ecmwf_param, wn2_name in SURFACE_PARAM_MAP.items()}


def _fetch_static(ekd, date: datetime.datetime, cache_path: Path) -> dict:
    """Time-invariant fields — fetched once from any recent date, cached forever."""
    if cache_path.exists():
        return dict(np.load(str(cache_path)))

    raw = _fetch_params(ekd, date, list(STATIC_PARAM_MAP), levtype="sfc")
    static = {wn2_name: raw[ecmwf_param] for ecmwf_param, wn2_name in STATIC_PARAM_MAP.items()}
    cache_path.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(str(cache_path), **static)
    return static


# ── Public API ──────────────────────────────────────────────────────────────────

def load_ics(cache_dir: Path | str = DEFAULT_STATIC_CACHE, date: datetime.datetime | None = None):
    """
    Generator: fetches the two input frames (t-6h, t) plus static fields, and
    assembles the exact ``state`` dict WeatherNext2FeatureExtractor expects
    (see its documented contract): time-varying vars shaped
    ``[1, 2, (levels,) 721, 1440]``, static vars shaped ``[721, 1440]``,
    plus the 4 calendar-forcing variables computed for both frames (see
    :func:`_forcings_for_frame` — the model's config lists these as regular
    inputs, not just target-step forcings).

    Yields
    ------
    ("log", str)                          -- progress messages
    ("result", (state, date))             -- final payload
    """
    import earthkit.data as ekd

    from ecmwf.opendata import Client as OpendataClient

    ekd.config.set({"cache-policy": "user"})

    if date is None:
        date = OpendataClient("ecmwf").latest()
        yield "log", f"📅  Latest ECMWF run: {date}"
    else:
        yield "log", f"📅  ECMWF run: {date}"

    t0 = date - datetime.timedelta(hours=STEP_HOURS)

    yield "log", "⬇️   Fetching static fields (orography, land-sea mask)…"
    static = _fetch_static(ekd, date, Path(cache_dir) if isinstance(cache_dir, (str, Path)) else DEFAULT_STATIC_CACHE)

    yield "log", f"⬇️   Fetching pressure-level fields @ t-6h ({t0})…"
    pressure_t0 = _fetch_pressure_frame(ekd, t0)
    yield "log", f"⬇️   Fetching pressure-level fields @ t ({date})…"
    pressure_t1 = _fetch_pressure_frame(ekd, date)

    yield "log", f"⬇️   Fetching surface fields @ t-6h ({t0})…"
    surface_t0 = _fetch_surface_frame(ekd, t0)
    yield "log", f"⬇️   Fetching surface fields @ t ({date})…"
    surface_t1 = _fetch_surface_frame(ekd, date)

    forcings_t0 = _forcings_for_frame(t0)
    forcings_t1 = _forcings_for_frame(date)

    state = {}
    for name in PRESSURE_PARAM_MAP.values():
        state[name] = np.stack([pressure_t0[name], pressure_t1[name]])[np.newaxis, ...]  # [1, 2, levels, 721, 1440]
    for name in SURFACE_PARAM_MAP.values():
        state[name] = np.stack([surface_t0[name], surface_t1[name]])[np.newaxis, ...]    # [1, 2, 721, 1440]
    for name in STATIC_PARAM_MAP.values():
        state[name] = static[name]                                                       # [721, 1440]
    for name in ("year_progress_sin", "year_progress_cos", "day_progress_sin", "day_progress_cos"):
        state[name] = np.stack([forcings_t0[name], forcings_t1[name]], axis=1)           # [1, 2] or [1, 2, 1440]

    yield "log", "✓  All fields fetched — ready for WeatherNext2."
    yield "result", (state, date)


_MODEL_CACHE: dict = {}


def _load_model(device: str):
    if "model" not in _MODEL_CACHE:
        from transformers import WeatherNext2FeatureExtractor, WeatherNext2ForWeatherForecasting

        model = WeatherNext2ForWeatherForecasting.from_pretrained(CHECKPOINT, device_map=device).eval()
        processor = WeatherNext2FeatureExtractor.from_pretrained(CHECKPOINT)
        _MODEL_CACHE["model"] = model
        _MODEL_CACHE["processor"] = processor
    return _MODEL_CACHE["model"], _MODEL_CACHE["processor"]


def run_forecast(state: dict, date: datetime.datetime, num_steps: int = 1, device: str = "cpu", seed: int = 0):
    """
    Generator: runs `num_steps` autoregressive 6h WeatherNext2 steps from
    `state` (as built by :func:`load_ics`), advancing the conditioning state
    between steps via the feature extractor's own ``advance_state`` — drops
    the oldest frame, appends the just-predicted one, recomputes calendar
    forcings for the new valid time.

    Yields ("log", str) then ("result", states) — `states` is a list (one
    per step, matching AIFS's forecast state list) of forecast states on
    WeatherNext2's native regular grid — NOT flattened, unlike AIFS's
    irregular-grid states:
    ``{"date", "fields": {name: (levels, 721, 1440) or (721, 1440) array}}``.
    Grid coordinates are the module-level ``GRID_LATITUDES``/``GRID_LONGITUDES``
    (both 1-D, since it's a regular grid) — use :func:`plot_field` to render.

    `device` defaults to "cpu" deliberately — the model needs ~50 GB for a
    single member at fp32, comfortably more than this app's default ZeroGPU
    allocation (48 GB); CPU sidesteps that (the model card's own benchmark:
    ~96s per forward pass on CPU) without needing a bigger, costlier GPU tier.

    NOT verified end-to-end in development (the dev sandbox this was written
    in has 11 GB RAM, far short of what real inference needs). What HAS been
    verified, using the real, lightweight feature extractor (no model weights
    needed) fed synthetic-but-correctly-shaped data: the fetch/state-assembly
    steps against real ECMWF data, the calendar-forcing math against the
    real compute_forcings (bit-for-bit match), a full processor(...) call
    with zero missing/mismatched inputs, and the whole advance_state rollout
    loop staying shape-consistent across 3 autoregressive steps. NOT
    verified: the actual model forward pass and postprocess() — those need
    the real ~50GB weights.
    """
    import torch

    yield "log", "🤖  Loading WeatherNext2 (first run downloads weights + builds the mesh — a few minutes)…"
    model, processor = _load_model(device)

    states = []
    current_state = state
    current_date = date
    for step in range(1, num_steps + 1):
        target_date = current_date + datetime.timedelta(hours=STEP_HOURS)
        seconds_since_epoch = np.array([target_date.replace(tzinfo=datetime.timezone.utc).timestamp()])

        yield "log", f"🌩️  Step {step}/{num_steps} — target valid time {target_date} UTC (seed={seed})…"
        inputs = processor(current_state, seconds_since_epoch=seconds_since_epoch).to(model.device)
        with torch.no_grad():
            outputs = model(**inputs, generator=torch.Generator().manual_seed(seed))
        forecast = processor.postprocess(outputs.prediction, current_state)

        fields = {}
        for name in PLOTTABLE_ATMOSPHERIC + PLOTTABLE_SURFACE:
            if name not in forecast:
                continue
            fields[name] = np.squeeze(np.asarray(forecast[name]))
        states.append({"date": target_date, "fields": fields})
        yield "log", f"✅  Step {step}/{num_steps} done ({target_date} UTC)."

        if step < num_steps:
            current_state = processor.advance_state(current_state, forecast, seconds_since_epoch)
            current_date = target_date

    yield "log", f"✅  WeatherNext2 rollout complete — {len(states)} step(s)."
    yield "result", states


# ── Plotting (regular grid — a plain pcolormesh, no triangulation needed) ─────

def plot_field(result: dict, field_name: str, level: int | None = None):
    """Plot one output field from :func:`run_forecast` on its native regular grid."""
    import cartopy.crs as ccrs
    import cartopy.feature as cfeature
    import matplotlib.pyplot as plt

    if field_name not in result["fields"]:
        raise KeyError(f"'{field_name}' not in forecast output. Available: {sorted(result['fields'])}")

    data = result["fields"][field_name]
    if data.ndim == 3:  # (levels, lat, lon) — pick one level to plot
        if level is None:
            raise ValueError(f"'{field_name}' is a pressure-level field — pass `level`, one of {PRESSURE_LEVELS}.")
        data = data[PRESSURE_LEVELS.index(level)]

    lons_plot = np.where(GRID_LONGITUDES > 180, GRID_LONGITUDES - 360, GRID_LONGITUDES)
    order = np.argsort(lons_plot)

    fig, ax = plt.subplots(figsize=(11, 6), subplot_kw={"projection": ccrs.PlateCarree()})
    ax.coastlines()
    ax.add_feature(cfeature.BORDERS, linestyle=":")
    mesh = ax.pcolormesh(
        lons_plot[order], GRID_LATITUDES, data[:, order],
        transform=ccrs.PlateCarree(), cmap="RdBu_r", shading="auto",
    )
    fig.colorbar(mesh, ax=ax, orientation="vertical", shrink=0.7, label=field_name)
    title = f"WeatherNext2 — {field_name}" + (f" @ {level}hPa" if level else "") + f" — {result['date']}"
    ax.set_title(title, fontsize=11)
    fig.tight_layout()
    return fig