|
| 1 | +# /// script |
| 2 | +# requires-python = ">=3.10" |
| 3 | +# dependencies = [ |
| 4 | +# "xarray_sql", |
| 5 | +# "xarray[io]", |
| 6 | +# "gcsfs", |
| 7 | +# "numpy", |
| 8 | +# ] |
| 9 | +# |
| 10 | +# [tool.uv.sources] |
| 11 | +# xarray_sql = { path = "..", editable = true } |
| 12 | +# /// |
| 13 | +"""Differentiable SQL over ARCO-ERA5. |
| 14 | +
|
| 15 | +A minimal demonstration of xarray-sql's autograd: take a real climate archive |
| 16 | +(ARCO-ERA5, read anonymously from GCS), express a physical quantity as an |
| 17 | +*analytic* SQL formula over its variables, and let ``grad(...)`` differentiate |
| 18 | +that formula symbolically — evaluated per grid cell, which is the relational |
| 19 | +equivalent of ``jax.vmap(jax.grad(f))`` (each row is an independent point). |
| 20 | +
|
| 21 | +Note this is *symbolic* differentiation of an expression, not a finite- |
| 22 | +difference spatial gradient: ``grad(f(u, v), u)`` is the exact partial |
| 23 | +derivative of the formula ``f``, evaluated at every cell's values. |
| 24 | +
|
| 25 | +Two cases: |
| 26 | +
|
| 27 | +1. Wind-speed magnitude ``speed = sqrt(u^2 + v^2)``. Its sensitivity to the |
| 28 | + eastward wind is ``d(speed)/du = u / speed`` — checked exactly. |
| 29 | +
|
| 30 | +2. Saturation vapour pressure ``e_s(T)`` (August-Roche-Magnus form of the |
| 31 | + Clausius-Clapeyron relation). ``d(e_s)/dT`` governs how fast the atmosphere's |
| 32 | + moisture capacity grows with temperature — checked against the closed-form |
| 33 | + slope. |
| 34 | +
|
| 35 | +Run standalone (builds the local extension on first use): |
| 36 | +
|
| 37 | + uv run benchmarks/grad_era5.py |
| 38 | +""" |
| 39 | + |
| 40 | +from __future__ import annotations |
| 41 | + |
| 42 | +import time |
| 43 | + |
| 44 | +import numpy as np |
| 45 | +import xarray as xr |
| 46 | + |
| 47 | +import xarray_sql as xql |
| 48 | + |
| 49 | +ARCO_ERA5 = ( |
| 50 | + "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3" |
| 51 | +) |
| 52 | + |
| 53 | +# The SQL result comes back with ascending coordinates; ERA5's native latitude |
| 54 | +# is descending. Sort both sides before comparing so equality is by label. |
| 55 | +_SORT = ["latitude", "longitude"] |
| 56 | + |
| 57 | +# ERA5 variable names start with a digit, so they must be double-quoted in SQL. |
| 58 | +U = '"10m_u_component_of_wind"' |
| 59 | +V = '"10m_v_component_of_wind"' |
| 60 | +T = '"2m_temperature"' |
| 61 | + |
| 62 | + |
| 63 | +def load_era5_block() -> xr.Dataset: |
| 64 | + """Open ARCO-ERA5 and pull one timestamp over a small region. |
| 65 | +
|
| 66 | + Lazy open of the whole archive; only the requested block is read. We keep |
| 67 | + it to a few thousand cells so the demo runs in seconds. |
| 68 | + """ |
| 69 | + full = xr.open_zarr( |
| 70 | + ARCO_ERA5, chunks=None, storage_options={"token": "anon"} |
| 71 | + ) |
| 72 | + block = ( |
| 73 | + full[ |
| 74 | + [ |
| 75 | + "10m_u_component_of_wind", |
| 76 | + "10m_v_component_of_wind", |
| 77 | + "2m_temperature", |
| 78 | + ] |
| 79 | + ] |
| 80 | + .sel(time="2020-01-01T00") |
| 81 | + # A ~North-America box (index-based to avoid lat-orientation pitfalls). |
| 82 | + .isel(latitude=slice(120, 200), longitude=slice(900, 1000)) |
| 83 | + .load() |
| 84 | + ) |
| 85 | + return block.chunk({"latitude": 40}) |
| 86 | + |
| 87 | + |
| 88 | +def wind_speed_sensitivity(ctx: xql.XarrayContext, ref: xr.Dataset) -> None: |
| 89 | + """grad(sqrt(u^2 + v^2)) checked against the exact u / speed, v / speed.""" |
| 90 | + speed = f"sqrt(power({U}, 2) + power({V}, 2))" |
| 91 | + out = ( |
| 92 | + ctx.sql( |
| 93 | + f""" |
| 94 | + SELECT |
| 95 | + latitude, |
| 96 | + longitude, |
| 97 | + {speed} AS wind_speed, |
| 98 | + grad({speed}, {U}) AS d_speed_d_u, |
| 99 | + grad({speed}, {V}) AS d_speed_d_v |
| 100 | + FROM era5 |
| 101 | + """ |
| 102 | + ) |
| 103 | + .to_dataset(dims=["latitude", "longitude"]) |
| 104 | + .sortby(_SORT) |
| 105 | + ) |
| 106 | + |
| 107 | + u = ref["10m_u_component_of_wind"] |
| 108 | + v = ref["10m_v_component_of_wind"] |
| 109 | + speed_ref = np.sqrt(u**2 + v**2).sortby(_SORT) |
| 110 | + |
| 111 | + xr.testing.assert_allclose( |
| 112 | + out["wind_speed"], speed_ref.rename("wind_speed") |
| 113 | + ) |
| 114 | + xr.testing.assert_allclose( |
| 115 | + out["d_speed_d_u"], (u / speed_ref).sortby(_SORT).rename("d_speed_d_u") |
| 116 | + ) |
| 117 | + xr.testing.assert_allclose( |
| 118 | + out["d_speed_d_v"], (v / speed_ref).sortby(_SORT).rename("d_speed_d_v") |
| 119 | + ) |
| 120 | + print(" wind-speed sensitivity matches u/|w|, v/|w| exactly") |
| 121 | + print(out) |
| 122 | + |
| 123 | + |
| 124 | +def clausius_clapeyron(ctx: xql.XarrayContext, ref: xr.Dataset) -> None: |
| 125 | + """grad(e_s(T)) checked against the closed-form Clausius-Clapeyron slope.""" |
| 126 | + # August-Roche-Magnus: e_s(T) = A * exp(B * tc / (tc + C)), tc = T - 273.15. |
| 127 | + a, b, c = 6.1094, 17.625, 243.04 |
| 128 | + tc = f"({T} - 273.15)" |
| 129 | + es = f"{a} * exp({b} * {tc} / ({tc} + {c}))" |
| 130 | + out = ( |
| 131 | + ctx.sql( |
| 132 | + f""" |
| 133 | + SELECT |
| 134 | + latitude, |
| 135 | + longitude, |
| 136 | + {es} AS e_s, |
| 137 | + grad({es}, {T}) AS de_s_dt |
| 138 | + FROM era5 |
| 139 | + """ |
| 140 | + ) |
| 141 | + .to_dataset(dims=["latitude", "longitude"]) |
| 142 | + .sortby(_SORT) |
| 143 | + ) |
| 144 | + |
| 145 | + # Reference in float64 (the columns are float32): the exact derivative is |
| 146 | + # d(e_s)/dT = e_s * B*C / (tc + C)^2. |
| 147 | + temp = ref["2m_temperature"].astype("float64") |
| 148 | + tc_ref = temp - 273.15 |
| 149 | + es_ref = a * np.exp(b * tc_ref / (tc_ref + c)) |
| 150 | + des_dt_ref = es_ref * (b * c) / (tc_ref + c) ** 2 |
| 151 | + |
| 152 | + xr.testing.assert_allclose( |
| 153 | + out["e_s"], es_ref.sortby(_SORT).rename("e_s"), rtol=1e-5 |
| 154 | + ) |
| 155 | + xr.testing.assert_allclose( |
| 156 | + out["de_s_dt"], des_dt_ref.sortby(_SORT).rename("de_s_dt"), rtol=1e-5 |
| 157 | + ) |
| 158 | + print(" d(e_s)/dT matches the closed-form Clausius-Clapeyron slope") |
| 159 | + print(out) |
| 160 | + |
| 161 | + |
| 162 | +def main() -> None: |
| 163 | + t0 = time.time() |
| 164 | + ds = load_era5_block() |
| 165 | + print(f"loaded ERA5 block {dict(ds.sizes)} in {time.time() - t0:.1f}s") |
| 166 | + |
| 167 | + ctx = xql.XarrayContext() |
| 168 | + ctx.from_dataset("era5", ds) |
| 169 | + |
| 170 | + print("\n== wind-speed sensitivity: grad(sqrt(u^2 + v^2)) ==") |
| 171 | + wind_speed_sensitivity(ctx, ds) |
| 172 | + |
| 173 | + print("\n== Clausius-Clapeyron: grad(e_s(T)) ==") |
| 174 | + clausius_clapeyron(ctx, ds) |
| 175 | + |
| 176 | + print("\nOK: symbolic SQL gradients match the analytic references.") |
| 177 | + |
| 178 | + |
| 179 | +if __name__ == "__main__": |
| 180 | + main() |
0 commit comments