aifs/era5_verify.py

"""
aifs.era5_verify
=================
Score an AIFS forecast (a list of state dicts from :func:`aifs.forecast.run_forecast`,
initialised from a *historical* date via ``aifs.initial_conditions.load_ics(date=...)``)
against EarthMover's public ERA5 reanalysis, plus a persistence baseline (ERA5 held
fixed at the initial time — "assume nothing changes").

Only forecasts initialised from a past date can be scored this way: a forecast's
valid times are in the future relative to when it was run, and ERA5 (reanalysis)
for those times doesn't exist until well after the fact.
"""

from __future__ import annotations

import datetime

import numpy as np

from aifs.era5_env import fetch_era5_fields

STEP_HOURS = 6  # AIFS advances in fixed 6-hour steps (see aifs.forecast)
_EPOCH = datetime.datetime(1940, 1, 1)  # ERA5's valid_time units: "hours since 1940-01-01"

#: AIFS field name -> (ERA5 group, ERA5 variable name, pressure level or None).
#: Names/groups confirmed against the live store; wave fields (swh, mwp, ...) are
#: NOT included here — EarthMover's ERA5 mirror doesn't carry the wave stream.
EVAL_FIELD_MAP = {
    "2t":    ("single",   "t2m", None),
    "10u":   ("single",   "u10", None),
    "10v":   ("single",   "v10", None),
    "100u":  ("single",   "u100", None),
    "100v":  ("single",   "v100", None),
    "msl":   ("single",   "msl", None),
    "sp":    ("single",   "sp",  None),
    "tcw":   ("single",   "tcw", None),
    "t_850": ("pressure", "t",   850),
    "t_500": ("pressure", "t",   500),
    "u_850": ("pressure", "u",   850),
    "v_850": ("pressure", "v",   850),
    "z_500": ("pressure", "z",   500),
    "q_700": ("pressure", "q",   700),
}

#: ERA5's regular 0.25 degree grid (see aifs.era5_worker probing): lat descends
#: 90 -> -90, lon ascends 0 -> 359.75 — same [0, 360) convention AIFS uses.
_ERA5_LAT0, _ERA5_DLAT, _ERA5_NLAT = 90.0, -0.25, 721
_ERA5_LON0, _ERA5_DLON, _ERA5_NLON = 0.0, 0.25, 1440


