Spaces:
Running on Zero
Running on Zero
Commit Β·
786225b
1
Parent(s): 71701b2
add comparision
Browse files- .gitignore +2 -0
- aifs/era5_env.py +128 -0
- aifs/era5_verify.py +337 -0
- aifs/era5_worker.py +130 -0
- aifs/initial_conditions.py +61 -7
- app.py +322 -3
.gitignore
CHANGED
|
@@ -3,4 +3,6 @@ __pycache__/
|
|
| 3 |
*.pyc
|
| 4 |
.npz
|
| 5 |
ic_cache/
|
|
|
|
|
|
|
| 6 |
.idea/
|
|
|
|
| 3 |
*.pyc
|
| 4 |
.npz
|
| 5 |
ic_cache/
|
| 6 |
+
era5_env_isolated/
|
| 7 |
+
era5_static_cache/
|
| 8 |
.idea/
|
aifs/era5_env.py
ADDED
|
@@ -0,0 +1,128 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
aifs.era5_env
|
| 3 |
+
=============
|
| 4 |
+
Bootstraps and drives the isolated environment that talks to EarthMover's
|
| 5 |
+
ERA5 Icechunk store (see :mod:`aifs.era5_worker`).
|
| 6 |
+
|
| 7 |
+
Why isolated: ``icechunk`` requires ``zarr>=3``, but ``anemoi-datasets``
|
| 8 |
+
(already required for AIFS inference in this Space) pins ``zarr<=2.18``.
|
| 9 |
+
pip cannot satisfy both in one environment. Instead we ``pip install
|
| 10 |
+
--target=`` a private directory on first use, and run the actual ERA5
|
| 11 |
+
reads in a subprocess whose ``PYTHONPATH`` is prepended with that
|
| 12 |
+
directory β the subprocess resolves ``zarr`` to the isolated v3 install
|
| 13 |
+
regardless of what's importable from the main env's site-packages
|
| 14 |
+
(verified: this is plain CPython import-order semantics, not a hack
|
| 15 |
+
specific to zarr).
|
| 16 |
+
"""
|
| 17 |
+
|
| 18 |
+
from __future__ import annotations
|
| 19 |
+
|
| 20 |
+
import json
|
| 21 |
+
import os
|
| 22 |
+
import platform
|
| 23 |
+
import shutil
|
| 24 |
+
import subprocess
|
| 25 |
+
import sys
|
| 26 |
+
import tempfile
|
| 27 |
+
from pathlib import Path
|
| 28 |
+
|
| 29 |
+
import numpy as np
|
| 30 |
+
|
| 31 |
+
_HERE = Path(__file__).resolve().parent
|
| 32 |
+
WORKER_SCRIPT = _HERE / "era5_worker.py"
|
| 33 |
+
ISOLATED_DIR = _HERE.parent / "era5_env_isolated"
|
| 34 |
+
SENTINEL = ISOLATED_DIR / ".bootstrap_ok"
|
| 35 |
+
|
| 36 |
+
ISOLATED_PACKAGES = ["icechunk>=2.1", "zarr>=3", "pcodec"]
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
def _fingerprint() -> str:
|
| 40 |
+
"""
|
| 41 |
+
Identifies the machine/interpreter the isolated install's compiled
|
| 42 |
+
wheels (numpy, icechunk, pcodec) are built for. Compared against the
|
| 43 |
+
sentinel on every call so a directory built on one machine (e.g. a
|
| 44 |
+
dev sandbox) is never reused on another (e.g. the actual Space host)
|
| 45 |
+
β compiled extensions are platform- and Python-version-specific, and
|
| 46 |
+
a mismatch fails with a confusing "numpy C-extensions" ImportError
|
| 47 |
+
rather than anything that points at the real cause.
|
| 48 |
+
"""
|
| 49 |
+
return f"{platform.system()}-{platform.machine()}-py{sys.version_info.major}.{sys.version_info.minor}"
|
| 50 |
+
|
| 51 |
+
|
| 52 |
+
def ensure_bootstrapped(log=lambda msg: None) -> None:
|
| 53 |
+
"""Install the isolated zarr>=3 / icechunk stack if not already present for this machine."""
|
| 54 |
+
fingerprint = _fingerprint()
|
| 55 |
+
if SENTINEL.exists() and SENTINEL.read_text().strip() == fingerprint:
|
| 56 |
+
return
|
| 57 |
+
|
| 58 |
+
if ISOLATED_DIR.exists():
|
| 59 |
+
log("π Isolated ERA5 environment was built for a different machine β reinstallingβ¦")
|
| 60 |
+
shutil.rmtree(ISOLATED_DIR)
|
| 61 |
+
|
| 62 |
+
ISOLATED_DIR.mkdir(parents=True, exist_ok=True)
|
| 63 |
+
log(f"π¦ Setting up isolated ERA5 environment (one-time, ~30s)β¦")
|
| 64 |
+
result = subprocess.run(
|
| 65 |
+
[sys.executable, "-m", "pip", "install", "-q", "--target", str(ISOLATED_DIR), *ISOLATED_PACKAGES],
|
| 66 |
+
capture_output=True, text=True,
|
| 67 |
+
)
|
| 68 |
+
if result.returncode != 0:
|
| 69 |
+
raise RuntimeError(
|
| 70 |
+
f"Failed to bootstrap the isolated ERA5 environment:\n{result.stderr[-2000:]}"
|
| 71 |
+
)
|
| 72 |
+
SENTINEL.write_text(fingerprint)
|
| 73 |
+
log("β
Isolated ERA5 environment ready.")
|
| 74 |
+
|
| 75 |
+
|
| 76 |
+
def _subprocess_env() -> dict:
|
| 77 |
+
env = dict(os.environ)
|
| 78 |
+
existing = env.get("PYTHONPATH", "")
|
| 79 |
+
env["PYTHONPATH"] = f"{ISOLATED_DIR}{os.pathsep}{existing}" if existing else str(ISOLATED_DIR)
|
| 80 |
+
return env
|
| 81 |
+
|
| 82 |
+
|
| 83 |
+
def fetch_era5_fields(requests: list[dict], log=lambda msg: None, timeout: int = 600) -> tuple[dict, dict]:
|
| 84 |
+
"""
|
| 85 |
+
Run a batch of ERA5 reads in the isolated subprocess.
|
| 86 |
+
|
| 87 |
+
``requests`` β list of ``{"group", "var", "level", "time_idx"}`` dicts.
|
| 88 |
+
Returns ``(arrays, meta)`` where ``arrays`` maps request index (int) to
|
| 89 |
+
a ``(721, 1440)`` float32 grid, and ``meta`` has ``"errors"`` (index ->
|
| 90 |
+
message, for requests that failed) and ``"resolved_levels"``.
|
| 91 |
+
"""
|
| 92 |
+
ensure_bootstrapped(log)
|
| 93 |
+
|
| 94 |
+
with tempfile.TemporaryDirectory() as tmp:
|
| 95 |
+
request_path = Path(tmp) / "request.json"
|
| 96 |
+
response_prefix = Path(tmp) / "response"
|
| 97 |
+
request_path.write_text(json.dumps(requests))
|
| 98 |
+
|
| 99 |
+
proc = subprocess.Popen(
|
| 100 |
+
[sys.executable, str(WORKER_SCRIPT), str(request_path), str(response_prefix)],
|
| 101 |
+
env=_subprocess_env(),
|
| 102 |
+
stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1,
|
| 103 |
+
)
|
| 104 |
+
stderr_tail = []
|
| 105 |
+
try:
|
| 106 |
+
for line in proc.stdout:
|
| 107 |
+
line = line.rstrip()
|
| 108 |
+
stderr_tail.append(line)
|
| 109 |
+
if line.startswith("PROGRESS") or line.startswith("RETRY"):
|
| 110 |
+
log(f"π‘ {line}")
|
| 111 |
+
proc.wait(timeout=timeout)
|
| 112 |
+
except subprocess.TimeoutExpired:
|
| 113 |
+
proc.kill()
|
| 114 |
+
raise RuntimeError("Timed out waiting for the ERA5 worker subprocess.")
|
| 115 |
+
|
| 116 |
+
if proc.returncode != 0:
|
| 117 |
+
raise RuntimeError(
|
| 118 |
+
"ERA5 worker subprocess failed:\n" + "\n".join(stderr_tail[-30:])
|
| 119 |
+
)
|
| 120 |
+
|
| 121 |
+
npz_path = f"{response_prefix}.npz"
|
| 122 |
+
meta_path = f"{response_prefix}.meta.json"
|
| 123 |
+
with np.load(npz_path) as npz:
|
| 124 |
+
arrays = {int(k): npz[k] for k in npz.files}
|
| 125 |
+
meta = json.loads(Path(meta_path).read_text())
|
| 126 |
+
meta["errors"] = {int(k): v for k, v in meta["errors"].items()}
|
| 127 |
+
meta["resolved_levels"] = {int(k): v for k, v in meta["resolved_levels"].items()}
|
| 128 |
+
return arrays, meta
|
aifs/era5_verify.py
ADDED
|
@@ -0,0 +1,337 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
aifs.era5_verify
|
| 3 |
+
=================
|
| 4 |
+
Score an AIFS forecast (a list of state dicts from :func:`aifs.forecast.run_forecast`,
|
| 5 |
+
initialised from a *historical* date via ``aifs.initial_conditions.load_ics(date=...)``)
|
| 6 |
+
against EarthMover's public ERA5 reanalysis, plus a persistence baseline (ERA5 held
|
| 7 |
+
fixed at the initial time β "assume nothing changes").
|
| 8 |
+
|
| 9 |
+
Only forecasts initialised from a past date can be scored this way: a forecast's
|
| 10 |
+
valid times are in the future relative to when it was run, and ERA5 (reanalysis)
|
| 11 |
+
for those times doesn't exist until well after the fact.
|
| 12 |
+
"""
|
| 13 |
+
|
| 14 |
+
from __future__ import annotations
|
| 15 |
+
|
| 16 |
+
import datetime
|
| 17 |
+
|
| 18 |
+
import numpy as np
|
| 19 |
+
|
| 20 |
+
from aifs.era5_env import fetch_era5_fields
|
| 21 |
+
|
| 22 |
+
STEP_HOURS = 6 # AIFS advances in fixed 6-hour steps (see aifs.forecast)
|
| 23 |
+
_EPOCH = datetime.datetime(1940, 1, 1) # ERA5's valid_time units: "hours since 1940-01-01"
|
| 24 |
+
|
| 25 |
+
#: AIFS field name -> (ERA5 group, ERA5 variable name, pressure level or None).
|
| 26 |
+
#: Names/groups confirmed against the live store; wave fields (swh, mwp, ...) are
|
| 27 |
+
#: NOT included here β EarthMover's ERA5 mirror doesn't carry the wave stream.
|
| 28 |
+
EVAL_FIELD_MAP = {
|
| 29 |
+
"2t": ("single", "t2m", None),
|
| 30 |
+
"10u": ("single", "u10", None),
|
| 31 |
+
"10v": ("single", "v10", None),
|
| 32 |
+
"msl": ("single", "msl", None),
|
| 33 |
+
"sp": ("single", "sp", None),
|
| 34 |
+
"tcw": ("single", "tcw", None),
|
| 35 |
+
"t_850": ("pressure", "t", 850),
|
| 36 |
+
"t_500": ("pressure", "t", 500),
|
| 37 |
+
"u_850": ("pressure", "u", 850),
|
| 38 |
+
"v_850": ("pressure", "v", 850),
|
| 39 |
+
"z_500": ("pressure", "z", 500),
|
| 40 |
+
"q_700": ("pressure", "q", 700),
|
| 41 |
+
}
|
| 42 |
+
|
| 43 |
+
#: ERA5's regular 0.25 degree grid (see aifs.era5_worker probing): lat descends
|
| 44 |
+
#: 90 -> -90, lon ascends 0 -> 359.75 β same [0, 360) convention AIFS uses.
|
| 45 |
+
_ERA5_LAT0, _ERA5_DLAT, _ERA5_NLAT = 90.0, -0.25, 721
|
| 46 |
+
_ERA5_LON0, _ERA5_DLON, _ERA5_NLON = 0.0, 0.25, 1440
|
| 47 |
+
|
| 48 |
+
|
| 49 |
+
def _time_idx(dt: datetime.datetime) -> int:
|
| 50 |
+
return int((dt - _EPOCH).total_seconds() // 3600)
|
| 51 |
+
|
| 52 |
+
|
| 53 |
+
def _nearest_indices(lats: np.ndarray, lons: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
|
| 54 |
+
"""Map each (lat, lon) point to its nearest cell in ERA5's regular grid."""
|
| 55 |
+
lat_idx = np.clip(np.round((lats - _ERA5_LAT0) / _ERA5_DLAT).astype(int), 0, _ERA5_NLAT - 1)
|
| 56 |
+
lon_idx = np.round((lons - _ERA5_LON0) / _ERA5_DLON).astype(int) % _ERA5_NLON
|
| 57 |
+
return lat_idx, lon_idx
|
| 58 |
+
|
| 59 |
+
|
| 60 |
+
def sample_at_points(grid: np.ndarray, lat_idx: np.ndarray, lon_idx: np.ndarray) -> np.ndarray:
|
| 61 |
+
"""Sample a (721, 1440) ERA5 grid at pre-computed nearest-neighbour indices."""
|
| 62 |
+
return grid[lat_idx, lon_idx]
|
| 63 |
+
|
| 64 |
+
|
| 65 |
+
def metrics(forecast: np.ndarray, truth: np.ndarray) -> dict:
|
| 66 |
+
"""RMSE / MAE / bias / correlation between two same-shape point arrays, NaN-safe."""
|
| 67 |
+
mask = ~(np.isnan(forecast) | np.isnan(truth))
|
| 68 |
+
f, t = np.asarray(forecast)[mask], np.asarray(truth)[mask]
|
| 69 |
+
if len(f) < 2:
|
| 70 |
+
return {"rmse": float("nan"), "mae": float("nan"), "bias": float("nan"), "corr": float("nan"), "n": int(len(f))}
|
| 71 |
+
diff = f - t
|
| 72 |
+
corr = float(np.corrcoef(f, t)[0, 1])
|
| 73 |
+
return {
|
| 74 |
+
"rmse": float(np.sqrt(np.mean(diff ** 2))),
|
| 75 |
+
"mae": float(np.mean(np.abs(diff))),
|
| 76 |
+
"bias": float(np.mean(diff)),
|
| 77 |
+
"corr": corr,
|
| 78 |
+
"n": int(len(f)),
|
| 79 |
+
}
|
| 80 |
+
|
| 81 |
+
|
| 82 |
+
def _fetch_era5_bounds(group: str, extra_requests: list[dict], log):
|
| 83 |
+
"""
|
| 84 |
+
Fetch a group's live ``valid_time`` coverage bounds together with
|
| 85 |
+
``extra_requests`` in one batched call β avoids paying icechunk's
|
| 86 |
+
session-open cost twice. EarthMover's free ERA5 archive is NOT current
|
| 87 |
+
to "now" (it lags real-time by several months, and how far isn't
|
| 88 |
+
documented reliably), so callers check this live rather than trusting
|
| 89 |
+
a hardcoded date.
|
| 90 |
+
|
| 91 |
+
Returns ``(era5_start, era5_end, arrays, meta)`` β ``arrays``/``meta``
|
| 92 |
+
cover ALL requests: bounds are indices 0 and 1, ``extra_requests``
|
| 93 |
+
start at index 2.
|
| 94 |
+
"""
|
| 95 |
+
bounds_requests = [
|
| 96 |
+
{"group": group, "var": "valid_time", "level": None, "time_idx": 0},
|
| 97 |
+
{"group": group, "var": "valid_time", "level": None, "time_idx": -1},
|
| 98 |
+
]
|
| 99 |
+
arrays, meta = fetch_era5_fields(bounds_requests + extra_requests, log=log)
|
| 100 |
+
if 0 in meta["errors"] or 1 in meta["errors"]:
|
| 101 |
+
raise RuntimeError(
|
| 102 |
+
f"Could not read ERA5's time coverage for group '{group}': "
|
| 103 |
+
f"{meta['errors'].get(0) or meta['errors'].get(1)}"
|
| 104 |
+
)
|
| 105 |
+
era5_start = _EPOCH + datetime.timedelta(hours=float(arrays[0]))
|
| 106 |
+
era5_end = _EPOCH + datetime.timedelta(hours=float(arrays[1]))
|
| 107 |
+
return era5_start, era5_end, arrays, meta
|
| 108 |
+
|
| 109 |
+
|
| 110 |
+
def evaluate_forecast(states: list[dict], field_name: str, log=lambda msg: None) -> dict:
|
| 111 |
+
"""
|
| 112 |
+
Score every step of a historical-init forecast against ERA5, plus a
|
| 113 |
+
persistence baseline (ERA5 held fixed at t0).
|
| 114 |
+
|
| 115 |
+
Returns a dict: ``{"field", "lats", "lons", "per_step": [...]}`` β each
|
| 116 |
+
per_step entry has ``date``, ``lead_h``, ``forecast``, ``truth`` (both
|
| 117 |
+
sampled onto the AIFS grid), and ``model_*``/``persistence_*`` metrics.
|
| 118 |
+
"""
|
| 119 |
+
if not states:
|
| 120 |
+
raise ValueError("No forecast states to evaluate β run a forecast first.")
|
| 121 |
+
if field_name not in EVAL_FIELD_MAP:
|
| 122 |
+
raise ValueError(
|
| 123 |
+
f"'{field_name}' has no ERA5 counterpart configured. "
|
| 124 |
+
f"Available: {sorted(EVAL_FIELD_MAP)}"
|
| 125 |
+
)
|
| 126 |
+
|
| 127 |
+
group, var, level = EVAL_FIELD_MAP[field_name]
|
| 128 |
+
lats = np.asarray(states[0]["latitudes"]).ravel()
|
| 129 |
+
lons = np.asarray(states[0]["longitudes"]).ravel()
|
| 130 |
+
lat_idx, lon_idx = _nearest_indices(lats, lons)
|
| 131 |
+
|
| 132 |
+
t0 = states[0]["date"] - datetime.timedelta(hours=STEP_HOURS)
|
| 133 |
+
valid_times = [t0] + [s["date"] for s in states]
|
| 134 |
+
|
| 135 |
+
# De-duplicate in case any two steps land on the same valid time.
|
| 136 |
+
unique_times = sorted(set(valid_times))
|
| 137 |
+
data_requests = [
|
| 138 |
+
{"group": group, "var": var, "level": level, "time_idx": _time_idx(t)}
|
| 139 |
+
for t in unique_times
|
| 140 |
+
]
|
| 141 |
+
|
| 142 |
+
log(f"π‘ Fetching ERA5 '{var}'" + (f"@{level}hPa" if level else "") +
|
| 143 |
+
f" for {len(unique_times)} valid time(s)β¦")
|
| 144 |
+
era5_start, era5_end, arrays, meta = _fetch_era5_bounds(group, data_requests, log)
|
| 145 |
+
|
| 146 |
+
earliest_needed, latest_needed = unique_times[0], unique_times[-1]
|
| 147 |
+
if latest_needed > era5_end or earliest_needed < era5_start:
|
| 148 |
+
raise ValueError(
|
| 149 |
+
f"This forecast needs ERA5 truth from {earliest_needed} to {latest_needed}, but "
|
| 150 |
+
f"EarthMover's free ERA5 archive currently only covers {era5_start} to {era5_end}. "
|
| 151 |
+
f"Pick an earlier historical init date β with enough room before {era5_end} to fit "
|
| 152 |
+
f"your full lead time β and run the forecast again."
|
| 153 |
+
)
|
| 154 |
+
|
| 155 |
+
data_errors = {i - 2: err for i, err in meta["errors"].items() if i >= 2}
|
| 156 |
+
if data_errors:
|
| 157 |
+
first_err = next(iter(data_errors.values()))
|
| 158 |
+
raise RuntimeError(f"Could not read '{field_name}' from EarthMover's ERA5 store: {first_err}")
|
| 159 |
+
|
| 160 |
+
truth_by_time = {
|
| 161 |
+
unique_times[i]: sample_at_points(arrays[i + 2], lat_idx, lon_idx)
|
| 162 |
+
for i in range(len(unique_times))
|
| 163 |
+
}
|
| 164 |
+
|
| 165 |
+
truth_t0 = truth_by_time[t0]
|
| 166 |
+
per_step = []
|
| 167 |
+
for i, state in enumerate(states, start=1):
|
| 168 |
+
truth = truth_by_time[state["date"]]
|
| 169 |
+
forecast = np.asarray(state["fields"][field_name])
|
| 170 |
+
|
| 171 |
+
m = metrics(forecast, truth)
|
| 172 |
+
p = metrics(truth_t0, truth)
|
| 173 |
+
|
| 174 |
+
per_step.append({
|
| 175 |
+
"date": str(state["date"]),
|
| 176 |
+
"lead_h": i * STEP_HOURS,
|
| 177 |
+
"forecast": forecast,
|
| 178 |
+
"truth": truth,
|
| 179 |
+
**{f"model_{k}": v for k, v in m.items()},
|
| 180 |
+
**{f"persistence_{k}": v for k, v in p.items()},
|
| 181 |
+
})
|
| 182 |
+
|
| 183 |
+
log(f"β
Scored {len(per_step)} step(s) against ERA5 for '{field_name}'.")
|
| 184 |
+
return {"field": field_name, "lats": lats, "lons": lons, "per_step": per_step}
|
| 185 |
+
|
| 186 |
+
|
| 187 |
+
def compute_climatology(state: dict, field_name: str, num_years: int = 10, log=lambda msg: None) -> dict:
|
| 188 |
+
"""
|
| 189 |
+
Compare ONE forecast state's field against ERA5's climatology for that
|
| 190 |
+
calendar day/hour β the mean (and spread) over ``num_years`` years of
|
| 191 |
+
ERA5 history on the same month/day/hour.
|
| 192 |
+
|
| 193 |
+
Unlike :func:`evaluate_forecast`, this works for a LIVE (today's or a
|
| 194 |
+
future) forecast: no ERA5 truth is needed for the forecast's own valid
|
| 195 |
+
time, only for many *past* years on the same calendar day β sidestepping
|
| 196 |
+
ERA5's coverage-ceiling problem entirely. It answers "is this forecast
|
| 197 |
+
unusual for the time of year?", not "is this forecast correct?".
|
| 198 |
+
"""
|
| 199 |
+
if field_name not in EVAL_FIELD_MAP:
|
| 200 |
+
raise ValueError(
|
| 201 |
+
f"'{field_name}' has no ERA5 counterpart configured. "
|
| 202 |
+
f"Available: {sorted(EVAL_FIELD_MAP)}"
|
| 203 |
+
)
|
| 204 |
+
if num_years < 2:
|
| 205 |
+
raise ValueError("num_years must be at least 2 to compute a meaningful climatology.")
|
| 206 |
+
|
| 207 |
+
group, var, level = EVAL_FIELD_MAP[field_name]
|
| 208 |
+
lats = np.asarray(state["latitudes"]).ravel()
|
| 209 |
+
lons = np.asarray(state["longitudes"]).ravel()
|
| 210 |
+
lat_idx, lon_idx = _nearest_indices(lats, lons)
|
| 211 |
+
|
| 212 |
+
valid_date: datetime.datetime = state["date"]
|
| 213 |
+
|
| 214 |
+
# Request a few extra candidate years beyond num_years so that skipped
|
| 215 |
+
# years (Feb 29 in a non-leap year, or anything ERA5 doesn't have yet
|
| 216 |
+
# for the most recent year or two) still leave enough successful
|
| 217 |
+
# samples β cheaper than pre-checking the archive's bounds separately.
|
| 218 |
+
BUFFER = 5
|
| 219 |
+
candidate_years = []
|
| 220 |
+
year = valid_date.year - 1 # climatology uses PAST years only
|
| 221 |
+
while len(candidate_years) < num_years + BUFFER and year >= 1940:
|
| 222 |
+
try:
|
| 223 |
+
candidate_years.append(valid_date.replace(year=year))
|
| 224 |
+
except ValueError:
|
| 225 |
+
pass # Feb 29 in a non-leap year β skip it
|
| 226 |
+
year -= 1
|
| 227 |
+
|
| 228 |
+
requests = [
|
| 229 |
+
{"group": group, "var": var, "level": level, "time_idx": _time_idx(t)}
|
| 230 |
+
for t in candidate_years
|
| 231 |
+
]
|
| 232 |
+
log(f"π‘ Fetching up to {num_years} years of ERA5 '{var}'" + (f"@{level}hPa" if level else "") +
|
| 233 |
+
f" for {valid_date.strftime('%m-%d %H:%M')} UTCβ¦")
|
| 234 |
+
arrays, meta = fetch_era5_fields(requests, log=log)
|
| 235 |
+
|
| 236 |
+
successful = [(candidate_years[i], arrays[i]) for i in range(len(candidate_years)) if i not in meta["errors"]]
|
| 237 |
+
if len(successful) < 2:
|
| 238 |
+
first_err = next(iter(meta["errors"].values()), "no data available")
|
| 239 |
+
raise RuntimeError(
|
| 240 |
+
f"Could not build a climatology for '{field_name}' β only {len(successful)} of "
|
| 241 |
+
f"{len(candidate_years)} candidate years had ERA5 data ({first_err})."
|
| 242 |
+
)
|
| 243 |
+
successful = successful[:num_years]
|
| 244 |
+
used_years = [t.year for t, _ in successful]
|
| 245 |
+
|
| 246 |
+
samples = np.stack([sample_at_points(grid, lat_idx, lon_idx) for _, grid in successful])
|
| 247 |
+
clim_mean = np.nanmean(samples, axis=0)
|
| 248 |
+
clim_std = np.nanstd(samples, axis=0)
|
| 249 |
+
|
| 250 |
+
forecast = np.asarray(state["fields"][field_name])
|
| 251 |
+
anomaly = forecast - clim_mean
|
| 252 |
+
zscore = np.full_like(anomaly, np.nan)
|
| 253 |
+
valid_std = clim_std > 1e-6
|
| 254 |
+
zscore[valid_std] = anomaly[valid_std] / clim_std[valid_std]
|
| 255 |
+
|
| 256 |
+
log(f"β
Built a {len(used_years)}-year climatology ({min(used_years)}β{max(used_years)}) for '{field_name}'.")
|
| 257 |
+
|
| 258 |
+
return {
|
| 259 |
+
"field": field_name, "lats": lats, "lons": lons,
|
| 260 |
+
"date": str(valid_date), "years": used_years,
|
| 261 |
+
"forecast": forecast, "climatology_mean": clim_mean, "climatology_std": clim_std,
|
| 262 |
+
"anomaly": anomaly, "zscore": zscore,
|
| 263 |
+
}
|
| 264 |
+
|
| 265 |
+
|
| 266 |
+
# ββ Plotting (matches aifs.plot.plot_field's Cartopy tricontourf style) ββββββββ
|
| 267 |
+
|
| 268 |
+
def _map_figure(lats, lons, data, title: str, cmap: str = "RdBu_r"):
|
| 269 |
+
import cartopy.crs as ccrs
|
| 270 |
+
import cartopy.feature as cfeature
|
| 271 |
+
import matplotlib.pyplot as plt
|
| 272 |
+
import matplotlib.tri as tri
|
| 273 |
+
|
| 274 |
+
lons_plot = np.where(lons > 180, lons - 360, lons)
|
| 275 |
+
triangulation = tri.Triangulation(lons_plot, lats)
|
| 276 |
+
|
| 277 |
+
fig, ax = plt.subplots(figsize=(9, 5), subplot_kw={"projection": ccrs.PlateCarree()})
|
| 278 |
+
ax.coastlines()
|
| 279 |
+
ax.add_feature(cfeature.BORDERS, linestyle=":")
|
| 280 |
+
contour = ax.tricontourf(triangulation, data, levels=20, transform=ccrs.PlateCarree(), cmap=cmap)
|
| 281 |
+
fig.colorbar(contour, ax=ax, orientation="vertical", shrink=0.7)
|
| 282 |
+
ax.set_title(title, fontsize=11)
|
| 283 |
+
fig.tight_layout()
|
| 284 |
+
return fig
|
| 285 |
+
|
| 286 |
+
|
| 287 |
+
def plot_eval_maps(eval_data: dict, step_index: int):
|
| 288 |
+
"""Forecast / ERA5 truth / (forecast - truth) maps for one evaluated step."""
|
| 289 |
+
step = eval_data["per_step"][step_index]
|
| 290 |
+
lats, lons, field = eval_data["lats"], eval_data["lons"], eval_data["field"]
|
| 291 |
+
|
| 292 |
+
fig_fc = _map_figure(lats, lons, step["forecast"], f"AIFS forecast β {field} @ {step['date']}")
|
| 293 |
+
fig_truth = _map_figure(lats, lons, step["truth"], f"ERA5 truth β {field} @ {step['date']}")
|
| 294 |
+
fig_diff = _map_figure(
|
| 295 |
+
lats, lons, step["forecast"] - step["truth"],
|
| 296 |
+
f"Forecast β ERA5 β {field} @ {step['date']}",
|
| 297 |
+
)
|
| 298 |
+
return fig_fc, fig_truth, fig_diff
|
| 299 |
+
|
| 300 |
+
|
| 301 |
+
def plot_climatology_maps(clim_data: dict):
|
| 302 |
+
"""Forecast / ERA5 climatology-mean / anomaly maps for one field."""
|
| 303 |
+
lats, lons, field = clim_data["lats"], clim_data["lons"], clim_data["field"]
|
| 304 |
+
year_range = f"{min(clim_data['years'])}β{max(clim_data['years'])}"
|
| 305 |
+
n = len(clim_data["years"])
|
| 306 |
+
|
| 307 |
+
fig_fc = _map_figure(lats, lons, clim_data["forecast"], f"AIFS forecast β {field} @ {clim_data['date']}")
|
| 308 |
+
fig_clim = _map_figure(
|
| 309 |
+
lats, lons, clim_data["climatology_mean"],
|
| 310 |
+
f"ERA5 climatology mean β {field} ({year_range}, n={n})",
|
| 311 |
+
)
|
| 312 |
+
fig_anom = _map_figure(
|
| 313 |
+
lats, lons, clim_data["anomaly"],
|
| 314 |
+
f"Anomaly (forecast β climatology) β {field}",
|
| 315 |
+
)
|
| 316 |
+
return fig_fc, fig_clim, fig_anom
|
| 317 |
+
|
| 318 |
+
|
| 319 |
+
def plot_skill_curve(eval_data: dict):
|
| 320 |
+
"""RMSE vs lead time: AIFS forecast vs a persistence (ERA5-held-at-t0) baseline."""
|
| 321 |
+
import matplotlib.pyplot as plt
|
| 322 |
+
|
| 323 |
+
steps = eval_data["per_step"]
|
| 324 |
+
lead_h = [s["lead_h"] for s in steps]
|
| 325 |
+
model_rmse = [s["model_rmse"] for s in steps]
|
| 326 |
+
persistence_rmse = [s["persistence_rmse"] for s in steps]
|
| 327 |
+
|
| 328 |
+
fig, ax = plt.subplots(figsize=(7, 4))
|
| 329 |
+
ax.plot(lead_h, model_rmse, marker="o", label="AIFS forecast", color="#1f6feb")
|
| 330 |
+
ax.plot(lead_h, persistence_rmse, marker="o", linestyle="--", label="Persistence (t0 held fixed)", color="#8b949e")
|
| 331 |
+
ax.set_xlabel("Lead time (hours)")
|
| 332 |
+
ax.set_ylabel("RMSE")
|
| 333 |
+
ax.set_title(f"Skill vs ERA5 β {eval_data['field']}")
|
| 334 |
+
ax.legend()
|
| 335 |
+
ax.grid(alpha=0.3)
|
| 336 |
+
fig.tight_layout()
|
| 337 |
+
return fig
|
aifs/era5_worker.py
ADDED
|
@@ -0,0 +1,130 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
aifs.era5_worker
|
| 3 |
+
================
|
| 4 |
+
Standalone subprocess entry point β reads EarthMover's public ERA5
|
| 5 |
+
Icechunk/Zarr-v3 store on S3 (anonymous access, no account needed).
|
| 6 |
+
|
| 7 |
+
This script is deliberately NOT imported by the rest of the app. It is
|
| 8 |
+
invoked via ``subprocess`` with ``PYTHONPATH`` pointed at an isolated
|
| 9 |
+
``pip install --target`` directory (see :mod:`aifs.era5_env`), because
|
| 10 |
+
``icechunk`` requires ``zarr>=3`` while ``anemoi-datasets`` (already a
|
| 11 |
+
hard dependency of this Space, for AIFS inference) pins ``zarr<=2.18``.
|
| 12 |
+
The two cannot coexist in one interpreter's import path, so this script
|
| 13 |
+
only ever runs in a separate process with its own isolated zarr install.
|
| 14 |
+
|
| 15 |
+
Only stdlib + numpy + icechunk + zarr are imported here β keep it that
|
| 16 |
+
way so it never accidentally picks up the main env's (incompatible)
|
| 17 |
+
zarr via a transitive import.
|
| 18 |
+
|
| 19 |
+
Usage
|
| 20 |
+
-----
|
| 21 |
+
python3 era5_worker.py <request.json> <response_prefix>
|
| 22 |
+
|
| 23 |
+
``request.json`` is a list of ``{"group", "var", "level", "time_idx"}``
|
| 24 |
+
dicts. Writes ``<response_prefix>.npz`` (arrays keyed by request index,
|
| 25 |
+
one per successfully-read request) and ``<response_prefix>.meta.json``
|
| 26 |
+
(``{"errors": {...}, "resolved_levels": {...}}``). Progress is reported
|
| 27 |
+
as ``PROGRESS <i>/<n> <group>/<var>`` lines on stdout.
|
| 28 |
+
"""
|
| 29 |
+
|
| 30 |
+
import json
|
| 31 |
+
import sys
|
| 32 |
+
import time
|
| 33 |
+
|
| 34 |
+
import numpy as np
|
| 35 |
+
|
| 36 |
+
BUCKET = "earthmover-icechunk-era5"
|
| 37 |
+
PREFIX = "icechunkV2"
|
| 38 |
+
REGION = "us-east-1"
|
| 39 |
+
|
| 40 |
+
MAX_RETRIES = 6
|
| 41 |
+
_RETRIABLE_KEYWORDS = (
|
| 42 |
+
"429", "rate limit", "too many requests", "timeout", "connection reset",
|
| 43 |
+
"503", "service unavailable", "throughput", "streaming error", "i/o error",
|
| 44 |
+
)
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
def _retriable(exc: Exception) -> bool:
|
| 48 |
+
msg = str(exc).lower()
|
| 49 |
+
return any(k in msg for k in _RETRIABLE_KEYWORDS)
|
| 50 |
+
|
| 51 |
+
|
| 52 |
+
def _with_retry(fn, label: str):
|
| 53 |
+
for attempt in range(MAX_RETRIES):
|
| 54 |
+
try:
|
| 55 |
+
return fn()
|
| 56 |
+
except Exception as exc:
|
| 57 |
+
if attempt < MAX_RETRIES - 1 and _retriable(exc):
|
| 58 |
+
wait = min(3 * (2 ** attempt), 30)
|
| 59 |
+
print(f"RETRY {label}: {exc} (attempt {attempt + 2}/{MAX_RETRIES}, waiting {wait}s)", flush=True)
|
| 60 |
+
time.sleep(wait)
|
| 61 |
+
else:
|
| 62 |
+
raise
|
| 63 |
+
|
| 64 |
+
|
| 65 |
+
def _open_store():
|
| 66 |
+
import icechunk
|
| 67 |
+
|
| 68 |
+
def _open():
|
| 69 |
+
storage = icechunk.s3_storage(bucket=BUCKET, prefix=PREFIX, region=REGION, anonymous=True)
|
| 70 |
+
repo = icechunk.Repository.open(storage)
|
| 71 |
+
session = repo.readonly_session("main")
|
| 72 |
+
return session.store
|
| 73 |
+
|
| 74 |
+
return _with_retry(_open, "open icechunk repo")
|
| 75 |
+
|
| 76 |
+
|
| 77 |
+
def main():
|
| 78 |
+
request_path, response_prefix = sys.argv[1], sys.argv[2]
|
| 79 |
+
with open(request_path) as f:
|
| 80 |
+
requests = json.load(f)
|
| 81 |
+
|
| 82 |
+
import zarr
|
| 83 |
+
|
| 84 |
+
store = _open_store()
|
| 85 |
+
groups = {}
|
| 86 |
+
level_coords = {}
|
| 87 |
+
|
| 88 |
+
results = {}
|
| 89 |
+
errors = {}
|
| 90 |
+
resolved_levels = {}
|
| 91 |
+
|
| 92 |
+
n = len(requests)
|
| 93 |
+
for i, req in enumerate(requests):
|
| 94 |
+
group, var, level, time_idx = req["group"], req["var"], req.get("level"), req["time_idx"]
|
| 95 |
+
print(f"PROGRESS {i + 1}/{n} {group}/{var}" + (f"@{level}hPa" if level is not None else ""), flush=True)
|
| 96 |
+
try:
|
| 97 |
+
if group not in groups:
|
| 98 |
+
groups[group] = _with_retry(
|
| 99 |
+
lambda g=group: zarr.open_group(store, mode="r", path=f"{g}/spatial"),
|
| 100 |
+
f"open group {group}",
|
| 101 |
+
)
|
| 102 |
+
g = groups[group]
|
| 103 |
+
|
| 104 |
+
if var not in g.array_keys():
|
| 105 |
+
raise KeyError(f"'{var}' not found in ERA5 group '{group}' (have: {sorted(g.array_keys())})")
|
| 106 |
+
|
| 107 |
+
arr = g[var]
|
| 108 |
+
|
| 109 |
+
def _read():
|
| 110 |
+
if level is not None:
|
| 111 |
+
if group not in level_coords:
|
| 112 |
+
level_coords[group] = np.asarray(g["pressure_level"][:])
|
| 113 |
+
levels = level_coords[group]
|
| 114 |
+
pos = int(np.argmin(np.abs(levels - level)))
|
| 115 |
+
resolved_levels[str(i)] = float(levels[pos])
|
| 116 |
+
return arr[time_idx, pos]
|
| 117 |
+
return arr[time_idx]
|
| 118 |
+
|
| 119 |
+
data = _with_retry(_read, f"read {group}/{var}")
|
| 120 |
+
results[str(i)] = np.asarray(data, dtype=np.float32)
|
| 121 |
+
except Exception as exc:
|
| 122 |
+
errors[str(i)] = str(exc)
|
| 123 |
+
|
| 124 |
+
np.savez_compressed(f"{response_prefix}.npz", **results)
|
| 125 |
+
with open(f"{response_prefix}.meta.json", "w") as f:
|
| 126 |
+
json.dump({"errors": errors, "resolved_levels": resolved_levels}, f)
|
| 127 |
+
|
| 128 |
+
|
| 129 |
+
if __name__ == "__main__":
|
| 130 |
+
main()
|
aifs/initial_conditions.py
CHANGED
|
@@ -35,6 +35,32 @@ LEVELS = [1000, 925, 850, 700, 600, 500, 400, 300, 250, 200, 150, 100, 50, 10]
|
|
| 35 |
|
| 36 |
SOURCE = "ecmwf"
|
| 37 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 38 |
# ββ Cache helpers βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 39 |
|
| 40 |
DEFAULT_CACHE_DIR = Path("ic_cache")
|
|
@@ -97,7 +123,7 @@ def _fetch_with_retry(ekd, ekr, date, param, max_retries: int = 6, **kwargs):
|
|
| 97 |
raise
|
| 98 |
|
| 99 |
|
| 100 |
-
def _fetch_fields(ekd, ekr, date, param, levelist=None, **kwargs) -> dict:
|
| 101 |
"""
|
| 102 |
Download ``param`` for two time-steps (t-6h, t) and return a dict
|
| 103 |
``{variable_name: np.ndarray shape (2, N320_nodes)}``.
|
|
@@ -111,7 +137,7 @@ def _fetch_fields(ekd, ekr, date, param, levelist=None, **kwargs) -> dict:
|
|
| 111 |
date=t,
|
| 112 |
param=param,
|
| 113 |
levelist=levelist,
|
| 114 |
-
source=
|
| 115 |
**kwargs,
|
| 116 |
)
|
| 117 |
for field in dataset:
|
|
@@ -134,7 +160,7 @@ def _fetch_fields(ekd, ekr, date, param, levelist=None, **kwargs) -> dict:
|
|
| 134 |
|
| 135 |
|
| 136 |
|
| 137 |
-
def _build_fields(ekd, ekr, date: datetime.datetime):
|
| 138 |
"""Download and transform all required fields for ``date``."""
|
| 139 |
fields: dict = {}
|
| 140 |
log_lines: list[str] = []
|
|
@@ -148,6 +174,7 @@ def _build_fields(ekd, ekr, date: datetime.datetime):
|
|
| 148 |
"""Yield log messages then return the result dict."""
|
| 149 |
yield log(label)
|
| 150 |
result = None
|
|
|
|
| 151 |
for kind, payload in _fetch_with_retry(ekd, ekr, *args, **kwargs):
|
| 152 |
if kind == "log":
|
| 153 |
yield log(f" {payload}")
|
|
@@ -216,10 +243,21 @@ def _build_fields(ekd, ekr, date: datetime.datetime):
|
|
| 216 |
def load_ics(
|
| 217 |
cache_dir: Path | str = DEFAULT_CACHE_DIR,
|
| 218 |
force: bool = False,
|
|
|
|
| 219 |
):
|
| 220 |
"""
|
| 221 |
Generator version of load_ics.
|
| 222 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 223 |
Yields
|
| 224 |
------
|
| 225 |
("log", str) -- progress messages
|
|
@@ -227,13 +265,29 @@ def load_ics(
|
|
| 227 |
"""
|
| 228 |
import earthkit.data as ekd
|
| 229 |
import earthkit.regrid as ekr
|
| 230 |
-
from ecmwf.opendata import Client as OpendataClient
|
| 231 |
|
| 232 |
ekd.config.set({"cache-policy": "user"})
|
| 233 |
cache_dir = Path(cache_dir)
|
| 234 |
|
| 235 |
-
date
|
| 236 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 237 |
|
| 238 |
if not force:
|
| 239 |
cached, path = _try_load(date, cache_dir)
|
|
@@ -246,7 +300,7 @@ def load_ics(
|
|
| 246 |
yield "log", "β¬οΈ Downloading initial conditions β¦"
|
| 247 |
|
| 248 |
fields = None
|
| 249 |
-
for kind, payload in _build_fields(ekd, ekr, date):
|
| 250 |
if kind == "log":
|
| 251 |
yield "log", payload
|
| 252 |
else: # "result"
|
|
|
|
| 35 |
|
| 36 |
SOURCE = "ecmwf"
|
| 37 |
|
| 38 |
+
#: ECMWF Open Data's primary endpoint ("ecmwf") only keeps a rolling ~4-day
|
| 39 |
+
#: window. The "aws" named source (an S3 mirror, still via the same
|
| 40 |
+
#: ecmwf-opendata/earthkit-data client) retains a much deeper archive
|
| 41 |
+
#: (observed back to 2023-01-18) β used for historical initial conditions.
|
| 42 |
+
HISTORICAL_SOURCE = "aws"
|
| 43 |
+
|
| 44 |
+
#: Valid ECMWF synoptic run hours (UTC) β the archive only has data at these.
|
| 45 |
+
VALID_RUN_HOURS = (0, 6, 12, 18)
|
| 46 |
+
|
| 47 |
+
#: Earliest date the "aws" historical archive has been observed to serve.
|
| 48 |
+
EARLIEST_HISTORICAL_DATE = datetime.date(2023, 1, 18)
|
| 49 |
+
|
| 50 |
+
#: Before this date, ECMWF's 06/18 UTC runs used a reduced product ("scda" for
|
| 51 |
+
#: pressure levels, "scwv" for waves) instead of the full "oper"/"wave" stream
|
| 52 |
+
#: β confirmed missing: all pressure-level vars at 10 hPa, and everything in
|
| 53 |
+
#: PARAM_WAVE except mwd/mwp/swh. 00/12 UTC always used the full stream, even
|
| 54 |
+
#: before this cutover. This was unified across all run hours starting on this
|
| 55 |
+
#: date, but that's *after* EarthMover's free ERA5 archive's coverage ends (see
|
| 56 |
+
#: aifs.era5_verify) β so for this app's actual use case (historical init +
|
| 57 |
+
#: ERA5 verification), 06/18 UTC is never a usable choice anyway.
|
| 58 |
+
REDUCED_PRODUCT_CUTOVER = datetime.date(2026, 5, 12)
|
| 59 |
+
REDUCED_PRODUCT_HOURS = (6, 18)
|
| 60 |
+
|
| 61 |
+
#: Run hours guaranteed to carry the full field set, at any historical date.
|
| 62 |
+
FULL_FIELD_RUN_HOURS = (0, 12)
|
| 63 |
+
|
| 64 |
# ββ Cache helpers βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 65 |
|
| 66 |
DEFAULT_CACHE_DIR = Path("ic_cache")
|
|
|
|
| 123 |
raise
|
| 124 |
|
| 125 |
|
| 126 |
+
def _fetch_fields(ekd, ekr, date, param, levelist=None, source=SOURCE, **kwargs) -> dict:
|
| 127 |
"""
|
| 128 |
Download ``param`` for two time-steps (t-6h, t) and return a dict
|
| 129 |
``{variable_name: np.ndarray shape (2, N320_nodes)}``.
|
|
|
|
| 137 |
date=t,
|
| 138 |
param=param,
|
| 139 |
levelist=levelist,
|
| 140 |
+
source=source,
|
| 141 |
**kwargs,
|
| 142 |
)
|
| 143 |
for field in dataset:
|
|
|
|
| 160 |
|
| 161 |
|
| 162 |
|
| 163 |
+
def _build_fields(ekd, ekr, date: datetime.datetime, source: str = SOURCE):
|
| 164 |
"""Download and transform all required fields for ``date``."""
|
| 165 |
fields: dict = {}
|
| 166 |
log_lines: list[str] = []
|
|
|
|
| 174 |
"""Yield log messages then return the result dict."""
|
| 175 |
yield log(label)
|
| 176 |
result = None
|
| 177 |
+
kwargs.setdefault("source", source)
|
| 178 |
for kind, payload in _fetch_with_retry(ekd, ekr, *args, **kwargs):
|
| 179 |
if kind == "log":
|
| 180 |
yield log(f" {payload}")
|
|
|
|
| 243 |
def load_ics(
|
| 244 |
cache_dir: Path | str = DEFAULT_CACHE_DIR,
|
| 245 |
force: bool = False,
|
| 246 |
+
date: datetime.datetime | None = None,
|
| 247 |
):
|
| 248 |
"""
|
| 249 |
Generator version of load_ics.
|
| 250 |
|
| 251 |
+
Parameters
|
| 252 |
+
----------
|
| 253 |
+
date:
|
| 254 |
+
If ``None`` (default), fetch the latest available run from
|
| 255 |
+
ECMWF Open Data's primary endpoint β unchanged live behaviour.
|
| 256 |
+
If given, fetch that specific historical run instead, from the
|
| 257 |
+
"aws" S3 mirror (a much deeper archive than the ~4-day rolling
|
| 258 |
+
window of the primary endpoint β observed back to 2023-01-18).
|
| 259 |
+
Must fall on a synoptic run hour (00/06/12/18 UTC).
|
| 260 |
+
|
| 261 |
Yields
|
| 262 |
------
|
| 263 |
("log", str) -- progress messages
|
|
|
|
| 265 |
"""
|
| 266 |
import earthkit.data as ekd
|
| 267 |
import earthkit.regrid as ekr
|
|
|
|
| 268 |
|
| 269 |
ekd.config.set({"cache-policy": "user"})
|
| 270 |
cache_dir = Path(cache_dir)
|
| 271 |
|
| 272 |
+
if date is None:
|
| 273 |
+
from ecmwf.opendata import Client as OpendataClient
|
| 274 |
+
|
| 275 |
+
date = OpendataClient(SOURCE).latest()
|
| 276 |
+
yield "log", f"π
Latest ECMWF run: {date}"
|
| 277 |
+
source = SOURCE
|
| 278 |
+
else:
|
| 279 |
+
if date.hour not in VALID_RUN_HOURS or date.minute or date.second or date.microsecond:
|
| 280 |
+
raise ValueError(
|
| 281 |
+
f"date must fall on a synoptic run hour {VALID_RUN_HOURS} UTC, got {date}"
|
| 282 |
+
)
|
| 283 |
+
if date.hour in REDUCED_PRODUCT_HOURS and date.date() < REDUCED_PRODUCT_CUTOVER:
|
| 284 |
+
raise ValueError(
|
| 285 |
+
f"{date} is a 06/18 UTC run before {REDUCED_PRODUCT_CUTOVER.isoformat()} β "
|
| 286 |
+
"ECMWF served a reduced product at those hours back then (missing 10 hPa "
|
| 287 |
+
"pressure levels and most wave fields AIFS needs). Use a 00 or 12 UTC run instead."
|
| 288 |
+
)
|
| 289 |
+
yield "log", f"π
Historical ECMWF run: {date} (via '{HISTORICAL_SOURCE}' archive)"
|
| 290 |
+
source = HISTORICAL_SOURCE
|
| 291 |
|
| 292 |
if not force:
|
| 293 |
cached, path = _try_load(date, cache_dir)
|
|
|
|
| 300 |
yield "log", "β¬οΈ Downloading initial conditions β¦"
|
| 301 |
|
| 302 |
fields = None
|
| 303 |
+
for kind, payload in _build_fields(ekd, ekr, date, source=source):
|
| 304 |
if kind == "log":
|
| 305 |
yield "log", payload
|
| 306 |
else: # "result"
|
app.py
CHANGED
|
@@ -1,3 +1,4 @@
|
|
|
|
|
| 1 |
import os
|
| 2 |
import ssl
|
| 3 |
import warnings
|
|
@@ -19,6 +20,8 @@ import gradio as gr
|
|
| 19 |
import tempfile
|
| 20 |
|
| 21 |
from aifs.device import device_label
|
|
|
|
|
|
|
| 22 |
|
| 23 |
import shutil
|
| 24 |
|
|
@@ -59,17 +62,40 @@ def _run_forecast_gpu(fields, date, lead_time, num_chunks):
|
|
| 59 |
yield from _run_forecast(fields, date, lead_time=lead_time, num_chunks=num_chunks)
|
| 60 |
|
| 61 |
|
| 62 |
-
def
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 63 |
from aifs.initial_conditions import load_ics
|
| 64 |
|
| 65 |
def emit(status, phase, dd=gr.update(choices=[]), states=[], btn=_BTN_RUNNING):
|
| 66 |
return status, _phase(phase), dd, states, btn
|
| 67 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 68 |
yield emit("Starting upβ¦", "π₯ Downloading initial conditions from ECMWFβ¦")
|
| 69 |
|
| 70 |
try:
|
| 71 |
fields = date = None
|
| 72 |
-
for kind, payload in load_ics(cache_dir="ic_cache"):
|
| 73 |
if kind == "log":
|
| 74 |
yield emit(payload, "π₯ Downloading initial conditions from ECMWFβ¦")
|
| 75 |
else:
|
|
@@ -137,6 +163,193 @@ def plot_forecast_handler(field_name: str, timestamp: str, states: list):
|
|
| 137 |
)
|
| 138 |
return path, stats
|
| 139 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 140 |
# ββ UI ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 141 |
DARK_CSS = """
|
| 142 |
body, .gradio-container {
|
|
@@ -188,6 +401,27 @@ with gr.Blocks(css=DARK_CSS, title="AIFS Single v2 Forecast") as demo:
|
|
| 188 |
with gr.Row():
|
| 189 |
with gr.Column(scale=1, elem_classes="panel"):
|
| 190 |
gr.Markdown("### βοΈ Forecast Settings")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 191 |
lead_time_sl = gr.Slider(
|
| 192 |
minimum=6, maximum=96, step=6, value=6,
|
| 193 |
label="Lead time (hours)",
|
|
@@ -223,9 +457,79 @@ with gr.Blocks(css=DARK_CSS, title="AIFS Single v2 Forecast") as demo:
|
|
| 223 |
map_img = gr.Image(label="Global Map", type="filepath")
|
| 224 |
stats_md = gr.Markdown()
|
| 225 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 226 |
run_btn.click(
|
| 227 |
fn=run_forecast,
|
| 228 |
-
inputs=[lead_time_sl, num_chunks_sl],
|
| 229 |
outputs=[status_box, phase_md, timestamp_dd, forecast_state, run_btn],
|
| 230 |
)
|
| 231 |
plot_btn.click(
|
|
@@ -233,6 +537,21 @@ with gr.Blocks(css=DARK_CSS, title="AIFS Single v2 Forecast") as demo:
|
|
| 233 |
inputs=[field_dd, timestamp_dd, forecast_state],
|
| 234 |
outputs=[map_img, stats_md],
|
| 235 |
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 236 |
|
| 237 |
gr.Markdown(
|
| 238 |
"""
|
|
|
|
| 1 |
+
import datetime
|
| 2 |
import os
|
| 3 |
import ssl
|
| 4 |
import warnings
|
|
|
|
| 20 |
import tempfile
|
| 21 |
|
| 22 |
from aifs.device import device_label
|
| 23 |
+
from aifs.initial_conditions import VALID_RUN_HOURS, EARLIEST_HISTORICAL_DATE, FULL_FIELD_RUN_HOURS
|
| 24 |
+
from aifs.era5_verify import EVAL_FIELD_MAP
|
| 25 |
|
| 26 |
import shutil
|
| 27 |
|
|
|
|
| 62 |
yield from _run_forecast(fields, date, lead_time=lead_time, num_chunks=num_chunks)
|
| 63 |
|
| 64 |
|
| 65 |
+
def _parse_historical_date(date_str: str, hour_str: str) -> datetime.datetime:
|
| 66 |
+
"""Parse and validate the historical-date UI inputs, or raise ValueError."""
|
| 67 |
+
try:
|
| 68 |
+
year, month, day = (int(p) for p in date_str.strip().split("-"))
|
| 69 |
+
picked = datetime.date(year, month, day)
|
| 70 |
+
except Exception:
|
| 71 |
+
raise ValueError(f"could not parse date '{date_str}' β use YYYY-MM-DD.")
|
| 72 |
+
|
| 73 |
+
if picked < EARLIEST_HISTORICAL_DATE or picked > datetime.date.today():
|
| 74 |
+
raise ValueError(
|
| 75 |
+
f"date must be between {EARLIEST_HISTORICAL_DATE.isoformat()} and today, got {picked}."
|
| 76 |
+
)
|
| 77 |
+
return datetime.datetime(picked.year, picked.month, picked.day, int(hour_str))
|
| 78 |
+
|
| 79 |
+
|
| 80 |
+
def run_forecast(lead_time: int, num_chunks: int, ic_mode: str, hist_date_str: str, hist_hour: str):
|
| 81 |
from aifs.initial_conditions import load_ics
|
| 82 |
|
| 83 |
def emit(status, phase, dd=gr.update(choices=[]), states=[], btn=_BTN_RUNNING):
|
| 84 |
return status, _phase(phase), dd, states, btn
|
| 85 |
|
| 86 |
+
historical_dt = None
|
| 87 |
+
if ic_mode == "Historical date":
|
| 88 |
+
try:
|
| 89 |
+
historical_dt = _parse_historical_date(hist_date_str, hist_hour)
|
| 90 |
+
except ValueError as exc:
|
| 91 |
+
yield emit(f"β {exc}", "β Invalid historical date", btn=_BTN_READY)
|
| 92 |
+
return
|
| 93 |
+
|
| 94 |
yield emit("Starting upβ¦", "π₯ Downloading initial conditions from ECMWFβ¦")
|
| 95 |
|
| 96 |
try:
|
| 97 |
fields = date = None
|
| 98 |
+
for kind, payload in load_ics(cache_dir="ic_cache", date=historical_dt):
|
| 99 |
if kind == "log":
|
| 100 |
yield emit(payload, "π₯ Downloading initial conditions from ECMWFβ¦")
|
| 101 |
else:
|
|
|
|
| 163 |
)
|
| 164 |
return path, stats
|
| 165 |
|
| 166 |
+
# ββ ERA5 evaluation (historical runs only) βββββββββββββββββββββββββββββββββββββ
|
| 167 |
+
|
| 168 |
+
_EVAL_BTN_RUNNING = gr.update(interactive=False, value="β³ Evaluatingβ¦")
|
| 169 |
+
_EVAL_BTN_READY = gr.update(interactive=True, value="π‘ Evaluate vs ERA5")
|
| 170 |
+
|
| 171 |
+
|
| 172 |
+
def run_evaluation(field_name: str, states: list):
|
| 173 |
+
"""
|
| 174 |
+
Thread + queue wrapper around aifs.era5_verify.evaluate_forecast (a plain
|
| 175 |
+
blocking call with a log callback) so progress streams to the UI the same
|
| 176 |
+
way run_forecast's GPU step does.
|
| 177 |
+
"""
|
| 178 |
+
import queue
|
| 179 |
+
import threading
|
| 180 |
+
|
| 181 |
+
from aifs.era5_verify import evaluate_forecast
|
| 182 |
+
|
| 183 |
+
def emit(status, dd=gr.update(choices=[]), eval_data=None, skill_path=None, btn=_EVAL_BTN_RUNNING):
|
| 184 |
+
return status, dd, eval_data if eval_data is not None else {}, skill_path, btn
|
| 185 |
+
|
| 186 |
+
if not states:
|
| 187 |
+
yield emit("β οΈ Run a forecast from a historical date first, then evaluate it.", btn=_EVAL_BTN_READY)
|
| 188 |
+
return
|
| 189 |
+
|
| 190 |
+
result_queue: queue.Queue = queue.Queue()
|
| 191 |
+
|
| 192 |
+
def _worker():
|
| 193 |
+
try:
|
| 194 |
+
data = evaluate_forecast(states, field_name, log=lambda msg: result_queue.put(("log", msg)))
|
| 195 |
+
result_queue.put(("done", data))
|
| 196 |
+
except Exception as exc:
|
| 197 |
+
result_queue.put(("error", exc))
|
| 198 |
+
|
| 199 |
+
thread = threading.Thread(target=_worker, daemon=True)
|
| 200 |
+
thread.start()
|
| 201 |
+
|
| 202 |
+
log_lines: list[str] = []
|
| 203 |
+
eval_data = None
|
| 204 |
+
while True:
|
| 205 |
+
kind, payload = result_queue.get()
|
| 206 |
+
if kind == "log":
|
| 207 |
+
log_lines.append(payload)
|
| 208 |
+
yield emit("\n".join(log_lines))
|
| 209 |
+
elif kind == "done":
|
| 210 |
+
eval_data = payload
|
| 211 |
+
break
|
| 212 |
+
elif kind == "error":
|
| 213 |
+
yield emit("\n".join(log_lines) + f"\nβ Error: {payload}", btn=_EVAL_BTN_READY)
|
| 214 |
+
return
|
| 215 |
+
|
| 216 |
+
thread.join(timeout=5.0)
|
| 217 |
+
|
| 218 |
+
from aifs.era5_verify import plot_skill_curve
|
| 219 |
+
|
| 220 |
+
skill_path = os.path.join(tempfile.gettempdir(), "aifs_eval_skill.png")
|
| 221 |
+
plot_skill_curve(eval_data).savefig(skill_path, dpi=150, bbox_inches="tight")
|
| 222 |
+
|
| 223 |
+
timestamps = [s["date"] for s in eval_data["per_step"]]
|
| 224 |
+
yield (
|
| 225 |
+
f"β
Evaluated {len(timestamps)} step(s) against ERA5 for '{field_name}'.",
|
| 226 |
+
gr.update(choices=timestamps, value=timestamps[-1]),
|
| 227 |
+
eval_data,
|
| 228 |
+
skill_path,
|
| 229 |
+
_EVAL_BTN_READY,
|
| 230 |
+
)
|
| 231 |
+
|
| 232 |
+
|
| 233 |
+
def plot_eval_handler(timestamp: str, eval_data: dict):
|
| 234 |
+
from aifs.era5_verify import plot_eval_maps
|
| 235 |
+
|
| 236 |
+
if not eval_data or not eval_data.get("per_step"):
|
| 237 |
+
return None, None, None, "Run an evaluation first."
|
| 238 |
+
|
| 239 |
+
per_step = eval_data["per_step"]
|
| 240 |
+
idx = next((i for i, s in enumerate(per_step) if s["date"] == timestamp), len(per_step) - 1)
|
| 241 |
+
step = per_step[idx]
|
| 242 |
+
|
| 243 |
+
fig_fc, fig_truth, fig_diff = plot_eval_maps(eval_data, idx)
|
| 244 |
+
paths = []
|
| 245 |
+
for name, fig in zip(("fc", "truth", "diff"), (fig_fc, fig_truth, fig_diff)):
|
| 246 |
+
path = os.path.join(tempfile.gettempdir(), f"aifs_eval_{name}.png")
|
| 247 |
+
fig.savefig(path, dpi=150, bbox_inches="tight")
|
| 248 |
+
paths.append(path)
|
| 249 |
+
|
| 250 |
+
units = UNITS_MAP.get(eval_data["field"], "")
|
| 251 |
+
stats = (
|
| 252 |
+
f"**{eval_data['field']}** @ {step['date']} (lead +{step['lead_h']}h)\n\n"
|
| 253 |
+
f"| | AIFS forecast | Persistence (t0) |\n|---|---|---|\n"
|
| 254 |
+
f"| RMSE | `{step['model_rmse']:.4g}` {units} | `{step['persistence_rmse']:.4g}` {units} |\n"
|
| 255 |
+
f"| MAE | `{step['model_mae']:.4g}` {units} | `{step['persistence_mae']:.4g}` {units} |\n"
|
| 256 |
+
f"| Bias | `{step['model_bias']:.4g}` {units} | `{step['persistence_bias']:.4g}` {units} |\n"
|
| 257 |
+
f"| Corr. w/ ERA5 | `{step['model_corr']:.3f}` | `{step['persistence_corr']:.3f}` |\n"
|
| 258 |
+
)
|
| 259 |
+
return paths[0], paths[1], paths[2], stats
|
| 260 |
+
|
| 261 |
+
# ββ ERA5 climatology (works for any forecast β live or historical) ββββββββββββ
|
| 262 |
+
|
| 263 |
+
_CLIM_BTN_RUNNING = gr.update(interactive=False, value="β³ Computingβ¦")
|
| 264 |
+
_CLIM_BTN_READY = gr.update(interactive=True, value="π‘οΈ Compare vs Climatology")
|
| 265 |
+
|
| 266 |
+
|
| 267 |
+
def run_climatology(field_name: str, timestamp: str, num_years: int, states: list):
|
| 268 |
+
"""
|
| 269 |
+
Compares one forecast step's field to its ERA5 climatology (mean over
|
| 270 |
+
many past years on the same calendar day). Unlike the ERA5-truth
|
| 271 |
+
evaluation above, this works for ANY forecast β including a live one
|
| 272 |
+
initialised from "Latest" β since it never needs ERA5 truth for the
|
| 273 |
+
forecast's own (possibly future) valid time, only for past years.
|
| 274 |
+
"""
|
| 275 |
+
import queue
|
| 276 |
+
import threading
|
| 277 |
+
|
| 278 |
+
import numpy as np
|
| 279 |
+
|
| 280 |
+
from aifs.era5_verify import compute_climatology
|
| 281 |
+
|
| 282 |
+
def emit(status, fc_path=None, clim_path=None, anom_path=None, stats="", btn=_CLIM_BTN_RUNNING):
|
| 283 |
+
return status, fc_path, clim_path, anom_path, stats, btn
|
| 284 |
+
|
| 285 |
+
if not states:
|
| 286 |
+
yield emit("β οΈ Run a forecast first, then compare it to climatology.", btn=_CLIM_BTN_READY)
|
| 287 |
+
return
|
| 288 |
+
|
| 289 |
+
state = next((s for s in states if str(s["date"]) == timestamp), states[-1])
|
| 290 |
+
if field_name not in state.get("fields", {}):
|
| 291 |
+
yield emit(f"β οΈ Field '{field_name}' not available in this forecast state.", btn=_CLIM_BTN_READY)
|
| 292 |
+
return
|
| 293 |
+
|
| 294 |
+
result_queue: queue.Queue = queue.Queue()
|
| 295 |
+
|
| 296 |
+
def _worker():
|
| 297 |
+
try:
|
| 298 |
+
data = compute_climatology(
|
| 299 |
+
state, field_name, num_years=int(num_years),
|
| 300 |
+
log=lambda msg: result_queue.put(("log", msg)),
|
| 301 |
+
)
|
| 302 |
+
result_queue.put(("done", data))
|
| 303 |
+
except Exception as exc:
|
| 304 |
+
result_queue.put(("error", exc))
|
| 305 |
+
|
| 306 |
+
thread = threading.Thread(target=_worker, daemon=True)
|
| 307 |
+
thread.start()
|
| 308 |
+
|
| 309 |
+
log_lines: list[str] = []
|
| 310 |
+
clim_data = None
|
| 311 |
+
while True:
|
| 312 |
+
kind, payload = result_queue.get()
|
| 313 |
+
if kind == "log":
|
| 314 |
+
log_lines.append(payload)
|
| 315 |
+
yield emit("\n".join(log_lines))
|
| 316 |
+
elif kind == "done":
|
| 317 |
+
clim_data = payload
|
| 318 |
+
break
|
| 319 |
+
elif kind == "error":
|
| 320 |
+
yield emit("\n".join(log_lines) + f"\nβ Error: {payload}", btn=_CLIM_BTN_READY)
|
| 321 |
+
return
|
| 322 |
+
|
| 323 |
+
thread.join(timeout=5.0)
|
| 324 |
+
|
| 325 |
+
from aifs.era5_verify import plot_climatology_maps
|
| 326 |
+
|
| 327 |
+
fig_fc, fig_clim, fig_anom = plot_climatology_maps(clim_data)
|
| 328 |
+
paths = []
|
| 329 |
+
for name, fig in zip(("fc", "clim", "anom"), (fig_fc, fig_clim, fig_anom)):
|
| 330 |
+
path = os.path.join(tempfile.gettempdir(), f"aifs_clim_{name}.png")
|
| 331 |
+
fig.savefig(path, dpi=150, bbox_inches="tight")
|
| 332 |
+
paths.append(path)
|
| 333 |
+
|
| 334 |
+
units = UNITS_MAP.get(field_name, "")
|
| 335 |
+
anomaly, zscore = clim_data["anomaly"], clim_data["zscore"]
|
| 336 |
+
z_mask = ~np.isnan(zscore)
|
| 337 |
+
pct_unusual = float(np.mean(np.abs(zscore[z_mask]) > 2) * 100) if z_mask.any() else float("nan")
|
| 338 |
+
years = clim_data["years"]
|
| 339 |
+
stats = (
|
| 340 |
+
f"**{field_name}** @ {clim_data['date']} Β· {len(years)}-year climatology "
|
| 341 |
+
f"({min(years)}β{max(years)})\n\n"
|
| 342 |
+
f"- Mean anomaly: `{np.nanmean(anomaly):.4g}` {units}\n"
|
| 343 |
+
f"- Mean |z-score|: `{np.nanmean(np.abs(zscore)):.3g}`\n"
|
| 344 |
+
f"- Area with |z| > 2 (\"unusual for the time of year\"): `{pct_unusual:.1f}%`\n"
|
| 345 |
+
)
|
| 346 |
+
|
| 347 |
+
yield (
|
| 348 |
+
f"β
Built a {len(years)}-year climatology for '{field_name}'.",
|
| 349 |
+
paths[0], paths[1], paths[2], stats,
|
| 350 |
+
_CLIM_BTN_READY,
|
| 351 |
+
)
|
| 352 |
+
|
| 353 |
# ββ UI ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 354 |
DARK_CSS = """
|
| 355 |
body, .gradio-container {
|
|
|
|
| 401 |
with gr.Row():
|
| 402 |
with gr.Column(scale=1, elem_classes="panel"):
|
| 403 |
gr.Markdown("### βοΈ Forecast Settings")
|
| 404 |
+
ic_mode_radio = gr.Radio(
|
| 405 |
+
["Latest", "Historical date"], value="Latest",
|
| 406 |
+
label="Initial conditions",
|
| 407 |
+
info="Historical dates pull from ECMWF's deeper S3 archive (from "
|
| 408 |
+
f"{EARLIEST_HISTORICAL_DATE.isoformat()}) instead of the live feed.",
|
| 409 |
+
)
|
| 410 |
+
with gr.Row(visible=False) as historical_row:
|
| 411 |
+
hist_date_tb = gr.Textbox(
|
| 412 |
+
label="Date (UTC)", placeholder="YYYY-MM-DD",
|
| 413 |
+
value=(datetime.date.today() - datetime.timedelta(days=7)).isoformat(),
|
| 414 |
+
)
|
| 415 |
+
hist_hour_dd = gr.Dropdown(
|
| 416 |
+
[f"{h:02d}" for h in FULL_FIELD_RUN_HOURS], value="00",
|
| 417 |
+
label="Run hour (UTC)",
|
| 418 |
+
info="Limited to 00/12 UTC β 06/18 UTC used a reduced ECMWF product "
|
| 419 |
+
"before 2026-05-12 that's missing fields AIFS needs.",
|
| 420 |
+
)
|
| 421 |
+
ic_mode_radio.change(
|
| 422 |
+
fn=lambda mode: gr.update(visible=(mode == "Historical date")),
|
| 423 |
+
inputs=ic_mode_radio, outputs=historical_row,
|
| 424 |
+
)
|
| 425 |
lead_time_sl = gr.Slider(
|
| 426 |
minimum=6, maximum=96, step=6, value=6,
|
| 427 |
label="Lead time (hours)",
|
|
|
|
| 457 |
map_img = gr.Image(label="Global Map", type="filepath")
|
| 458 |
stats_md = gr.Markdown()
|
| 459 |
|
| 460 |
+
gr.Markdown("---")
|
| 461 |
+
with gr.Row():
|
| 462 |
+
with gr.Column(scale=1, elem_classes="panel"):
|
| 463 |
+
gr.Markdown(
|
| 464 |
+
"### π Evaluate vs ERA5\n"
|
| 465 |
+
"Only meaningful for a forecast initialised from a **historical date** "
|
| 466 |
+
"above β ERA5 (EarthMover's public reanalysis) is the ground truth being "
|
| 467 |
+
"compared against, and it only exists for times already in the past.\n\n"
|
| 468 |
+
"EarthMover's free ERA5 archive also isn't current to *now* β it lags "
|
| 469 |
+
"real-time by several months. If a date is too recent (init date + lead "
|
| 470 |
+
"time falls after the archive's latest coverage), evaluation will fail "
|
| 471 |
+
"with a message telling you the exact usable window β pick an earlier "
|
| 472 |
+
"historical date and re-run the forecast above."
|
| 473 |
+
)
|
| 474 |
+
eval_field_dd = gr.Dropdown(
|
| 475 |
+
choices=sorted(EVAL_FIELD_MAP), value="2t",
|
| 476 |
+
label="Field to evaluate",
|
| 477 |
+
info="Wave fields aren't available in EarthMover's ERA5 mirror, so they're excluded here.",
|
| 478 |
+
)
|
| 479 |
+
eval_btn = gr.Button("π‘ Evaluate vs ERA5", variant="primary")
|
| 480 |
+
eval_status = gr.Textbox(
|
| 481 |
+
label="Detailed log", lines=6, interactive=False,
|
| 482 |
+
placeholder="Run a historical forecast above, then evaluate it hereβ¦",
|
| 483 |
+
)
|
| 484 |
+
eval_step_dd = gr.Dropdown(choices=[], label="Forecast step to inspect")
|
| 485 |
+
eval_show_btn = gr.Button("π Show Maps for Step", variant="secondary")
|
| 486 |
+
|
| 487 |
+
with gr.Column(scale=2, elem_classes="panel"):
|
| 488 |
+
skill_img = gr.Image(label="Skill vs lead time (RMSE)", type="filepath")
|
| 489 |
+
with gr.Row():
|
| 490 |
+
eval_fc_img = gr.Image(label="AIFS forecast", type="filepath")
|
| 491 |
+
eval_truth_img = gr.Image(label="ERA5 truth", type="filepath")
|
| 492 |
+
eval_diff_img = gr.Image(label="Forecast β ERA5", type="filepath")
|
| 493 |
+
eval_stats_md = gr.Markdown()
|
| 494 |
+
|
| 495 |
+
eval_state = gr.State({})
|
| 496 |
+
|
| 497 |
+
gr.Markdown("---")
|
| 498 |
+
with gr.Row():
|
| 499 |
+
with gr.Column(scale=1, elem_classes="panel"):
|
| 500 |
+
gr.Markdown(
|
| 501 |
+
"### π‘οΈ Compare vs ERA5 Climatology\n"
|
| 502 |
+
"Works for **any** forecast, including a live one initialised from "
|
| 503 |
+
"\"Latest\" β this compares a forecast step to the ERA5 average for "
|
| 504 |
+
"that calendar day over many *past* years, so it never needs ERA5 "
|
| 505 |
+
"truth for the forecast's own (possibly future) valid time. It answers "
|
| 506 |
+
"*\"is this forecast unusual for the time of year?\"*, not \"is it correct?\"."
|
| 507 |
+
)
|
| 508 |
+
clim_field_dd = gr.Dropdown(
|
| 509 |
+
choices=sorted(EVAL_FIELD_MAP), value="2t",
|
| 510 |
+
label="Field to compare",
|
| 511 |
+
)
|
| 512 |
+
clim_years_sl = gr.Slider(
|
| 513 |
+
minimum=3, maximum=30, step=1, value=10,
|
| 514 |
+
label="Climatology years",
|
| 515 |
+
info="How many past years of ERA5 to average β more is slower but smoother.",
|
| 516 |
+
)
|
| 517 |
+
clim_btn = gr.Button("π‘οΈ Compare vs Climatology", variant="primary")
|
| 518 |
+
clim_status = gr.Textbox(
|
| 519 |
+
label="Detailed log", lines=6, interactive=False,
|
| 520 |
+
placeholder="Run a forecast above, pick a step in \"Forecast step\", then compare hereβ¦",
|
| 521 |
+
)
|
| 522 |
+
|
| 523 |
+
with gr.Column(scale=2, elem_classes="panel"):
|
| 524 |
+
with gr.Row():
|
| 525 |
+
clim_fc_img = gr.Image(label="AIFS forecast", type="filepath")
|
| 526 |
+
clim_mean_img = gr.Image(label="ERA5 climatology mean", type="filepath")
|
| 527 |
+
clim_anom_img = gr.Image(label="Anomaly (forecast β climatology)", type="filepath")
|
| 528 |
+
clim_stats_md = gr.Markdown()
|
| 529 |
+
|
| 530 |
run_btn.click(
|
| 531 |
fn=run_forecast,
|
| 532 |
+
inputs=[lead_time_sl, num_chunks_sl, ic_mode_radio, hist_date_tb, hist_hour_dd],
|
| 533 |
outputs=[status_box, phase_md, timestamp_dd, forecast_state, run_btn],
|
| 534 |
)
|
| 535 |
plot_btn.click(
|
|
|
|
| 537 |
inputs=[field_dd, timestamp_dd, forecast_state],
|
| 538 |
outputs=[map_img, stats_md],
|
| 539 |
)
|
| 540 |
+
eval_btn.click(
|
| 541 |
+
fn=run_evaluation,
|
| 542 |
+
inputs=[eval_field_dd, forecast_state],
|
| 543 |
+
outputs=[eval_status, eval_step_dd, eval_state, skill_img, eval_btn],
|
| 544 |
+
)
|
| 545 |
+
eval_show_btn.click(
|
| 546 |
+
fn=plot_eval_handler,
|
| 547 |
+
inputs=[eval_step_dd, eval_state],
|
| 548 |
+
outputs=[eval_fc_img, eval_truth_img, eval_diff_img, eval_stats_md],
|
| 549 |
+
)
|
| 550 |
+
clim_btn.click(
|
| 551 |
+
fn=run_climatology,
|
| 552 |
+
inputs=[clim_field_dd, timestamp_dd, clim_years_sl, forecast_state],
|
| 553 |
+
outputs=[clim_status, clim_fc_img, clim_mean_img, clim_anom_img, clim_stats_md, clim_btn],
|
| 554 |
+
)
|
| 555 |
|
| 556 |
gr.Markdown(
|
| 557 |
"""
|