add storage

aifs/archive.py

@@ -0,0 +1,153 @@
+"""
+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

app.py

@@ -187,6 +187,72 @@ def toggle_model_controls(model: str):
     )
 
 
+def format_status(all_states: dict) -> str:
+    """One line per model: whether it's been run this session and, if so, its step range."""
+    lines = []
+    for model in MODELS:
+        states = all_states.get(model)
+        if states:
+            lines.append(f"- **{model}**: ✅ {len(states)} step(s) ready — {states[0]['date']} → {states[-1]['date']}")
+        else:
+            lines.append(f"- **{model}**: ⬜ not run yet")
+    return "\n".join(lines)
+
+
+# ── Forecast archive (save/load runs to the HF dataset in aifs.archive) ───────
+
+def save_current_run(model: str, all_states: dict) -> str:
+    from aifs.archive import save_run
+
+    states = all_states.get(model)
+    if not states:
+        return f"⚠️  Run {model} first, then save it."
+
+    log_lines: list[str] = []
+    try:
+        save_run(model, states, log=log_lines.append)
+        return "\n".join(log_lines)
+    except Exception as exc:
+        return "\n".join(log_lines) + f"\n❌  Error: {exc}"
+
+
+def refresh_saved_runs():
+    from aifs.archive import list_saved_runs
+
+    log_lines: list[str] = []
+    try:
+        runs = list_saved_runs(log=log_lines.append)
+    except Exception as exc:
+        return gr.update(choices=[], value=None), "\n".join(log_lines) + f"\n❌  Error: {exc}"
+
+    choices = [
+        (
+            f"{r['model']} — init {r['init_date']:%Y-%m-%d %H:%M} — {r['num_steps']} step(s) "
+            f"— saved {r['saved_at']:%Y-%m-%d %H:%M}",
+            r["filename"],
+        )
+        for r in runs
+    ]
+    return gr.update(choices=choices, value=(choices[0][1] if choices else None)), "\n".join(log_lines)
+
+
+def load_saved_run(filename: str, all_states: dict):
+    from aifs.archive import load_run
+
+    if not filename:
+        return all_states, "⚠️  Pick a saved run to load first."
+
+    log_lines: list[str] = []
+    try:
+        model, states = load_run(filename, log=log_lines.append)
+    except Exception as exc:
+        return all_states, "\n".join(log_lines) + f"\n❌  Error: {exc}"
+
+    new_all_states = dict(all_states)
+    new_all_states[model] = states
+    return new_all_states, "\n".join(log_lines)
+
+
 # ── Visualize (any one model/step/field) ───────────────────────────────────────
 
 def plot_selected(model: str, step_str: str, field: str, all_states: dict):
@@ -273,6 +339,29 @@ textarea, input, select { background: #1c2128!important; color: #cdd9e5!importan
 textarea::placeholder, input::placeholder { color: #a8b3c0!important; }
 .output-markdown { color: #cdd9e5!important; }
 
+/* === Dropdown / radio internals ===
+   Gradio's Dropdown/Radio are custom components, not a plain <select>/
+   <input> — the rule above doesn't reach their popup list or selected-pill
+   styling, which otherwise falls back to Gradio's default (light) theme. */
+ul[class*="options"], li[class*="item"] {
+    background: #1c2128!important;
+    color: #cdd9e5!important;
+}
+li[class*="item"]:hover, li[class*="item"][aria-selected="true"] {
+    background: #30363d!important;
+    color: #cdd9e5!important;
+}
+.wrap label {
+    background: #1c2128!important;
+    color: #cdd9e5!important;
+    border-color: #30363d!important;
+}
+label.selected, label[class*="selected"] {
+    background: #1f6feb!important;
+    color: #ffffff!important;
+    border-color: #1f6feb!important;
+}
+
 /* === Notes / Markdown prose === */
 .prose,
 .prose p,
@@ -302,6 +391,10 @@ with gr.Blocks(css=DARK_CSS, title="Weather Model Comparison") as demo:
 
     all_states = gr.State({})
 
+    with gr.Group(elem_classes="panel"):
+        gr.Markdown("### 📋 Session Status")
+        status_md = gr.Markdown(format_status({}))
+
     with gr.Row():
         with gr.Column(scale=1, elem_classes="panel"):
             gr.Markdown("### ⚙️ Run a Forecast")
@@ -359,6 +452,8 @@ with gr.Blocks(css=DARK_CSS, title="Weather Model Comparison") as demo:
                 label="Detailed log", lines=8, interactive=False,
                 placeholder="Progress details will appear here…",
             )
+            save_btn = gr.Button("💾 Save Current Run to Archive", variant="secondary")
+            save_status = gr.Textbox(label="Save log", lines=2, interactive=False)
 
         with gr.Column(scale=2, elem_classes="panel"):
             gr.Markdown("### 🗺️ Visualize")
@@ -372,6 +467,21 @@ with gr.Blocks(css=DARK_CSS, title="Weather Model Comparison") as demo:
             viz_img = gr.Image(label="Map", type="filepath")
             viz_stats_md = gr.Markdown()
 
+    gr.Markdown("---")
+    with gr.Row():
+        with gr.Column(elem_classes="panel"):
+            gr.Markdown(
+                "### 📂 Saved Runs\n"
+                "Forecasts saved above persist in the "
+                "[weather-forecast-archive](https://huggingface.co/datasets/EmmaScharfmann/weather-forecast-archive) "
+                "dataset — load one back here to plot/compare it without re-running the model."
+            )
+            with gr.Row():
+                saved_runs_dd = gr.Dropdown(choices=[], label="Saved run", scale=3)
+                refresh_saved_btn = gr.Button("🔄 Refresh", scale=1)
+                load_saved_btn = gr.Button("📥 Load", variant="secondary", scale=1)
+            saved_runs_status = gr.Textbox(label="Archive log", lines=2, interactive=False)
+
     gr.Markdown("---")
     with gr.Row():
         with gr.Column(scale=1, elem_classes="panel"):
@@ -417,6 +527,26 @@ with gr.Blocks(css=DARK_CSS, title="Weather Model Comparison") as demo:
         inputs=[cmp_model_a_dd, cmp_step_a_dd, cmp_model_b_dd, cmp_step_b_dd, cmp_field_dd, all_states],
         outputs=[cmp_img_a, cmp_img_b, cmp_img_diff, cmp_stats_md],
     )
+    save_btn.click(
+        fn=save_current_run,
+        inputs=[model_dd, all_states],
+        outputs=[save_status],
+    )
+    refresh_saved_btn.click(
+        fn=refresh_saved_runs,
+        inputs=[],
+        outputs=[saved_runs_dd, saved_runs_status],
+    )
+    load_saved_btn.click(
+        fn=load_saved_run,
+        inputs=[saved_runs_dd, all_states],
+        outputs=[all_states, saved_runs_status],
+    )
+    all_states.change(
+        fn=format_status,
+        inputs=all_states,
+        outputs=status_md,
+    )
 
     gr.Markdown(
         """

requirements.txt

@@ -23,6 +23,11 @@ cartopy
 git+https://github.com/kashif/transformers.git@add-weathernext2
 accelerate
 
+# Forecast archive (save/load runs to the EmmaScharfmann/weather-forecast-archive
+# HF dataset) — already an implicit dependency of transformers, listed
+# explicitly since aifs/archive.py imports it directly.
+huggingface_hub
+
 # HF Space environment config:
 # Python: 3.12.12 (main, Feb 24 2026, 21:49:09) [GCC 14.2.0]
 #Platform: Linux-6.12.80-106.156.amzn2023.x86_64-x86_64-with-glibc2.41