From 6949f079337decdbd19fab6fbb403280e0805dd5 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sun, 25 Oct 2026 09:47:06 -0700 Subject: [PATCH 1/5] Restore the GraphCast demo, moved here from #269 Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h --- benchmarks/ml/03_graphcast.py | 724 ++++++++++++++++++++++++++++++++++ benchmarks/ml/README.md | 35 +- 2 files changed, 752 insertions(+), 7 deletions(-) create mode 100644 benchmarks/ml/03_graphcast.py diff --git a/benchmarks/ml/03_graphcast.py b/benchmarks/ml/03_graphcast.py new file mode 100644 index 00000000..c2704518 --- /dev/null +++ b/benchmarks/ml/03_graphcast.py @@ -0,0 +1,724 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "xarray-sql[ml,duckdb]", +# "jax", +# "gcsfs", +# "zarr>=3", +# ] +# +# [tool.uv.sources] +# xarray-sql = { path = "../../", editable = true } +# /// +"""Demo 3: a GraphCast-shaped weather model in SQL, over ERA5 in Zarr. + +GraphCast (Lam et al., 2023) forecasts the atmosphere 6 hours ahead from +the two previous states. It is a graph neural network in three stages: + +* an **encoder** maps the latitude–longitude grid onto a "multi-mesh", + a refined icosahedron that keeps the edges of every refinement level; +* a **processor** passes messages along the mesh's edges, round after + round; +* a **decoder** maps the mesh back onto the grid, as the 6-hour change. + +Here it is in miniature, as SQL over ERA5, ECMWF's reanalysis of the +atmosphere since 1940. xarray-sql reads it from the Zarr store of +WeatherBench 2, a benchmark for weather models (or a synthetic stand-in, +offline). Every +graph is a table of edges ``(edge, sender, receiver)``, every state a +table of ``(id, f, val)``, and every stage the same relational steps: + +================= ============================================ +GNN step In SQL +================= ============================================ +Gather join the nodes' states to the edges, by sender + and by receiver +Concatenate ``UNION ALL``, offsetting each part's features +MLP einsums with ``params`` around swish, ``x / (1 + exp(-x))`` +Aggregate ``SUM`` of the edges' messages, ``GROUP BY`` receiver +================= ============================================ + +The geometry (the mesh, and which grid cells and mesh nodes are joined) +is computed once with numpy; the model, its inputs and its forecast are +SQL. The script checks the forecast, and the gradient of the forecast +error, against the same model in JAX, Google's library for differentiable +array programs. ddx (https://github.com/xqlsystems/ddx), which +differentiates SQL queries, computes the gradient. + +Then it measures what the third target demo of einfold +(https://github.com/xqlsystems/einfold, a query optimizer that rewrites +SQL plans into array algorithms; the demo is in its ``docs/demos.md``) is +about: how much less a regional forecast needs than a global one. It +reports the receptive field: how many mesh nodes and grid cells each +stage needs for the region. + +Simplifications from GraphCast: random weights, a 64×32 grid and six +variables, a mesh refined 2 times rather than 6, 2 processor rounds +rather than 16, no biases or layer norms, edge features of 3-D +displacement and length, decoder edges from each cell's 3 nearest mesh +nodes, and edge states that are embedded once rather than updated. + +Run ``uv run benchmarks/ml/03_graphcast.py --help`` for the options. +""" + +from __future__ import annotations + +import jax +import jax.numpy as jnp +import numpy as np +import pandas as pd +import xarray as xr +from ddxdb import DdxError +from _bench import Report, assert_close, cli, hosts + +jax.config.update("jax_enable_x64", True) + +ERA5 = ( + "gs://weatherbench2/datasets/era5/" + "1959-2023_01_10-6h-64x32_equiangular_conservative.zarr" +) +ERA5_TABLES = { + ("time", "longitude", "latitude"): "surface", + ("time", "level", "longitude", "latitude"): "atmosphere", +} +# The forecast's variables, as (table, variable, pressure level). +VARIABLES = [ + ("surface", "2m_temperature", None), + ("surface", "mean_sea_level_pressure", None), + ("surface", "10m_u_component_of_wind", None), + ("surface", "10m_v_component_of_wind", None), + ("atmosphere", "geopotential", 500), + ("atmosphere", "temperature", 850), +] +# The two input states and the state to forecast, 6 hours apart. +TIMES = pd.to_datetime(["2020-01-01T00", "2020-01-01T06", "2020-01-01T12"]) + +HIDDEN = 16 # channels of every node and edge state +INPUTS = 2 * len(VARIABLES) + 3 # two states, cos(lat), sin(lon), cos(lon) +EDGE_FEATURES = 4 # receiver minus sender in 3-D, and the edge's length +# A region to forecast: the contiguous United States, roughly. +BOUNDS = {"latitude": (25, 50), "longitude": (235, 295)} +REGION = " AND ".join( + f"c.{coord} BETWEEN {low} AND {high}" + for coord, (low, high) in BOUNDS.items() +) + +# -------------------------------------------------------------------------- +# The model, in SQL. +# -------------------------------------------------------------------------- + + +def mlp(out: str, x: str, net: str) -> str: + """A two-layer MLP over the features of each ``id`` of ``x``.""" + return f""" +{out}_hidden AS ( + SELECT x.id, p.g AS f, SUM(x.val * p.w) AS val + FROM {x} x JOIN params p ON p.net = '{net}.0' AND x.f = p.f + GROUP BY x.id, p.g), +{out} AS ( + SELECT h.id, p.g AS f, SUM(h.val / (1 + exp(-h.val)) * p.w) AS val + FROM {out}_hidden h JOIN params p ON p.net = '{net}.1' AND h.f = p.f + GROUP BY h.id, p.g),""" + + +def message_passing(out, net, edges, edge_h, senders, receivers) -> str: + """One round of an interaction network on the graph ``edges``: each + edge's message is an MLP of the edge, its sender and its receiver; + each receiver adds an MLP of itself and its messages' sum to itself.""" + return f""" +{out}_edge_in AS ( + SELECT id, f, val FROM {edge_h} + UNION ALL + SELECT e.edge AS id, s.f + {HIDDEN} AS f, s.val + FROM {edges} e JOIN {senders} s ON e.sender = s.id + UNION ALL + SELECT e.edge AS id, r.f + {2 * HIDDEN} AS f, r.val + FROM {edges} e JOIN {receivers} r ON e.receiver = r.id), +{mlp(f"{out}_message", f"{out}_edge_in", f"{net}.edge")} +{out}_sum AS ( + SELECT e.receiver AS id, m.f, SUM(m.val) AS val + FROM {out}_message m JOIN {edges} e ON m.id = e.edge + GROUP BY e.receiver, m.f), +{out}_node_in AS ( + SELECT id, f, val FROM {receivers} + UNION ALL + SELECT id, f + {HIDDEN} AS f, val FROM {out}_sum), +{mlp(f"{out}_update", f"{out}_node_in", f"{net}.node")} +{out} AS ( + SELECT r.id, r.f, r.val + u.val AS val + FROM {receivers} r JOIN {out}_update u ON r.id = u.id AND r.f = u.f),""" + + +def era5_state(time, first_feature: int) -> list[str]: + """Each variable at ``time``, normalized, as (cell, feature, value) + rows. Only these variables, levels and times are read from ERA5.""" + parts = [] + for var, (table, name, level) in enumerate(VARIABLES): + on_level = f" AND x.level = {level}" if level else "" + parts.append( + f""" SELECT c.cell AS id, {first_feature + var} AS f, + (x."{name}" - n.mean) / n.std AS val + FROM era5.{table} x + JOIN cells c ON x.latitude = c.latitude AND x.longitude = c.longitude + JOIN scale n ON n.var = {var} + WHERE x.time = TIMESTAMP '{time}'{on_level}""" + ) + return parts + + +def position(table: str, key: str, first_feature: int) -> list[str]: + """Where each row of ``table`` is on the sphere, as 3 features.""" + features = [ + "cos(radians(latitude))", + "sin(radians(longitude))", + "cos(radians(longitude))", + ] + return [ + f" SELECT {key} AS id, {first_feature + i} AS f, {feature} AS val " + f"FROM {table}" + for i, feature in enumerate(features) + ] + + +def union(name: str, parts: list[str]) -> str: + return f"\n{name} AS (\n" + "\n UNION ALL\n".join(parts) + ")," + + +def edge_features(graph: str) -> str: + return f"(SELECT edge AS id, f, val FROM {graph}.features)" + + +def model(rounds: int) -> str: + """The CTEs of the whole model, ending in ``forecast``: the change of + each normalized variable over the next 6 hours, at each cell.""" + n = len(VARIABLES) + processor = "".join( + message_passing( + f"mesh{r + 1}", + f"processor.{r}", + "multimesh.edges", + "multimesh_h", + f"mesh{r}", + f"mesh{r}", + ) + for r in range(rounds) + ) + return f"""{ + union( + "grid_in", + era5_state(TIMES[0], 0) + + era5_state(TIMES[1], n) + + position("cells", "cell", 2 * n), + ) + } +{union("mesh_in", position("mesh", "node", 0))} +{mlp("grid_h", "grid_in", "embed.grid")} +{mlp("mesh_h", "mesh_in", "embed.mesh")} +{mlp("g2m_h", edge_features("g2m"), "embed.g2m")} +{mlp("multimesh_h", edge_features("multimesh"), "embed.multimesh")} +{mlp("m2g_h", edge_features("m2g"), "embed.m2g")} +{message_passing("mesh0", "encoder", "g2m.edges", "g2m_h", "grid_h", "mesh_h")} +{processor} +{ + message_passing( + "grid_out", + "decoder", + "m2g.edges", + "m2g_h", + f"mesh{rounds}", + "grid_h", + ) + } +forecast AS ( + SELECT g.id, p.g AS f, SUM(g.val * p.w) AS val + FROM grid_out g JOIN params p ON p.net = 'output' AND g.f = p.f + GROUP BY g.id, p.g)""" + + +def forecast_statement(rounds: int, where: str = "TRUE") -> str: + """The forecast change of each variable at each cell in ``where``.""" + return f"""WITH {model(rounds)} +SELECT c.latitude, c.longitude, f.f AS var, f.val AS change +FROM forecast f JOIN cells c ON f.id = c.cell +WHERE {where} +ORDER BY c.latitude, c.longitude, f.f""" + + +def gradient_statement(rounds: int) -> str: + """Each weight's gradient of the forecast's mean squared error + against what ERA5 recorded 6 hours later.""" + n = len(VARIABLES) + return f"""WITH {model(rounds)}, +{union("later", era5_state(TIMES[2], 0))} +error AS ( + SELECT AVG(power(f.val - (y.val - x.val), 2)) AS l + FROM forecast f + JOIN later y ON f.id = y.id AND f.f = y.f + JOIN grid_in x ON f.id = x.id AND f.f + {n} = x.f) +SELECT weight, grad(error, w) AS w FROM params ORDER BY weight""" + + +def receptive_field(host, rounds: int) -> list[tuple[str, int]]: + """How many cells and mesh nodes each stage needs for the region, + found by walking the edges backward from it: semi-joins in SQL.""" + + def count(sql): + return host.query(f"SELECT COUNT(*) AS n FROM ({sql}) t")["n"][ + 0 + ].as_py() + + needed = f"SELECT cell AS id FROM cells c WHERE {REGION}" + stages = [("decoder: grid cells forecast", count(needed))] + needed = ( + f"SELECT DISTINCT sender AS id FROM m2g.edges " + f"WHERE receiver IN ({needed})" + ) + stages.append((f"processor round {rounds}: mesh nodes", count(needed))) + for r in reversed(range(rounds)): + needed = ( + f"SELECT DISTINCT sender AS id FROM multimesh.edges " + f"WHERE receiver IN ({needed}) UNION SELECT id FROM ({needed}) n" + ) + label = f"processor round {r}" if r else "encoder" + stages.append((f"{label}: mesh nodes", count(needed))) + needed = ( + f"SELECT DISTINCT sender AS id FROM g2m.edges " + f"WHERE receiver IN ({needed})" + ) + stages.append(("encoder: grid cells read", count(needed))) + return stages + + +# -------------------------------------------------------------------------- +# The same model in JAX: the reference. +# -------------------------------------------------------------------------- + + +def jax_mlp(p, net, x): + h = x @ p[f"{net}.0"] + return (h * jax.nn.sigmoid(h)) @ p[f"{net}.1"] + + +def jax_message_passing(p, net, edges, edge_h, senders, receivers): + sender, receiver = edges + edge_in = jnp.concatenate( + [edge_h, senders[sender], receivers[receiver]], axis=1 + ) + message = jax_mlp(p, f"{net}.edge", edge_in) + total = jax.ops.segment_sum(message, receiver, len(receivers)) + node_in = jnp.concatenate([receivers, total], axis=1) + return receivers + jax_mlp(p, f"{net}.node", node_in) + + +def jax_forecast(p, graph, grid_in, rounds): + grid_h = jax_mlp(p, "embed.grid", grid_in) + mesh_h = jax_mlp(p, "embed.mesh", graph["mesh_in"]) + edge_h = { + g: jax_mlp(p, f"embed.{g}", graph[g]["features"]) + for g in ("g2m", "multimesh", "m2g") + } + mesh = jax_message_passing( + p, "encoder", graph["g2m"]["edges"], edge_h["g2m"], grid_h, mesh_h + ) + for r in range(rounds): + mesh = jax_message_passing( + p, + f"processor.{r}", + graph["multimesh"]["edges"], + edge_h["multimesh"], + mesh, + mesh, + ) + grid_out = jax_message_passing( + p, "decoder", graph["m2g"]["edges"], edge_h["m2g"], mesh, grid_h + ) + return grid_out @ p["output"] + + +def jax_error(p, graph, grid_in, target, rounds): + return jnp.mean((jax_forecast(p, graph, grid_in, rounds) - target) ** 2) + + +# -------------------------------------------------------------------------- +# Geometry, with numpy: the multi-mesh and the edges between it and the grid. +# -------------------------------------------------------------------------- + + +def icosahedron(): + t = (1 + 5**0.5) / 2 + vertices = [ + (-1, t, 0), + (1, t, 0), + (-1, -t, 0), + (1, -t, 0), + (0, -1, t), + (0, 1, t), + (0, -1, -t), + (0, 1, -t), + (t, 0, -1), + (t, 0, 1), + (-t, 0, -1), + (-t, 0, 1), + ] + faces = [ + (0, 11, 5), + (0, 5, 1), + (0, 1, 7), + (0, 7, 10), + (0, 10, 11), + (1, 5, 9), + (5, 11, 4), + (11, 10, 2), + (10, 7, 6), + (7, 1, 8), + (3, 9, 4), + (3, 4, 2), + (3, 2, 6), + (3, 6, 8), + (3, 8, 9), + (4, 9, 5), + (2, 4, 11), + (6, 2, 10), + (8, 6, 7), + (9, 8, 1), + ] + vertices = [np.array(v) / np.linalg.norm(v) for v in vertices] + return vertices, faces + + +def refine(vertices: list, faces: list) -> list: + """Split each triangle into four, adding its edges' midpoints, on the + sphere, to ``vertices``.""" + midpoints = {} + + def midpoint(a, b): + key = (min(a, b), max(a, b)) + if key not in midpoints: + m = vertices[a] + vertices[b] + vertices.append(m / np.linalg.norm(m)) + midpoints[key] = len(vertices) - 1 + return midpoints[key] + + finer = [] + for a, b, c in faces: + ab, bc, ca = midpoint(a, b), midpoint(b, c), midpoint(c, a) + finer += [(a, ab, ca), (b, bc, ab), (c, ca, bc), (ab, bc, ca)] + return finer + + +def multimesh(refinements: int): + """The vertices of an icosahedron refined ``refinements`` times, the + edges of every level of refinement in both directions, and the length + of the finest level's longest edge.""" + vertices, faces = icosahedron() + levels = [faces] + for _ in range(refinements): + levels.append(refine(vertices, levels[-1])) + edges = { + (u, v) + for level in levels + for face in level + for u, v in zip(face, face[1:] + face[:1]) + } + edges |= {(v, u) for u, v in edges} + vertices = np.array(vertices) + finest = max( + np.linalg.norm(vertices[u] - vertices[v]) + for face in levels[-1] + for u, v in zip(face, face[1:] + face[:1]) + ) + return vertices, np.array(sorted(edges)).T, finest + + +def xyz(latitude, longitude): + lat, lon = np.radians(latitude), np.radians(longitude) + return np.stack( + [np.cos(lat) * np.cos(lon), np.cos(lat) * np.sin(lon), np.sin(lat)], -1 + ) + + +def graph_table(senders, receivers, points_from, points_to) -> xr.Dataset: + """A graph as a Dataset: the table ``edges(edge, sender, receiver)`` + and the table ``features(edge, f, val)``.""" + offset = points_to[receivers] - points_from[senders] + features = np.concatenate( + [offset, np.linalg.norm(offset, axis=1, keepdims=True)], axis=1 + ) + return xr.Dataset( + { + "sender": ("edge", senders), + "receiver": ("edge", receivers), + "val": (("edge", "f"), features), + }, + coords={"edge": np.arange(len(senders)), "f": np.arange(EDGE_FEATURES)}, + ) + + +GRAPH_TABLES = {("edge",): "edges", ("edge", "f"): "features"} + + +def geometry(era5: xr.Dataset, refinements: int) -> dict[str, xr.Dataset]: + """The tables of the model's fixed structure, as Datasets.""" + lat, lon = xr.broadcast(era5.latitude, era5.longitude) + lat, lon = lat.values.ravel(), lon.values.ravel() + cell_xyz = xyz(lat, lon) + node_xyz, (senders, receivers), finest = multimesh(refinements) + node_lat = np.degrees(np.arcsin(node_xyz[:, 2])) + node_lon = np.degrees(np.arctan2(node_xyz[:, 1], node_xyz[:, 0])) % 360 + + # Grid to mesh: each cell to the mesh nodes within 0.6 of the finest + # mesh's longest edge, GraphCast's radius. + distance = np.linalg.norm(cell_xyz[:, None] - node_xyz[None], axis=-1) + g2m_cell, g2m_node = np.nonzero(distance < 0.6 * finest) + # Mesh to grid: each cell from its 3 nearest mesh nodes. + nearest = np.argsort(distance, axis=1)[:, :3] + m2g_cell = np.repeat(np.arange(len(lat)), 3) + m2g_node = nearest.ravel() + + return { + "cells": xr.Dataset( + {"latitude": ("cell", lat), "longitude": ("cell", lon)}, + coords={"cell": np.arange(len(lat))}, + ), + "mesh": xr.Dataset( + {"latitude": ("node", node_lat), "longitude": ("node", node_lon)}, + coords={"node": np.arange(len(node_xyz))}, + ), + "g2m": graph_table(g2m_cell, g2m_node, cell_xyz, node_xyz), + "multimesh": graph_table(senders, receivers, node_xyz, node_xyz), + "m2g": graph_table(m2g_node, m2g_cell, node_xyz, cell_xyz), + } + + +# -------------------------------------------------------------------------- +# Data and parameters. +# -------------------------------------------------------------------------- + + +def open_era5() -> xr.Dataset: + """WeatherBench 2's ERA5, lazily; a synthetic stand-in offline.""" + try: + return xr.open_zarr( + ERA5, chunks=None, storage_options={"token": "anon"} + ) + except Exception as e: # noqa: BLE001 - any failure to reach the store + print(f" ERA5 unavailable ({type(e).__name__}); using synthetic data") + lon = np.arange(64) * 5.625 + lat = -87.1875 + np.arange(32) * 5.625 + rng = np.random.default_rng(3) + surface = rng.normal(size=(len(TIMES), 64, 32)) + atmosphere = rng.normal(size=(len(TIMES), 2, 64, 32)) + data: dict[str, tuple[tuple[str, ...], np.ndarray]] = { + name: ( + ("time", "longitude", "latitude"), + surface + i, + ) + for i, (table, name, _) in enumerate(VARIABLES) + if table == "surface" + } + data["geopotential"] = ( + ("time", "level", "longitude", "latitude"), + atmosphere, + ) + data["temperature"] = ( + ("time", "level", "longitude", "latitude"), + -atmosphere, + ) + return xr.Dataset( + data, + coords={ + "time": TIMES, + "level": [500, 850], + "longitude": lon, + "latitude": lat, + }, + ) + + +def era5_arrays(era5: xr.Dataset, time) -> np.ndarray: + """Each variable at ``time`` as a column, one row per cell, read with + xarray: the reference's input, independent of the SQL's.""" + columns = [] + for _, name, level in VARIABLES: + field = era5[name].sel(time=time) + if level: + field = field.sel(level=level) + columns.append(field.transpose("latitude", "longitude").values.ravel()) + return np.stack(columns, axis=1) + + +def init_params(rounds: int, rng) -> dict[str, np.ndarray]: + """Each network's weight matrices, by name, with fan-in scaling.""" + shapes = {"output": (HIDDEN, len(VARIABLES))} + + def two_layer(net, inputs): + shapes[f"{net}.0"] = (inputs, HIDDEN) + shapes[f"{net}.1"] = (HIDDEN, HIDDEN) + + two_layer("embed.grid", INPUTS) + two_layer("embed.mesh", 3) + for graph in ("g2m", "multimesh", "m2g"): + two_layer(f"embed.{graph}", EDGE_FEATURES) + for net in ["encoder", "decoder"] + [ + f"processor.{r}" for r in range(rounds) + ]: + two_layer(f"{net}.edge", 3 * HIDDEN) + two_layer(f"{net}.node", 2 * HIDDEN) + return { + net: rng.normal(0, 1 / np.sqrt(shape[0]), shape) + for net, shape in shapes.items() + } + + +def params_table(params: dict[str, np.ndarray]) -> xr.Dataset: + """Every weight as a row of ``params(weight, net, f, g, w)``.""" + rows = [ + (net, f, g, w) + for net, matrix in params.items() + for (f, g), w in np.ndenumerate(matrix) + ] + net, f, g, w = map(np.array, zip(*rows)) + return xr.Dataset( + { + "net": ("weight", net), + "f": ("weight", f), + "g": ("weight", g), + "w": ("weight", w), + }, + coords={"weight": np.arange(len(rows))}, + ) + + +def main(): + parser = cli(__doc__.splitlines()[0], sizes=[2], size_name="refinements") + parser.add_argument( + "--rounds", type=int, default=2, help="processor rounds" + ) + args = parser.parse_args() + report = Report("graphcast", "refinements", args.reps) + rounds = args.rounds + era5 = open_era5() + + # The normalization: each variable's mean and spread at the inputs. + now, before, later = (era5_arrays(era5, t) for t in TIMES[[1, 0, 2]]) + mean = np.concatenate([before, now]).mean(axis=0) + std = np.concatenate([before, now]).std(axis=0) + scale = xr.Dataset( + {"mean": ("var", mean), "std": ("var", std)}, + coords={"var": np.arange(len(VARIABLES))}, + ) + params = init_params(rounds, np.random.default_rng(0)) + + for refinements in args.sizes: + tables = geometry(era5, refinements) + cells = tables["cells"] + print( + f" mesh refined {refinements}×: {tables['mesh'].sizes['node']} " + f"nodes, {tables['multimesh'].sizes['edge']} mesh edges, " + f"{tables['g2m'].sizes['edge']} grid-to-mesh, " + f"{tables['m2g'].sizes['edge']} mesh-to-grid; " + f"{len(params_table(params).weight)} weights" + ) + + # The reference, in JAX, from arrays read with xarray. + graph = { + name: { + "edges": ( + tables[name].sender.values, + tables[name].receiver.values, + ), + "features": tables[name].val.values, + } + for name in ("g2m", "multimesh", "m2g") + } + graph["mesh_in"] = np.stack( + [ + np.cos(np.radians(tables["mesh"].latitude.values)), + np.sin(np.radians(tables["mesh"].longitude.values)), + np.cos(np.radians(tables["mesh"].longitude.values)), + ], + axis=1, + ) + position = np.stack( + [ + np.cos(np.radians(cells.latitude.values)), + np.sin(np.radians(cells.longitude.values)), + np.cos(np.radians(cells.longitude.values)), + ], + axis=1, + ) + grid_in = np.concatenate( + [(before - mean) / std, (now - mean) / std, position], axis=1 + ) + target = (later - now) / std + want = jax_forecast(params, graph, grid_in, rounds) + + want_grad = jax.grad(jax_error)(params, graph, grid_in, target, rounds) + want_grad = np.concatenate([want_grad[net].ravel() for net in params]) + + for host in hosts(args.hosts): + host.load("era5", era5, table_names=ERA5_TABLES) + for name, table in tables.items(): + options = ( + {"table_names": GRAPH_TABLES} + if "edge" in table.dims + else {} + ) + host.load(name, table, **options) + host.load("scale", scale) + host.load("params", params_table(params)) + + statement = forecast_statement(rounds) + rows = report.time( + "forecast, global", host, refinements, host.query, statement + ) + got = rows["change"].to_numpy().reshape(len(cells.cell), -1) + # The rows are ordered by latitude then longitude, as cells are. + report.check(assert_close("forecast", got, want, rtol=1e-9)) + report.note("operators", host.operators(statement)) + + statement = forecast_statement(rounds, REGION) + rows = report.time( + "forecast, regional", host, refinements, host.query, statement + ) + in_region = np.logical_and.reduce( + [ + (cells[c] >= lo) & (cells[c] <= hi) + for c, (lo, hi) in BOUNDS.items() + ] + ) + regional = rows["change"].to_numpy().reshape(-1, len(VARIABLES)) + report.check(assert_close("regional", regional, want[in_region])) + report.note("operators", host.operators(statement)) + + step = "gradient of the error" + if not host.differentiates: + report.skip(step, host, refinements, "no ddx backend") + continue + try: + rows = report.time( + step, + host, + refinements, + host.query, + gradient_statement(rounds), + ) + except DdxError as e: + # Recorded rather than raised: a gap in ddx, which this + # row tracks until ddx closes it. + reason = str(e).removeprefix("not supported by ddx-ad yet: ") + report.skip(step, host, refinements, f"ddx: {reason}") + continue + report.check( + assert_close("gradient", rows["w"].to_numpy(), want_grad) + ) + + print("\n receptive field of the regional forecast:") + for stage, n in receptive_field(host, rounds): + print(f" {stage:<34} {n:>6}") + + report.finish(args.csv) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/ml/README.md b/benchmarks/ml/README.md index ae5b5950..e0e080d2 100644 --- a/benchmarks/ml/README.md +++ b/benchmarks/ml/README.md @@ -1,14 +1,14 @@ # Machine learning in SQL -These scripts write machine-learning models as SQL over tables that +These scripts write three machine-learning models as SQL over tables that [xarray-sql](../../README.md) registers from xarray Datasets, and time the same SQL on several engines. Two related projects shape them: - [ddx](https://github.com/xqlsystems/ddx) differentiates SQL queries: it turns the gradient of a query into more SQL, which trains the models. - [einfold](https://github.com/xqlsystems/einfold) is a query optimizer - that rewrites SQL plans into array algorithms. Two of these models - are its [target demos](https://github.com/xqlsystems/einfold/blob/main/docs/demos.md), + that rewrites SQL plans into array algorithms. These models are its + three [target demos](https://github.com/xqlsystems/einfold/blob/main/docs/demos.md), and this SQL is the baseline its plans must beat. Each script checks its results against the same model in @@ -19,6 +19,7 @@ array programs. |---|---|---|---| | 01 | [`01_attention.py`](01_attention.py) | causal softmax attention, the core of a transformer, as a view (the model) in einfold's window and `GROUP BY` spellings, and a loss that reads it | the naive plan, which stores every score, for [FlashAttention](https://arxiv.org/abs/2205.14135)'s tiled plan to beat; gradients with respect to Q, K and V; masked pairs as absent rows; a control that isn't attention, the Nadaraya–Watson kernel average | | 02 | [`02_nanogpt.py`](02_nanogpt.py) | [nanoGPT](https://github.com/karpathy/nanoGPT)'s GPT-2 model, at 1 layer and 4,416 parameters, as a view, and a cross-entropy loss that reads it, on Tiny Shakespeare | the loss on each host; training steps with SGD or, as nanoGPT trains, AdamW, every parameter table checked after every step | +| 03 | [`03_graphcast.py`](03_graphcast.py) | a miniature [GraphCast](https://arxiv.org/abs/2212.12794), DeepMind's graph neural network for weather forecasting, over [ERA5](https://www.ecmwf.int/en/forecasts/dataset/ecmwf-reanalysis-v5) reanalysis in Zarr | a global against a regional forecast; the gradient of the forecast error; the receptive field of a region | The SQL favors legibility over speed. Each model reads like its definition, and the slow parts are left slow on purpose: they are what ddx, xarray-sql @@ -38,15 +39,16 @@ mismatch fails the run. ```shell uv run benchmarks/ml/01_attention.py uv run benchmarks/ml/02_nanogpt.py --steps 5 +uv run benchmarks/ml/03_graphcast.py --rounds 1 ``` Each script carries [inline `uv` metadata](https://docs.astral.sh/uv/guides/scripts/) and runs against the xarray-sql in this checkout, with the `ml` extra. The -options are the same in each: +options are the same in all three: ```shell --hosts datafusion duckdb # the engines to run on (default: all) ---sizes 64 256 # problem sizes: tokens, or batch size +--sizes 64 256 # problem sizes: tokens, batch size, or mesh refinements --reps 3 # timed repetitions; the median is reported --csv results.csv # also write the results table ``` @@ -55,7 +57,8 @@ Demo 2 adds `--steps`, the number of training steps, and `--optimizer`: `sgd` (the default), or `adamw` to train as nanoGPT does, and the model's size: `--layers`, `--embed`, `--heads` and `--context`, nanoGPT's `n_layer`, `n_embd`, `n_head` and `block_size` (1, 16, 2 and 16 by -default; nanoGPT's Shakespeare model is 6, 384, 6 and 256). +default; nanoGPT's Shakespeare model is 6, 384, 6 and 256). Demo 3 adds +`--rounds`, the number of processor rounds. The output ends with one table: a row per step, host and size, with its median time, its error against JAX, and the number of operators in the @@ -130,13 +133,20 @@ transcription matches PyTorch exactly. **DataFusion computes a CTE once for every place that reads it.** A model reads its activations more than once: a layer norm reads its input -three times, a residual connection twice. DataFusion inlines each read, +three times, a residual connection twice, a message-passing round reads +the mesh as sender, receiver and residual. DataFusion inlines each read, so the plan grows with the model's depth; DuckDB computes a CTE that is read more than once a single time. | Query | DataFusion operators | DuckDB operators | DataFusion | DuckDB | |---|---:|---:|---:|---:| | GPT loss, 1 layer, batch 4 | 13,756 | 126 | 3.0 s | 0.063 s | +| GraphCast forecast, 1 round | 1,816 | 243 | 43 s | 17 s | +| GraphCast forecast, 2 rounds | 5,948 | 282 | 159 s | 13 s | + +The GraphCast rows are from 2026-10-17, before in-memory tables. + +The GraphCast rows are from 2026-10-17, before in-memory tables. Setting up each lazy read of a Dataset takes xarray-sql about 1.5 ms, so the GPT's 2,989 reads in one forward pass took 6.7 s of its 9.4 s. A small @@ -193,12 +203,23 @@ against 0.20 s for one step), but copying Q, K and V into native tables only brings the gradient to 1.7 s, and slows the forward pass (measured on 2026-10-24, with ddx's development branch and DuckDB 1.5.6). +**A regional forecast costs what a global one does.** For the contiguous +United States, the forecast covers 55 of 2,048 grid cells. Walking the +graph back from them, the decoder needs 15 mesh nodes, the processor's +two rounds need 36 and then 82, and the encoder reads 1,103 grid cells +(54%); with one round, 576 (28%). Neither host uses any of this: the +regional forecast takes as long as the global one on both. Pushing the +region's filter back through the graph is einfold's Demo 3. + **What ddx cannot differentiate yet.** Each of these raised a typed error naming the construct, never a wrong number: - A `GROUP BY` key that depends on a differentiated column. einfold's window spelling groups by the normalizer `z`; the demos divide inside the sum instead, `SUM(w / z * v)`, which einfold also lists. +- `UNION ALL` between a parameter table and the loss. GraphCast + concatenates features with `UNION ALL`, so its gradient row records the + error. It passes once ddx differentiates a union. - On DuckDB, a view whose CTE computes a window and is read more than once, as the softmax's weights are in attention's window spelling and in the GPT, so neither trains on DuckDB yet. ddx reads DuckDB's plans From 2df5307654bee84dca4a4b25dd97d99c55c02df1 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sat, 17 Oct 2026 16:31:26 -0700 Subject: [PATCH 2/5] WIP: GraphCast_small in SQL with DeepMind's weights Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h --- benchmarks/ml/03_graphcast.py | 1007 ++++++++++++--------------------- benchmarks/ml/_bench.py | 9 +- benchmarks/ml/_graphcast.py | 453 +++++++++++++++ benchmarks/ml/_hosts.py | 27 +- 4 files changed, 850 insertions(+), 646 deletions(-) create mode 100644 benchmarks/ml/_graphcast.py diff --git a/benchmarks/ml/03_graphcast.py b/benchmarks/ml/03_graphcast.py index c2704518..379aaf08 100644 --- a/benchmarks/ml/03_graphcast.py +++ b/benchmarks/ml/03_graphcast.py @@ -2,720 +2,447 @@ # requires-python = ">=3.11" # dependencies = [ # "xarray-sql[ml,duckdb]", -# "jax", +# "weathernext @ git+https://github.com/google-deepmind/graphcast", # "gcsfs", -# "zarr>=3", # ] # # [tool.uv.sources] # xarray-sql = { path = "../../", editable = true } # /// -"""Demo 3: a GraphCast-shaped weather model in SQL, over ERA5 in Zarr. +"""Demo 3: GraphCast in SQL, with DeepMind's published weights. GraphCast (Lam et al., 2023) forecasts the atmosphere 6 hours ahead from -the two previous states. It is a graph neural network in three stages: - -* an **encoder** maps the latitude–longitude grid onto a "multi-mesh", - a refined icosahedron that keeps the edges of every refinement level; -* a **processor** passes messages along the mesh's edges, round after - round; -* a **decoder** maps the mesh back onto the grid, as the 6-hour change. - -Here it is in miniature, as SQL over ERA5, ECMWF's reanalysis of the -atmosphere since 1940. xarray-sql reads it from the Zarr store of -WeatherBench 2, a benchmark for weather models (or a synthetic stand-in, -offline). Every -graph is a table of edges ``(edge, sender, receiver)``, every state a -table of ``(id, f, val)``, and every stage the same relational steps: - -================= ============================================ -GNN step In SQL -================= ============================================ -Gather join the nodes' states to the edges, by sender - and by receiver -Concatenate ``UNION ALL``, offsetting each part's features -MLP einsums with ``params`` around swish, ``x / (1 + exp(-x))`` -Aggregate ``SUM`` of the edges' messages, ``GROUP BY`` receiver -================= ============================================ - -The geometry (the mesh, and which grid cells and mesh nodes are joined) -is computed once with numpy; the model, its inputs and its forecast are -SQL. The script checks the forecast, and the gradient of the forecast -error, against the same model in JAX, Google's library for differentiable -array programs. ddx (https://github.com/xqlsystems/ddx), which -differentiates SQL queries, computes the gradient. - -Then it measures what the third target demo of einfold -(https://github.com/xqlsystems/einfold, a query optimizer that rewrites -SQL plans into array algorithms; the demo is in its ``docs/demos.md``) is -about: how much less a regional forecast needs than a global one. It -reports the receptive field: how many mesh nodes and grid cells each -stage needs for the region. - -Simplifications from GraphCast: random weights, a 64×32 grid and six -variables, a mesh refined 2 times rather than 6, 2 processor rounds -rather than 16, no biases or layer norms, edge features of 3-D -displacement and length, decoder edges from each cell's 3 nearest mesh -nodes, and edge states that are embedded once rather than updated. +the two previous states. This runs GraphCast_small, as DeepMind publishes +it (1° grid, 13 pressure levels, a mesh refined 5 times, 16 rounds of +message passing, 36 million weights), as a program of SQL statements over +tables xarray-sql registers: + +* ``era5``: DeepMind's example inputs (two states, 6 hours apart) and the + forcings at the forecast's time; +* ``weights``, ``biases``, ``norms`` and ``nets``: every parameter in + DeepMind's checkpoint, one row each; +* ``grid2mesh``, ``mesh`` and ``mesh2grid``: the three graphs, as tables of + edges and of edge features, built by DeepMind's code. + +The model has three stages: + +* an **encoder** embeds each grid cell's 183 input channels, and passes + one round of messages from the grid to the mesh; +* a **processor** passes 16 rounds of messages along the multi-mesh, the + icosahedron's edges at every level of refinement; +* a **decoder** passes one round back to the grid, and turns each cell's + state into 83 output channels: the 6-hour change of each variable. + +Every network is the same SQL: an MLP (linear, swish, linear, layer norm) +over a table of ``(id, f, val)``, one row per node or edge and feature. +Each stage is one statement, and its result is stored as a table, so no +host computes a stage twice. + +The check is DeepMind's own model. Its JAX code, run on the same inputs, +gives the stages' outputs; the SQL must match them. A plain JAX +transcription of the model (``_graphcast.mirror``), checked against +DeepMind's, also gives each processor round, so a run can stop after a +few rounds and still be checked. A full forecast is about two trillion +multiply-adds, which hash joins take hours over; DeepMind's JAX takes 15 +seconds. Run ``uv run benchmarks/ml/03_graphcast.py --help`` for the options. +``--sizes`` is the number of processor rounds to run; with 16, the +forecast is complete. """ from __future__ import annotations import jax -import jax.numpy as jnp import numpy as np -import pandas as pd -import xarray as xr -from ddxdb import DdxError + +import _graphcast as gc from _bench import Report, assert_close, cli, hosts jax.config.update("jax_enable_x64", True) -ERA5 = ( - "gs://weatherbench2/datasets/era5/" - "1959-2023_01_10-6h-64x32_equiangular_conservative.zarr" -) -ERA5_TABLES = { - ("time", "longitude", "latitude"): "surface", - ("time", "level", "longitude", "latitude"): "atmosphere", -} -# The forecast's variables, as (table, variable, pressure level). -VARIABLES = [ - ("surface", "2m_temperature", None), - ("surface", "mean_sea_level_pressure", None), - ("surface", "10m_u_component_of_wind", None), - ("surface", "10m_v_component_of_wind", None), - ("atmosphere", "geopotential", 500), - ("atmosphere", "temperature", 850), -] -# The two input states and the state to forecast, 6 hours apart. -TIMES = pd.to_datetime(["2020-01-01T00", "2020-01-01T06", "2020-01-01T12"]) - -HIDDEN = 16 # channels of every node and edge state -INPUTS = 2 * len(VARIABLES) + 3 # two states, cos(lat), sin(lon), cos(lon) -EDGE_FEATURES = 4 # receiver minus sender in 3-D, and the edge's length -# A region to forecast: the contiguous United States, roughly. -BOUNDS = {"latitude": (25, 50), "longitude": (235, 295)} -REGION = " AND ".join( - f"c.{coord} BETWEEN {low} AND {high}" - for coord, (low, high) in BOUNDS.items() -) +INPUTS = 183 # input channels: the variables at each time and level +ROUNDS = 16 # processor rounds in GraphCast_small # -------------------------------------------------------------------------- -# The model, in SQL. +# The networks, in SQL. A state is a table of (id, f, val). # -------------------------------------------------------------------------- -def mlp(out: str, x: str, net: str) -> str: - """A two-layer MLP over the features of each ``id`` of ``x``.""" - return f""" -{out}_hidden AS ( - SELECT x.id, p.g AS f, SUM(x.val * p.w) AS val - FROM {x} x JOIN params p ON p.net = '{net}.0' AND x.f = p.f - GROUP BY x.id, p.g), -{out} AS ( - SELECT h.id, p.g AS f, SUM(h.val / (1 + exp(-h.val)) * p.w) AS val - FROM {out}_hidden h JOIN params p ON p.net = '{net}.1' AND h.f = p.f - GROUP BY h.id, p.g),""" - - -def message_passing(out, net, edges, edge_h, senders, receivers) -> str: - """One round of an interaction network on the graph ``edges``: each - edge's message is an MLP of the edge, its sender and its receiver; - each receiver adds an MLP of itself and its messages' sum to itself.""" - return f""" -{out}_edge_in AS ( - SELECT id, f, val FROM {edge_h} - UNION ALL - SELECT e.edge AS id, s.f + {HIDDEN} AS f, s.val - FROM {edges} e JOIN {senders} s ON e.sender = s.id - UNION ALL - SELECT e.edge AS id, r.f + {2 * HIDDEN} AS f, r.val - FROM {edges} e JOIN {receivers} r ON e.receiver = r.id), -{mlp(f"{out}_message", f"{out}_edge_in", f"{net}.edge")} -{out}_sum AS ( - SELECT e.receiver AS id, m.f, SUM(m.val) AS val - FROM {out}_message m JOIN {edges} e ON m.id = e.edge - GROUP BY e.receiver, m.f), -{out}_node_in AS ( - SELECT id, f, val FROM {receivers} - UNION ALL - SELECT id, f + {HIDDEN} AS f, val FROM {out}_sum), -{mlp(f"{out}_update", f"{out}_node_in", f"{net}.node")} -{out} AS ( - SELECT r.id, r.f, r.val + u.val AS val - FROM {receivers} r JOIN {out}_update u ON r.id = u.id AND r.f = u.f),""" +def weights(name: str) -> str: + """The weight matrix of one linear layer, by DeepMind's name.""" + return f"""( + SELECT w.f, w.g, w.w FROM nets n JOIN weights w ON w.net = n.net + WHERE n.name = '{name}')""" -def era5_state(time, first_feature: int) -> list[str]: - """Each variable at ``time``, normalized, as (cell, feature, value) - rows. Only these variables, levels and times are read from ERA5.""" - parts = [] - for var, (table, name, level) in enumerate(VARIABLES): - on_level = f" AND x.level = {level}" if level else "" - parts.append( - f""" SELECT c.cell AS id, {first_feature + var} AS f, - (x."{name}" - n.mean) / n.std AS val - FROM era5.{table} x - JOIN cells c ON x.latitude = c.latitude AND x.longitude = c.longitude - JOIN scale n ON n.var = {var} - WHERE x.time = TIMESTAMP '{time}'{on_level}""" - ) - return parts +def biases(name: str) -> str: + """The bias vector of one linear layer, by DeepMind's name.""" + return f"""( + SELECT b.g, b.b FROM nets n JOIN biases b ON b.net = n.net + WHERE n.name = '{name}')""" -def position(table: str, key: str, first_feature: int) -> list[str]: - """Where each row of ``table`` is on the sphere, as 3 features.""" - features = [ - "cos(radians(latitude))", - "sin(radians(longitude))", - "cos(radians(longitude))", - ] - return [ - f" SELECT {key} AS id, {first_feature + i} AS f, {feature} AS val " - f"FROM {table}" - for i, feature in enumerate(features) - ] +def mlp(out: str, x: str, net: str, layer_norm: bool = True) -> str: + """DeepMind's MLP over the features of each ``id`` of ``x``: linear, + swish, linear, and (unless it is the decoder's last) a layer norm with + a learned scale and offset. A linear layer is a matrix product, a + join and a sum, and then its bias.""" + first, second = f"{net}_mlp/linear_0", f"{net}_mlp/linear_1" + sql = f""" +{out}_product_0 AS ( + SELECT x.id, w.g AS f, SUM(x.val * w.w) AS val + FROM {x} x JOIN {weights(first)} w ON x.f = w.f + GROUP BY x.id, w.g), +{out}_hidden AS ( + SELECT h.id, h.f, (h.val + b.b) / (1 + exp(-(h.val + b.b))) AS val + FROM {out}_product_0 h JOIN {biases(first)} b ON h.f = b.g), +{out}_product_1 AS ( + SELECT h.id, w.g AS f, SUM(h.val * w.w) AS val + FROM {out}_hidden h JOIN {weights(second)} w ON h.f = w.f + GROUP BY h.id, w.g), +{out}_linear AS ( + SELECT y.id, y.f, y.val + b.b AS val + FROM {out}_product_1 y JOIN {biases(second)} b ON y.f = b.g),""" + if not layer_norm: + return sql + f"\n{out} AS (SELECT * FROM {out}_linear)," + return ( + sql + + f""" +{out} AS ( + SELECT y.id, y.f, + (y.val - AVG(y.val) OVER (PARTITION BY y.id)) + / sqrt(VAR_POP(y.val) OVER (PARTITION BY y.id) + 1e-5) + * n.scale + n.offset AS val + FROM {out}_linear y + JOIN norms n ON n.g = y.f + JOIN nets m ON m.net = n.net + WHERE m.name = '{net}_layer_norm'),""" + ) -def union(name: str, parts: list[str]) -> str: - return f"\n{name} AS (\n" + "\n UNION ALL\n".join(parts) + ")," +def select(ctes: str, out: str) -> str: + return f"WITH {ctes.rstrip(',')}\nSELECT id, f, val FROM {out}" -def edge_features(graph: str) -> str: - return f"(SELECT edge AS id, f, val FROM {graph}.features)" +def embed(x: str, net: str) -> str: + """A statement: an MLP of every row of ``x``.""" + return select(mlp("embedded", x, net), "embedded") -def model(rounds: int) -> str: - """The CTEs of the whole model, ending in ``forecast``: the change of - each normalized variable over the next 6 hours, at each cell.""" - n = len(VARIABLES) - processor = "".join( - message_passing( - f"mesh{r + 1}", - f"processor.{r}", - "multimesh.edges", - "multimesh_h", - f"mesh{r}", - f"mesh{r}", - ) - for r in range(rounds) - ) - return f"""{ - union( - "grid_in", - era5_state(TIMES[0], 0) - + era5_state(TIMES[1], n) - + position("cells", "cell", 2 * n), - ) - } -{union("mesh_in", position("mesh", "node", 0))} -{mlp("grid_h", "grid_in", "embed.grid")} -{mlp("mesh_h", "mesh_in", "embed.mesh")} -{mlp("g2m_h", edge_features("g2m"), "embed.g2m")} -{mlp("multimesh_h", edge_features("multimesh"), "embed.multimesh")} -{mlp("m2g_h", edge_features("m2g"), "embed.m2g")} -{message_passing("mesh0", "encoder", "g2m.edges", "g2m_h", "grid_h", "mesh_h")} -{processor} -{ - message_passing( - "grid_out", - "decoder", - "m2g.edges", - "m2g_h", - f"mesh{rounds}", - "grid_h", - ) - } -forecast AS ( - SELECT g.id, p.g AS f, SUM(g.val * p.w) AS val - FROM grid_out g JOIN params p ON p.net = 'output' AND g.f = p.f - GROUP BY g.id, p.g)""" - - -def forecast_statement(rounds: int, where: str = "TRUE") -> str: - """The forecast change of each variable at each cell in ``where``.""" - return f"""WITH {model(rounds)} -SELECT c.latitude, c.longitude, f.f AS var, f.val AS change -FROM forecast f JOIN cells c ON f.id = c.cell -WHERE {where} -ORDER BY c.latitude, c.longitude, f.f""" - - -def gradient_statement(rounds: int) -> str: - """Each weight's gradient of the forecast's mean squared error - against what ERA5 recorded 6 hours later.""" - n = len(VARIABLES) - return f"""WITH {model(rounds)}, -{union("later", era5_state(TIMES[2], 0))} -error AS ( - SELECT AVG(power(f.val - (y.val - x.val), 2)) AS l - FROM forecast f - JOIN later y ON f.id = y.id AND f.f = y.f - JOIN grid_in x ON f.id = x.id AND f.f + {n} = x.f) -SELECT weight, grad(error, w) AS w FROM params ORDER BY weight""" - - -def receptive_field(host, rounds: int) -> list[tuple[str, int]]: - """How many cells and mesh nodes each stage needs for the region, - found by walking the edges backward from it: semi-joins in SQL.""" - - def count(sql): - return host.query(f"SELECT COUNT(*) AS n FROM ({sql}) t")["n"][ - 0 - ].as_py() - - needed = f"SELECT cell AS id FROM cells c WHERE {REGION}" - stages = [("decoder: grid cells forecast", count(needed))] - needed = ( - f"SELECT DISTINCT sender AS id FROM m2g.edges " - f"WHERE receiver IN ({needed})" +def messages(net: str, graph: str, edges: str, senders: str, receivers: str): + """A statement: each edge's message, an MLP of the edge's state, its + sender's and its receiver's, concatenated.""" + return select( + f""" +edge_in AS ( + SELECT x.id, x.f, x.val FROM {edges} x JOIN {graph} e ON x.id = e.edge + UNION ALL + SELECT e.edge AS id, s.f + 512 AS f, s.val + FROM {graph} e JOIN {senders} s ON e.sender = s.id + UNION ALL + SELECT e.edge AS id, r.f + 1024 AS f, r.val + FROM {graph} e JOIN {receivers} r ON e.receiver = r.id),""" + + mlp("message", "edge_in", net), + "message", ) - stages.append((f"processor round {rounds}: mesh nodes", count(needed))) - for r in reversed(range(rounds)): - needed = ( - f"SELECT DISTINCT sender AS id FROM multimesh.edges " - f"WHERE receiver IN ({needed}) UNION SELECT id FROM ({needed}) n" - ) - label = f"processor round {r}" if r else "encoder" - stages.append((f"{label}: mesh nodes", count(needed))) - needed = ( - f"SELECT DISTINCT sender AS id FROM g2m.edges " - f"WHERE receiver IN ({needed})" + + +def update(net: str, graph: str, nodes: str, messages: str | None) -> str: + """A statement: each node plus an MLP of itself and the sum of the + messages it receives (none, for the encoder's grid cells).""" + node_in = f"SELECT id, f, val FROM {nodes}" + if messages: + node_in += f""" + UNION ALL + SELECT e.receiver AS id, m.f + 512 AS f, SUM(m.val) AS val + FROM {messages} m JOIN {graph} e ON m.id = e.edge + GROUP BY e.receiver, m.f""" + return select( + f"\nnode_in AS (\n {node_in})," + + mlp("updated", "node_in", net) + + f""" +new AS ( + SELECT x.id, x.f, x.val + u.val AS val + FROM {nodes} x JOIN updated u ON x.id = u.id AND x.f = u.f),""", + "new", ) - stages.append(("encoder: grid cells read", count(needed))) - return stages -# -------------------------------------------------------------------------- -# The same model in JAX: the reference. -# -------------------------------------------------------------------------- +def residual(state: str, change: str) -> str: + """A statement: ``state + change``, row by row.""" + return f"""SELECT x.id, x.f, x.val + d.val AS val +FROM {state} x JOIN {change} d ON x.id = d.id AND x.f = d.f""" -def jax_mlp(p, net, x): - h = x @ p[f"{net}.0"] - return (h * jax.nn.sigmoid(h)) @ p[f"{net}.1"] +ERA5_TABLES = { + ("time", "lat", "lon"): "surface", + ("time", "level", "lat", "lon"): "atmosphere", + ("lat", "lon"): "static", + ("time", "lon"): "day", + ("time",): "year", +} +# How each era5 table's rows reach the grid's cells. +CELLS = { + "surface": "JOIN cells c ON x.lat = c.lat AND x.lon = c.lon", + "atmosphere": "JOIN cells c ON x.lat = c.lat AND x.lon = c.lon", + "static": "JOIN cells c ON x.lat = c.lat AND x.lon = c.lon", + "day": "JOIN cells c ON x.lon = c.lon", + "year": "CROSS JOIN cells c", +} -def jax_message_passing(p, net, edges, edge_h, senders, receivers): - sender, receiver = edges - edge_in = jnp.concatenate( - [edge_h, senders[sender], receivers[receiver]], axis=1 - ) - message = jax_mlp(p, f"{net}.edge", edge_in) - total = jax.ops.segment_sum(message, receiver, len(receivers)) - node_in = jnp.concatenate([receivers, total], axis=1) - return receivers + jax_mlp(p, f"{net}.node", node_in) - - -def jax_forecast(p, graph, grid_in, rounds): - grid_h = jax_mlp(p, "embed.grid", grid_in) - mesh_h = jax_mlp(p, "embed.mesh", graph["mesh_in"]) - edge_h = { - g: jax_mlp(p, f"embed.{g}", graph[g]["features"]) - for g in ("g2m", "multimesh", "m2g") - } - mesh = jax_message_passing( - p, "encoder", graph["g2m"]["edges"], edge_h["g2m"], grid_h, mesh_h - ) - for r in range(rounds): - mesh = jax_message_passing( - p, - f"processor.{r}", - graph["multimesh"]["edges"], - edge_h["multimesh"], - mesh, - mesh, +def grid_inputs(era5) -> str: + """A statement: each cell's input channels, normalized, then its + position. Every variable is read from its ``era5`` table at the times + and levels the channels ask for.""" + parts = [] + for name in sorted(era5.data_vars): + dims = era5[name].dims + table = ERA5_TABLES[dims] + time = "ch.time = x.time" if "time" in dims else "ch.time = -1" + level = "ch.level = x.level" if "level" in dims else "ch.level = 0" + parts.append( + f"""SELECT c.cell AS id, ch.channel AS f, + (x."{name}" - ch.mean) / ch.std AS val +FROM era5.{table} x +{CELLS[table]} +JOIN channels ch ON ch.variable = '{name}' AND {time} AND {level}""" ) - grid_out = jax_message_passing( - p, "decoder", graph["m2g"]["edges"], edge_h["m2g"], mesh, grid_h + parts.append(f"SELECT cell AS id, f + {INPUTS} AS f, val FROM cell_features") + return "\nUNION ALL\n".join(parts) + + +def forecast(era5) -> str: + """A statement: the forecast of each variable at each cell and level, + the last input plus the output's normalized 6-hour change.""" + targets = sorted(set(gc.outputs_table().variable.values)) + last = "\n UNION ALL\n".join( + f""" SELECT c.cell AS id, '{name}' AS variable, + {"x.level" if "level" in era5[name].dims else "0"} AS level, + x."{name}" AS value + FROM era5.{ERA5_TABLES[era5[name].dims]} x + {CELLS[ERA5_TABLES[era5[name].dims]]} + WHERE x.time = 0""" + for name in targets + ) + return f"""WITH last_input AS ( +{last}) +SELECT c.lat, c.lon, o.variable, o.level, + x.value + y.val * o.diff_std AS value +FROM output y +JOIN outputs o ON o.channel = y.f +JOIN cells c ON c.cell = y.id +JOIN last_input x + ON x.id = y.id AND x.variable = o.variable AND x.level = o.level +ORDER BY o.variable, o.level, c.lat, c.lon""" + + +def flatten(ds) -> np.ndarray: + """A forecast Dataset in ``forecast``'s row order: by variable, level, + latitude and longitude.""" + return np.concatenate( + [ds[name].transpose(..., "lat", "lon").values.ravel() + for name in sorted(ds.data_vars)] ) - return grid_out @ p["output"] - - -def jax_error(p, graph, grid_in, target, rounds): - return jnp.mean((jax_forecast(p, graph, grid_in, rounds) - target) ** 2) # -------------------------------------------------------------------------- -# Geometry, with numpy: the multi-mesh and the edges between it and the grid. +# The program: one stored table per stage. # -------------------------------------------------------------------------- -def icosahedron(): - t = (1 + 5**0.5) / 2 - vertices = [ - (-1, t, 0), - (1, t, 0), - (-1, -t, 0), - (1, -t, 0), - (0, -1, t), - (0, 1, t), - (0, -1, -t), - (0, 1, -t), - (t, 0, -1), - (t, 0, 1), - (-t, 0, -1), - (-t, 0, 1), - ] - faces = [ - (0, 11, 5), - (0, 5, 1), - (0, 1, 7), - (0, 7, 10), - (0, 10, 11), - (1, 5, 9), - (5, 11, 4), - (11, 10, 2), - (10, 7, 6), - (7, 1, 8), - (3, 9, 4), - (3, 4, 2), - (3, 2, 6), - (3, 6, 8), - (3, 8, 9), - (4, 9, 5), - (2, 4, 11), - (6, 2, 10), - (8, 6, 7), - (9, 8, 1), - ] - vertices = [np.array(v) / np.linalg.norm(v) for v in vertices] - return vertices, faces - - -def refine(vertices: list, faces: list) -> list: - """Split each triangle into four, adding its edges' midpoints, on the - sphere, to ``vertices``.""" - midpoints = {} - - def midpoint(a, b): - key = (min(a, b), max(a, b)) - if key not in midpoints: - m = vertices[a] + vertices[b] - vertices.append(m / np.linalg.norm(m)) - midpoints[key] = len(vertices) - 1 - return midpoints[key] - - finer = [] - for a, b, c in faces: - ab, bc, ca = midpoint(a, b), midpoint(b, c), midpoint(c, a) - finer += [(a, ab, ca), (b, bc, ab), (c, ca, bc), (ab, bc, ca)] - return finer - - -def multimesh(refinements: int): - """The vertices of an icosahedron refined ``refinements`` times, the - edges of every level of refinement in both directions, and the length - of the finest level's longest edge.""" - vertices, faces = icosahedron() - levels = [faces] - for _ in range(refinements): - levels.append(refine(vertices, levels[-1])) - edges = { - (u, v) - for level in levels - for face in level - for u, v in zip(face, face[1:] + face[:1]) - } - edges |= {(v, u) for u, v in edges} - vertices = np.array(vertices) - finest = max( - np.linalg.norm(vertices[u] - vertices[v]) - for face in levels[-1] - for u, v in zip(face, face[1:] + face[:1]) - ) - return vertices, np.array(sorted(edges)).T, finest +# A region to forecast: the contiguous United States, roughly. +BOUNDS = {"lat": (25, 50), "lon": (235, 295)} +REGION = " AND ".join( + f"c.{coord} BETWEEN {low} AND {high}" for coord, (low, high) in BOUNDS.items() +) -def xyz(latitude, longitude): - lat, lon = np.radians(latitude), np.radians(longitude) - return np.stack( - [np.cos(lat) * np.cos(lon), np.cos(lat) * np.sin(lon), np.sin(lat)], -1 +def receptive_field(rounds: int): + """The cells and mesh nodes a regional forecast needs, as stored + ``(table, statement)``s, found by walking the graphs backward from the + region: semi-joins of each graph's edges with what the next stage + needs.""" + yield "needed_cells", f"SELECT cell AS id FROM cells c WHERE {REGION}" + yield f"needed_mesh_{rounds}", """SELECT DISTINCT sender AS id +FROM mesh2grid.edges WHERE receiver IN (SELECT id FROM needed_cells)""" + for r in reversed(range(rounds)): + yield f"needed_mesh_{r}", f"""SELECT id FROM needed_mesh_{r + 1} +UNION +SELECT sender AS id FROM mesh.edges +WHERE receiver IN (SELECT id FROM needed_mesh_{r + 1})""" + yield "needed_grid", """SELECT id FROM needed_cells +UNION +SELECT sender AS id FROM grid2mesh.edges +WHERE receiver IN (SELECT id FROM needed_mesh_0)""" + + +def program(era5, rounds: int, regional: bool = False): + """The stages in order, as ``(table, statement)``; ``(table, None)`` + marks a table the reference checks. A regional program runs the same + stages on the rows its receptive field needs.""" + + def within(table: str, column: str, needed: str) -> str: + if not regional: + return table + return f"(SELECT * FROM {table} WHERE {column} IN (SELECT id FROM {needed}))" + + def features(graph: str, edges: str) -> str: + return f"""(SELECT x.edge AS id, x.f, x.val FROM {graph}.features x + WHERE x.edge IN (SELECT edge FROM {edges}))""" + + checked = not regional + yield "grid_in", ( + f"SELECT * FROM ({grid_inputs(era5)}) g\n" + "WHERE id IN (SELECT id FROM needed_grid)" + if regional + else grid_inputs(era5) ) - - -def graph_table(senders, receivers, points_from, points_to) -> xr.Dataset: - """A graph as a Dataset: the table ``edges(edge, sender, receiver)`` - and the table ``features(edge, f, val)``.""" - offset = points_to[receivers] - points_from[senders] - features = np.concatenate( - [offset, np.linalg.norm(offset, axis=1, keepdims=True)], axis=1 + if checked: + yield "grid_in", None + # Encoder: embed the grid, the mesh and the grid-to-mesh edges, then + # one round of messages from the grid to the mesh. + g2m = "grid2mesh_gnn" + g2m_graph = within("grid2mesh.edges", "receiver", "needed_mesh_0") + yield "grid", embed("grid_in", f"{g2m}/encoder_nodes_grid_nodes") + # DeepMind gives the mesh nodes 183 zero inputs; zeros are absent rows. + yield "mesh", embed( + f"""(SELECT node AS id, f + {INPUTS} AS f, val + FROM {within("node_features", "node", "needed_mesh_0")} n)""", + f"{g2m}/encoder_nodes_mesh_nodes", ) - return xr.Dataset( - { - "sender": ("edge", senders), - "receiver": ("edge", receivers), - "val": (("edge", "f"), features), - }, - coords={"edge": np.arange(len(senders)), "f": np.arange(EDGE_FEATURES)}, + yield "g2m_edges", embed( + features("grid2mesh", g2m_graph), f"{g2m}/encoder_edges_grid2mesh" ) - - -GRAPH_TABLES = {("edge",): "edges", ("edge", "f"): "features"} - - -def geometry(era5: xr.Dataset, refinements: int) -> dict[str, xr.Dataset]: - """The tables of the model's fixed structure, as Datasets.""" - lat, lon = xr.broadcast(era5.latitude, era5.longitude) - lat, lon = lat.values.ravel(), lon.values.ravel() - cell_xyz = xyz(lat, lon) - node_xyz, (senders, receivers), finest = multimesh(refinements) - node_lat = np.degrees(np.arcsin(node_xyz[:, 2])) - node_lon = np.degrees(np.arctan2(node_xyz[:, 1], node_xyz[:, 0])) % 360 - - # Grid to mesh: each cell to the mesh nodes within 0.6 of the finest - # mesh's longest edge, GraphCast's radius. - distance = np.linalg.norm(cell_xyz[:, None] - node_xyz[None], axis=-1) - g2m_cell, g2m_node = np.nonzero(distance < 0.6 * finest) - # Mesh to grid: each cell from its 3 nearest mesh nodes. - nearest = np.argsort(distance, axis=1)[:, :3] - m2g_cell = np.repeat(np.arange(len(lat)), 3) - m2g_node = nearest.ravel() - - return { - "cells": xr.Dataset( - {"latitude": ("cell", lat), "longitude": ("cell", lon)}, - coords={"cell": np.arange(len(lat))}, - ), - "mesh": xr.Dataset( - {"latitude": ("node", node_lat), "longitude": ("node", node_lon)}, - coords={"node": np.arange(len(node_xyz))}, - ), - "g2m": graph_table(g2m_cell, g2m_node, cell_xyz, node_xyz), - "multimesh": graph_table(senders, receivers, node_xyz, node_xyz), - "m2g": graph_table(m2g_node, m2g_cell, node_xyz, cell_xyz), - } - - -# -------------------------------------------------------------------------- -# Data and parameters. -# -------------------------------------------------------------------------- - - -def open_era5() -> xr.Dataset: - """WeatherBench 2's ERA5, lazily; a synthetic stand-in offline.""" - try: - return xr.open_zarr( - ERA5, chunks=None, storage_options={"token": "anon"} + yield "g2m_messages", messages( + f"{g2m}/processor_edges_0_grid2mesh", g2m_graph, + "g2m_edges", "grid", "mesh", + ) + yield "mesh_0", update( + f"{g2m}/processor_nodes_0_mesh_nodes", g2m_graph, + "mesh", "g2m_messages", + ) + yield "grid_0", update( + f"{g2m}/processor_nodes_0_grid_nodes", g2m_graph, + within("grid", "id", "needed_cells"), None, + ) + if checked: + yield "mesh_0", None + yield "grid_0", None + # Processor: rounds of messages along the multi-mesh, edges and nodes + # each updated with a residual. + graphs = [ + within("mesh.edges", "receiver", f"needed_mesh_{r + 1}") + for r in range(rounds) + ] + yield "mesh_edges_0", embed( + features("mesh", graphs[0] if rounds else "mesh.edges"), + "mesh_gnn/encoder_edges_mesh", + ) + for r in range(rounds): + yield f"mesh_messages_{r}", messages( + f"mesh_gnn/processor_edges_{r}_mesh", graphs[r], + f"mesh_edges_{r}", f"mesh_{r}", f"mesh_{r}", ) - except Exception as e: # noqa: BLE001 - any failure to reach the store - print(f" ERA5 unavailable ({type(e).__name__}); using synthetic data") - lon = np.arange(64) * 5.625 - lat = -87.1875 + np.arange(32) * 5.625 - rng = np.random.default_rng(3) - surface = rng.normal(size=(len(TIMES), 64, 32)) - atmosphere = rng.normal(size=(len(TIMES), 2, 64, 32)) - data: dict[str, tuple[tuple[str, ...], np.ndarray]] = { - name: ( - ("time", "longitude", "latitude"), - surface + i, + yield f"mesh_{r + 1}", update( + f"mesh_gnn/processor_nodes_{r}_mesh_nodes", graphs[r], + within(f"mesh_{r}", "id", f"needed_mesh_{r + 1}"), + f"mesh_messages_{r}", ) - for i, (table, name, _) in enumerate(VARIABLES) - if table == "surface" - } - data["geopotential"] = ( - ("time", "level", "longitude", "latitude"), - atmosphere, + yield f"mesh_edges_{r + 1}", residual( + f"mesh_edges_{r}", f"mesh_messages_{r}" + ) + if checked: + yield f"mesh_{r + 1}", None + if rounds < ROUNDS: + return + # Decoder: one round of messages from the mesh to the grid, then each + # cell's 83 output channels. + m2g = "mesh2grid_gnn" + m2g_graph = within("mesh2grid.edges", "receiver", "needed_cells") + yield "m2g_edges", embed( + features("mesh2grid", m2g_graph), f"{m2g}/encoder_edges_mesh2grid" ) - data["temperature"] = ( - ("time", "level", "longitude", "latitude"), - -atmosphere, + yield "m2g_messages", messages( + f"{m2g}/processor_edges_0_mesh2grid", m2g_graph, + "m2g_edges", f"mesh_{rounds}", "grid_0", ) - return xr.Dataset( - data, - coords={ - "time": TIMES, - "level": [500, 850], - "longitude": lon, - "latitude": lat, - }, + yield "grid_1", update( + f"{m2g}/processor_nodes_0_grid_nodes", m2g_graph, + "grid_0", "m2g_messages", ) + yield "output", select( + mlp("out", "grid_1", f"{m2g}/decoder_nodes_grid_nodes", layer_norm=False), + "out", + ) + if checked: + yield "output", None -def era5_arrays(era5: xr.Dataset, time) -> np.ndarray: - """Each variable at ``time`` as a column, one row per cell, read with - xarray: the reference's input, independent of the SQL's.""" - columns = [] - for _, name, level in VARIABLES: - field = era5[name].sel(time=time) - if level: - field = field.sel(level=level) - columns.append(field.transpose("latitude", "longitude").values.ravel()) - return np.stack(columns, axis=1) - - -def init_params(rounds: int, rng) -> dict[str, np.ndarray]: - """Each network's weight matrices, by name, with fan-in scaling.""" - shapes = {"output": (HIDDEN, len(VARIABLES))} - - def two_layer(net, inputs): - shapes[f"{net}.0"] = (inputs, HIDDEN) - shapes[f"{net}.1"] = (HIDDEN, HIDDEN) - - two_layer("embed.grid", INPUTS) - two_layer("embed.mesh", 3) - for graph in ("g2m", "multimesh", "m2g"): - two_layer(f"embed.{graph}", EDGE_FEATURES) - for net in ["encoder", "decoder"] + [ - f"processor.{r}" for r in range(rounds) - ]: - two_layer(f"{net}.edge", 3 * HIDDEN) - two_layer(f"{net}.node", 2 * HIDDEN) - return { - net: rng.normal(0, 1 / np.sqrt(shape[0]), shape) - for net, shape in shapes.items() - } - - -def params_table(params: dict[str, np.ndarray]) -> xr.Dataset: - """Every weight as a row of ``params(weight, net, f, g, w)``.""" - rows = [ - (net, f, g, w) - for net, matrix in params.items() - for (f, g), w in np.ndenumerate(matrix) - ] - net, f, g, w = map(np.array, zip(*rows)) - return xr.Dataset( - { - "net": ("weight", net), - "f": ("weight", f), - "g": ("weight", g), - "w": ("weight", w), - }, - coords={"weight": np.arange(len(rows))}, - ) +def matrix(host, table: str) -> np.ndarray: + """A stored ``(id, f, val)`` table as an (id, f) matrix.""" + rows = host.query(f"SELECT id, f, val FROM {table} ORDER BY id, f") + n = rows["id"][-1].as_py() + 1 + return rows["val"].to_numpy().reshape(n, -1) def main(): - parser = cli(__doc__.splitlines()[0], sizes=[2], size_name="refinements") - parser.add_argument( - "--rounds", type=int, default=2, help="processor rounds" - ) + parser = cli(__doc__.splitlines()[0], sizes=[1], size_name="rounds") args = parser.parse_args() - report = Report("graphcast", "refinements", args.reps) - rounds = args.rounds - era5 = open_era5() - - # The normalization: each variable's mean and spread at the inputs. - now, before, later = (era5_arrays(era5, t) for t in TIMES[[1, 0, 2]]) - mean = np.concatenate([before, now]).mean(axis=0) - std = np.concatenate([before, now]).std(axis=0) - scale = xr.Dataset( - {"mean": ("var", mean), "std": ("var", std)}, - coords={"var": np.arange(len(VARIABLES))}, - ) - params = init_params(rounds, np.random.default_rng(0)) - - for refinements in args.sizes: - tables = geometry(era5, refinements) - cells = tables["cells"] - print( - f" mesh refined {refinements}×: {tables['mesh'].sizes['node']} " - f"nodes, {tables['multimesh'].sizes['edge']} mesh edges, " - f"{tables['g2m'].sizes['edge']} grid-to-mesh, " - f"{tables['m2g'].sizes['edge']} mesh-to-grid; " - f"{len(params_table(params).weight)} weights" - ) + report = Report("graphcast", "rounds", args.reps) + + print(" loading DeepMind's checkpoint, example and graphs") + reference = gc.deepmind_stages() + era5 = gc.era5() + tables = { + "era5": era5, + "channels": gc.channels_table(), + "outputs": gc.outputs_table(), + **gc.graph_tables(), + **gc.param_tables(), + } - # The reference, in JAX, from arrays read with xarray. - graph = { - name: { - "edges": ( - tables[name].sender.values, - tables[name].receiver.values, - ), - "features": tables[name].val.values, - } - for name in ("g2m", "multimesh", "m2g") + for rounds in args.sizes: + mirror = gc.mirror(reference["features"], rounds=rounds) + expected = { + "grid_in": np.concatenate( + [reference["features"], gc.graphs()["grid_nodes"]], axis=1 + ), + "grid_0": reference["grid"], + "mesh_0": reference["mesh"], + **{f"mesh_{r + 1}": mirror[f"round_{r}"] for r in range(rounds)}, + "output": reference["output"], } - graph["mesh_in"] = np.stack( - [ - np.cos(np.radians(tables["mesh"].latitude.values)), - np.sin(np.radians(tables["mesh"].longitude.values)), - np.cos(np.radians(tables["mesh"].longitude.values)), - ], - axis=1, - ) - position = np.stack( - [ - np.cos(np.radians(cells.latitude.values)), - np.sin(np.radians(cells.longitude.values)), - np.cos(np.radians(cells.longitude.values)), - ], - axis=1, - ) - grid_in = np.concatenate( - [(before - mean) / std, (now - mean) / std, position], axis=1 - ) - target = (later - now) / std - want = jax_forecast(params, graph, grid_in, rounds) - - want_grad = jax.grad(jax_error)(params, graph, grid_in, target, rounds) - want_grad = np.concatenate([want_grad[net].ravel() for net in params]) - for host in hosts(args.hosts): - host.load("era5", era5, table_names=ERA5_TABLES) for name, table in tables.items(): - options = ( - {"table_names": GRAPH_TABLES} - if "edge" in table.dims - else {} - ) + options = {} + if name == "era5": + options = {"table_names": ERA5_TABLES} + elif "edge" in table.dims: + options = {"table_names": gc.GRAPH_TABLES} host.load(name, table, **options) - host.load("scale", scale) - host.load("params", params_table(params)) - - statement = forecast_statement(rounds) - rows = report.time( - "forecast, global", host, refinements, host.query, statement - ) - got = rows["change"].to_numpy().reshape(len(cells.cell), -1) - # The rows are ordered by latitude then longitude, as cells are. - report.check(assert_close("forecast", got, want, rtol=1e-9)) - report.note("operators", host.operators(statement)) - - statement = forecast_statement(rounds, REGION) - rows = report.time( - "forecast, regional", host, refinements, host.query, statement - ) - in_region = np.logical_and.reduce( - [ - (cells[c] >= lo) & (cells[c] <= hi) - for c, (lo, hi) in BOUNDS.items() - ] - ) - regional = rows["change"].to_numpy().reshape(-1, len(VARIABLES)) - report.check(assert_close("regional", regional, want[in_region])) - report.note("operators", host.operators(statement)) - - step = "gradient of the error" - if not host.differentiates: - report.skip(step, host, refinements, "no ddx backend") - continue - try: + + for table, statement in program(era5, rounds): + if statement is None: + got = matrix(host, table) + # The reference runs in float32, the SQL in float64. + error = assert_close( + table, got, expected[table], rtol=1e-4, atol=1e-3 + ) + report.check(error, step=table) + continue + report.time(table, host, rounds, host.store, table, statement) + + if rounds == ROUNDS: rows = report.time( - step, - host, - refinements, - host.query, - gradient_statement(rounds), + "forecast", host, rounds, host.query, forecast(era5) + ) + want = flatten(gc.deepmind_forecast().isel(time=0)) + error = assert_close( + "forecast", rows["value"].to_numpy(), want, rtol=1e-4 ) - except DdxError as e: - # Recorded rather than raised: a gap in ddx, which this - # row tracks until ddx closes it. - reason = str(e).removeprefix("not supported by ddx-ad yet: ") - report.skip(step, host, refinements, f"ddx: {reason}") - continue - report.check( - assert_close("gradient", rows["w"].to_numpy(), want_grad) - ) - - print("\n receptive field of the regional forecast:") - for stage, n in receptive_field(host, rounds): - print(f" {stage:<34} {n:>6}") + report.check(error) report.finish(args.csv) diff --git a/benchmarks/ml/_bench.py b/benchmarks/ml/_bench.py index 5a3235e0..6898aafd 100644 --- a/benchmarks/ml/_bench.py +++ b/benchmarks/ml/_bench.py @@ -116,10 +116,11 @@ def attempt(self, step: str, host: Host, size: int, run: Callable, *args): self.skip(step, host, size, f"refused: {reason}") return None - def check(self, error: float) -> None: - """Mark the last row's result as checked, with its error (from - :func:`assert_close`).""" - self.rows[-1]["check"] = f"ok ({error:.0e})" + def check(self, error: float, step: str | None = None) -> None: + """Mark a row's result as checked, with its error (from + :func:`assert_close`): the last row, or the last row of ``step``.""" + row = next(r for r in reversed(self.rows) if step in (None, r["step"])) + row["check"] = f"ok ({error:.0e})" def skip(self, step: str, host: Host, size: int, why: str) -> None: """Record a step the host cannot run, and why.""" diff --git a/benchmarks/ml/_graphcast.py b/benchmarks/ml/_graphcast.py new file mode 100644 index 00000000..7dbf1ecb --- /dev/null +++ b/benchmarks/ml/_graphcast.py @@ -0,0 +1,453 @@ +"""GraphCast_small, as DeepMind publishes it, for the SQL demo to match. + +Everything here is DeepMind's code or data, or a direct transcription of +it into tables: + +* :func:`checkpoint`, :func:`example` and :func:`stats` fetch the + published weights, a 1° ERA5 example and the normalization statistics + from DeepMind's public bucket, cached in ``~/.cache/xarray-sql``; +* :func:`graphs` builds the three graphs with DeepMind's own code: which + cells and mesh nodes are joined, and their geometric features; +* :func:`deepmind_stages` runs DeepMind's model, in float32, and returns + the output of each stage: the reference the SQL must match; +* :func:`mirror` is the same model written plainly in JAX, from the same + weights. It is checked against :func:`deepmind_stages`, and it can stop + after any processor round, which DeepMind's code cannot, so each SQL + stage can be checked on its own; +* :func:`params_table`, :func:`channels_table` and :func:`graph_tables` + turn the weights, the input channels and the graphs into Datasets that + xarray-sql registers as tables. +""" + +from __future__ import annotations + +import dataclasses +import functools +import pathlib + +import gcsfs +import haiku as hk +import jax +import jax.numpy as jnp +import numpy as np +import xarray as xr +from weathernext.utils import checkpoint as checkpoint_lib +from weathernext.utils import data_utils, normalization +from weathernext.weathernext1_graph import graphcast + +CACHE = pathlib.Path.home() / ".cache" / "xarray-sql" / "graphcast" +BUCKET = "dm_graphcast/graphcast" +PARAMS = ( + "params/GraphCast_small - ERA5 1979-2015 - resolution 1.0 - " + "pressure levels 13 - mesh 2to5 - precipitation input and output.npz" +) +EXAMPLE = "dataset/source-era5_date-2022-01-01_res-1.0_levels-13_steps-01.nc" +STATS = { + "diffs_stddev": "stats/diffs_stddev_by_level.nc", + "mean": "stats/mean_by_level.nc", + "stddev": "stats/stddev_by_level.nc", +} + + +def fetch(path: str) -> pathlib.Path: + """A file from DeepMind's bucket, downloaded once.""" + local = CACHE / pathlib.Path(path).name + if not local.exists(): + CACHE.mkdir(parents=True, exist_ok=True) + gcsfs.GCSFileSystem(token="anon").get(f"{BUCKET}/{path}", str(local)) + return local + + +@functools.cache +def checkpoint() -> graphcast.CheckPoint: + with open(fetch(PARAMS), "rb") as f: + return checkpoint_lib.load(f, graphcast.CheckPoint) + + +@functools.cache +def stats() -> dict[str, xr.Dataset]: + return { + name: xr.load_dataset(fetch(path)).compute() + for name, path in STATS.items() + } + + +@functools.cache +def example() -> tuple[xr.Dataset, xr.Dataset, xr.Dataset]: + """The inputs, the 6-hour target and the forcings, as DeepMind's + notebook extracts them from the example, without the batch dim.""" + data = xr.load_dataset(fetch(EXAMPLE), decode_timedelta=True).compute() + task = dataclasses.asdict(checkpoint().task_config) + inputs, targets, forcings = data_utils.extract_inputs_targets_forcings( + data, target_lead_times=slice("6h", "6h"), **task + ) + return tuple( + ds.isel(batch=0).drop_vars("datetime", errors="ignore") + for ds in (inputs, targets, forcings) + ) + + +# -------------------------------------------------------------------------- +# DeepMind's model: the graphs it builds and the output of each stage. +# -------------------------------------------------------------------------- + + +@functools.cache +def _model() -> graphcast.GraphCast: + """A GraphCast with its graphs built for the example's grid.""" + ckpt = checkpoint() + built = {} + + def build(): + model = graphcast.GraphCast(ckpt.model_config, ckpt.task_config) + model._maybe_init(example()[0]) + built["model"] = model + + hk.transform(build).init(None) + return built["model"] + + +@functools.cache +def graphs() -> dict: + """The grid-to-mesh, mesh and mesh-to-grid graphs: their edges' + senders and receivers, and the nodes' and edges' features.""" + model = _model() + out = {} + for name, graph in ( + ("grid2mesh", model._grid2mesh_graph_structure), + ("mesh", model._mesh_graph_structure), + ("mesh2grid", model._mesh2grid_graph_structure), + ): + (edges,) = graph.edges.values() + out[name] = { + "senders": np.asarray(edges.indices.senders), + "receivers": np.asarray(edges.indices.receivers), + "features": np.asarray(edges.features), + } + nodes = model._grid2mesh_graph_structure.nodes + out["grid_nodes"] = np.asarray(nodes["grid_nodes"].features) + out["mesh_nodes"] = np.asarray(nodes["mesh_nodes"].features) + out["grid_lat"] = model._grid_nodes_lat + out["grid_lon"] = model._grid_nodes_lon + out["mesh_lat"] = model._mesh_nodes_lat + out["mesh_lon"] = model._mesh_nodes_lon + return out + + +def _normalized(inputs, forcings): + s = stats() + return ( + normalization.normalize(inputs, s["stddev"], s["mean"]), + normalization.normalize(forcings, s["stddev"], s["mean"]), + ) + + +@functools.cache +def deepmind_stages() -> dict[str, np.ndarray]: + """DeepMind's model, in float32, stage by stage: the grid features it + reads, the latent grid and mesh after the encoder, the mesh after the + processor, and the decoder's normalized output.""" + ckpt = checkpoint() + inputs, _, forcings = example() + + @hk.transform_with_state + def stages(inputs, forcings): + model = graphcast.GraphCast(ckpt.model_config, ckpt.task_config) + inputs, forcings = _normalized(inputs, forcings) + model._maybe_init(inputs) + features = model._inputs_to_grid_node_features( + inputs.expand_dims("batch"), forcings.expand_dims("batch") + ) + mesh, grid = model._run_grid2mesh_gnn(features) + processed = model._run_mesh_gnn(mesh) + output = model._run_mesh2grid_gnn(processed, grid) + return dict( + features=features, grid=grid, mesh=mesh, + processed=processed, output=output, + ) + + out, _ = jax.jit(stages.apply)( + ckpt.params, {}, None, + inputs.astype(np.float32), forcings.astype(np.float32), + ) + return {k: np.asarray(v)[:, 0] for k, v in out.items()} + + +@functools.cache +def deepmind_forecast() -> xr.Dataset: + """DeepMind's 6-hour forecast, in float32, with its normalization and + residual: what the notebook's ``run_forward`` returns, without the + bfloat16 cast.""" + ckpt = checkpoint() + inputs, targets, forcings = example() + s = stats() + + @hk.transform_with_state + def forward(inputs, targets, forcings): + model = normalization.InputsAndResiduals( + graphcast.GraphCast(ckpt.model_config, ckpt.task_config), + diffs_stddev_by_level=s["diffs_stddev"], + mean_by_level=s["mean"], + stddev_by_level=s["stddev"], + ) + return model(inputs, targets_template=targets, forcings=forcings) + + batch = [ds.expand_dims("batch").astype(np.float32) + for ds in (inputs, targets * np.nan, forcings)] + out, _ = forward.apply(ckpt.params, {}, None, *batch) + return out.isel(batch=0).compute() + + +# -------------------------------------------------------------------------- +# The same model, written plainly: the oracle for each stage. +# -------------------------------------------------------------------------- + + +def weights(net: str) -> dict[str, np.ndarray]: + """One network's parameters, by DeepMind's name for it, e.g. + ``mesh_gnn/processor_edges_3_mesh``.""" + params = checkpoint().params + gnn, block = net.split("/") + prefix = f"{gnn}/~_networks_builder/{block}" + out = {} + for layer in ("linear_0", "linear_1"): + for p, v in params[f"{prefix}_mlp/~/{layer}"].items(): + out[f"{layer}:{p}"] = np.asarray(v, np.float64) + if f"{prefix}_layer_norm" in params: + for p, v in params[f"{prefix}_layer_norm"].items(): + out[f"layer_norm:{p}"] = np.asarray(v, np.float64) + return out + + +def mlp(net: str, x): + """Linear, swish, linear, then a layer norm, as DeepMind's MLPs are.""" + p = weights(net) + h = jax.nn.swish(x @ p["linear_0:w"] + p["linear_0:b"]) + y = h @ p["linear_1:w"] + p["linear_1:b"] + if "layer_norm:scale" not in p: + return y + mean = y.mean(-1, keepdims=True) + var = ((y - mean) ** 2).mean(-1, keepdims=True) + return (y - mean) / jnp.sqrt(var + 1e-5) * p["layer_norm:scale"] + p[ + "layer_norm:offset" + ] + + +def interaction(gnn, step, edge_set, graph, edges, senders, receivers, node_set): + """One message-passing step: new edges, then the receivers' update, + each with its residual. Returns the new edges and receivers.""" + s, r = graph["senders"], graph["receivers"] + message = mlp( + f"{gnn}/processor_edges_{step}_{edge_set}", + jnp.concatenate([edges, senders[s], receivers[r]], axis=1), + ) + total = jax.ops.segment_sum(message, r, len(receivers)) + update = mlp( + f"{gnn}/processor_nodes_{step}_{node_set}", + jnp.concatenate([receivers, total], axis=1), + ) + return edges + message, receivers + update + + +def mirror(features: np.ndarray, rounds: int | None = None) -> dict: + """The model from the grid features DeepMind's model reads, in float64. + With ``rounds``, it stops after that many processor rounds.""" + g = graphs() + x = jnp.asarray(features, jnp.float64) + grid_in = jnp.concatenate([x, g["grid_nodes"]], axis=1) + mesh_in = jnp.concatenate( + [jnp.zeros((len(g["mesh_nodes"]), x.shape[1])), g["mesh_nodes"]], axis=1 + ) + grid = mlp("grid2mesh_gnn/encoder_nodes_grid_nodes", grid_in) + mesh = mlp("grid2mesh_gnn/encoder_nodes_mesh_nodes", mesh_in) + edges = mlp("grid2mesh_gnn/encoder_edges_grid2mesh", g["grid2mesh"]["features"]) + _, mesh = interaction( + "grid2mesh_gnn", 0, "grid2mesh", g["grid2mesh"], edges, grid, mesh, + "mesh_nodes", + ) + grid = grid + mlp("grid2mesh_gnn/processor_nodes_0_grid_nodes", grid) + out = {"grid": grid, "mesh": mesh} + + edges = mlp("mesh_gnn/encoder_edges_mesh", g["mesh"]["features"]) + total = checkpoint().model_config.gnn_msg_steps + for step in range(total if rounds is None else rounds): + edges, mesh = interaction( + "mesh_gnn", step, "mesh", g["mesh"], edges, mesh, mesh, "mesh_nodes" + ) + out[f"round_{step}"] = mesh + if rounds is not None and rounds < total: + return out + out["processed"] = mesh + + edges = mlp("mesh2grid_gnn/encoder_edges_mesh2grid", g["mesh2grid"]["features"]) + _, grid = interaction( + "mesh2grid_gnn", 0, "mesh2grid", g["mesh2grid"], edges, mesh, grid, + "grid_nodes", + ) + out["output"] = mlp("mesh2grid_gnn/decoder_nodes_grid_nodes", grid) + return out + + +# -------------------------------------------------------------------------- +# Tables: the example, the input channels, the graphs and the weights. +# -------------------------------------------------------------------------- + + +def _hours(ds: xr.Dataset) -> xr.Dataset: + """``time`` as whole hours from the forecast's start.""" + return ds.assign_coords(time=(ds.time / np.timedelta64(1, "h")).astype(int)) + + +def era5() -> xr.Dataset: + """The model's inputs at -6 h and 0 h, and the forcings at +6 h.""" + inputs, _, forcings = example() + return xr.merge( + [_hours(inputs), _hours(forcings)], compat="no_conflicts", join="outer" + ) + + +def _stacking(ds: xr.Dataset, time_offset: int = 0) -> list[tuple]: + """DeepMind's channel order for ``ds``: its variables sorted by name, + each over its own dimensions other than lat and lon, as + ``(variable, time, level)``; ``level`` is 0 for a surface variable.""" + channels = [] + for name in sorted(ds.data_vars): + var = ds[name] + times = var.time.values if "time" in var.dims else [None] + levels = var.level.values if "level" in var.dims else [0] + for t in times: + for level in levels: + time = None if t is None else int(t / np.timedelta64(1, "h")) + channels.append((name, time, int(level))) + return channels + + +def channels_table() -> xr.Dataset: + """Each input channel: its variable, time and level, and the mean and + standard deviation that normalize it.""" + inputs, _, forcings = example() + rows = _stacking(inputs) + _stacking(forcings) + s = stats() + + def stat(kind, name, level): + value = s[kind][name] + return float(value.sel(level=level) if "level" in value.dims else value) + + return xr.Dataset( + { + "variable": ("channel", [r[0] for r in rows]), + # Static variables have no time; -1 never matches a row's time. + "time": ("channel", [-1 if r[1] is None else r[1] for r in rows]), + "level": ("channel", [r[2] for r in rows]), + "mean": ("channel", [stat("mean", r[0], r[2]) for r in rows]), + "std": ("channel", [stat("stddev", r[0], r[2]) for r in rows]), + }, + coords={"channel": np.arange(len(rows))}, + ) + + +def outputs_table() -> xr.Dataset: + """Each output channel: its variable and level, and the standard + deviation of its 6-hour change, which turns the output into one.""" + _, targets, _ = example() + rows = _stacking(targets.isel(time=0, drop=True)) + diffs = stats()["diffs_stddev"] + + def std(name, level): + value = diffs[name] + return float(value.sel(level=level) if "level" in value.dims else value) + + return xr.Dataset( + { + "variable": ("channel", [r[0] for r in rows]), + "level": ("channel", [r[2] for r in rows]), + "diff_std": ("channel", [std(r[0], r[2]) for r in rows]), + }, + coords={"channel": np.arange(len(rows))}, + ) + + +def _long(name: str, values: np.ndarray, key: str) -> xr.Dataset: + """A (key, feature) matrix as the table ``(key, f, val)``.""" + return xr.Dataset( + {"val": ((key, "f"), values.astype(np.float64))}, + coords={key: np.arange(len(values)), "f": np.arange(values.shape[1])}, + ) + + +GRAPH_TABLES = {("edge",): "edges", ("edge", "f"): "features"} + + +def graph_tables() -> dict[str, xr.Dataset]: + """The grid, the mesh, and the three graphs, as Datasets. Grid cells + are numbered as DeepMind numbers them: by latitude, then longitude.""" + g = graphs() + tables = { + "cells": xr.Dataset( + {"lat": ("cell", g["grid_lat"]), "lon": ("cell", g["grid_lon"])}, + coords={"cell": np.arange(len(g["grid_lat"]))}, + ), + "cell_features": _long("cell_features", g["grid_nodes"], "cell"), + "node_features": _long("node_features", g["mesh_nodes"], "node"), + } + for name in ("grid2mesh", "mesh", "mesh2grid"): + graph = g[name] + tables[name] = xr.Dataset( + { + "sender": ("edge", graph["senders"]), + "receiver": ("edge", graph["receivers"]), + "val": (("edge", "f"), graph["features"].astype(np.float64)), + }, + coords={ + "edge": np.arange(len(graph["senders"])), + "f": np.arange(graph["features"].shape[1]), + }, + ) + return tables + + +def _short(module: str) -> str: + """DeepMind's module path without Haiku's scaffolding: + ``mesh_gnn/processor_edges_3_mesh_mlp/linear_0``.""" + return module.replace("~_networks_builder/", "").replace("~/", "") + + +def param_tables() -> dict[str, xr.Dataset]: + """Every parameter, long form: ``weights(net, f, g, w)``, + ``biases(net, g, b)``, ``norms(net, g, scale, offset)``, and + ``nets(net, name)`` naming each network by DeepMind's path.""" + params = checkpoint().params + names = sorted(_short(m) for m in params) + ids = {name: i for i, name in enumerate(names)} + weights, biases, norms = [], [], [] + for module, p in params.items(): + net = ids[_short(module)] + if "w" in p: + f, g = np.indices(p["w"].shape).reshape(2, -1) + weights.append((np.full(f.size, net), f, g, p["w"].ravel())) + biases.append((np.full(p["b"].size, net), np.arange(p["b"].size), p["b"])) + else: + norms.append( + (np.full(p["scale"].size, net), np.arange(p["scale"].size), + p["scale"], p["offset"]) + ) + + def table(dim, columns, parts): + arrays = [np.concatenate(c) for c in zip(*parts)] + return xr.Dataset( + { + name: (dim, a.astype(np.float64 if a.dtype.kind == "f" else np.int32)) + for name, a in zip(columns, arrays) + }, + coords={dim: np.arange(len(arrays[0]))}, + ) + + return { + "nets": xr.Dataset( + {"name": ("net", np.array(names, dtype=object))}, + coords={"net": np.arange(len(names), dtype=np.int32)}, + ), + "weights": table("weight", ("net", "f", "g", "w"), weights), + "biases": table("bias", ("net", "g", "b"), biases), + "norms": table("norm", ("net", "g", "scale", "offset"), norms), + } diff --git a/benchmarks/ml/_hosts.py b/benchmarks/ml/_hosts.py index d5bcadcf..44937581 100644 --- a/benchmarks/ml/_hosts.py +++ b/benchmarks/ml/_hosts.py @@ -5,7 +5,7 @@ plan its *host*. Every demo writes its model once, as SQL, and runs that same text on each host, so a difference in time is a difference between engines, never between -queries. A host does five things: +queries. A host does these things: * ``load(name, data, **options)`` registers an ``xr.Dataset`` as the table ``name`` through xarray-sql, replacing any table of that name @@ -17,7 +17,10 @@ differentiates, shares the backward pass between those that take ``grad`` of the same loss; * ``operators(sql)`` counts the operators in the physical plan it would - run, which shows, for one, whether a CTE read twice is computed twice. + run, which shows, for one, whether a CTE read twice is computed twice; +* ``store(name, sql)`` runs a statement and keeps its rows as the table + ``name``, and ``drop(name)`` removes it: a model too large for one query + runs as a program of stored stages. Forward queries run everywhere. Gradients (``grad``, ``vjp``, ``jvp``) need ddx (https://github.com/xqlsystems/ddx), which differentiates SQL @@ -67,6 +70,12 @@ def query_all(self, sqls: Sequence[str]) -> list[pa.Table]: def operators(self, sql: str) -> int: raise NotImplementedError + def store(self, name: str, sql: str) -> None: + raise NotImplementedError + + def drop(self, name: str) -> None: + raise NotImplementedError + class DataFusion(Host): """xarray-sql's ``Context``, with ddx's gradients (the ``ml`` extra).""" @@ -95,6 +104,14 @@ def operators(self, sql): plan = self.ctx.sql(sql).execution_plan().display_indent() return len(plan.splitlines()) + def store(self, name, sql): + partitions = self.ctx.sql(sql).collect_partitioned() + self.ctx.deregister_table(name) + self.ctx.register_record_batches(name, partitions) + + def drop(self, name): + self.ctx.deregister_table(name) + class DuckDB(Host): """DuckDB, reading the same Datasets through xarray-sql's adapter, with @@ -133,6 +150,12 @@ def operators(self, sql): (_, plan), *_ = self.con.sql(f"EXPLAIN (FORMAT JSON) {sql}").fetchall() return sum(_count(node) for node in json.loads(plan)) + def store(self, name, sql): + self.con.sql(f"CREATE OR REPLACE TABLE {name} AS {sql}") + + def drop(self, name): + self.con.sql(f"DROP TABLE IF EXISTS {name}") + def _count(node: dict) -> int: return 1 + sum(_count(child) for child in node.get("children", [])) From a55e5a34a12aa3571476d12ac038b586b20ff75e Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sat, 17 Oct 2026 16:52:06 -0700 Subject: [PATCH 3/5] Give each host a memory limit and a spill directory Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h --- benchmarks/ml/01_attention.py | 2 +- benchmarks/ml/02_nanogpt.py | 2 +- benchmarks/ml/03_graphcast.py | 2 +- benchmarks/ml/_bench.py | 12 +++++++++--- benchmarks/ml/_hosts.py | 32 ++++++++++++++++++++++++++++---- 5 files changed, 40 insertions(+), 10 deletions(-) diff --git a/benchmarks/ml/01_attention.py b/benchmarks/ml/01_attention.py index 30ede227..5c435adf 100644 --- a/benchmarks/ml/01_attention.py +++ b/benchmarks/ml/01_attention.py @@ -185,7 +185,7 @@ def main(): want = attention(q, k, v) want_grads = jax.grad(loss, argnums=(0, 1, 2))(q, k, v) - for host in hosts(args.hosts): + for host in hosts(args.hosts, args.memory): for name, data in zip("qkv", (q, k, v)): host.load(name, table(data)) diff --git a/benchmarks/ml/02_nanogpt.py b/benchmarks/ml/02_nanogpt.py index b0cd15f3..cf7c203e 100644 --- a/benchmarks/ml/02_nanogpt.py +++ b/benchmarks/ml/02_nanogpt.py @@ -569,7 +569,7 @@ def main(): ) ) - for host in hosts(args.hosts): + for host in hosts(args.hosts, args.memory): reference = dict(init) moments = { n: (np.zeros_like(w), np.zeros_like(w)) for n, w in init.items() diff --git a/benchmarks/ml/03_graphcast.py b/benchmarks/ml/03_graphcast.py index 379aaf08..0ef62ac4 100644 --- a/benchmarks/ml/03_graphcast.py +++ b/benchmarks/ml/03_graphcast.py @@ -414,7 +414,7 @@ def main(): **{f"mesh_{r + 1}": mirror[f"round_{r}"] for r in range(rounds)}, "output": reference["output"], } - for host in hosts(args.hosts): + for host in hosts(args.hosts, args.memory): for name, table in tables.items(): options = {} if name == "era5": diff --git a/benchmarks/ml/_bench.py b/benchmarks/ml/_bench.py index 6898aafd..af6b209d 100644 --- a/benchmarks/ml/_bench.py +++ b/benchmarks/ml/_bench.py @@ -50,13 +50,19 @@ def cli(description: str, sizes: Sequence[int], size_name: str = "n"): default=1, help="timed repetitions per step; the median is reported", ) + parser.add_argument( + "--memory", + type=float, + default=6, + help="each host's memory limit, in GB, past which it spills to disk", + ) parser.add_argument("--csv", help="also write the results table here") return parser -def hosts(names: Sequence[str]) -> list[Host]: - """A fresh instance of each named host.""" - return [HOSTS[name]() for name in names] +def hosts(names: Sequence[str], memory: float = 6) -> list[Host]: + """A fresh instance of each named host, limited to ``memory`` GB.""" + return [HOSTS[name](int(memory * 2**30)) for name in names] def assert_close(name: str, got, want, *, rtol=1e-9, atol=1e-9) -> float: diff --git a/benchmarks/ml/_hosts.py b/benchmarks/ml/_hosts.py index 44937581..3929b24b 100644 --- a/benchmarks/ml/_hosts.py +++ b/benchmarks/ml/_hosts.py @@ -30,6 +30,10 @@ A gradient that ddx (or the engine under it) refuses to plan raises one of the host's ``refusals``, and the demos report that step with the reason. +Every host gets the same memory limit, in bytes, and spills to +``~/.cache/xarray-sql/spill`` past it, so a stage too large for memory +slows down rather than taking the machine down. + Adding a host (einfold's optimizer on DataFusion, or DuckDB with a GPU extension) is one class here and one entry in ``HOSTS``. """ @@ -37,17 +41,22 @@ from __future__ import annotations import json +import pathlib from collections.abc import Sequence import duckdb import pyarrow as pa import xarray as xr +from datafusion import RuntimeEnvBuilder from ddxdb import DdxError from ddxdb import duckdb as ddx_duckdb import xarray_sql as xql +SPILL = pathlib.Path.home() / ".cache" / "xarray-sql" / "spill" + + class Host: """An engine that runs the demos' SQL. Subclasses fill in the rest.""" @@ -55,6 +64,9 @@ class Host: differentiates = False refusals: tuple[type[Exception], ...] = () + def __init__(self, memory: int): + SPILL.mkdir(parents=True, exist_ok=True) + def load(self, name: str, data: xr.Dataset, **options) -> None: raise NotImplementedError @@ -84,8 +96,14 @@ class DataFusion(Host): differentiates = True refusals = (DdxError,) - def __init__(self): - self.ctx = xql.Context() + def __init__(self, memory): + super().__init__(memory) + runtime = ( + RuntimeEnvBuilder() + .with_disk_manager_specified(str(SPILL)) + .with_fair_spill_pool(memory) + ) + self.ctx = xql.Context(runtime=runtime) def load(self, name, data, **options): self.ctx.deregister_table(name) @@ -124,8 +142,14 @@ class DuckDB(Host): # plan it can't write. refusals = (DdxError, duckdb.NotImplementedException) - def __init__(self): - self.con = duckdb.connect() + def __init__(self, memory): + super().__init__(memory) + self.con = duckdb.connect( + config={ + "memory_limit": f"{memory // 2**20}MB", + "temp_directory": str(SPILL), + } + ) def load(self, name, data, **options): xql.register(self.con, name, data, **options) From b3d380bc2b88a1514186e68fb40fcce65656d098 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sat, 17 Oct 2026 19:12:47 -0700 Subject: [PATCH 4/5] Introduce GraphCast, ERA5, JAX and swish in the GraphCast docs Also fix the docstring's name for param_tables, and say that DeepMind's 15-second forecast includes JAX's compilation. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h --- benchmarks/ml/03_graphcast.py | 31 ++++++++++++++++++------------- benchmarks/ml/_graphcast.py | 8 +++++--- 2 files changed, 23 insertions(+), 16 deletions(-) diff --git a/benchmarks/ml/03_graphcast.py b/benchmarks/ml/03_graphcast.py index 0ef62ac4..9ed09594 100644 --- a/benchmarks/ml/03_graphcast.py +++ b/benchmarks/ml/03_graphcast.py @@ -11,14 +11,17 @@ # /// """Demo 3: GraphCast in SQL, with DeepMind's published weights. -GraphCast (Lam et al., 2023) forecasts the atmosphere 6 hours ahead from -the two previous states. This runs GraphCast_small, as DeepMind publishes -it (1° grid, 13 pressure levels, a mesh refined 5 times, 16 rounds of -message passing, 36 million weights), as a program of SQL statements over -tables xarray-sql registers: - -* ``era5``: DeepMind's example inputs (two states, 6 hours apart) and the - forcings at the forecast's time; +GraphCast (Lam et al., 2023, https://arxiv.org/abs/2212.12794) is +DeepMind's graph neural network for weather forecasting. It forecasts the +atmosphere 6 hours ahead from the two previous states. This runs +GraphCast_small, as DeepMind publishes it (1° grid, 13 pressure levels, a +mesh refined 5 times, 16 rounds of message passing, 36 million weights), +as a program of SQL statements over tables xarray-sql registers: + +* ``era5``: DeepMind's example inputs, two states 6 hours apart from ERA5 + (ECMWF's reanalysis of the atmosphere since 1940), and the forcings at + the forecast's time: inputs known in advance, such as the sunlight + reaching the top of the atmosphere; * ``weights``, ``biases``, ``norms`` and ``nets``: every parameter in DeepMind's checkpoint, one row each; * ``grid2mesh``, ``mesh`` and ``mesh2grid``: the three graphs, as tables of @@ -35,16 +38,18 @@ Every network is the same SQL: an MLP (linear, swish, linear, layer norm) over a table of ``(id, f, val)``, one row per node or edge and feature. -Each stage is one statement, and its result is stored as a table, so no -host computes a stage twice. +Swish is the activation ``x · sigmoid(x)``. Each stage is one statement, +and its result is stored as a table, so no host (an engine that runs the +SQL; see ``_hosts.py``) computes a stage twice. -The check is DeepMind's own model. Its JAX code, run on the same inputs, +The check is DeepMind's own model. Its code, in JAX (Google's library for +differentiable array programs), run on the same inputs, gives the stages' outputs; the SQL must match them. A plain JAX transcription of the model (``_graphcast.mirror``), checked against DeepMind's, also gives each processor round, so a run can stop after a few rounds and still be checked. A full forecast is about two trillion -multiply-adds, which hash joins take hours over; DeepMind's JAX takes 15 -seconds. +multiply-adds, which hash joins take hours over; DeepMind's JAX code +takes about 15 seconds on a 6-core CPU, including its compilation. Run ``uv run benchmarks/ml/03_graphcast.py --help`` for the options. ``--sizes`` is the number of processor rounds to run; with 16, the diff --git a/benchmarks/ml/_graphcast.py b/benchmarks/ml/_graphcast.py index 7dbf1ecb..90cbd1d3 100644 --- a/benchmarks/ml/_graphcast.py +++ b/benchmarks/ml/_graphcast.py @@ -4,8 +4,10 @@ it into tables: * :func:`checkpoint`, :func:`example` and :func:`stats` fetch the - published weights, a 1° ERA5 example and the normalization statistics - from DeepMind's public bucket, cached in ``~/.cache/xarray-sql``; + published weights, an example of the model's inputs from ERA5 (ECMWF's + reanalysis of the atmosphere) on a 1° grid, and the normalization + statistics from DeepMind's public bucket, cached in + ``~/.cache/xarray-sql``; * :func:`graphs` builds the three graphs with DeepMind's own code: which cells and mesh nodes are joined, and their geometric features; * :func:`deepmind_stages` runs DeepMind's model, in float32, and returns @@ -14,7 +16,7 @@ weights. It is checked against :func:`deepmind_stages`, and it can stop after any processor round, which DeepMind's code cannot, so each SQL stage can be checked on its own; -* :func:`params_table`, :func:`channels_table` and :func:`graph_tables` +* :func:`param_tables`, :func:`channels_table` and :func:`graph_tables` turn the weights, the input channels and the graphs into Datasets that xarray-sql registers as tables. """ From 5c39c17692d1dfb55658311a0c0487e2f0e062a7 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Sat, 17 Oct 2026 22:02:54 -0700 Subject: [PATCH 5/5] Format the GraphCast demo and its reference Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h --- benchmarks/ml/03_graphcast.py | 179 ++++++++++++++++++++++++---------- benchmarks/ml/_graphcast.py | 65 +++++++++--- 2 files changed, 177 insertions(+), 67 deletions(-) diff --git a/benchmarks/ml/03_graphcast.py b/benchmarks/ml/03_graphcast.py index 9ed09594..08dfc661 100644 --- a/benchmarks/ml/03_graphcast.py +++ b/benchmarks/ml/03_graphcast.py @@ -214,7 +214,9 @@ def grid_inputs(era5) -> str: {CELLS[table]} JOIN channels ch ON ch.variable = '{name}' AND {time} AND {level}""" ) - parts.append(f"SELECT cell AS id, f + {INPUTS} AS f, val FROM cell_features") + parts.append( + f"SELECT cell AS id, f + {INPUTS} AS f, val FROM cell_features" + ) return "\nUNION ALL\n".join(parts) @@ -247,8 +249,10 @@ def flatten(ds) -> np.ndarray: """A forecast Dataset in ``forecast``'s row order: by variable, level, latitude and longitude.""" return np.concatenate( - [ds[name].transpose(..., "lat", "lon").values.ravel() - for name in sorted(ds.data_vars)] + [ + ds[name].transpose(..., "lat", "lon").values.ravel() + for name in sorted(ds.data_vars) + ] ) @@ -260,7 +264,8 @@ def flatten(ds) -> np.ndarray: # A region to forecast: the contiguous United States, roughly. BOUNDS = {"lat": (25, 50), "lon": (235, 295)} REGION = " AND ".join( - f"c.{coord} BETWEEN {low} AND {high}" for coord, (low, high) in BOUNDS.items() + f"c.{coord} BETWEEN {low} AND {high}" + for coord, (low, high) in BOUNDS.items() ) @@ -270,17 +275,26 @@ def receptive_field(rounds: int): region: semi-joins of each graph's edges with what the next stage needs.""" yield "needed_cells", f"SELECT cell AS id FROM cells c WHERE {REGION}" - yield f"needed_mesh_{rounds}", """SELECT DISTINCT sender AS id -FROM mesh2grid.edges WHERE receiver IN (SELECT id FROM needed_cells)""" + yield ( + f"needed_mesh_{rounds}", + """SELECT DISTINCT sender AS id +FROM mesh2grid.edges WHERE receiver IN (SELECT id FROM needed_cells)""", + ) for r in reversed(range(rounds)): - yield f"needed_mesh_{r}", f"""SELECT id FROM needed_mesh_{r + 1} + yield ( + f"needed_mesh_{r}", + f"""SELECT id FROM needed_mesh_{r + 1} UNION SELECT sender AS id FROM mesh.edges -WHERE receiver IN (SELECT id FROM needed_mesh_{r + 1})""" - yield "needed_grid", """SELECT id FROM needed_cells +WHERE receiver IN (SELECT id FROM needed_mesh_{r + 1})""", + ) + yield ( + "needed_grid", + """SELECT id FROM needed_cells UNION SELECT sender AS id FROM grid2mesh.edges -WHERE receiver IN (SELECT id FROM needed_mesh_0)""" +WHERE receiver IN (SELECT id FROM needed_mesh_0)""", + ) def program(era5, rounds: int, regional: bool = False): @@ -298,11 +312,14 @@ def features(graph: str, edges: str) -> str: WHERE x.edge IN (SELECT edge FROM {edges}))""" checked = not regional - yield "grid_in", ( - f"SELECT * FROM ({grid_inputs(era5)}) g\n" - "WHERE id IN (SELECT id FROM needed_grid)" - if regional - else grid_inputs(era5) + yield ( + "grid_in", + ( + f"SELECT * FROM ({grid_inputs(era5)}) g\n" + "WHERE id IN (SELECT id FROM needed_grid)" + if regional + else grid_inputs(era5) + ), ) if checked: yield "grid_in", None @@ -312,25 +329,47 @@ def features(graph: str, edges: str) -> str: g2m_graph = within("grid2mesh.edges", "receiver", "needed_mesh_0") yield "grid", embed("grid_in", f"{g2m}/encoder_nodes_grid_nodes") # DeepMind gives the mesh nodes 183 zero inputs; zeros are absent rows. - yield "mesh", embed( - f"""(SELECT node AS id, f + {INPUTS} AS f, val + yield ( + "mesh", + embed( + f"""(SELECT node AS id, f + {INPUTS} AS f, val FROM {within("node_features", "node", "needed_mesh_0")} n)""", - f"{g2m}/encoder_nodes_mesh_nodes", + f"{g2m}/encoder_nodes_mesh_nodes", + ), ) - yield "g2m_edges", embed( - features("grid2mesh", g2m_graph), f"{g2m}/encoder_edges_grid2mesh" + yield ( + "g2m_edges", + embed( + features("grid2mesh", g2m_graph), f"{g2m}/encoder_edges_grid2mesh" + ), ) - yield "g2m_messages", messages( - f"{g2m}/processor_edges_0_grid2mesh", g2m_graph, - "g2m_edges", "grid", "mesh", + yield ( + "g2m_messages", + messages( + f"{g2m}/processor_edges_0_grid2mesh", + g2m_graph, + "g2m_edges", + "grid", + "mesh", + ), ) - yield "mesh_0", update( - f"{g2m}/processor_nodes_0_mesh_nodes", g2m_graph, - "mesh", "g2m_messages", + yield ( + "mesh_0", + update( + f"{g2m}/processor_nodes_0_mesh_nodes", + g2m_graph, + "mesh", + "g2m_messages", + ), ) - yield "grid_0", update( - f"{g2m}/processor_nodes_0_grid_nodes", g2m_graph, - within("grid", "id", "needed_cells"), None, + yield ( + "grid_0", + update( + f"{g2m}/processor_nodes_0_grid_nodes", + g2m_graph, + within("grid", "id", "needed_cells"), + None, + ), ) if checked: yield "mesh_0", None @@ -341,22 +380,36 @@ def features(graph: str, edges: str) -> str: within("mesh.edges", "receiver", f"needed_mesh_{r + 1}") for r in range(rounds) ] - yield "mesh_edges_0", embed( - features("mesh", graphs[0] if rounds else "mesh.edges"), - "mesh_gnn/encoder_edges_mesh", + yield ( + "mesh_edges_0", + embed( + features("mesh", graphs[0] if rounds else "mesh.edges"), + "mesh_gnn/encoder_edges_mesh", + ), ) for r in range(rounds): - yield f"mesh_messages_{r}", messages( - f"mesh_gnn/processor_edges_{r}_mesh", graphs[r], - f"mesh_edges_{r}", f"mesh_{r}", f"mesh_{r}", - ) - yield f"mesh_{r + 1}", update( - f"mesh_gnn/processor_nodes_{r}_mesh_nodes", graphs[r], - within(f"mesh_{r}", "id", f"needed_mesh_{r + 1}"), + yield ( f"mesh_messages_{r}", + messages( + f"mesh_gnn/processor_edges_{r}_mesh", + graphs[r], + f"mesh_edges_{r}", + f"mesh_{r}", + f"mesh_{r}", + ), + ) + yield ( + f"mesh_{r + 1}", + update( + f"mesh_gnn/processor_nodes_{r}_mesh_nodes", + graphs[r], + within(f"mesh_{r}", "id", f"needed_mesh_{r + 1}"), + f"mesh_messages_{r}", + ), ) - yield f"mesh_edges_{r + 1}", residual( - f"mesh_edges_{r}", f"mesh_messages_{r}" + yield ( + f"mesh_edges_{r + 1}", + residual(f"mesh_edges_{r}", f"mesh_messages_{r}"), ) if checked: yield f"mesh_{r + 1}", None @@ -366,20 +419,42 @@ def features(graph: str, edges: str) -> str: # cell's 83 output channels. m2g = "mesh2grid_gnn" m2g_graph = within("mesh2grid.edges", "receiver", "needed_cells") - yield "m2g_edges", embed( - features("mesh2grid", m2g_graph), f"{m2g}/encoder_edges_mesh2grid" + yield ( + "m2g_edges", + embed( + features("mesh2grid", m2g_graph), f"{m2g}/encoder_edges_mesh2grid" + ), ) - yield "m2g_messages", messages( - f"{m2g}/processor_edges_0_mesh2grid", m2g_graph, - "m2g_edges", f"mesh_{rounds}", "grid_0", + yield ( + "m2g_messages", + messages( + f"{m2g}/processor_edges_0_mesh2grid", + m2g_graph, + "m2g_edges", + f"mesh_{rounds}", + "grid_0", + ), ) - yield "grid_1", update( - f"{m2g}/processor_nodes_0_grid_nodes", m2g_graph, - "grid_0", "m2g_messages", + yield ( + "grid_1", + update( + f"{m2g}/processor_nodes_0_grid_nodes", + m2g_graph, + "grid_0", + "m2g_messages", + ), ) - yield "output", select( - mlp("out", "grid_1", f"{m2g}/decoder_nodes_grid_nodes", layer_norm=False), - "out", + yield ( + "output", + select( + mlp( + "out", + "grid_1", + f"{m2g}/decoder_nodes_grid_nodes", + layer_norm=False, + ), + "out", + ), ) if checked: yield "output", None diff --git a/benchmarks/ml/_graphcast.py b/benchmarks/ml/_graphcast.py index 90cbd1d3..cef6de7d 100644 --- a/benchmarks/ml/_graphcast.py +++ b/benchmarks/ml/_graphcast.py @@ -164,13 +164,19 @@ def stages(inputs, forcings): processed = model._run_mesh_gnn(mesh) output = model._run_mesh2grid_gnn(processed, grid) return dict( - features=features, grid=grid, mesh=mesh, - processed=processed, output=output, + features=features, + grid=grid, + mesh=mesh, + processed=processed, + output=output, ) out, _ = jax.jit(stages.apply)( - ckpt.params, {}, None, - inputs.astype(np.float32), forcings.astype(np.float32), + ckpt.params, + {}, + None, + inputs.astype(np.float32), + forcings.astype(np.float32), ) return {k: np.asarray(v)[:, 0] for k, v in out.items()} @@ -194,8 +200,10 @@ def forward(inputs, targets, forcings): ) return model(inputs, targets_template=targets, forcings=forcings) - batch = [ds.expand_dims("batch").astype(np.float32) - for ds in (inputs, targets * np.nan, forcings)] + batch = [ + ds.expand_dims("batch").astype(np.float32) + for ds in (inputs, targets * np.nan, forcings) + ] out, _ = forward.apply(ckpt.params, {}, None, *batch) return out.isel(batch=0).compute() @@ -235,7 +243,9 @@ def mlp(net: str, x): ] -def interaction(gnn, step, edge_set, graph, edges, senders, receivers, node_set): +def interaction( + gnn, step, edge_set, graph, edges, senders, receivers, node_set +): """One message-passing step: new edges, then the receivers' update, each with its residual. Returns the new edges and receivers.""" s, r = graph["senders"], graph["receivers"] @@ -262,9 +272,17 @@ def mirror(features: np.ndarray, rounds: int | None = None) -> dict: ) grid = mlp("grid2mesh_gnn/encoder_nodes_grid_nodes", grid_in) mesh = mlp("grid2mesh_gnn/encoder_nodes_mesh_nodes", mesh_in) - edges = mlp("grid2mesh_gnn/encoder_edges_grid2mesh", g["grid2mesh"]["features"]) + edges = mlp( + "grid2mesh_gnn/encoder_edges_grid2mesh", g["grid2mesh"]["features"] + ) _, mesh = interaction( - "grid2mesh_gnn", 0, "grid2mesh", g["grid2mesh"], edges, grid, mesh, + "grid2mesh_gnn", + 0, + "grid2mesh", + g["grid2mesh"], + edges, + grid, + mesh, "mesh_nodes", ) grid = grid + mlp("grid2mesh_gnn/processor_nodes_0_grid_nodes", grid) @@ -281,9 +299,17 @@ def mirror(features: np.ndarray, rounds: int | None = None) -> dict: return out out["processed"] = mesh - edges = mlp("mesh2grid_gnn/encoder_edges_mesh2grid", g["mesh2grid"]["features"]) + edges = mlp( + "mesh2grid_gnn/encoder_edges_mesh2grid", g["mesh2grid"]["features"] + ) _, grid = interaction( - "mesh2grid_gnn", 0, "mesh2grid", g["mesh2grid"], edges, mesh, grid, + "mesh2grid_gnn", + 0, + "mesh2grid", + g["mesh2grid"], + edges, + mesh, + grid, "grid_nodes", ) out["output"] = mlp("mesh2grid_gnn/decoder_nodes_grid_nodes", grid) @@ -427,18 +453,27 @@ def param_tables() -> dict[str, xr.Dataset]: if "w" in p: f, g = np.indices(p["w"].shape).reshape(2, -1) weights.append((np.full(f.size, net), f, g, p["w"].ravel())) - biases.append((np.full(p["b"].size, net), np.arange(p["b"].size), p["b"])) + biases.append( + (np.full(p["b"].size, net), np.arange(p["b"].size), p["b"]) + ) else: norms.append( - (np.full(p["scale"].size, net), np.arange(p["scale"].size), - p["scale"], p["offset"]) + ( + np.full(p["scale"].size, net), + np.arange(p["scale"].size), + p["scale"], + p["offset"], + ) ) def table(dim, columns, parts): arrays = [np.concatenate(c) for c in zip(*parts)] return xr.Dataset( { - name: (dim, a.astype(np.float64 if a.dtype.kind == "f" else np.int32)) + name: ( + dim, + a.astype(np.float64 if a.dtype.kind == "f" else np.int32), + ) for name, a in zip(columns, arrays) }, coords={dim: np.arange(len(arrays[0]))},