|
| 1 | +#!/usr/bin/env python3 |
| 2 | +# /// script |
| 3 | +# requires-python = ">=3.11" |
| 4 | +# dependencies = [ |
| 5 | +# "xarray-sql", |
| 6 | +# "xarray", |
| 7 | +# "numpy", |
| 8 | +# "pyproj", |
| 9 | +# "pyarrow", |
| 10 | +# "scipy", |
| 11 | +# "xee", |
| 12 | +# "earthengine-api", |
| 13 | +# "shapely", |
| 14 | +# ] |
| 15 | +# |
| 16 | +# [tool.uv.sources] |
| 17 | +# xarray-sql = { path = "../../", editable = true } |
| 18 | +# /// |
| 19 | +"""Warp — reprojecting *and* resampling a raster is case 07's UDF + case 08's JOIN. |
| 20 | +
|
| 21 | +A *warp* moves a raster from one CRS onto a grid in another — the everyday GDAL/ |
| 22 | +rasterio ``reproject`` that GIS runs constantly. It is exactly the composition of |
| 23 | +the two "hard" cases: |
| 24 | +
|
| 25 | +* **case 07** — reproject coordinates with a scalar PROJ **UDF**; and |
| 26 | +* **case 08** — resample values with a sparse-weight **JOIN**. |
| 27 | +
|
| 28 | +The pipeline reads as that composition, and it shows the division of labor cleanly: |
| 29 | +
|
| 30 | +1. **SQL reprojects the target grid** (the 07 UDF): for every target ``(lon, lat)`` |
| 31 | + cell, ``reproject()`` returns where it falls in the source CRS (UTM ``x``/``y``). |
| 32 | +2. **Arrays build the bilinear weights** (the geometry): each reprojected target |
| 33 | + point lands between four source pixels; we compute its four bilinear weights — |
| 34 | + the genuinely geometric step the array world owns. (This is the same |
| 35 | + "arrays compute the weights, SQL applies them" boundary as case 08, except the |
| 36 | + target points are *scattered* in source space because they were reprojected, |
| 37 | + so the weights are a per-point stencil rather than a separable lat×lon grid.) |
| 38 | +3. **SQL applies the weights** (the 08 JOIN): join the source values onto the |
| 39 | + weight table and ``SUM(value * weight)`` per target cell. |
| 40 | +
|
| 41 | +**References.** The exact check is the array paradigm doing the same warp — plain |
| 42 | +``pyproj`` + xarray ``.interp`` at the reprojected points — which the SQL result |
| 43 | +matches to floating-point tolerance. As an *independent* real-world cross-check we |
| 44 | +also open the **same SRTM** directly on the lon/lat grid through Xee (Earth |
| 45 | +Engine's own warp) and report the agreement; it is close (a few metres median) but |
| 46 | +not bit-exact, because EE resamples from native 30 m while our source is the coarse |
| 47 | +UTM grid — which is exactly why the deterministic warp, not EE, is the tolerance |
| 48 | +reference. |
| 49 | +
|
| 50 | +Data: real **SRTM elevation** (Northern California terrain) via [Xee](https://github.com/google/Xee), |
| 51 | +opened once on a UTM grid (the source) and once on a lon/lat grid (the EE |
| 52 | +cross-check). Requires Earth Engine access; skips cleanly otherwise. |
| 53 | +""" |
| 54 | + |
| 55 | +from __future__ import annotations |
| 56 | + |
| 57 | +import numpy as np |
| 58 | +import pyarrow as pa |
| 59 | +import pyproj |
| 60 | +import shapely.geometry as sgeom |
| 61 | +import xarray as xr |
| 62 | +from datafusion import udf |
| 63 | + |
| 64 | +import xarray_sql as xql |
| 65 | + |
| 66 | +from _harness import ( |
| 67 | + CaseSkipped, |
| 68 | + assert_grid_close, |
| 69 | + initialize_earth_engine, |
| 70 | + measured, |
| 71 | + run_case, |
| 72 | + show_result, |
| 73 | + show_sql, |
| 74 | + timed, |
| 75 | +) |
| 76 | + |
| 77 | +_SRC_CRS = "EPSG:32610" # UTM zone 10N — the source raster's CRS |
| 78 | +_DST_CRS = "EPSG:4326" # lon/lat — the target grid's CRS |
| 79 | +_AOI = (-122.6, 37.4, -121.6, 38.4) # ~1° box of Northern California terrain |
| 80 | +_SRC_SCALE_M = 2_000.0 # ~2 km source pixels |
| 81 | +_DST_SCALE_DEG = 0.02 # ~2 km target cells |
| 82 | + |
| 83 | + |
| 84 | +def _register_reproject_udf(ctx, src_crs, dst_crs, name="reproject"): |
| 85 | + """Register ``reproject(a, b) -> {x, y}`` — case 07's PROJ scalar UDF. |
| 86 | +
|
| 87 | + Vectorized over each Arrow batch; ``always_xy=True`` keeps (easting, northing) |
| 88 | + /(lon, lat) order. Returns both output coordinates from one struct-returning |
| 89 | + call (PROJ contexts are not thread-safe, so one UDF, evaluated serially). |
| 90 | + """ |
| 91 | + ret = pa.struct([("x", pa.float64()), ("y", pa.float64())]) |
| 92 | + |
| 93 | + def _fn(a: pa.Array, b: pa.Array) -> pa.Array: |
| 94 | + transformer = pyproj.Transformer.from_crs( |
| 95 | + src_crs, dst_crs, always_xy=True |
| 96 | + ) |
| 97 | + xs = np.asarray(a.to_numpy(zero_copy_only=False), dtype="float64") |
| 98 | + ys = np.asarray(b.to_numpy(zero_copy_only=False), dtype="float64") |
| 99 | + ox, oy = transformer.transform(xs, ys) |
| 100 | + return pa.StructArray.from_arrays( |
| 101 | + [ |
| 102 | + pa.array(np.asarray(ox, "float64")), |
| 103 | + pa.array(np.asarray(oy, "float64")), |
| 104 | + ], |
| 105 | + names=["x", "y"], |
| 106 | + ) |
| 107 | + |
| 108 | + ctx.register_udf( |
| 109 | + udf(_fn, [pa.float64(), pa.float64()], ret, "immutable", name) |
| 110 | + ) |
| 111 | + |
| 112 | + |
| 113 | +def _open_srtm( |
| 114 | + grid_crs: str, scale: tuple[float, float], xy_names |
| 115 | +) -> xr.DataArray: |
| 116 | + """Open SRTM elevation over the AOI on the requested grid via Xee (lazy).""" |
| 117 | + try: |
| 118 | + from xee import helpers |
| 119 | + except ImportError as exc: # pragma: no cover |
| 120 | + raise CaseSkipped( |
| 121 | + "Earth Engine support needs `pip install earthengine-api xee`" |
| 122 | + ) from exc |
| 123 | + |
| 124 | + ee = initialize_earth_engine() |
| 125 | + grid = helpers.fit_geometry( |
| 126 | + sgeom.box(*_AOI), |
| 127 | + geometry_crs="EPSG:4326", |
| 128 | + grid_crs=grid_crs, |
| 129 | + grid_scale=scale, |
| 130 | + ) |
| 131 | + ic = ee.ImageCollection([ee.Image("USGS/SRTMGL1_003")]) |
| 132 | + da = xr.open_dataset(ic, engine="ee", **grid)["elevation"].isel(time=0) |
| 133 | + a, b = xy_names |
| 134 | + rename = {} |
| 135 | + for d in da.dims: |
| 136 | + dl = d.lower() |
| 137 | + if dl in ("y", "lat", "latitude"): |
| 138 | + rename[d] = a |
| 139 | + elif dl in ("x", "lon", "longitude"): |
| 140 | + rename[d] = b |
| 141 | + da = da.rename(rename).sortby(a).sortby(b) |
| 142 | + return da.assign_coords( |
| 143 | + {a: da[a].astype("float64"), b: da[b].astype("float64")} |
| 144 | + ) |
| 145 | + |
| 146 | + |
| 147 | +def _warp_weight_table( |
| 148 | + sx: np.ndarray, |
| 149 | + sy: np.ndarray, |
| 150 | + dst_lon: np.ndarray, |
| 151 | + dst_lat: np.ndarray, |
| 152 | + px: np.ndarray, |
| 153 | + py: np.ndarray, |
| 154 | +) -> xr.Dataset: |
| 155 | + """Bilinear weights for reprojected target points — the geometry step. |
| 156 | +
|
| 157 | + Each target cell ``(dst_lat, dst_lon)`` was reprojected to source coordinates |
| 158 | + ``(px, py)``; here we find the four surrounding source pixels and their |
| 159 | + bilinear weights. One row per (target cell, source corner). Targets that fall |
| 160 | + outside the source footprint contribute no rows (and are dropped). |
| 161 | + """ |
| 162 | + dst_lats, dst_lons, src_xs, src_ys, weights = [], [], [], [], [] |
| 163 | + for k in range(len(px)): |
| 164 | + x, y = px[k], py[k] |
| 165 | + if not (sx[0] <= x <= sx[-1] and sy[0] <= y <= sy[-1]): |
| 166 | + continue |
| 167 | + i = int(np.clip(np.searchsorted(sx, x) - 1, 0, len(sx) - 2)) |
| 168 | + j = int(np.clip(np.searchsorted(sy, y) - 1, 0, len(sy) - 2)) |
| 169 | + tx = (x - sx[i]) / (sx[i + 1] - sx[i]) |
| 170 | + ty = (y - sy[j]) / (sy[j + 1] - sy[j]) |
| 171 | + for ii, wx in ((i, 1.0 - tx), (i + 1, tx)): |
| 172 | + for jj, wy in ((j, 1.0 - ty), (j + 1, ty)): |
| 173 | + dst_lats.append(dst_lat[k]) |
| 174 | + dst_lons.append(dst_lon[k]) |
| 175 | + src_xs.append(sx[ii]) |
| 176 | + src_ys.append(sy[jj]) |
| 177 | + weights.append(wx * wy) |
| 178 | + n = len(weights) |
| 179 | + return xr.Dataset( |
| 180 | + { |
| 181 | + "dst_lat": (["pair"], np.array(dst_lats, "float64")), |
| 182 | + "dst_lon": (["pair"], np.array(dst_lons, "float64")), |
| 183 | + "src_x": (["pair"], np.array(src_xs, "float64")), |
| 184 | + "src_y": (["pair"], np.array(src_ys, "float64")), |
| 185 | + "weight": (["pair"], np.array(weights, "float64")), |
| 186 | + }, |
| 187 | + coords={"pair": np.arange(n)}, |
| 188 | + ).chunk({"pair": n}) |
| 189 | + |
| 190 | + |
| 191 | +def main() -> None: |
| 192 | + with timed("open SRTM on UTM + lon/lat grids via Xee"): |
| 193 | + src = _open_srtm(_SRC_CRS, (_SRC_SCALE_M, _SRC_SCALE_M), ("y", "x")) |
| 194 | + ref_ee = _open_srtm( |
| 195 | + _DST_CRS, (_DST_SCALE_DEG, _DST_SCALE_DEG), ("lat", "lon") |
| 196 | + ) |
| 197 | + sx, sy = src.x.values, src.y.values |
| 198 | + |
| 199 | + # Target lon/lat grid strictly inside the source UTM footprint, so every |
| 200 | + # target cell reprojects to a point with four source corners (no edge cells |
| 201 | + # to drop). Inscribe a lon/lat box in the reprojected UTM rectangle. |
| 202 | + inv = pyproj.Transformer.from_crs(_SRC_CRS, _DST_CRS, always_xy=True) |
| 203 | + cx = [sx[0], sx[-1], sx[0], sx[-1]] |
| 204 | + cy = [sy[0], sy[0], sy[-1], sy[-1]] |
| 205 | + clon, clat = inv.transform(cx, cy) |
| 206 | + lon0, lon1 = max(clon[0], clon[2]) + 0.01, min(clon[1], clon[3]) - 0.01 |
| 207 | + lat0, lat1 = max(clat[0], clat[1]) + 0.01, min(clat[2], clat[3]) - 0.01 |
| 208 | + tlon = np.linspace(lon0, lon1, 60) |
| 209 | + tlat = np.linspace(lat0, lat1, 60) |
| 210 | + print( |
| 211 | + f" source UTM grid {len(sy)}×{len(sx)} → target lon/lat grid " |
| 212 | + f"{len(tlat)}×{len(tlon)} ({_SRC_CRS} → {_DST_CRS})" |
| 213 | + ) |
| 214 | + |
| 215 | + ctx = xql.XarrayContext() |
| 216 | + _register_reproject_udf(ctx, _DST_CRS, _SRC_CRS) |
| 217 | + |
| 218 | + # The target grid as a (dst_lat, dst_lon) table. |
| 219 | + LON, LAT = np.meshgrid(tlon, tlat) |
| 220 | + target = xr.Dataset( |
| 221 | + { |
| 222 | + "dst_lon": (["cell"], LON.ravel()), |
| 223 | + "dst_lat": (["cell"], LAT.ravel()), |
| 224 | + }, |
| 225 | + coords={"cell": np.arange(LON.size)}, |
| 226 | + ).chunk({"cell": LON.size}) |
| 227 | + ctx.from_dataset("target", target, chunks={"cell": LON.size}) |
| 228 | + |
| 229 | + # 1) SQL reprojects the target grid into the source CRS (case 07's UDF). |
| 230 | + reproj_sql = """ |
| 231 | + SELECT dst_lat, dst_lon, |
| 232 | + reproject(dst_lon, dst_lat)['x'] AS sx, |
| 233 | + reproject(dst_lon, dst_lat)['y'] AS sy |
| 234 | + FROM target |
| 235 | + """ |
| 236 | + show_sql(reproj_sql, label="SQL — reproject target grid (PROJ UDF)") |
| 237 | + rp = ctx.sql(reproj_sql).to_pandas() |
| 238 | + px, py = rp["sx"].to_numpy(), rp["sy"].to_numpy() |
| 239 | + |
| 240 | + # 2) Arrays turn the reprojected points into a bilinear weight table. |
| 241 | + weights = _warp_weight_table( |
| 242 | + sx, sy, rp["dst_lon"].to_numpy(), rp["dst_lat"].to_numpy(), px, py |
| 243 | + ) |
| 244 | + ctx.from_dataset( |
| 245 | + "src", src.to_dataset(name="value"), chunks={"y": len(sy), "x": len(sx)} |
| 246 | + ) |
| 247 | + ctx.from_dataset("weights", weights, chunks={"pair": weights.sizes["pair"]}) |
| 248 | + |
| 249 | + # 3) SQL applies the weights (case 08's JOIN). |
| 250 | + apply_sql = """ |
| 251 | + SELECT w.dst_lat AS lat, w.dst_lon AS lon, |
| 252 | + SUM(s.value * w.weight) AS warped |
| 253 | + FROM weights w |
| 254 | + JOIN src s ON s.x = w.src_x AND s.y = w.src_y |
| 255 | + GROUP BY w.dst_lat, w.dst_lon |
| 256 | + ORDER BY w.dst_lat, w.dst_lon |
| 257 | + """ |
| 258 | + show_sql(apply_sql, label="SQL — apply bilinear weights (JOIN)") |
| 259 | + for _ in measured("SQL warp (reproject UDF + regrid JOIN)"): |
| 260 | + got = ctx.sql(apply_sql).to_dataset(dims=["lat", "lon"]).warped |
| 261 | + |
| 262 | + # Reference: the array paradigm doing the same warp — pyproj reproject of the |
| 263 | + # target grid, then xarray's own bilinear .interp at those source points. |
| 264 | + for _ in measured("xarray reference (pyproj + .interp)"): |
| 265 | + tr = pyproj.Transformer.from_crs(_DST_CRS, _SRC_CRS, always_xy=True) |
| 266 | + rx, ry = tr.transform(LON.ravel(), LAT.ravel()) |
| 267 | + warped = src.interp( |
| 268 | + x=xr.DataArray(rx, dims="cell"), |
| 269 | + y=xr.DataArray(ry, dims="cell"), |
| 270 | + method="linear", |
| 271 | + ).values.reshape(len(tlat), len(tlon)) |
| 272 | + ref = xr.DataArray( |
| 273 | + warped, dims=["lat", "lon"], coords={"lat": tlat, "lon": tlon} |
| 274 | + ) |
| 275 | + |
| 276 | + assert_grid_close("warped elevation (m)", got, ref, rtol=1e-6, atol=1e-4) |
| 277 | + show_result(got) |
| 278 | + |
| 279 | + # Independent cross-check: EE's own SRTM on the lon/lat grid (a real warp). |
| 280 | + ee_on_grid = ref_ee.interp(lat=got.lat, lon=got.lon, method="linear").values |
| 281 | + a, b = got.values.ravel(), ee_on_grid.ravel() |
| 282 | + m = np.isfinite(a) & np.isfinite(b) |
| 283 | + corr = float(np.corrcoef(a[m], b[m])[0, 1]) |
| 284 | + print( |
| 285 | + f"\n vs Earth Engine's own lon/lat SRTM: median |Δ| " |
| 286 | + f"{np.nanmedian(np.abs(a[m] - b[m])):.1f} m, correlation {corr:.4f} " |
| 287 | + f"(EE resamples native 30 m; ours warps the {_SRC_SCALE_M:.0f} m UTM grid)" |
| 288 | + ) |
| 289 | + |
| 290 | + |
| 291 | +if __name__ == "__main__": |
| 292 | + raise SystemExit(run_case(main, "Warp: reproject UDF + regrid JOIN (SRTM)")) |
0 commit comments