aifs/archive.py

"""
aifs.archive
============
Save/load forecast runs to/from a Hugging Face Hub dataset repo
(EmmaScharfmann/weather-forecast-archive), so a run from one session can be
recalled and re-plotted later without re-running the model.

Serialization is generic across all three "models" in aifs.compare: AIFS
states carry their own (irregular-grid) latitudes/longitudes and are stored
alongside the fields; WeatherNext2/climatology states don't (their grid is
the fixed module-level constants in aifs.weathernext2), and are stored
without them, unchanged on load.
"""

from __future__ import annotations

import datetime
import tempfile
from pathlib import Path

import numpy as np

REPO_ID = "EmmaScharfmann/weather-forecast-archive"
REPO_TYPE = "dataset"

# Short filename slugs for the three models in aifs.compare — kept as a
# local mapping (not importing aifs.compare's exact strings) so archive
# filenames stay stable even if a model's display name changes later.
_SLUG_BY_MODEL = {
    "AIFS": "AIFS",
    "WeatherNext2": "WeatherNext2",
    "Climatology (ERA5 baseline)": "Climatology",
}
_MODEL_BY_SLUG = {slug: model for model, slug in _SLUG_BY_MODEL.items()}

_DATETIME_FMT = "%Y%m%dT%H%M%S"


def _serialize_states(states: list[dict]) -> dict:
    """Flattens a states list into one dict of numpy arrays, npz-savable."""
    payload = {"dates": np.array([s["date"].isoformat() for s in states])}

    if "latitudes" in states[0]:
        payload["latitudes"] = np.asarray(states[0]["latitudes"])
        payload["longitudes"] = np.asarray(states[0]["longitudes"])

    field_names = sorted(states[0]["fields"].keys())
    payload["field_names"] = np.array(field_names)
    for i, state in enumerate(states):
        for name in field_names:
            payload[f"field__{i}__{name}"] = np.asarray(state["fields"][name])
    return payload


def _deserialize_states(npz) -> list[dict]:
    """Inverse of _serialize_states — reconstructs the original states list."""
    dates = [datetime.datetime.fromisoformat(str(d)) for d in npz["dates"]]
    field_names = [str(n) for n in npz["field_names"]]
    has_grid = "latitudes" in npz

    states = []
    for i, date in enumerate(dates):
        fields = {name: npz[f"field__{i}__{name}"] for name in field_names}
        state = {"date": date, "fields": fields}
        if has_grid:
            state["latitudes"] = npz["latitudes"]
            state["longitudes"] = npz["longitudes"]
        states.append(state)
    return states


def save_run(model: str, states: list[dict], log=lambda msg: None) -> str:
    """Uploads one model's forecast states as a single .npz to the archive dataset.

    Returns the remote filename (also encodes model/init-date/step-count/
    saved-at, so :func:`list_saved_runs` can describe it without downloading).
    """
    from huggingface_hub import HfApi

    if not states:
        raise ValueError("No states to save — run a forecast first.")
    if model not in _SLUG_BY_MODEL:
        raise ValueError(f"Unknown model {model!r}.")

    now = datetime.datetime.now(datetime.timezone.utc)
    slug = _SLUG_BY_MODEL[model]
    filename = (
        f"{slug}__init-{states[0]['date'].strftime(_DATETIME_FMT)}"
        f"__{len(states)}steps__saved-{now.strftime(_DATETIME_FMT)}.npz"
    )

    payload = _serialize_states(states)
    with tempfile.TemporaryDirectory() as tmp:
        local_path = Path(tmp) / filename
        np.savez_compressed(str(local_path), **payload)

        log(f"☁️  Uploading {filename} to {REPO_ID}…")
        HfApi().upload_file(
            path_or_fileobj=str(local_path),
            path_in_repo=filename,
            repo_id=REPO_ID,
            repo_type=REPO_TYPE,
        )
    log(f"✅  Saved to {REPO_ID}/{filename}")
    return filename


def list_saved_runs(log=lambda msg: None) -> list[dict]:
    """
    Lists archived runs, newest first — metadata parsed from filenames only
    (no downloads), so this stays fast even with many saved runs.
    """
    from huggingface_hub import HfApi

    log(f"☁️  Listing saved runs in {REPO_ID}…")
    files = HfApi().list_repo_files(repo_id=REPO_ID, repo_type=REPO_TYPE)

    runs = []
    for fname in files:
        if not fname.endswith(".npz"):
            continue
        try:
            slug, rest = fname[:-len(".npz")].split("__init-", 1)
            init_str, rest = rest.split("__", 1)
            steps_str, saved_str = rest.split("__saved-", 1)
            runs.append({
                "filename": fname,
                "model": _MODEL_BY_SLUG.get(slug, slug),
                "init_date": datetime.datetime.strptime(init_str, _DATETIME_FMT),
                "num_steps": int(steps_str.replace("steps", "")),
                "saved_at": datetime.datetime.strptime(saved_str, _DATETIME_FMT),
            })
        except Exception:
            continue  # not one of our files — skip rather than fail the whole listing

    runs.sort(key=lambda r: r["saved_at"], reverse=True)
    log(f"✅  Found {len(runs)} saved run(s).")
    return runs


def load_run(filename: str, log=lambda msg: None) -> tuple[str, list[dict]]:
    """Downloads and deserializes one saved run. Returns (model, states)."""
    from huggingface_hub import hf_hub_download

    log(f"☁️  Downloading {filename} from {REPO_ID}…")
    local_path = hf_hub_download(repo_id=REPO_ID, repo_type=REPO_TYPE, filename=filename)
    with np.load(local_path) as npz:
        states = _deserialize_states(npz)

    slug = filename.split("__init-", 1)[0]
    model = _MODEL_BY_SLUG.get(slug, slug)
    log(f"✅  Loaded {len(states)} step(s) for {model}.")
    return model, states