Add differentiable-SQL demos: ARCO-ERA5 and gradient descent

claude · claude · commit fdb17fb39cb3 · 2026-06-28T15:17:32.000Z
Stacked demo branch (on the autograd feature) holding the runnable benchmark scripts, kept out of the core branch so it stays reviewable. * grad_era5.py: symbolic grad over real ARCO-ERA5 data (wind-speed sensitivity checked exactly; saturation vapour pressure checked against the closed-form Clausius-Clapeyron slope). * grad_descent.py: gradient descent in SQL. The parameters live in a DataFusion in-memory table; each step appends the next generation with a SQL INSERT whose values are the SQL-computed update (a - lr*AVG(grad(loss, a))). The whole loss curve is one GROUP BY over the history. Fit matches numpy least-squares. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_017mDoFJgsm9kS7SicGoCVF6
diff --git a/benchmarks/README.md b/benchmarks/README.md
@@ -0,0 +1,54 @@
+# Benchmarks & demos
+
+Standalone scripts that exercise xarray-sql against real data. Each declares its
+own dependencies inline (PEP 723) and points `xarray_sql` at this checkout, so
+they run with no setup:
+
+```bash
+uv run benchmarks/grad_era5.py
+```
+
+## `grad_era5.py` — differentiable SQL over ARCO-ERA5
+
+Demonstrates the autograd feature on a real climate archive
+([ARCO-ERA5](https://github.com/google-research/arco-era5), read anonymously
+from GCS — needs `gcsfs` and network access).
+
+The key idea: a physical quantity is written as an **analytic SQL formula** over
+ERA5 variables, and `grad(...)` differentiates that formula **symbolically**,
+evaluated at every grid cell. Because each row is an independent point, this is
+the relational equivalent of `jax.vmap(jax.grad(f))`. It is *not* a finite-
+difference spatial gradient — `grad(f(u, v), u)` is the exact partial derivative
+of `f`.
+
+Two worked cases, each checked against an analytic reference:
+
+| Quantity | SQL | Derivative | Check |
+| --- | --- | --- | --- |
+| Wind speed | `sqrt(power(u,2) + power(v,2))` | `grad(speed, u) = u/speed` | exact |
+| Saturation vapour pressure | `A*exp(B*tc/(tc+C))` | `grad(e_s, T)` | closed-form Clausius-Clapeyron slope |
+
+Each query round-trips back to an `xarray.Dataset` via `.to_dataset(...)`.
+
+## `grad_descent.py` — gradient descent in SQL
+
+Fits a line `y ~= a*x + b` by minimising the mean squared error, with the
+gradients w.r.t. the parameters computed in SQL. Differentiating through an
+aggregate is just linearity:
+
+```
+d/dθ AVG(loss) = AVG(grad(loss, θ))
+```
+
+so the gradient is an ordinary aggregate of a `grad` expression — no special
+"differentiate through GROUP BY" machinery, since `grad` becomes plain SQL
+before the aggregate runs.
+
+The optimiser trajectory is itself a relation: `params(step, a, b)` is a
+DataFusion **in-memory table**, and each step **appends the next generation with
+a SQL `INSERT`** computed by descending from the current row
+(`new_a = a - lr * AVG(grad(loss, a))`). A training step is literally "add a row
+to a table", and the whole loss curve is a single `GROUP BY` over the history
+joined to the data. The Python loop only sequences the steps. The fit matches
+numpy's least-squares solution. Self-contained (no network).
+
diff --git a/benchmarks/grad_descent.py b/benchmarks/grad_descent.py
@@ -0,0 +1,120 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#   "xarray_sql",
+#   "xarray",
+#   "numpy",
+#   "pyarrow",
+# ]
+#
+# [tool.uv.sources]
+# xarray_sql = { path = "..", editable = true }
+# ///
+"""Gradient descent in SQL, with the optimiser trajectory as a table.
+
+Fits a line ``y ~= a*x + b`` by minimising the mean squared error. Two ideas:
+
+1. **The gradient is a SQL aggregate.** Differentiating through ``AVG`` is just
+   linearity, so ``d/dθ AVG(loss) = AVG(grad(loss, θ))`` — an ordinary aggregate
+   of a ``grad`` expression.
+
+2. **The parameters are a real (mutable) SQL table.** ``params(step, a, b)`` is
+   a DataFusion in-memory table, and each gradient step **appends the next
+   generation with a SQL ``INSERT``** computed by descending from the current
+   row:
+
+       INSERT INTO params
+       SELECT step+1, a - lr*AVG(grad(loss, a)), b - lr*AVG(grad(loss, b)) ...
+
+   So a training step is literally "add a row to a table", and the whole
+   optimisation history is a relation you can query — e.g. the loss curve over
+   every generation in a single ``GROUP BY`` over the history joined to the
+   data. The Python ``for`` loop only sequences the steps; the gradient, the
+   update, and the append all happen in SQL. (A fully in-engine loop would be a
+   recursive CTE — see issue #194 — which doesn't round-trip Substrait yet.)
+
+Run standalone:
+
+    uv run benchmarks/grad_descent.py
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pyarrow as pa
+import xarray as xr
+
+import xarray_sql as xql
+
+# Per-row loss r^2 with residual r = y - (a*x + b).
+RESIDUAL = "(y - (a * x + b))"
+LOSS = f"{RESIDUAL} * {RESIDUAL}"
+LR = 0.4
+STEPS = 200
+
+
+def main() -> None:
+    rng = np.random.default_rng(0)
+    n = 500
+    x = rng.uniform(0.0, 1.0, n)
+    a_true, b_true = 2.0, -1.0
+    y = a_true * x + b_true + rng.normal(0.0, 0.01, n)
+
+    ctx = xql.XarrayContext()
+    ctx.from_dataset(
+        "d",
+        xr.Dataset(
+            {"x": (("i",), x), "y": (("i",), y)}, coords={"i": np.arange(n)}
+        ),
+        chunks={"i": n},
+    )
+
+    # The parameters live in a mutable DataFusion in-memory table, seeded with
+    # generation 0 at (a, b) = (0, 0).
+    ctx.register_record_batches(
+        "params",
+        [[pa.RecordBatch.from_pydict({"step": [0], "a": [0.0], "b": [0.0]})]],
+    )
+
+    for k in range(STEPS):
+        # Compute the next generation in SQL: descend from the current row along
+        # the SQL-computed gradient...
+        row = ctx.sql(
+            f"""
+            WITH cur AS (SELECT a, b FROM params WHERE step = {k})
+            SELECT cur.a - {LR} * AVG(grad({LOSS}, a)) AS a,
+                   cur.b - {LR} * AVG(grad({LOSS}, b)) AS b
+            FROM d CROSS JOIN cur
+            GROUP BY cur.a, cur.b
+            """
+        ).to_pandas()
+        # ...and append it to the parameter table with a SQL INSERT.
+        na, nb = float(row["a"][0]), float(row["b"][0])
+        ctx.sql(
+            f"INSERT INTO params VALUES ({k + 1}, {na!r}, {nb!r})"
+        ).collect()
+
+    # The optimiser trajectory is now a table: compute the loss at every step in
+    # a single query over the parameter history joined to the data.
+    curve = ctx.sql(
+        f"SELECT p.step AS step, AVG({LOSS}) AS loss "
+        f"FROM d CROSS JOIN params p GROUP BY p.step ORDER BY p.step"
+    ).to_pandas()
+    print("loss curve (every 40th generation):")
+    print(curve.iloc[::40].to_string(index=False))
+
+    final = ctx.sql(f"SELECT a, b FROM params WHERE step = {STEPS}").to_pandas()
+    a, b = float(final["a"][0]), float(final["b"][0])
+    a_ols, b_ols = np.polyfit(x, y, 1)
+    print(
+        f"\nSQL gradient descent:  a={a:.4f}  b={b:.4f}  ({STEPS + 1} generations)"
+    )
+    print(f"least-squares (numpy): a={a_ols:.4f}  b={b_ols:.4f}")
+    assert abs(a - a_ols) < 1e-2 and abs(b - b_ols) < 1e-2
+    print(
+        "\nOK: SQL-computed gradients fit the line to the least-squares solution."
+    )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/benchmarks/grad_era5.py b/benchmarks/grad_era5.py
@@ -0,0 +1,180 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#   "xarray_sql",
+#   "xarray[io]",
+#   "gcsfs",
+#   "numpy",
+# ]
+#
+# [tool.uv.sources]
+# xarray_sql = { path = "..", editable = true }
+# ///
+"""Differentiable SQL over ARCO-ERA5.
+
+A minimal demonstration of xarray-sql's autograd: take a real climate archive
+(ARCO-ERA5, read anonymously from GCS), express a physical quantity as an
+*analytic* SQL formula over its variables, and let ``grad(...)`` differentiate
+that formula symbolically — evaluated per grid cell, which is the relational
+equivalent of ``jax.vmap(jax.grad(f))`` (each row is an independent point).
+
+Note this is *symbolic* differentiation of an expression, not a finite-
+difference spatial gradient: ``grad(f(u, v), u)`` is the exact partial
+derivative of the formula ``f``, evaluated at every cell's values.
+
+Two cases:
+
+1. Wind-speed magnitude ``speed = sqrt(u^2 + v^2)``. Its sensitivity to the
+   eastward wind is ``d(speed)/du = u / speed`` — checked exactly.
+
+2. Saturation vapour pressure ``e_s(T)`` (August-Roche-Magnus form of the
+   Clausius-Clapeyron relation). ``d(e_s)/dT`` governs how fast the atmosphere's
+   moisture capacity grows with temperature — checked against the closed-form
+   slope.
+
+Run standalone (builds the local extension on first use):
+
+    uv run benchmarks/grad_era5.py
+"""
+
+from __future__ import annotations
+
+import time
+
+import numpy as np
+import xarray as xr
+
+import xarray_sql as xql
+
+ARCO_ERA5 = (
+    "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3"
+)
+
+# The SQL result comes back with ascending coordinates; ERA5's native latitude
+# is descending. Sort both sides before comparing so equality is by label.
+_SORT = ["latitude", "longitude"]
+
+# ERA5 variable names start with a digit, so they must be double-quoted in SQL.
+U = '"10m_u_component_of_wind"'
+V = '"10m_v_component_of_wind"'
+T = '"2m_temperature"'
+
+
+def load_era5_block() -> xr.Dataset:
+    """Open ARCO-ERA5 and pull one timestamp over a small region.
+
+    Lazy open of the whole archive; only the requested block is read. We keep
+    it to a few thousand cells so the demo runs in seconds.
+    """
+    full = xr.open_zarr(
+        ARCO_ERA5, chunks=None, storage_options={"token": "anon"}
+    )
+    block = (
+        full[
+            [
+                "10m_u_component_of_wind",
+                "10m_v_component_of_wind",
+                "2m_temperature",
+            ]
+        ]
+        .sel(time="2020-01-01T00")
+        # A ~North-America box (index-based to avoid lat-orientation pitfalls).
+        .isel(latitude=slice(120, 200), longitude=slice(900, 1000))
+        .load()
+    )
+    return block.chunk({"latitude": 40})
+
+
+def wind_speed_sensitivity(ctx: xql.XarrayContext, ref: xr.Dataset) -> None:
+    """grad(sqrt(u^2 + v^2)) checked against the exact u / speed, v / speed."""
+    speed = f"sqrt(power({U}, 2) + power({V}, 2))"
+    out = (
+        ctx.sql(
+            f"""
+        SELECT
+          latitude,
+          longitude,
+          {speed}            AS wind_speed,
+          grad({speed}, {U}) AS d_speed_d_u,
+          grad({speed}, {V}) AS d_speed_d_v
+        FROM era5
+        """
+        )
+        .to_dataset(dims=["latitude", "longitude"])
+        .sortby(_SORT)
+    )
+
+    u = ref["10m_u_component_of_wind"]
+    v = ref["10m_v_component_of_wind"]
+    speed_ref = np.sqrt(u**2 + v**2).sortby(_SORT)
+
+    xr.testing.assert_allclose(
+        out["wind_speed"], speed_ref.rename("wind_speed")
+    )
+    xr.testing.assert_allclose(
+        out["d_speed_d_u"], (u / speed_ref).sortby(_SORT).rename("d_speed_d_u")
+    )
+    xr.testing.assert_allclose(
+        out["d_speed_d_v"], (v / speed_ref).sortby(_SORT).rename("d_speed_d_v")
+    )
+    print("  wind-speed sensitivity matches u/|w|, v/|w| exactly")
+    print(out)
+
+
+def clausius_clapeyron(ctx: xql.XarrayContext, ref: xr.Dataset) -> None:
+    """grad(e_s(T)) checked against the closed-form Clausius-Clapeyron slope."""
+    # August-Roche-Magnus: e_s(T) = A * exp(B * tc / (tc + C)), tc = T - 273.15.
+    a, b, c = 6.1094, 17.625, 243.04
+    tc = f"({T} - 273.15)"
+    es = f"{a} * exp({b} * {tc} / ({tc} + {c}))"
+    out = (
+        ctx.sql(
+            f"""
+            SELECT
+              latitude,
+              longitude,
+              {es}            AS e_s,
+              grad({es}, {T}) AS de_s_dt
+            FROM era5
+            """
+        )
+        .to_dataset(dims=["latitude", "longitude"])
+        .sortby(_SORT)
+    )
+
+    # Reference in float64 (the columns are float32): the exact derivative is
+    #   d(e_s)/dT = e_s * B*C / (tc + C)^2.
+    temp = ref["2m_temperature"].astype("float64")
+    tc_ref = temp - 273.15
+    es_ref = a * np.exp(b * tc_ref / (tc_ref + c))
+    des_dt_ref = es_ref * (b * c) / (tc_ref + c) ** 2
+
+    xr.testing.assert_allclose(
+        out["e_s"], es_ref.sortby(_SORT).rename("e_s"), rtol=1e-5
+    )
+    xr.testing.assert_allclose(
+        out["de_s_dt"], des_dt_ref.sortby(_SORT).rename("de_s_dt"), rtol=1e-5
+    )
+    print("  d(e_s)/dT matches the closed-form Clausius-Clapeyron slope")
+    print(out)
+
+
+def main() -> None:
+    t0 = time.time()
+    ds = load_era5_block()
+    print(f"loaded ERA5 block {dict(ds.sizes)} in {time.time() - t0:.1f}s")
+
+    ctx = xql.XarrayContext()
+    ctx.from_dataset("era5", ds)
+
+    print("\n== wind-speed sensitivity: grad(sqrt(u^2 + v^2)) ==")
+    wind_speed_sensitivity(ctx, ds)
+
+    print("\n== Clausius-Clapeyron: grad(e_s(T)) ==")
+    clausius_clapeyron(ctx, ds)
+
+    print("\nOK: symbolic SQL gradients match the analytic references.")
+
+
+if __name__ == "__main__":
+    main()
