add aifs wrapper

AIFS-tutorial

@@ -1 +0,0 @@
-Subproject commit 00f678c38f3d4222cf0e7cf0835732a754e02237

aifs/compat.py

@@ -0,0 +1,86 @@
+"""
+Import this module *before* importing anything from ``anemoi`` to allow the compatibility with available GPU / CPU
+"""
+
+import sys
+import time
+import types
+
+import torch
+import torch.nn.functional as F
+
+
+# ── SDPA-based attention replacement ─────────────────────────────────────────
+
+def _sdpa_compat(q, k, v, causal=False, window_size=(-1, -1), dropout_p=0.0, softcap=None, alibi_slopes=None):
+    """
+    Drop-in replacement for ``flash_attn_func``.
+
+    Parameters mirror the flash-attn 2.x signature that Anemoi calls.
+    Input tensors are shaped ``(batch, seq, heads, dim)``.
+    """
+    t0 = time.perf_counter()
+
+    # flash-attn layout: (B, S, H, D)  →  SDPA layout: (B, H, S, D)
+    q, k, v = (t.permute(0, 2, 1, 3) for t in (q, k, v))
+    ws = window_size[0] if isinstance(window_size, (tuple, list)) else int(window_size)
+
+    if q.device.type == "cuda":
+        # Full global attention; SDPA dispatches to flash-attn kernel when available
+        out = F.scaled_dot_product_attention(q, k, v, dropout_p=dropout_p)
+
+    elif ws > 0:
+        # MPS: chunked sliding-window attention (avoids OOM on large sequences)
+        B, H, S, D = q.shape
+        out = torch.zeros_like(q)
+        for i in range(0, S, ws):
+            k_start = max(0, i - ws)
+            k_end   = min(S, i + ws + ws)
+            out[:, :, i : i + ws] = F.scaled_dot_product_attention(
+                q[:, :, i : i + ws],
+                k[:, :, k_start:k_end],
+                v[:, :, k_start:k_end],
+                dropout_p=dropout_p,
+            )
+
+    else:
+        # CPU fallback — move to CPU in case tensors are on an unsupported device
+        out = F.scaled_dot_product_attention(
+            q.cpu(), k.cpu(), v.cpu(), dropout_p=dropout_p
+        ).to(q.device)
+
+    elapsed = time.perf_counter() - t0
+    print(f"  [compat] attn {elapsed:.2f}s  device={q.device.type}  ws={ws}")
+
+    return out.permute(0, 2, 1, 3)
+
+
+# ── Build stub modules ────────────────────────────────────────────────────────
+
+def _patch():
+    """Install the flash_attn stub into ``sys.modules``."""
+    if "flash_attn" in sys.modules:
+        return
+
+    flash_attn = types.ModuleType("flash_attn")
+    flash_attn.__version__ = "2.6.0"  # version Anemoi checks against
+
+    # flash_attn.layers.rotary  (imported but only used on specific GPU paths)
+    layers_mod = types.ModuleType("flash_attn.layers")
+    rotary_mod = types.ModuleType("flash_attn.layers.rotary")
+    rotary_mod.RotaryEmbedding = None
+    layers_mod.rotary = rotary_mod
+    flash_attn.layers = layers_mod
+
+    # flash_attn.flash_attn_interface  (the one Anemoi actually calls)
+    interface_mod = types.ModuleType("flash_attn.flash_attn_interface")
+    interface_mod.flash_attn_func = _sdpa_compat
+    flash_attn.flash_attn_interface = interface_mod
+
+    sys.modules["flash_attn"]                    = flash_attn
+    sys.modules["flash_attn.layers"]             = layers_mod
+    sys.modules["flash_attn.layers.rotary"]      = rotary_mod
+    sys.modules["flash_attn.flash_attn_interface"] = interface_mod
+
+
+_patch()

aifs/device.py

