From 48e65a160a5d172476ff79d96a21a779ce9b0eea Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Fri, 30 Oct 2026 19:23:31 -0700 Subject: [PATCH 1/2] Run GraphCast as views of its stages, on the current demos Each stage of GraphCast_small is now a view of the stages before it, and the checks and the forecast are plain queries of those views. This drops Host.store and Host.drop, which kept each stage as a table so no engine computed a CTE twice: the demos write the SQL a reader would, and making it fast is the engines' work. The hosts keep their memory limit and spill directory, which change no SQL. Rebased onto ml-demos at 6debe24: einfold is a host, with the memory limit like the others. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h --- benchmarks/ml/03_graphcast.py | 524 ++++++++++++++++++++++++++++++++++ benchmarks/ml/README.md | 16 +- benchmarks/ml/_bench.py | 12 +- benchmarks/ml/_graphcast.py | 490 +++++++++++++++++++++++++++++++ benchmarks/ml/_hosts.py | 35 ++- 5 files changed, 1063 insertions(+), 14 deletions(-) create mode 100644 benchmarks/ml/03_graphcast.py create mode 100644 benchmarks/ml/_graphcast.py diff --git a/benchmarks/ml/03_graphcast.py b/benchmarks/ml/03_graphcast.py new file mode 100644 index 00000000..0ddf3868 --- /dev/null +++ b/benchmarks/ml/03_graphcast.py @@ -0,0 +1,524 @@ +# /// script +# requires-python = ">=3.11" +# dependencies = [ +# "xarray-sql[ml,duckdb]", +# "weathernext @ git+https://github.com/google-deepmind/graphcast", +# "gcsfs", +# ] +# +# [tool.uv.sources] +# xarray-sql = { path = "../../", editable = true } +# /// +"""Demo 3: GraphCast in SQL, with DeepMind's published weights. + +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 SQL views 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 + 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. +Swish is the activation ``x · sigmoid(x)``. Each stage is a view of the +stages before it, and the forecast is one query of the last, which each +host (an engine that runs the SQL; see ``_hosts.py``) plans as a whole. + +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 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 +forecast is complete. +""" + +from __future__ import annotations + +import _graphcast as gc +import jax +import numpy as np +from _bench import Report, assert_close, cli, hosts + +jax.config.update("jax_enable_x64", True) + +INPUTS = 183 # input channels: the variables at each time and level +ROUNDS = 16 # processor rounds in GraphCast_small + +# -------------------------------------------------------------------------- +# The networks, in SQL. A state is a table of (id, f, val). +# -------------------------------------------------------------------------- + + +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 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 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 select(ctes: str, out: str) -> str: + return f"WITH {ctes.rstrip(',')}\nSELECT id, f, val FROM {out}" + + +def embed(x: str, net: str) -> str: + """A statement: an MLP of every row of ``x``.""" + return select(mlp("embedded", x, net), "embedded") + + +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", + ) + + +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", + ) + + +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""" + + +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 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}""" + ) + 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) + ] + ) + + +# -------------------------------------------------------------------------- +# The model: one view per stage. +# -------------------------------------------------------------------------- + + +# 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 receptive_field(rounds: int): + """The cells and mesh nodes a regional forecast needs, as + ``(view, query)``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 ``(view, query)``. 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}))""" + + 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) + ), + ) + # 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", + ), + ) + 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 ( + "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, + ), + ) + # 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}", + ), + ) + 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}"), + ) + 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" + ), + ) + 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 ( + "output", + select( + mlp( + "out", + "grid_1", + f"{m2g}/decoder_nodes_grid_nodes", + layer_norm=False, + ), + "out", + ), + ) + + +def matrix(host, view: str) -> np.ndarray: + """An ``(id, f, val)`` view as an (id, f) matrix.""" + rows = host.query(f"SELECT id, f, val FROM {view} 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=[1], size_name="rounds") + args = parser.parse_args() + 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(), + } + + 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"], + } + for host in hosts(args.hosts, args.memory): + for name, table in tables.items(): + 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) + + for view, query in program(era5, rounds): + host.execute(f"CREATE OR REPLACE VIEW {view} AS {query}") + + # Each check is one query of a stage's view: the whole model + # up to that stage. + checked = dict.fromkeys( + ["grid_in", "grid_0", "mesh_0", f"mesh_{rounds}"] + ) + for view in checked: + got = report.time(view, host, rounds, matrix, host, view) + # The reference runs in float32, the SQL in float64. + error = assert_close( + view, got, expected[view], rtol=1e-4, atol=1e-3 + ) + report.check(error) + + if rounds == ROUNDS: + rows = report.time( + "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 + ) + report.check(error) + + report.finish(args.csv) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/ml/README.md b/benchmarks/ml/README.md index ee632f96..96a0ec0d 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 the demos run on DataFusion with einfold as one of their engines. 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, against [FlashAttention](https://arxiv.org/abs/2205.14135)'s tiled plan, which einfold finds from the same SQL; gradients with respect to Q, K and V; masked pairs as absent rows; a control that isn't attention: the [Nadaraya–Watson estimator](https://en.wikipedia.org/wiki/Kernel_regression#Nadaraya%E2%80%93Watson_kernel_regression), an average of nearby points weighted by a Gaussian kernel, which einfold must fuse with the same rules | | 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) | [GraphCast](https://arxiv.org/abs/2212.12794)_small, DeepMind's graph neural network for weather forecasting, with DeepMind's published weights (36 million), as one view per stage, on DeepMind's example inputs from [ERA5](https://www.ecmwf.int/en/forecasts/dataset/ecmwf-reanalysis-v5), ECMWF's reanalysis of the atmosphere | each stage against DeepMind's own model in JAX, a processor round at a time; a 6-hour forecast at 16 rounds | 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,6 +39,7 @@ 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 --hosts einfold --sizes 1 ``` Each script carries [inline `uv` metadata](https://docs.astral.sh/uv/guides/scripts/) @@ -46,8 +48,9 @@ options are the same in each: ```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 rounds --reps 3 # timed repetitions; the median is reported +--memory 6 # each host's memory limit in GB, past which it spills --csv results.csv # also write the results table ``` @@ -55,7 +58,10 @@ 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). In demo 3, +`--sizes` is the number of processor rounds; at 16, the forecast is +complete. Its first run downloads DeepMind's checkpoint, example and +statistics (about 0.5 GB) into `~/.cache/xarray-sql/graphcast`. 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 diff --git a/benchmarks/ml/_bench.py b/benchmarks/ml/_bench.py index dcb044e8..519ce207 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/_graphcast.py b/benchmarks/ml/_graphcast.py new file mode 100644 index 00000000..cef6de7d --- /dev/null +++ b/benchmarks/ml/_graphcast.py @@ -0,0 +1,490 @@ +"""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, 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 + 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:`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 __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 b002b79b..3af94c57 100644 --- a/benchmarks/ml/_hosts.py +++ b/benchmarks/ml/_hosts.py @@ -27,6 +27,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 (DuckDB with a GPU extension, say) is one class here and one entry in ``HOSTS``. """ @@ -34,17 +38,21 @@ 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 from einfold.datafusion import EinFold as EinFoldExtension 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.""" @@ -53,6 +61,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 @@ -76,8 +87,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) @@ -111,8 +128,8 @@ class EinFold(DataFusion): name = "einfold" - def __init__(self): - super().__init__() + def __init__(self, memory): + super().__init__(memory) self.einfold = EinFoldExtension() self.ctx.with_extensions(self.einfold) @@ -128,8 +145,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 6505546ee2efb0c27c36203051b8b88385398412 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Fri, 30 Oct 2026 19:55:18 -0700 Subject: [PATCH 2/2] List the GraphCast demo, with no run small enough for CI Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h --- tests/test_ml_demos.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/test_ml_demos.py b/tests/test_ml_demos.py index b63ac452..44980ef5 100644 --- a/tests/test_ml_demos.py +++ b/tests/test_ml_demos.py @@ -24,6 +24,10 @@ ["--steps", "1"], ["--steps", "1", "--optimizer", "adamw"], ], + # GraphCast has no run small enough for CI: its smallest, the encoder + # and one processor round over the full grid, downloads DeepMind's + # checkpoint (about 0.5 GB) and takes minutes and over 10 GB. + "03_graphcast.py": [], }