EmmaScharfmann HF Staff commited on
Commit
786225b
Β·
1 Parent(s): 71701b2

add comparision

Browse files
Files changed (6) hide show
  1. .gitignore +2 -0
  2. aifs/era5_env.py +128 -0
  3. aifs/era5_verify.py +337 -0
  4. aifs/era5_worker.py +130 -0
  5. aifs/initial_conditions.py +61 -7
  6. 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=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: datetime.datetime = OpendataClient(SOURCE).latest()
236
- yield "log", f"πŸ“… Latest ECMWF run: {date}"
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 run_forecast(lead_time: int, num_chunks: int):
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
  """