Skip to content

Commit 1e5b057

Browse files
ghostiee-11ghostiee-11
authored
Add TEM benchmark (case 10): zonal means + eddy fluxes as a GROUP BY (#191)
1 parent b7c581c commit 1e5b057

2 files changed

Lines changed: 171 additions & 0 deletions

File tree

‎benchmarks/geospatial/10_tem.py‎

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,164 @@
1+
# /// script
2+
# requires-python = ">=3.11"
3+
# dependencies = [
4+
# "xarray-sql",
5+
# "xarray",
6+
# "gcsfs",
7+
# "zarr>=3",
8+
# ]
9+
#
10+
# [tool.uv.sources]
11+
# xarray-sql = { path = "../../", editable = true }
12+
# ///
13+
"""Transformed Eulerian Mean: zonal means and eddy fluxes are a GROUP BY.
14+
15+
The Transformed Eulerian Mean (TEM) is a standard atmospheric-circulation
16+
diagnostic: average the flow around each latitude circle, then measure how the
17+
departures from that average (the eddies) carry momentum and heat. In the array
18+
paradigm it is ``ds.mean("longitude")`` plus a few ``(x - x_bar)`` products
19+
averaged again over longitude.
20+
21+
Every piece of that is relational. A zonal mean is ``GROUP BY latitude``
22+
collapsing longitude. An eddy flux such as the momentum flux
23+
``u'v' = mean((u - u_bar)(v - v_bar))`` is, by the covariance identity, just
24+
``AVG(u*v) - AVG(u)*AVG(v)``: one grouped pass, no self-join. So the whole
25+
diagnostic is::
26+
27+
SELECT time, level, latitude,
28+
AVG(u) AS u_bar, ...,
29+
AVG(u*v) - AVG(u)*AVG(v) AS upvp, -- eddy momentum flux u'v'
30+
AVG(v*t) - AVG(v)*AVG(t) AS vptp, -- eddy heat flux v't'
31+
AVG(u*w) - AVG(u)*AVG(w) AS upwp -- vertical momentum flux u'w'
32+
FROM era5 GROUP BY time, level, latitude
33+
34+
This is the diagnostic dcherian raised in the Large Scale Geospatial Benchmarks
35+
discussion (coiled/benchmarks #1545); the SQL reads like its textbook definition.
36+
37+
Dataset: the full ARCO-ERA5 archive (0.25 degree, 37 pressure levels), opened
38+
lazily, so the query reads only u, v, T, w on the requested levels and timestep.
39+
Validated against the same diagnostic computed in pure xarray.
40+
"""
41+
42+
from __future__ import annotations
43+
44+
import datetime
45+
46+
import xarray as xr
47+
48+
import xarray_sql as xql
49+
50+
from _harness import (
51+
CaseSkipped,
52+
assert_grid_close,
53+
measured,
54+
run_case,
55+
show_result,
56+
show_sql,
57+
timed,
58+
)
59+
60+
_URL = "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3"
61+
_T = datetime.datetime(2020, 6, 1, 12)
62+
# Three representative pressure levels (hPa): upper jet, mid, lower troposphere.
63+
_LEVELS = (250, 500, 850)
64+
_VARS = [
65+
"u_component_of_wind",
66+
"v_component_of_wind",
67+
"temperature",
68+
"vertical_velocity",
69+
]
70+
# Timestep and levels are bound as query parameters, not formatted into the SQL.
71+
_PARAMS = {"t": _T, "l1": _LEVELS[0], "l2": _LEVELS[1], "l3": _LEVELS[2]}
72+
73+
74+
def main() -> None:
75+
try:
76+
import gcsfs # noqa: F401 (required by the gs:// protocol)
77+
78+
ds = xr.open_zarr(_URL, chunks=None, storage_options={"token": "anon"})
79+
except Exception as exc: # noqa: BLE001 (any failure skips, not crash)
80+
raise CaseSkipped(f"ARCO-ERA5 unavailable ({exc})") from exc
81+
82+
print(
83+
f" ARCO-ERA5: {ds.sizes['time']:,} timesteps x {ds.sizes['level']} levels "
84+
f"x {ds.sizes['latitude']}x{ds.sizes['longitude']} (no pre-slicing)"
85+
)
86+
87+
ctx = xql.XarrayContext()
88+
with timed("register full ERA5 (lazy)"):
89+
ctx.from_dataset(
90+
"era5",
91+
ds,
92+
chunks={"time": 1},
93+
table_names={
94+
("time", "latitude", "longitude"): "surface",
95+
("time", "level", "latitude", "longitude"): "atmosphere",
96+
},
97+
)
98+
99+
sql = """
100+
WITH f AS (
101+
SELECT time, level, latitude,
102+
"u_component_of_wind" AS u,
103+
"v_component_of_wind" AS v,
104+
"temperature" AS t,
105+
"vertical_velocity" AS w
106+
FROM era5.atmosphere
107+
WHERE time = $t
108+
AND level IN ($l1, $l2, $l3)
109+
)
110+
SELECT time, level, latitude,
111+
AVG(u) AS u_bar,
112+
AVG(v) AS v_bar,
113+
AVG(t) AS t_bar,
114+
AVG(w) AS w_bar,
115+
AVG(u * v) - AVG(u) * AVG(v) AS upvp,
116+
AVG(v * t) - AVG(v) * AVG(t) AS vptp,
117+
AVG(u * w) - AVG(u) * AVG(w) AS upwp
118+
FROM f
119+
GROUP BY time, level, latitude
120+
ORDER BY time, level, latitude
121+
"""
122+
show_sql(sql)
123+
124+
# Round-trip the diagnostic to a (time, level, latitude) Dataset.
125+
for _ in measured("SQL TEM (zonal means + eddy covariances, lazy read)"):
126+
got = ctx.sql(sql, param_values=_PARAMS).to_dataset(
127+
dims=["time", "level", "latitude"]
128+
)
129+
130+
# Array reference: the same TEM diagnostic in pure xarray.
131+
for _ in measured("xarray reference"):
132+
sub = ds[_VARS].sel(time=[_T], level=list(_LEVELS))
133+
u, v, t, w = (sub[n] for n in _VARS)
134+
135+
def zm(x: xr.DataArray) -> xr.DataArray:
136+
return x.mean("longitude")
137+
138+
u_bar, v_bar, t_bar, w_bar = zm(u), zm(v), zm(t), zm(w)
139+
ref_upvp = zm((u - u_bar) * (v - v_bar))
140+
ref_vptp = zm((v - v_bar) * (t - t_bar))
141+
ref_upwp = zm((u - u_bar) * (w - w_bar))
142+
143+
# Tolerance covers the SQL one-pass covariance (AVG(u*v) - AVG(u)*AVG(v))
144+
# against the two-pass xarray reference on float32 ERA5 fields.
145+
assert_grid_close(
146+
"zonal-mean u (u_bar)", got.u_bar, u_bar, rtol=1e-3, atol=1e-2
147+
)
148+
assert_grid_close(
149+
"eddy momentum flux u'v'", got.upvp, ref_upvp, rtol=1e-3, atol=1e-2
150+
)
151+
assert_grid_close(
152+
"eddy heat flux v't'", got.vptp, ref_vptp, rtol=1e-3, atol=1e-2
153+
)
154+
assert_grid_close(
155+
"vertical flux u'w'", got.upwp, ref_upwp, rtol=1e-3, atol=1e-2
156+
)
157+
158+
show_result(got)
159+
160+
161+
if __name__ == "__main__":
162+
raise SystemExit(
163+
run_case(main, "TEM: zonal means + eddy fluxes as GROUP BY (ARCO-ERA5)")
164+
)

‎benchmarks/geospatial/README.md‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,13 +24,17 @@ plain-English definition of the operation, and computes the same numbers.
2424
| 07 | `07_reproject_udf.py` | per-pixel CRS transform | scalar **UDF** (`reproject()`), à la PostGIS `ST_Transform` |
2525
| 08 | `08_regrid_weights.py` | interpolation to a new grid | sparse-weight table `JOIN` + weighted `GROUP BY` |
2626
| 09 | `09_warp.py` | reproject **and** resample (warp) | reproject **UDF** (07) → weight table `JOIN` (08) |
27+
| 10 | `10_tem.py` | zonal means + eddy fluxes (TEM diagnostic) | `GROUP BY (time, level, lat)` + covariance `AVG(u*v) - AVG(u)*AVG(v)` |
2728

2829
Cases 01–06 show operations that are *natively* relational. Cases 07–08 are the
2930
"hardest" array operations — reprojection and regridding — and show where a UDF
3031
fits (a per-row coordinate transform) versus where the operation is really a
3132
sparse matrix multiply expressed as a `JOIN`. Case 09 composes the two into a full
3233
**warp** (GDAL/rasterio `reproject`): the 07 UDF reprojects the target grid, arrays
3334
turn the reprojected points into bilinear weights, and the 08 `JOIN` applies them.
35+
Case 10 returns to a pure reduction: the Transformed Eulerian Mean, where the
36+
zonal means and the eddy momentum and heat fluxes fall out of a single grouped
37+
aggregate via the covariance identity.
3438
See
3539
[`docs/geospatial.md`](../../docs/geospatial.md) for the full narrative,
3640
including *where the array paradigm still earns its keep* (generating the
@@ -60,6 +64,9 @@ interpolation weights — the geometry — which SQL applies but does not comput
6064
validates against xarray's `.interp()` at the reprojected points, with Earth
6165
Engine's own lon/lat SRTM as a second, cross-CRS check. All three run against
6266
Earth Engine using your existing `gcloud` login, and skip cleanly without it.
67+
- **10 TEM** uses the same **ARCO-ERA5** archive as 02-06, reading the atmospheric
68+
wind and temperature fields (u, v, T, w) on a few pressure levels for one
69+
timestep. Network-backed; skips cleanly offline.
6370

6471
## Running
6572

0 commit comments

Comments
 (0)