@@ -0,0 +1,31 @@
+import os
+import ssl
+
+import certifi
+import torch
+
+
+def configure_ssl():
+    """Point OpenSSL to the certifi bundle (needed on some macOS setups)."""
+    os.environ.setdefault("SSL_CERT_FILE", certifi.where())
+    os.environ.setdefault("REQUESTS_CA_BUNDLE", certifi.where())
+    ssl._create_default_https_context = ssl.create_default_context
+
+
+def get_device() -> str:
+    """Return ``'cuda'``, ``'mps'``, or ``'cpu'`` in priority order."""
+    if torch.cuda.is_available():
+        return "cuda"
+    if torch.backends.mps.is_available():
+        return "mps"
+    return "cpu"
+
+
+def device_label() -> str:
+    """Human-readable description of the active compute device."""
+    device = get_device()
+    if device == "cuda":
+        return f"CUDA — {torch.cuda.get_device_name(0)}"
+    if device == "mps":
+        return "Apple MPS (Metal)"
+    return "CPU (no GPU detected — inference will be slow)"

aifs/forecast.py

@@ -0,0 +1,125 @@
+import os
+import datetime
+from typing import Generator
+
+# Apply flash-attn shim before any anemoi import
+import aifs.compat  # noqa: F401
+from aifs.device import get_device, device_label
+
+DEFAULT_CHECKPOINT = "aifs-single-2.0"
+
+CHECKPOINTS = {
+    DEFAULT_CHECKPOINT: {"huggingface": f"ecmwf/{DEFAULT_CHECKPOINT}"},
+}
+
+def run_forecast(
+    fields: dict,
+    date: datetime.datetime,
+    lead_time: int = 24,
+    num_chunks: int = 16,
+    checkpoint: str = DEFAULT_CHECKPOINT,
+    verbose: bool = True,
+) -> list[dict]:
+    """
+    Run an AIFS forecast and return all output states.
+
+    Parameters
+    ----------
+    fields:
+        Initial-condition field dict as returned by ``load_ics()``.
+        Shape of each array: ``(2, N320_nodes)``.
+    date:
+        Forecast initialisation datetime.
+    lead_time:
+        Forecast horizon in hours. Must be a multiple of 6.
+    num_chunks:
+        Number of chunks for the attention computation.
+        Increase if you run out of memory; decrease for speed.
+        16 is a safe default for 16 GB RAM / VRAM.
+    checkpoint:
+        Key into ``CHECKPOINTS`` dict, or a raw ``{"huggingface": "..."}``
+        dict you can pass directly.
+    verbose:
+        Print step-by-step progress.
+
+    Returns
+    -------
+    list of state dicts, one per 6-hour output step.
+    Each state dict has at minimum:
+        ``state["date"]``   — output datetime
+        ``state["fields"]`` — ``{variable: np.ndarray}``
+    """
+    if lead_time % 6 != 0:
+        raise ValueError(f"lead_time must be a multiple of 6, got {lead_time}")
+
+    from anemoi.inference.runners.simple import SimpleRunner
+
+    device = get_device()
+
+    # Environment knobs consumed by anemoi-inference internals
+    if device == "cuda":
+        os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+    os.environ["ANEMOI_INFERENCE_NUM_CHUNKS"] = str(num_chunks)
+
+    if verbose:
+        print(f"🖥️   Device  : {device_label()}")
+        print(f"📦  Checkpoint: {checkpoint}")
+        print(f"⏱️   Lead time : {lead_time} h  ({lead_time // 6} steps)")
+
+    ckpt = CHECKPOINTS.get(checkpoint, checkpoint)
+
+    if verbose:
+        print("🤖  Loading model …")
+
+    runner = SimpleRunner(ckpt)
+
+    if verbose:
+        print(f"🌍  Running inference …")
+
+    states: list[dict] = []
+    input_state = {"fields": fields, "date": date}
+    for state in runner.run(input_states=input_state, lead_time=lead_time):
+        states.append(state)
+        if verbose:
+            print(f"    ✓  {state['date']}")
+
+    if verbose:
+        print(f"✅  Done — {len(states)} steps produced.")
+
+    return states
+
+
+def run_forecast_streaming(
+    fields: dict,
+    date: datetime.datetime,
+    lead_time: int = 24,
+    num_chunks: int = 16,
+    checkpoint: str = DEFAULT_CHECKPOINT,
+) -> Generator[dict, None, None]:
+    """
+    Generator variant of :func:`run_forecast`.
+
+    Yields each state dict as soon as it is computed, which is useful for
+    Gradio apps or notebooks that want to display results incrementally.
+
+    Example
+    -------
+        for state in run_forecast_streaming(fields, date, lead_time=48):
+            plot_field(state)
+    """
+    if lead_time % 6 != 0:
+        raise ValueError(f"lead_time must be a multiple of 6, got {lead_time}")
+
+    from anemoi.inference.runners.simple import SimpleRunner
+
+    device = get_device()
+    if device == "cuda":
+        os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+    os.environ["ANEMOI_INFERENCE_NUM_CHUNKS"] = str(num_chunks)
+
+    ckpt = CHECKPOINTS.get(checkpoint, checkpoint)
+    runner = SimpleRunner(ckpt)
+
+    input_state = {"fields": fields, "date": date}
+    for state in runner.run(input_states=input_state, lead_time=lead_time):
+        yield state