def _time_idx(dt: datetime.datetime) -> int:
    return int((dt - _EPOCH).total_seconds() // 3600)


def _nearest_indices(lats: np.ndarray, lons: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Map each (lat, lon) point to its nearest cell in ERA5's regular grid."""
    lat_idx = np.clip(np.round((lats - _ERA5_LAT0) / _ERA5_DLAT).astype(int), 0, _ERA5_NLAT - 1)
    lon_idx = np.round((lons - _ERA5_LON0) / _ERA5_DLON).astype(int) % _ERA5_NLON
    return lat_idx, lon_idx


def sample_at_points(grid: np.ndarray, lat_idx: np.ndarray, lon_idx: np.ndarray) -> np.ndarray:
    """Sample a (721, 1440) ERA5 grid at pre-computed nearest-neighbour indices."""
    return grid[lat_idx, lon_idx]


def metrics(forecast: np.ndarray, truth: np.ndarray) -> dict:
    """RMSE / MAE / bias / correlation between two same-shape point arrays, NaN-safe."""
    mask = ~(np.isnan(forecast) | np.isnan(truth))
    f, t = np.asarray(forecast)[mask], np.asarray(truth)[mask]
    if len(f) < 2:
        return {"rmse": float("nan"), "mae": float("nan"), "bias": float("nan"), "corr": float("nan"), "n": int(len(f))}
    diff = f - t
    corr = float(np.corrcoef(f, t)[0, 1])
    return {
        "rmse": float(np.sqrt(np.mean(diff ** 2))),
        "mae": float(np.mean(np.abs(diff))),
        "bias": float(np.mean(diff)),
        "corr": corr,
        "n": int(len(f)),
    }


def _fetch_era5_bounds(group: str, extra_requests: list[dict], log):
    """
    Fetch a group's live ``valid_time`` coverage bounds together with
    ``extra_requests`` in one batched call — avoids paying icechunk's
    session-open cost twice. EarthMover's free ERA5 archive is NOT current
    to "now" (it lags real-time by several months, and how far isn't
    documented reliably), so callers check this live rather than trusting
    a hardcoded date.

    Returns ``(era5_start, era5_end, arrays, meta)`` — ``arrays``/``meta``
    cover ALL requests: bounds are indices 0 and 1, ``extra_requests``
    start at index 2.
    """
    bounds_requests = [
        {"group": group, "var": "valid_time", "level": None, "time_idx": 0},
        {"group": group, "var": "valid_time", "level": None, "time_idx": -1},
    ]
    arrays, meta = fetch_era5_fields(bounds_requests + extra_requests, log=log)
    if 0 in meta["errors"] or 1 in meta["errors"]:
        raise RuntimeError(
            f"Could not read ERA5's time coverage for group '{group}': "
            f"{meta['errors'].get(0) or meta['errors'].get(1)}"
        )
    era5_start = _EPOCH + datetime.timedelta(hours=float(arrays[0]))
    era5_end = _EPOCH + datetime.timedelta(hours=float(arrays[1]))
    return era5_start, era5_end, arrays, meta


def evaluate_forecast(states: list[dict], field_name: str, log=lambda msg: None) -> dict:
    """
    Score every step of a historical-init forecast against ERA5, plus a
    persistence baseline (ERA5 held fixed at t0).

    Returns a dict: ``{"field", "lats", "lons", "per_step": [...]}`` — each
    per_step entry has ``date``, ``lead_h``, ``forecast``, ``truth`` (both
    sampled onto the AIFS grid), and ``model_*``/``persistence_*`` metrics.
    """
    if not states:
        raise ValueError("No forecast states to evaluate — run a forecast first.")
    if field_name not in EVAL_FIELD_MAP:
        raise ValueError(
            f"'{field_name}' has no ERA5 counterpart configured. "
            f"Available: {sorted(EVAL_FIELD_MAP)}"
        )

    group, var, level = EVAL_FIELD_MAP[field_name]
    lats = np.asarray(states[0]["latitudes"]).ravel()
    lons = np.asarray(states[0]["longitudes"]).ravel()
    lat_idx, lon_idx = _nearest_indices(lats, lons)

    t0 = states[0]["date"] - datetime.timedelta(hours=STEP_HOURS)
    valid_times = [t0] + [s["date"] for s in states]

    # De-duplicate in case any two steps land on the same valid time.
    unique_times = sorted(set(valid_times))
    data_requests = [
        {"group": group, "var": var, "level": level, "time_idx": _time_idx(t)}
        for t in unique_times
    ]

    log(f"📡  Fetching ERA5 '{var}'" + (f"@{level}hPa" if level else "") +
        f" for {len(unique_times)} valid time(s)…")
    era5_start, era5_end, arrays, meta = _fetch_era5_bounds(group, data_requests, log)

    earliest_needed, latest_needed = unique_times[0], unique_times[-1]
    if latest_needed > era5_end or earliest_needed < era5_start:
        raise ValueError(
            f"This forecast needs ERA5 truth from {earliest_needed} to {latest_needed}, but "
            f"EarthMover's free ERA5 archive currently only covers {era5_start} to {era5_end}. "
            f"Pick an earlier historical init date — with enough room before {era5_end} to fit "
            f"your full lead time — and run the forecast again."
        )

    data_errors = {i - 2: err for i, err in meta["errors"].items() if i >= 2}
    if data_errors:
        first_err = next(iter(data_errors.values()))
        raise RuntimeError(f"Could not read '{field_name}' from EarthMover's ERA5 store: {first_err}")

    truth_by_time = {
        unique_times[i]: sample_at_points(arrays[i + 2], lat_idx, lon_idx)
        for i in range(len(unique_times))
    }

    truth_t0 = truth_by_time[t0]
    per_step = []
    for i, state in enumerate(states, start=1):
        truth = truth_by_time[state["date"]]
        forecast = np.asarray(state["fields"][field_name])

        m = metrics(forecast, truth)
        p = metrics(truth_t0, truth)

        per_step.append({
            "date": str(state["date"]),
            "lead_h": i * STEP_HOURS,
            "forecast": forecast,
            "truth": truth,
            **{f"model_{k}": v for k, v in m.items()},
            **{f"persistence_{k}": v for k, v in p.items()},
        })

    log(f"✅  Scored {len(per_step)} step(s) against ERA5 for '{field_name}'.")
    return {"field": field_name, "lats": lats, "lons": lons, "per_step": per_step}


def _climatology_grid(group: str, var: str, level: int | None, valid_date: datetime.datetime,
                       num_years: int, log) -> tuple[np.ndarray, np.ndarray, list[int]]:
    """
    Core climatology fetch: the ERA5 mean/std grid (both (721, 1440), NOT
    sampled onto any particular points yet) for one field on one calendar
    day/hour, averaged over up to `num_years` past years. Shared by
    :func:`compute_climatology` (vs. an existing forecast state) and
    :func:`run_climatology_baseline` (a standalone multi-step "model").
    """
    # Request a few extra candidate years beyond num_years so that skipped
    # years (Feb 29 in a non-leap year, or anything ERA5 doesn't have yet
    # for the most recent year or two) still leave enough successful
    # samples — cheaper than pre-checking the archive's bounds separately.
    BUFFER = 5
    candidate_years = []
    year = valid_date.year - 1  # climatology uses PAST years only
    while len(candidate_years) < num_years + BUFFER and year >= 1940:
        try:
            candidate_years.append(valid_date.replace(year=year))
        except ValueError:
            pass  # Feb 29 in a non-leap year — skip it
        year -= 1

    requests = [
        {"group": group, "var": var, "level": level, "time_idx": _time_idx(t)}
        for t in candidate_years
    ]
    log(f"📡  Fetching up to {num_years} years of ERA5 '{var}'" + (f"@{level}hPa" if level else "") +
        f" for {valid_date.strftime('%m-%d %H:%M')} UTC…")
    arrays, meta = fetch_era5_fields(requests, log=log)

    successful = [(candidate_years[i], arrays[i]) for i in range(len(candidate_years)) if i not in meta["errors"]]
    if len(successful) < 2:
        first_err = next(iter(meta["errors"].values()), "no data available")
        raise RuntimeError(
            f"Could not build a climatology for '{var}' — only {len(successful)} of "
            f"{len(candidate_years)} candidate years had ERA5 data ({first_err})."
        )
    successful = successful[:num_years]
    used_years = [t.year for t, _ in successful]

    samples = np.stack([grid for _, grid in successful])
    clim_mean = np.nanmean(samples, axis=0)
    clim_std = np.nanstd(samples, axis=0)
    return clim_mean, clim_std, used_years


def compute_climatology(state: dict, field_name: str, num_years: int = 10, log=lambda msg: None) -> dict:
    """
    Compare ONE forecast state's field against ERA5's climatology for that
    calendar day/hour — the mean (and spread) over ``num_years`` years of
    ERA5 history on the same month/day/hour.

    Unlike :func:`evaluate_forecast`, this works for a LIVE (today's or a
    future) forecast: no ERA5 truth is needed for the forecast's own valid
    time, only for many *past* years on the same calendar day — sidestepping
    ERA5's coverage-ceiling problem entirely. It answers "is this forecast
    unusual for the time of year?", not "is this forecast correct?".
    """
    if field_name not in EVAL_FIELD_MAP:
        raise ValueError(
            f"'{field_name}' has no ERA5 counterpart configured. "
            f"Available: {sorted(EVAL_FIELD_MAP)}"
        )
    if num_years < 2:
        raise ValueError("num_years must be at least 2 to compute a meaningful climatology.")

    group, var, level = EVAL_FIELD_MAP[field_name]
    lats = np.asarray(state["latitudes"]).ravel()
    lons = np.asarray(state["longitudes"]).ravel()
    lat_idx, lon_idx = _nearest_indices(lats, lons)

    valid_date: datetime.datetime = state["date"]
    clim_mean_grid, clim_std_grid, used_years = _climatology_grid(group, var, level, valid_date, num_years, log)
    clim_mean = sample_at_points(clim_mean_grid, lat_idx, lon_idx)
    clim_std = sample_at_points(clim_std_grid, lat_idx, lon_idx)

    forecast = np.asarray(state["fields"][field_name])
    anomaly = forecast - clim_mean
    zscore = np.full_like(anomaly, np.nan)
    valid_std = clim_std > 1e-6
    zscore[valid_std] = anomaly[valid_std] / clim_std[valid_std]

    log(f"✅  Built a {len(used_years)}-year climatology ({min(used_years)}–{max(used_years)}) for '{field_name}'.")

    return {
        "field": field_name, "lats": lats, "lons": lons,
        "date": str(valid_date), "years": used_years,
        "forecast": forecast, "climatology_mean": clim_mean, "climatology_std": clim_std,
        "anomaly": anomaly, "zscore": zscore,
    }


def run_climatology_baseline(date: datetime.datetime, num_steps: int, num_years: int = 10,
                              log=lambda msg: None) -> list[dict]:
    """
    A naive "forecast" built from pure ERA5 climatology: for each of
    `num_steps` 6h steps ahead of `date`, every canonical field's value is
    just the ERA5 climatological mean for that calendar day/hour (see
    :func:`_climatology_grid`) — zero model skill by construction, useful as
    a baseline to show whether AIFS/WeatherNext2 actually beat "just the
    climate normal".

    Returns a states list shaped like WeatherNext2's (regular grid, fields
    keyed by *canonical* name from :mod:`aifs.compare` — not AIFS's or
    WeatherNext2's own naming) — see aifs.compare.extract.
    """
    from aifs.compare import CANONICAL_FIELDS  # deferred: avoids a module-load cycle (compare imports this module)

    states = []
    for i in range(1, num_steps + 1):
        valid_time = date + datetime.timedelta(hours=STEP_HOURS * i)
        fields = {}
        for canonical_name, spec in CANONICAL_FIELDS.items():
            clim_mean, _, _ = _climatology_grid(
                spec["era5_group"], spec["era5_var"], spec["era5_level"], valid_time, num_years, log,
            )
            fields[canonical_name] = clim_mean
        states.append({"date": valid_time, "fields": fields})
        log(f"✅  Climatology baseline step {i}/{num_steps} ready ({valid_time} UTC).")

    return states


# ── Plotting (matches aifs.plot.plot_field's Cartopy tricontourf style) ────────

def _map_figure(lats, lons, data, title: str, cmap: str = "RdBu_r"):
    import cartopy.crs as ccrs
    import cartopy.feature as cfeature
    import matplotlib.pyplot as plt
    import matplotlib.tri as tri

    lons_plot = np.where(lons > 180, lons - 360, lons)
    triangulation = tri.Triangulation(lons_plot, lats)

    fig, ax = plt.subplots(figsize=(9, 5), subplot_kw={"projection": ccrs.PlateCarree()})
    ax.coastlines()
    ax.add_feature(cfeature.BORDERS, linestyle=":")
    contour = ax.tricontourf(triangulation, data, levels=20, transform=ccrs.PlateCarree(), cmap=cmap)
    fig.colorbar(contour, ax=ax, orientation="vertical", shrink=0.7)
    ax.set_title(title, fontsize=11)
    fig.tight_layout()
    return fig


def plot_eval_maps(eval_data: dict, step_index: int):
    """Forecast / ERA5 truth / (forecast - truth) maps for one evaluated step."""
    step = eval_data["per_step"][step_index]
    lats, lons, field = eval_data["lats"], eval_data["lons"], eval_data["field"]

    fig_fc = _map_figure(lats, lons, step["forecast"], f"AIFS forecast — {field} @ {step['date']}")
    fig_truth = _map_figure(lats, lons, step["truth"], f"ERA5 truth — {field} @ {step['date']}")
    fig_diff = _map_figure(
        lats, lons, step["forecast"] - step["truth"],
        f"Forecast − ERA5 — {field} @ {step['date']}",
    )
    return fig_fc, fig_truth, fig_diff


def plot_climatology_maps(clim_data: dict):
    """Forecast / ERA5 climatology-mean / anomaly maps for one field."""
    lats, lons, field = clim_data["lats"], clim_data["lons"], clim_data["field"]
    year_range = f"{min(clim_data['years'])}–{max(clim_data['years'])}"
    n = len(clim_data["years"])

    fig_fc = _map_figure(lats, lons, clim_data["forecast"], f"AIFS forecast — {field} @ {clim_data['date']}")
    fig_clim = _map_figure(
        lats, lons, clim_data["climatology_mean"],
        f"ERA5 climatology mean — {field} ({year_range}, n={n})",
    )
    fig_anom = _map_figure(
        lats, lons, clim_data["anomaly"],
        f"Anomaly (forecast − climatology) — {field}",
    )
    return fig_fc, fig_clim, fig_anom


def plot_skill_curve(eval_data: dict):
    """RMSE vs lead time: AIFS forecast vs a persistence (ERA5-held-at-t0) baseline."""
    import matplotlib.pyplot as plt

    steps = eval_data["per_step"]
    lead_h = [s["lead_h"] for s in steps]
    model_rmse = [s["model_rmse"] for s in steps]
    persistence_rmse = [s["persistence_rmse"] for s in steps]

    fig, ax = plt.subplots(figsize=(7, 4))
    ax.plot(lead_h, model_rmse, marker="o", label="AIFS forecast", color="#1f6feb")
    ax.plot(lead_h, persistence_rmse, marker="o", linestyle="--", label="Persistence (t0 held fixed)", color="#8b949e")
    ax.set_xlabel("Lead time (hours)")
    ax.set_ylabel("RMSE")
    ax.set_title(f"Skill vs ERA5 — {eval_data['field']}")
    ax.legend()
    ax.grid(alpha=0.3)
    fig.tight_layout()
    return fig