How GraphCast Works
Architecture, data flow, and the assumptions you'll inherit.
Traditional weather forecasting runs physics simulations on supercomputers. GraphCast replaces the simulation with a graph neural network that predicts the next weather state from the current one — and does it in under a minute on a single TPU. What's harder to see, and what this analysis surfaces, is what graphcast assumes about the world it runs in: The model was trained and evaluated at 1° latitude-longitude resolution; the authors state higher resolution would likely yield better results and is needed for operational use.
What graphcast Does
Runs and trains GraphCast/GenCast neural weather forecast models autoregressively
This package implements two deep learning weather forecasting models: GraphCast (deterministic, GNN-based) and GenCast (probabilistic, diffusion-based). Both models take ERA5 reanalysis atmospheric state data as input and predict future atmospheric states, with GraphCast producing single deterministic forecasts and GenCast producing ensemble samples via iterative denoising.
What graphcast Assumes But Doesn't Validate
Every system carries assumptions it never checks. CodeSea surfaced 14 in graphcast, 4 of them critical. The most consequential: The model was trained and evaluated at 1° latitude-longitude resolution; the authors state higher resolution would likely yield better results and is needed for operational use. The consequence when it breaks: Forecasts at 1° are coarser than current operational ensembles; running or interpreting outputs as operationally-ready high-resolution forecasts overstates fidelity relative to systems operating at finer resolutions.
Another the code takes for granted: GenCast does not generate its own analysis; it relies on initial conditions from a traditional NWP ensemble data assimilation system to seed the forecast. Without ensemble analysis (e.g. ERA5 EDA perturbations) as input, the model cannot represent initial-condition uncertainty and the ensemble spread/reliability results do not transfer.
These cluster around Scale, Contract, Domain, Temporal, Ordering, Environment, Resource: the dimensions most likely to shift as graphcast grows or its runtime changes. None are bugs today; they are where a future change surfaces one. See all 14 hidden assumptions →
Architecture Overview
graphcast is organized into 7 layers, with 8 components and 0 connections between them.
How Data Flows Through graphcast
Weather state data enters as xarray Datasets containing atmospheric variables (temperature, wind, geopotential, etc.) on a regular lat/lon grid at multiple pressure levels. The data is normalized by subtracting historical means and dividing by historical standard deviations. For multi-step predictions the autoregressive wrapper feeds outputs back as inputs via jax.lax.scan. Inside each one-step call, grid node features are assembled into flat vectors and projected onto icosahedral mesh nodes (grid→mesh encoding), then 16 rounds of GNN message passing propagate information across the mesh. The updated mesh node features are projected back to grid nodes (mesh→grid decoding), residually added to the input state, and denormalized to produce the final predicted atmospheric state. For GenCast, this entire encoder-decoder is used as a denoiser inside a DPM-Solver++ 2S reverse diffusion loop that iterates from max noise to near-zero noise to generate a probabilistic ensemble member.
1Load and normalize ERA5 inputs
Raw ERA5 xarray Datasets (inputs with 2 timesteps, targets, forcings like solar radiation) are passed to the Normalizer wrapper in normalization.py. Inputs are normalized by subtracting per-variable historical means and dividing by per-variable standard deviations. Target variables are normalized by their historical time-difference standard deviations (since the model predicts increments). Forcings and static variables use their own normalization statistics.
Config: TaskConfig.input_variables, TaskConfig.target_variables, TaskConfig.forcing_variables
2Cast inputs to bfloat16
The Bfloat16Cast wrapper in casting.py calls tree_map_cast to convert all floating-point arrays in the normalized input/target/forcing Datasets from float32 to bfloat16, and sets up a bfloat16_variable_view context so Haiku model weights are also viewed as bfloat16 during the forward pass, halving memory usage.
3Build icosahedral mesh and grid-mesh edges
At model initialization (not per-step), get_hierarchy_of_triangular_meshes_for_sphere in icosahedral_mesh.py creates a multi-resolution triangular mesh with mesh_size refinement splits (~40,962 vertices at split=5). radius_query_indices in grid_mesh_connectivity.py then uses a scipy cKDTree to find all grid lat/lon points within radius of each mesh vertex, producing the grid-to-mesh edge index arrays. in_mesh_triangle_indices finds which mesh triangle each grid point falls inside for mesh-to-grid edges.
Config: ModelConfig.mesh_size, ModelConfig.resolution, ModelConfig.radius_query_fraction_edge_length
4Assemble grid and mesh node features
In graphcast.py's _run_grid2mesh_gnn, model_utils.get_graph_spatial_features computes positional features for grid nodes (lat/lon as sin/cos, 3D unit-sphere position) and edge features (relative 3D displacement between connected nodes). xarray_jax.unwrap extracts the underlying JAX arrays from the xarray Dataset and stacks all input variables (across time and pressure levels) into a flat feature vector per grid node. These are assembled into a TypedGraph with grid-nodes, mesh-nodes, and grid-to-mesh edges.
Config: ModelConfig.latent_size, TaskConfig.pressure_levels
5Encode grid features to mesh nodes (Grid→Mesh GNN)
A DeepTypedGraphNet (in graphcast.py's _run_grid2mesh_gnn) applies a single round of message passing: each grid-to-mesh edge aggregates its sender (grid node) features, and each mesh node aggregates incoming edge messages to build an initial mesh node representation. This is a learned projection that maps the high-resolution grid signal onto the lower-dimensional mesh.
Config: ModelConfig.latent_size, ModelConfig.hidden_layers
6Run GNN message passing on mesh (Mesh→Mesh GNN)
A deeper DeepTypedGraphNet (_run_mesh_gnn in graphcast.py) with gnn_msg_steps=16 message-passing iterations operates on mesh-only nodes and edges. At each step, edge MLPs update edge features from sender+receiver node features, then node MLPs update node features by aggregating all incoming edge messages. With residual connections at each step, information propagates globally across the multi-resolution mesh.
Config: ModelConfig.gnn_msg_steps, ModelConfig.latent_size, ModelConfig.hidden_layers
7Decode mesh node outputs to grid (Mesh→Grid GNN)
A third DeepTypedGraphNet (_run_mesh2grid_gnn in graphcast.py) propagates updated mesh node features back to grid nodes via mesh-to-grid edges. The final grid node features are then split by variable and pressure level in model_utils.stacked_to_dataset, reshaping flat vectors back into the (lat, lon, level) structure of an xarray Dataset representing predicted variable increments.
Config: ModelConfig.latent_size, ModelConfig.hidden_layers, TaskConfig.target_variables
8Apply residual and denormalize predictions
The model predicts state increments (differences), not absolute values. The Normalizer wrapper in normalization.py adds the predicted increments back to the most recent input timestep and then multiplies by the historical time-difference standard deviations to denormalize, producing absolute atmospheric state values in the original physical units.
Config: TaskConfig.target_variables
9Autoregressive rollout (training) or GenCast diffusion sampling
For GraphCast training: autoregressive.py's Predictor uses jax.lax.scan to loop over target timesteps, feeding each step's predictions back as inputs via _get_flat_arrays_and_single_timestep_treedef and _unflatten_and_expand_time, accumulating per-timestep losses. For GenCast inference: dpm_solver_plus_plus_2s.py's Sampler iterates from max_noise_level down to min_noise_level across num_noise_levels steps, calling the denoiser twice per step (second-order correction) and optionally adding stochastic churn noise between steps.
Config: SamplerConfig.num_noise_levels, SamplerConfig.max_noise_level, SamplerConfig.min_noise_level
10Compute latitude-weighted MSE loss
losses.weighted_mse_per_level computes (prediction - target)^2, then multiplies by latitude weights (proportional to cos(lat) so equatorial grid cells with larger surface area contribute more) and by pressure-level weights. Per-variable losses are summed with per_variable_weights coefficients and averaged over all non-batch dimensions to produce a scalar loss per batch element.
Config: TaskConfig.target_variables, TaskConfig.pressure_levels
System Dynamics
Beyond the pipeline, graphcast has runtime behaviors that shape how it responds to load, failures, and configuration changes.
Data Pools
Model checkpoint (.npz file)
Stores serialized model parameters (Haiku parameter trees as nested dicts of numpy arrays) plus ModelConfig and TaskConfig dataclasses, flattened to colon-separated keys in numpy .npz format. Loaded once at startup to initialize the model.
Type: checkpoint
Normalization statistics
Stores per-variable historical means and standard deviations (and time-difference standard deviations for targets) as xarray Datasets, loaded from cloud storage. Used in every forward pass to normalize inputs and denormalize outputs.
Type: file-store
Autoregressive state buffer
A rolling window of the last N input timestep arrays, maintained as flat JAX arrays inside the jax.lax.scan carry state. New prediction outputs are concatenated and old inputs dropped at each step to feed the next step.
Type: in-memory
Graph topology (edges and spatial features)
Grid-to-mesh edge indices, mesh-to-mesh edge indices, mesh-to-grid edge indices, and their associated spatial features (relative 3D positions, lat/lon encodings) computed once at model init and reused in every forward pass.
Type: in-memory
Feedback Loops
Autoregressive prediction rollout
Trigger: Calling Predictor.__call__ or Predictor.loss with a targets_template spanning multiple timesteps → Runs one forward step of the wrapped predictor, appends predicted output to the input window, drops the oldest input frame, accumulates the per-step loss (exits when: All timesteps in targets_template have been predicted (loop count determined by targets_template.dims['time']))
Type: training-loop
GenCast reverse diffusion sampling loop
Trigger: Calling GenCast Predictor.__call__ at inference time → At each noise level step: optionally add stochastic churn noise, call denoiser to get a first-order estimate, compute a second-order correction via a second denoiser call, update the noisy sample using the DPM-Solver++ 2S update rule (exits when: num_noise_levels steps completed (noise level reaches min_noise_level))
Type: convergence
Gradient checkpointing in autoregressive training
Trigger: gradient_checkpointing=True in autoregressive.Predictor constructor → Uses hk.remat to rematerialize activations during the backward pass instead of storing them, trading ~2x compute for reduced peak memory (exits when: Per-step; applies at every autoregressive step)
Type: gradient-accumulation
Control Points
ModelConfig.mesh_size
ModelConfig.gnn_msg_steps
Bfloat16Cast.enabled
SamplerConfig.num_noise_levels
SamplerConfig.stochastic_churn_rate
autoregressive.Predictor.gradient_checkpointing
Delays
JAX JIT compilation on first call
Duration: Minutes (first call only)
Graph topology precomputation
Duration: Seconds to minutes depending on mesh_size
Technology Choices
graphcast is built with 7 key technologies. Each serves a specific role in the system.
Key Components
- GraphCast (Predictor) (processor): Implements the core one-step weather prediction: takes two consecutive atmospheric state snapshots (grid) plus forcings, encodes them to icosahedral mesh nodes via a learned MLP, runs 16 rounds of GNN message passing on the multi-resolution mesh, then decodes back to grid nodes to predict the next 6-hour state delta.
- Predictor (autoregressive.py) (orchestrator): Wraps a one-step predictor and drives multi-step rollouts for training: uses jax.lax.scan to loop over timesteps in a JAX-native way (keeping the loop differentiable), feeding each step's predictions back as the next step's inputs, maintaining a rolling window of the last N input frames, and averaging the one-step loss across all timesteps.
- DeepTypedGraphNet (processor): The core GNN engine: takes a TypedGraph with separate node/edge sets (grid-nodes, mesh-nodes, grid-to-mesh edges, mesh-to-mesh edges, mesh-to-grid edges), applies independent MLPs to each edge and node type at each message-passing step, and supports both shared-weight and unshared-weight configurations via num_message_passing_steps and num_processor_repetitions.
- Bfloat16Cast (adapter): Wraps any Predictor to transparently cast all floating-point inputs to bfloat16 before the forward pass (reducing memory and accelerating TPU/GPU computation), then casts outputs back to the original dtype (typically float32) — effectively a precision mode switch without changing any model code.
- Sampler (dpm_solver_plus_plus_2s.py) (executor): Implements the DPM-Solver++ 2S reverse diffusion sampling algorithm: starting from pure Gaussian noise at max_noise_level, iterates through a schedule of decreasing noise levels (spaced by rho parameter), calling the denoiser twice per step (second-order), optionally re-injecting noise (stochastic churn) to produce diverse ensemble members.
- checkpoint.dump / checkpoint.load (serializer): Serializes nested Python structures (dicts of numpy arrays, dataclasses with typed fields) to numpy .npz format by flattening the hierarchy into colon-separated flat keys (e.g. 'params:layer1:w'), and deserializes by unflattening and type-converting back using the schema type as a guide.
- get_hierarchy_of_triangular_meshes_for_sphere (factory): Builds the multi-resolution icosahedral mesh used as the GNN's working graph: starts from a 12-vertex icosahedron, subdivides each triangular face into 4 smaller triangles `splits` times (controlled by ModelConfig.mesh_size), projects all new vertices onto a unit sphere, and returns all resolution levels so edges from coarser levels can be included in the merged mesh.
- radius_query_indices (resolver): Computes which grid lat/lon points connect to which mesh vertices by building a KD-tree over mesh vertex 3D positions and querying all mesh nodes within a radius (derived from ModelConfig.radius_query_fraction_edge_length) of each grid point — these become the grid-to-mesh edges in the TypedGraph.
Who Should Read This
ML researchers interested in scientific applications, climate scientists, or engineers working with weather prediction systems.
This analysis was generated by CodeSea from the google-deepmind/graphcast source code. For the full interactive visualization — including pipeline graph, architecture diagram, and system behavior map — see the complete analysis.
Explore Further
Full Analysis
Interactive architecture map for graphcast
graphcast vs earth2studio
Side-by-side architecture comparison
Frequently Asked Questions
What is graphcast?
Runs and trains GraphCast/GenCast neural weather forecast models autoregressively
How does graphcast's pipeline work?
graphcast processes data through 10 stages: Load and normalize ERA5 inputs, Cast inputs to bfloat16, Build icosahedral mesh and grid-mesh edges, Assemble grid and mesh node features, Encode grid features to mesh nodes (Grid→Mesh GNN), and more. Weather state data enters as xarray Datasets containing atmospheric variables (temperature, wind, geopotential, etc.) on a regular lat/lon grid at multiple pressure levels. The data is normalized by subtracting historical means and dividing by historical standard deviations. For multi-step predictions the autoregressive wrapper feeds outputs back as inputs via jax.lax.scan. Inside each one-step call, grid node features are assembled into flat vectors and projected onto icosahedral mesh nodes (grid→mesh encoding), then 16 rounds of GNN message passing propagate information across the mesh. The updated mesh node features are projected back to grid nodes (mesh→grid decoding), residually added to the input state, and denormalized to produce the final predicted atmospheric state. For GenCast, this entire encoder-decoder is used as a denoiser inside a DPM-Solver++ 2S reverse diffusion loop that iterates from max noise to near-zero noise to generate a probabilistic ensemble member.
What tech stack does graphcast use?
graphcast is built with JAX (Provides JIT compilation, automatic differentiation, and vectorization for all model computation; jax.lax.scan drives the differentiable autoregressive loop), Haiku (hk) (Neural network parameter management for JAX — handles weight initialization, parameter trees, and module scoping for all MLPs, GNN layers, and transformers), xarray (Labeled multi-dimensional array container for weather data; all inputs, outputs, and normalization statistics are xarray Datasets with named dimensions (lat, lon, level, time, batch)), jraph (JAX-native graph neural network library providing GraphsTuple data structure and message-passing primitives used inside typed_graph_net.py), NumPy / SciPy (Used for non-JAX graph topology precomputation: building KD-trees for grid-mesh connectivity, icosahedron geometry, and sparse matrix operations in the transformer), and 2 more technologies.
How does graphcast handle errors and scaling?
graphcast uses 3 feedback loops, 6 control points, 4 data pools to manage its runtime behavior. These mechanisms handle error recovery, load distribution, and configuration changes.
How does graphcast compare to earth2studio?
CodeSea has detailed side-by-side architecture comparisons of graphcast with earth2studio. These cover tech stack differences, pipeline design, and system behavior.