[alxmrs]

Stacked on #269. This replaces the miniature 03_graphcast.py with GraphCast_small as DeepMind publishes it:

• a 1° grid with 13 pressure levels;
• a mesh refined 5 times;
• 16 message-passing rounds;
• 36M weights.

GraphCast is DeepMind's graph neural network for weather forecasting. Here it's written as SQL views over tables registered through xarray-sql, one view per stage.

Status: work in progress. As stored tables, the encoder and one processor round have run end to end on einfold and match DeepMind's model. As views, it doesn't run past the encoder on any engine yet, for the upstream reasons below.

What's here

• _graphcast.py: everything on DeepMind's side.
  • It fetches the checkpoint, the 1° example from ERA5 (ECMWF's reanalysis of the atmosphere) and the normalization statistics from gs://dm_graphcast, cached in ~/.cache/xarray-sql/graphcast.
  • It builds the three graphs with DeepMind's own code: grid→mesh with 101,892 edges, mesh with 81,900, and mesh→grid with 195,480. It 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() and deepmind_forecast() run DeepMind's model in float32.
  • mirror() is a plain JAX transcription of the model, checked against DeepMind's. It can stop after any processor round, so the SQL can be checked one round at a time.
• 03_graphcast.py: the model in SQL.
  • Every network is the same mlp(): a matrix product, bias plus swish (x · sigmoid(x)), a second matrix product and bias, then a layer norm through OVER (PARTITION BY id) windows.
  • Each stage is a CREATE VIEW of the stages before it.
  • The checks and the forecast are plain queries of those views.
  • Inputs are normalized and stacked into DeepMind's 183 channels in SQL. The mesh nodes' zero inputs are absent rows.
  • receptive_field() and program(..., regional=True) write an exact regional forecast: the same stages, each limited by semi-joins to what the region needs. It hasn't run yet.
  • --sizes is the number of processor rounds. At 16, the forecast is checked against DeepMind's.
• Hosts take a memory limit (--memory, in GB), and spill to ~/.cache/xarray-sql/spill.

Measured

Measured on an Intel i7-8700 with einfold 0.1.0rc2, DataFusion 55.0.0, DuckDB 1.5.6 and ddxdb 0.4.6. Details are in #285 and xqlsystems/einfold#36.

• Stored stages on einfold (the form before this rewrite):
  • The encoder plus one round took about 10 minutes.
  • Every checked stage matched DeepMind's to 1e-6 or better.
  • A full forecast extrapolates to 45–50 minutes, against 15 s for DeepMind's JAX.
• Grid embedding, 2,048 cells:
  • einfold: 1.8–2.2 s, with its matrix products running about 25× faster than DataFusion's;
  • DataFusion: 7.5 s;
  • DuckDB: out of memory at 6 and 9 GB.
• As views:
  • DataFusion's plan grows 4.7× per round: 50,078 operators and 150 s to plan at round 3 (DataFusion plans a CTE once per reference, so ML queries grow exponentially with depth #267).
  • DuckDB's plan stays small, but planning takes 45 s for one round and more than 10 minutes for two (DuckDB plans views that share reads in time exponential in their depth — draft for upstream #292).

Next

1. Run the views once DataFusion computes each shared read once (DataFusion plans a CTE once per reference, so ML queries grow exponentially with depth #267).
2. Run the full 16-round forecast and compare it with deepmind_forecast().
3. Run the exact regional forecast, then an 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

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

Share

• X
• Mastodon
• Reddit
• LinkedIn