Spaces:
Running on Zero
Running on Zero
File size: 5,666 Bytes
e69b4bb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 | """
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
|