[alxmrs]

Stacked on #269. Replaces the miniature 03_graphcast.py with GraphCast_small as DeepMind publishes it: a 1° grid, 13 pressure levels, a mesh refined 5 times, 16 message-passing rounds and 36M weights. It runs as a program of SQL statements over tables registered through xarray-sql.

Status: work in progress. A full forecast has not run end to end yet.

What's here

• _graphcast.py: everything on DeepMind's side.
  • It fetches the checkpoint, the 1° ERA5 example and the normalization stats from gs://dm_graphcast, cached in ~/.cache/xarray-sql/graphcast.
  • It builds the three graphs with DeepMind's own code (grid→mesh 101,892 edges, mesh 81,900, mesh→grid 195,480) and exports them as Datasets.
  • It turns the checkpoint into long-form tables: weights(net, f, g, w), biases, norms, and nets(net, name).
  • deepmind_stages() runs DeepMind's model in float32 and returns each stage's output, and deepmind_forecast() returns the full forecast.
  • mirror() is a plain JAX transcription of the model. It matches DeepMind's stages to about 1e-6 of each stage's largest value (float32 against float64), and it can stop after any processor round, so the SQL can be checked a round at a time.
• 03_graphcast.py: the model in SQL.
  • Every network is the same mlp(): matmul, bias plus swish, matmul, bias, then layer norm through OVER (PARTITION BY id) windows.
  • Each stage is one statement whose result the host stores as a table (Host.store), so no host computes a CTE twice (DataFusion plans a CTE once per reference, so ML queries grow exponentially with depth #267).
  • Inputs are normalized and stacked into DeepMind's 183 channels in SQL. The mesh nodes' zero inputs are simply absent rows.
  • receptive_field() and program(..., regional=True) add an exact regional forecast: the same stages, each limited by semi-joins to what the region needs. It's written but hasn't run yet.
  • --sizes is the number of processor rounds; at 16, the forecast is checked against DeepMind's.
• Hosts now take a memory limit (--memory, in GB) and spill to ~/.cache/xarray-sql/spill.

What's been checked

• Input channels: grid_in matches DeepMind's stacked inputs.
• Grid embedding (2,048-cell slice): matches the reference to 9e-6, which is the float32 input precision.

What went wrong

• The first full encoder run was OOM-killed at 13.9 GB. That took the agent session down with it. The memory limits above are the fix, and from now on long runs go inside a systemd-run --scope -p MemoryMax=… cgroup.
• Throughput is far below the einfold S21 figure. On a 2,048-cell slice of the grid embedding, DataFusion does 0.09 billion joined pairs/s after moving swish out of the join, about half of S21's 0.19. A full forecast is about 2.1 trillion pairs, so roughly 6–7 hours on DataFusion. DeepMind's JAX takes 15 s.
• DuckDB is slow on these aggregations. Most of its time is in HASH_GROUP_BY summing ~200M rows into ~1M groups, and the restructured MLP made it slower: 78 s against DataFusion's 8 s on the same slice. I'll look into that before filing it.

Next

1. Find the memory-limit settings at which the encoder completes, and check it against DeepMind's.
2. Check one processor round against mirror().
3. Run the full 16-round forecast and compare it with deepmind_forecast().
4. Run the exact regional forecast, then the approximate one (ORDER BY distance LIMIT k), and plot its error.

🤖 Generated with Claude Code

https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}