aifs/initial_conditions.py

@@ -0,0 +1,209 @@
+import datetime
+import warnings
+from collections import defaultdict
+from pathlib import Path
+
+import numpy as np
+
+warnings.filterwarnings("ignore", category=DeprecationWarning)
+warnings.filterwarnings("ignore", category=UserWarning)
+
+
+# ── Meteorological variable lists ─────────────────────────────────────────────
+
+#: Surface parameters (levtype=sfc)
+PARAM_SFC = [
+    "10u", "10v", "2d", "2t", "msl", "skt", "sp",
+    "tcw", "lsm", "z", "slor", "sdor", "sd",
+]
+
+#: Soil parameters (levtype=sfc, levelist=[1,2])
+PARAM_SOIL = ["vsw", "sot"]
+SOIL_LEVELS = [1, 2]
+
+#: Ocean-wave parameters (stream=wave)
+PARAM_WAVE = [
+    "wmb", "h1012", "h1214", "h1417", "h1721",
+    "h2125", "h2530", "mwd", "cdww", "mwp", "swh",
+]
+
+#: Pressure-level parameters
+PARAM_PL = ["gh", "t", "u", "v", "q"]
+LEVELS   = [1000, 925, 850, 700, 600, 500, 400, 300, 250, 200, 150, 100, 50, 10]
+
+SOURCE = "ecmwf"
+
+# ── Cache helpers ─────────────────────────────────────────────────────────────
+
+DEFAULT_CACHE_DIR = Path("ic_cache")
+
+
+def _cache_path(date: datetime.datetime, cache_dir: Path) -> Path:
+    return cache_dir / f"ic_{date.strftime('%Y%m%dT%H%M%S')}.npz"
+
+
+def _save(date: datetime.datetime, fields: dict, cache_dir: Path) -> Path:
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    path = _cache_path(date, cache_dir)
+    np.savez_compressed(str(path), **fields)
+    return path
+
+
+def _try_load(date: datetime.datetime, cache_dir: Path):
+    """Return ``(fields_dict, path)`` if cached, else ``(None, None)``."""
+    path = _cache_path(date, cache_dir)
+    if path.exists():
+        return dict(np.load(str(path))), path
+    return None, None
+
+
+def list_cached(cache_dir: Path = DEFAULT_CACHE_DIR) -> list[Path]:
+    """Return all cached .npz files, newest first."""
+    if not cache_dir.exists():
+        return []
+    return sorted(cache_dir.glob("ic_*.npz"), reverse=True)
+
+
+# ── Download helpers ──────────────────────────────────────────────────────────
+
+def _fetch_fields(ekd, ekr, date, param, levelist=None, **kwargs) -> dict:
+    """
+    Download ``param`` for two time-steps (t-6h, t) and return a dict
+    ``{variable_name: np.ndarray shape (2, N320_nodes)}``.
+    """
+    levelist = levelist or []
+    raw: dict[str, list] = defaultdict(list)
+
+    for t in [date - datetime.timedelta(hours=6), date]:
+        dataset = ekd.from_source(
+            "ecmwf-open-data",
+            date=t,
+            param=param,
+            levelist=levelist,
+            source=SOURCE,
+            **kwargs,
+        )
+        for field in dataset:
+            assert field.to_numpy().shape == (721, 1440), (
+                f"Unexpected grid shape for {field.metadata('param')}: "
+                f"{field.to_numpy().shape}"
+            )
+            # Shift lon from [0,360) to [-180,180) then regrid to N320 Gaussian
+            values = np.roll(field.to_numpy(), -field.shape[1] // 2, axis=1)
+            values = ekr.interpolate(values, {"grid": (0.25, 0.25)}, {"grid": "N320"})
+
+            if levelist:
+                name = f"{field.metadata('param')}_{field.metadata('levelist')}"
+            else:
+                name = field.metadata("param")
+            raw[name].append(values)
+
+    return {k: np.stack(v) for k, v in raw.items()}
+
+
+def _build_fields(ekd, ekr, date: datetime.datetime) -> dict:
+    """Download and transform all required fields for ``date``."""
+    fields: dict = {}
+
+    print("  ⬇  Surface fields …")
+    fields.update(_fetch_fields(ekd, ekr, date, PARAM_SFC, levtype="sfc"))
+
+    print("  ⬇  Wave fields …")
+    fields.update(_fetch_fields(ekd, ekr, date, PARAM_WAVE, stream="wave"))
+
+    print("  ⬇  Soil fields …")
+    soil = _fetch_fields(ekd, ekr, date, PARAM_SOIL, levelist=SOIL_LEVELS)
+
+    print("  ⬇  Pressure-level fields …")
+    fields.update(_fetch_fields(ekd, ekr, date, PARAM_PL, levelist=LEVELS))
+
+    # ── Transformations ───────────────────────────────────────────────────────
+
+    # Wave direction: decompose scalar angle into sin/cos components
+    mwd = fields.pop("mwd")
+    mwd_rad = np.deg2rad(mwd)
+    fields["cos_mwd"] = np.cos(mwd_rad)
+    fields["sin_mwd"] = np.sin(mwd_rad)
+
+    # Rename soil fields to ECMWF short-names expected by AIFS
+    _soil_rename = {
+        "sot_1": "stl1",  "sot_2": "stl2",
+        "vsw_1": "swvl1", "vsw_2": "swvl2",
+    }
+    for src, dst in _soil_rename.items():
+        fields[dst] = soil[src]
+
+    # Remove q levels that AIFS does not use
+    fields.pop("q_10", None)
+    fields.pop("q_50", None)
+
+    # Apply land-sea mask to snow depth and soil moisture (ocean → NaN)
+    try:
+        lsm = ekd.from_source("file", "lsm.grib")[0].to_numpy(flatten=True)
+        ocean_mask = np.equal(lsm, 0)
+        for var in ("sd", "swvl1", "swvl2"):
+            if var in fields:
+                fields[var][:, ocean_mask] = np.nan
+    except Exception:
+        pass  # lsm.grib not found; skip masking
+
+    # Convert geopotential height → geopotential  (Z = gh × g)
+    G = 9.80665
+    for level in LEVELS:
+        gh = fields.pop(f"gh_{level}", None)
+        if gh is not None:
+            fields[f"z_{level}"] = gh * G
+
+    return fields
+
+
+# ── Public API ────────────────────────────────────────────────────────────────
+
+def load_ics(
+    cache_dir: Path | str = DEFAULT_CACHE_DIR,
+    force: bool = False,
+) -> tuple[dict, datetime.datetime]:
+    """
+    Return ``(fields, date)`` for the latest available ECMWF Open Data run.
+
+    Parameters
+    ----------
+    cache_dir:
+        Directory where .npz caches are stored.
+    force:
+        Re-download even when a local cache exists.
+
+    Returns
+    -------
+    fields:
+        ``{variable_name: np.ndarray shape (2, N320_nodes)}``.
+        The first axis indexes the two input time-steps: ``[t-6h, t]``.
+    date:
+        The forecast initialisation date/time (the *later* of the two
+        time-steps).
+    """
+    import earthkit.data as ekd
+    import earthkit.regrid as ekr
+    from ecmwf.opendata import Client as OpendataClient
+
+    ekd.config.set({"cache-policy": "user"})
+    cache_dir = Path(cache_dir)
+
+    date: datetime.datetime = OpendataClient(SOURCE).latest()
+    print(f"📅  Latest ECMWF run: {date}")
+
+    if not force:
+        cached, path = _try_load(date, cache_dir)
+        if cached is not None:
+            sz_mb = path.stat().st_size / 1e6
+            print(f"✅  Loaded from cache  ({sz_mb:.0f} MB)  →  {path}")
+            return cached, date
+
+    print("⬇️   Downloading initial conditions …")
+    fields = _build_fields(ekd, ekr, date)
+
+    path = _save(date, fields, cache_dir)
+    sz_mb = path.stat().st_size / 1e6
+    print(f"💾  Saved to {path}  ({sz_mb:.0f} MB)")
+
+    return fields, date

aifs/plot.py

@@ -0,0 +1,338 @@
+"""
+aifs.plot
+=========
+Cartopy-based helpers for visualising AIFS forecast output on a global map.
+
+All functions return ``matplotlib.figure.Figure`` objects so they work in
+both notebooks (``plt.show()``) and scripts (``fig.savefig(...)``).
+
+Quickstart
+----------
+    from aifs.plot import plot_field, plot_field_sequence
+
+    # Single map
+    fig = plot_field(state, "2t", title="2-m Temperature — T+6h")
+    fig.savefig("t2m_T+6.png", dpi=150)
+
+    # Multi-panel sequence
+    fig = plot_field_sequence(states, "2t", max_steps=4)
+    fig.savefig("t2m_sequence.png", dpi=150)
+"""
+
+from __future__ import annotations
+
+import warnings
+import numpy as np
+
+warnings.filterwarnings("ignore", category=UserWarning)
+
+# ── Variable metadata ─────────────────────────────────────────────────────────
+
+#: Variables that can be extracted from forecast state dicts
+PLOTTABLE = [
+    "2t", "msl", "sp", "tcw", "10u", "10v", "swh", "mwp",
+    "t_850", "t_500", "u_850", "v_850", "z_500", "q_700",
+]
+
+_CMAP = {
+    "2t":    "RdBu_r", "t_850": "RdBu_r", "t_500": "RdBu_r",
+    "msl":   "viridis", "sp":    "viridis",
+    "10u":   "RdBu",    "10v":   "RdBu",
+    "u_850": "RdBu",    "v_850": "RdBu",
+    "swh":   "Blues",   "mwp":   "Blues",   "tcw":   "Blues",
+    "z_500": "plasma",  "q_700": "YlGn",
+}
+
+_UNITS = {
+    "2t":    "K",      "t_850": "K",      "t_500": "K",
+    "msl":   "Pa",     "sp":    "Pa",     "z_500": "m²/s²",
+    "10u":   "m/s",    "10v":   "m/s",    "u_850": "m/s",    "v_850": "m/s",
+    "swh":   "m",      "mwp":   "s",      "tcw":   "kg/m²",  "q_700": "kg/kg",
+}
+
+_LONG_NAME = {
+    "2t":    "2-m Temperature",
+    "msl":   "Mean Sea-Level Pressure",
+    "sp":    "Surface Pressure",
+    "tcw":   "Total Column Water",
+    "10u":   "10-m U Wind",
+    "10v":   "10-m V Wind",
+    "swh":   "Significant Wave Height",
+    "mwp":   "Mean Wave Period",
+    "t_850": "Temperature at 850 hPa",
+    "t_500": "Temperature at 500 hPa",
+    "u_850": "U Wind at 850 hPa",
+    "v_850": "V Wind at 850 hPa",
+    "z_500": "Geopotential at 500 hPa",
+    "q_700": "Specific Humidity at 700 hPa",
+}
+
+
+# ── Grid coordinate extraction ────────────────────────────────────────────────
+
+def _get_latlons(state: dict) -> tuple[np.ndarray, np.ndarray]:
+    """
+    Return (lats, lons) for the grid the forecast was run on.
+
+    The anemoi tensor handler injects ``state["latitudes"]`` and
+    ``state["longitudes"]`` from the checkpoint metadata before the first
+    inference step, and these are propagated to every output state via
+    ``new_states = input_states.copy()``.  We read them directly — no
+    separate grid-geometry lookup needed.
+
+    Longitudes are returned in the range [0, 360) as stored by anemoi;
+    callers that need [-180, 180) should call ``_to_180(lons)``.
+    """
+    lats = state.get("latitudes")
+    lons = state.get("longitudes")
+
+    if lats is None or lons is None:
+        raise KeyError(
+            "State dict does not contain 'latitudes'/'longitudes'. "
+            "Make sure you are passing a state returned by run_forecast() "
+            "and have not stripped those keys."
+        )
+
+    lats = np.asarray(lats).ravel()
+    lons = np.asarray(lons).ravel()
+
+    if len(lats) < 3 or len(lons) < 3:
+        raise ValueError(
+            f"Grid has only {len(lats)} points — expected ~542 080 for N320. "
+            "The state latitudes/longitudes may be corrupt."
+        )
+
+    return lats, lons
+
+
+def _to_180(lons: np.ndarray) -> np.ndarray:
+    """Normalise longitudes from [0, 360) to [-180, 180) for Cartopy."""
+    return np.where(lons > 180, lons - 360, lons)
+
+
+def _extract_field(state: dict, variable: str) -> np.ndarray | None:
+    """Pull ``variable`` out of ``state["fields"]``, return None if missing."""
+    return state.get("fields", {}).get(variable)
+
+
+# ── Public API ────────────────────────────────────────────────────────────────
+
+def plot_field(
+    state: dict,
+    variable: str,
+    title: str | None = None,
+    projection: str = "Robinson",
+    figsize: tuple[float, float] = (14, 7),
+    vmin: float | None = None,
+    vmax: float | None = None,
+) -> "matplotlib.figure.Figure":
+    """
+    Plot a single forecast field on a global map.
+
+    Parameters
+    ----------
+    state:
+        One element from the list returned by :func:`aifs.forecast.run_forecast`.
+    variable:
+        Short name of the field to plot (e.g. ``"2t"``, ``"msl"``).
+        See :data:`PLOTTABLE` for supported names.
+    title:
+        Figure title; defaults to ``"{long_name} — {state_date}"``.
+    projection:
+        Cartopy projection class name (e.g. ``"Robinson"``, ``"PlateCarree"``).
+    figsize:
+        Matplotlib figure size in inches.
+    vmin, vmax:
+        Colour-scale limits; auto-derived from percentiles if not provided.
+
+    Returns
+    -------
+    matplotlib.figure.Figure
+    """
+    import matplotlib
+    import matplotlib.pyplot as plt
+    import cartopy.crs as ccrs
+    import cartopy.feature as cfeature
+    import matplotlib.tri as tri
+
+    data = _extract_field(state, variable)
+    if data is None:
+        raise KeyError(
+            f"Variable '{variable}' not found in forecast state. "
+            f"Available: {sorted(state.get('fields', {}).keys())}"
+        )
+
+    data     = np.asarray(data).ravel()
+    lats, lons = _get_latlons(state)
+    lons_plot  = _to_180(lons)
+
+    units = _UNITS.get(variable, "")
+    lname = _LONG_NAME.get(variable, variable)
+    dt    = state.get("date", "")
+
+    # Default colormap: coolwarm gives blue=cold, red=warm; fall back per variable
+    cmap = _CMAP.get(variable, "coolwarm")
+
+    if vmin is None:
+        vmin = float(np.nanpercentile(data, 2))
+    if vmax is None:
+        vmax = float(np.nanpercentile(data, 98))
+
+    proj_cls = getattr(ccrs, projection, ccrs.Robinson)
+    proj     = proj_cls()
+    pc       = ccrs.PlateCarree()
+
+    # Pre-project coordinates into projection space before triangulating
+    xy    = proj.transform_points(pc, lons_plot, lats)   # (N, 3)
+    x, y  = xy[:, 0], xy[:, 1]
+
+    # Drop points that failed to project
+    valid = np.isfinite(x) & np.isfinite(y)
+    x, y, data = x[valid], y[valid], data[valid]
+
+    triangulation = tri.Triangulation(x, y)
+
+    # Mask triangles spanning the antimeridian (threshold in projection metres)
+    x_verts    = x[triangulation.triangles]
+    max_x_span = np.max(x_verts, axis=1) - np.min(x_verts, axis=1)
+    triangulation.set_mask(max_x_span > 1e6)
+
+    fig = plt.figure(figsize=figsize)
+    ax  = fig.add_subplot(1, 1, 1, projection=proj)
+    ax.set_global()
+    ax.add_feature(cfeature.COASTLINE, linewidth=0.5)
+    ax.add_feature(cfeature.BORDERS,   linewidth=0.3, alpha=0.5)
+
+    # No transform= — coordinates are already in projection space
+    pcm = ax.tripcolor(
+        triangulation, data,
+        cmap=cmap, vmin=vmin, vmax=vmax,
+        shading="gouraud",
+    )
+
+    cbar = fig.colorbar(pcm, ax=ax, orientation="horizontal",
+                        pad=0.04, fraction=0.03, shrink=0.8)
+    cbar.set_label(f"{lname}  [{units}]", fontsize=10)
+
+    if title is None:
+        title = f"{lname}  —  {dt}"
+    ax.set_title(title, fontsize=12, pad=10)
+
+    fig.tight_layout()
+
+    return fig
+
+
+def plot_field_sequence(
+    states: list[dict],
+    variable: str,
+    max_steps: int = 4,
+    figsize_per_panel: tuple[float, float] = (7, 3.5),
+    projection: str = "Robinson",
+    shared_colorscale: bool = True,
+) -> "matplotlib.figure.Figure":
+    """
+    Plot a sequence of forecast steps side-by-side in a single figure.
+
+    Parameters
+    ----------
+    states:
+        List of state dicts from :func:`aifs.forecast.run_forecast`.
+    variable:
+        Variable short name.
+    max_steps:
+        Maximum number of panels (steps) to show.
+    figsize_per_panel:
+        Width × height for each panel in inches.
+    projection:
+        Cartopy projection class name.
+    shared_colorscale:
+        Use the same colour limits for all panels (recommended).
+
+    Returns
+    -------
+    matplotlib.figure.Figure
+    """
+    import matplotlib.pyplot as plt
+    import cartopy.crs as ccrs
+    import cartopy.feature as cfeature
+    import matplotlib.tri as tri
+
+    steps = states[:max_steps]
+    n     = len(steps)
+    ncols = min(n, 2)
+    nrows = (n + ncols - 1) // ncols
+
+    fig_w    = figsize_per_panel[0] * ncols
+    fig_h    = figsize_per_panel[1] * nrows
+    proj_cls = getattr(ccrs, projection, ccrs.Robinson)
+    proj     = proj_cls()
+    pc       = ccrs.PlateCarree()
+
+    # Pre-project coordinates once from the first state
+    lats, lons = _get_latlons(steps[0])
+    lons_plot  = _to_180(lons)
+
+    xy    = proj.transform_points(pc, lons_plot, lats)
+    x, y  = xy[:, 0], xy[:, 1]
+    valid = np.isfinite(x) & np.isfinite(y)
+    x, y  = x[valid], y[valid]
+
+    triangulation = tri.Triangulation(x, y)
+    x_verts       = x[triangulation.triangles]
+    max_x_span    = np.max(x_verts, axis=1) - np.min(x_verts, axis=1)
+    triangulation.set_mask(max_x_span > 1e6)
+
+    cmap  = _CMAP.get(variable, "coolwarm")
+    units = _UNITS.get(variable, "")
+    lname = _LONG_NAME.get(variable, variable)
+
+    if shared_colorscale:
+        all_data = np.concatenate(
+            [np.asarray(_extract_field(s, variable)).ravel()[valid]
+             for s in steps if _extract_field(s, variable) is not None]
+        )
+        vmin = float(np.nanpercentile(all_data, 2))
+        vmax = float(np.nanpercentile(all_data, 98))
+    else:
+        vmin = vmax = None
+
+    fig, axes = plt.subplots(
+        nrows, ncols,
+        figsize=(fig_w, fig_h),
+        subplot_kw={"projection": proj},
+    )
+    axes_flat = np.array(axes).ravel()
+
+    for idx, (state, ax) in enumerate(zip(steps, axes_flat)):
+        data = _extract_field(state, variable)
+        if data is None:
+            ax.set_visible(False)
+            continue
+
+        data  = np.asarray(data).ravel()[valid]   # apply same valid mask
+        _vmin = vmin if shared_colorscale else float(np.nanpercentile(data, 2))
+        _vmax = vmax if shared_colorscale else float(np.nanpercentile(data, 98))
+
+        ax.set_global()
+        ax.add_feature(cfeature.COASTLINE, linewidth=0.4)
+
+        # No transform= — coordinates already in projection space
+        pcm = ax.tripcolor(
+            triangulation, data,
+            cmap=cmap, vmin=_vmin, vmax=_vmax,
+            shading="gouraud",
+        )
+        step_label = f"T+{(idx + 1) * 6}h  ({state.get('date', '')})"
+        ax.set_title(step_label, fontsize=9)
+
+        fig.colorbar(pcm, ax=ax, orientation="horizontal",
+                     pad=0.04, fraction=0.04, shrink=0.85,
+                     label=f"{units}")
+
+    for ax in axes_flat[len(steps):]:
+        ax.set_visible(False)
+
+    fig.suptitle(f"{lname} — AIFS Forecast", fontsize=13, y=1.01)
+    fig.tight_layout()
+    return fig