google-deepmind/graphcast
14 hidden assumptions · 10-stage pipeline · 8 components
Like any codebase, this ml inference makes assumptions it never checks — most are routine. The ones worth your attention are below, in plain language with what to do about each.
Runs and trains GraphCast/GenCast neural weather forecast models autoregressively
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.
Under the hood, the system uses 3 feedback loops, 4 data pools, 6 control points to manage its runtime behavior.
A 8-component ml inference. 37 files analyzed. Data flows through 10 distinct pipeline stages.
Hidden Assumptions
Most of what this code assumes is routine. These 3 are the ones most likely to cause trouble here — in plain terms, with what to do about each. The rest are minor; they're under "Show everything".
This forecasting system does not estimate the starting state of the atmosphere on its own. It depends on being handed starting conditions produced by a traditional physics-based weather analysis system, including a spread of slightly different starting states. If you feed it inputs that lack this kind of carefully prepared ensemble of starting conditions, the uncertainty estimates the paper reports will not hold.
What to do: Make sure your inputs come from a proper ensemble data assimilation analysis as the paper used, and do not trust the ensemble spread if you initialize from a single or improvised starting state.
“GenCast also relies on initial conditions from a traditional NWP ensemble data assimilation system, and therefore for operational use those systems must still be available.”
Read it in the paper · Conclusions ↗The authors tested this system only at a coarse global grid of roughly one degree. They explicitly say that for real operational use it should be trained and tested at finer resolution, and that finer resolution would likely give better results. The outputs you get are at the coarse scale the paper validated, not at the fine scale operational forecasts use.
What to do: Treat outputs as coarse-resolution forecasts and avoid presenting them as equivalent to finer-resolution operational ensemble products without re-validating at the resolution you need.
“For operational use, GenCast should be trained and tested at higher resolution, and would likely yield better results.”
Read it in the paper · Conclusions ↗The authors deliberately left precipitation out of their headline accuracy results because they are not fully confident in the quality of the training data for rainfall, and they did not tailor their evaluation to it. So any rain-related output should be treated with extra caution; it was not held to the same standard as the other weather variables.
What to do: Treat precipitation forecasts as provisional and check the paper's separate precipitation appendix and metrics before relying on rain predictions.
“we lack full confidence in the quality of ERA5 precipitation data, and that we have not tailored our evaluation to precipitation specifically”
Read it in the paper · Results ↗Show everything (11 more)
The system was built around a twelve-hour forecast step on purpose, because the underlying historical data mixes two different kinds of transitions depending on the timing window. By stepping twelve hours it always jumps cleanly between windows. If your data or step timing does not respect this structure, the model is being used outside the regime it was designed for.
What to do: Keep the twelve-hour step and align inputs to the same time-of-day convention the authors used, rather than re-timing steps in ways that break the assimilation-window structure.
“By choosing a 12 hour time step we avoid training on this bimodal distribution, and ensure that our model always predicts a target from the next assimilation window.”
Read it in the paper · Methods ↗autoregressive.py rollout window and 12h step configuration
The way the authors measured forecast reliability compares the spread of forecasts against a single best-guess past state, not against a spread of plausible past states. They warn that this scoring method unfairly favors forecasts that are too confident at short lead times, and can make a properly-spread forecast look overconfident. So short-range spread numbers should be read with that caveat in mind.
What to do: When judging short-lead-time spread or reliability, remember the paper's caution that the scoring method itself biases these numbers and do not over-interpret apparent over- or under-dispersion early in the forecast.
“this rewards under-dispersion at short lead times”
Read it in the paper · Verification ↗Verification / evaluation harness (latitude-weighted metrics against deterministic analysis)
The model relies on a separate statistics file to scale its inputs and outputs correctly. If that file doesn't exactly match the model you loaded — for example, if you downloaded the wrong statistics file, or mixed files from different model versions — the model will still run without any complaint, but every number it produces will be quietly wrong. The forecast will look like a forecast, but it won't be correct.
What to do: Before running, double-check that your statistics file, your model checkpoint, and your task configuration all come from the same model version — mismatching these is the single easiest way to get plausible-looking but incorrect forecasts.
graphcast/normalization.py:Normalizer
The model expects your weather data to be formatted in a very specific way: particular variable names, with latitude running from the North Pole down to the South Pole, and specific pressure altitudes. If you get data from any source other than ERA5 — or even ERA5 downloaded through a tool that reorders dimensions — the model will either crash or, more dangerously, silently use the wrong data in the wrong place and produce convincing-looking but wrong forecasts.
What to do: Before feeding in any data, confirm it uses ERA5 variable names, that latitude runs from 90 down to -90, longitude from 0 to 359.75, and that your pressure levels exactly match the list the model was configured with.
graphcast/graphcast.py:GraphCast.__init__
The model's internal map between the weather grid and its working graph was built assuming a specific grid resolution. If you feed it data at a different resolution — even a commonly available coarser resolution — the map is rebuilt with the wrong density, and the model silently operates on a completely different structure than it was trained on, producing bad forecasts.
What to do: Always use data at exactly the resolution the checkpoint was trained on; for the standard GraphCast model that means 0.25-degree global ERA5 data.
graphcast/grid_mesh_connectivity.py:radius_query_indices
When you load a saved model file, the code trusts that the file and the software version you're running are perfectly in sync. If the model file was saved with a newer or older version of the code, extra settings can be silently discarded or substituted with wrong defaults, causing the model to run differently than intended with no warning.
What to do: Always use the same version of the code to load a checkpoint as was used to save it, and keep track of which code version produced each checkpoint file.
graphcast/checkpoint.py:load
During multi-step forecasting, the model feeds its own predictions back as inputs for the next step. This only works correctly if the predictions come out in exactly the same format the model expects as input. If anything about the time ordering or variable structure is slightly off, either the model crashes part-way through a long forecast, or all steps after the first use misaligned data with no error shown.
What to do: When customizing inputs or variables, verify that every predicted output variable exactly matches the expected input structure before attempting a multi-step rollout.
graphcast/autoregressive.py:Predictor.__call__
The geometric structure of the model's internal graph is built using high-precision math on the CPU, then handed off to the accelerator. On certain hardware setups, that precision is silently cut in half before the model ever sees it, slightly but persistently corrupting the spatial encoding in every single forecast, with no warning.
What to do: If running on TPU or in a restricted-precision environment, verify that spatial edge features are explicitly cast to the intended precision before the first forward pass.
graphcast/graphcast.py:GraphCast.__init__
The training loss weights each location on the globe by its true surface area, which requires the latitude values to be in degrees. If a data pipeline inadvertently converts latitudes to radians before they reach the loss function, every location is weighted almost equally, the model gets the wrong training signal, and the resulting model performs worse — with no error or warning during training.
What to do: When writing custom data loaders, make sure latitude coordinates are kept in degrees all the way through to the loss function.
graphcast/losses.py:weighted_mse_per_level
The model always predicts one fixed time-step ahead regardless of what time labels you attach to your output template. If you accidentally ask for outputs at hourly or daily intervals instead of the model's native 6-hour step, the model runs without complaint but the time labels on the outputs are wrong — the forecast data corresponds to different times than indicated.
What to do: Make sure your output template uses time steps that exactly match the model's native forecast interval — 6 hours for GraphCast, 12 hours for GenCast.
graphcast/autoregressive.py:Predictor.__call__
Training the model over multiple forecast steps at once requires storing a large amount of intermediate data in memory — one copy for each time step in the rollout. By default, a setting that would cut this memory roughly in half is turned off. Trying to train over many steps without turning it on will simply crash with a confusing out-of-memory error.
What to do: When training with rollouts longer than a few steps, turn on the gradient checkpointing option to avoid running out of accelerator memory.
graphcast/autoregressive.py:Predictor.loss
Open the standalone hidden-assumptions report for graphcast →
How Data Flows Through the System
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.
- Load 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. [xarray.Dataset (weather state) → xarray.Dataset (weather state)] (config: TaskConfig.input_variables, TaskConfig.target_variables, TaskConfig.forcing_variables +1)
- Cast 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. [xarray.Dataset (weather state) → xarray.Dataset (weather state)]
- Build 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. [ModelConfig → TriangularMesh] (config: ModelConfig.mesh_size, ModelConfig.resolution, ModelConfig.radius_query_fraction_edge_length +1)
- Assemble 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. [xarray.Dataset (weather state) → TypedGraph] (config: ModelConfig.latent_size, TaskConfig.pressure_levels)
- Encode 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. [TypedGraph → TypedGraph] (config: ModelConfig.latent_size, ModelConfig.hidden_layers)
- Run 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. [TypedGraph → TypedGraph] (config: ModelConfig.gnn_msg_steps, ModelConfig.latent_size, ModelConfig.hidden_layers)
- Decode 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. [TypedGraph → xarray.Dataset (weather state)] (config: ModelConfig.latent_size, ModelConfig.hidden_layers, TaskConfig.target_variables +1)
- Apply 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. [xarray.Dataset (weather state) → xarray.Dataset (weather state)] (config: TaskConfig.target_variables)
- Autoregressive 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. [xarray.Dataset (weather state) → xarray.Dataset (weather state)] (config: SamplerConfig.num_noise_levels, SamplerConfig.max_noise_level, SamplerConfig.min_noise_level +2)
- Compute 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. [LossAndDiagnostics → LossAndDiagnostics] (config: TaskConfig.target_variables, TaskConfig.pressure_levels)
Data Models
The data structures that flow between stages — the contracts that hold the system together.
graphcast/graphcast.pyxarray.Dataset with variables like 'temperature', 'geopotential', 'u_component_of_wind' having dims (batch, time, lat, lon) for surface vars and (batch, time, level, lat, lon) for atmospheric vars; lat has 721 points (0.25° resolution), lon has 1440 points, level is one of 13/25/37 pressure levels in hPa
Created by loading ERA5 data slices, passed through normalization wrappers, split into grid node feature vectors for the GNN, then reconstructed from GNN output node features and written back as predicted weather state.
graphcast/graphcast.pychex.dataclass with input_variables: tuple[str,...], target_variables: tuple[str,...], forcing_variables: tuple[str,...], pressure_levels: tuple[int,...], input_duration: str (e.g. '12h' or '24h')
Defined as a constant (e.g. graphcast.TASK or gencast.TASK) at model definition time, used throughout to filter which variables are inputs vs. targets vs. forcings and which pressure levels to include.
graphcast/graphcast.pychex.dataclass with resolution: int, mesh_size: int (icosahedron split count 4-6), latent_size: int (e.g. 512), gnn_msg_steps: int (e.g. 16), hidden_layers: int, radius_query_fraction_edge_length: float, mesh2grid_edge_normalization_factor: Optional[float]
Loaded from checkpoint or constructed once at model initialization; controls the topology of the graph and the size of all neural network layers.
graphcast/typed_graph.pyNamedTuple with context: TypedGraphContext, nodes: Mapping[str, NodeSet], edges: Mapping[EdgeSetKey, EdgeSet]; NodeSet has features: ArrayLike[num_nodes, feature_dim]; EdgeSet has features: ArrayLike[num_edges, feature_dim] plus indices ArrayLike[num_edges] for senders and receivers
Constructed once per forward pass by embedding grid-node and mesh-node features into flat vectors, processed through multiple rounds of message passing in DeepTypedGraphNet, then the mesh-to-grid edge outputs are used to read off per-grid-node predictions.
graphcast/icosahedral_mesh.pyNamedTuple with vertices: np.ndarray[num_vertices, 3] (3D unit-sphere positions) and faces: np.ndarray[num_faces, 3] (integer vertex indices)
Built once at model initialization by starting with a 12-vertex icosahedron and subdividing each triangular face 4-6 times (controlled by mesh_size), projecting new vertices onto the unit sphere; the resulting vertices become mesh nodes in the GNN graph.
graphcast/gencast.pychex.dataclass with max_noise_level: float, min_noise_level: float, num_noise_levels: int, rho: float, stochastic_churn_rate: float, churn_min_noise_level: float, churn_max_noise_level: float, noise_level_inflation_factor: float
Created once when configuring GenCast inference, passed to the DPM-Solver++ 2S Sampler to define the noise schedule (number of steps, min/max noise, spacing via rho) and stochastic churn behavior for ensemble diversity.
graphcast/losses.pytuple[xarray.DataArray, xarray.Dataset] where DataArray has dims ('batch',) containing per-example scalar losses, and Dataset contains named diagnostic scalars (per-variable loss terms) also with dims ('batch',)
Produced by the loss function at each training step, averaged over batch dimension for the scalar that gets differentiated, while the diagnostic Dataset is logged separately for monitoring per-variable training progress.
System Behavior
How the system operates at runtime — where data accumulates, what loops, what waits, and what controls what.
Data Pools
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.
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.
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.
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.
Feedback Loops
- Autoregressive prediction rollout (training-loop, reinforcing) — Trigger: Calling Predictor.__call__ or Predictor.loss with a targets_template spanning multiple timesteps. Action: 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. Exit: All timesteps in targets_template have been predicted (loop count determined by targets_template.dims['time']).
- GenCast reverse diffusion sampling loop (convergence, balancing) — Trigger: Calling GenCast Predictor.__call__ at inference time. Action: 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. Exit: num_noise_levels steps completed (noise level reaches min_noise_level).
- Gradient checkpointing in autoregressive training (gradient-accumulation, balancing) — Trigger: gradient_checkpointing=True in autoregressive.Predictor constructor. Action: Uses hk.remat to rematerialize activations during the backward pass instead of storing them, trading ~2x compute for reduced peak memory. Exit: Per-step; applies at every autoregressive step.
Delays
- JAX JIT compilation on first call (compilation, ~Minutes (first call only)) — The first call to the JAX-jitted model traces and compiles the entire computation graph; subsequent calls use the compiled XLA code and run at full speed.
- Graph topology precomputation (warmup, ~Seconds to minutes depending on mesh_size) — Building the KD-tree and querying all grid-mesh connections at model init; for a 0.25° resolution grid (~1M grid points) and split-6 mesh this is a significant one-time cost before any inference can begin.
Control Points
- ModelConfig.mesh_size (architecture-switch) — Controls: Number of icosahedron subdivision splits (4-6), which determines the number of mesh nodes (~2,562 at split=4, ~40,962 at split=6) and hence model capacity and memory usage. Default: 6 for the full operational GraphCast model
- ModelConfig.gnn_msg_steps (hyperparameter) — Controls: Number of message-passing rounds in the Mesh→Mesh GNN, controlling how far information propagates across the mesh per forward pass and model depth. Default: 16 for the full operational model
- Bfloat16Cast.enabled (precision-mode) — Controls: Whether all model computation runs in bfloat16 (halved memory, faster on TPU/GPU) or float32 (higher precision but 2× memory). Default: True (enabled by default)
- SamplerConfig.num_noise_levels (hyperparameter) — Controls: Number of denoising steps in GenCast reverse diffusion; more steps improve sample quality at the cost of num_noise_levels × 2 denoiser forward passes
- SamplerConfig.stochastic_churn_rate (sampling-strategy) — Controls: S_churn parameter: 0 gives deterministic DDIM-like sampling; >0 re-injects noise at each step, increasing ensemble diversity at the cost of potentially noisier samples
- autoregressive.Predictor.gradient_checkpointing (feature-flag) — Controls: Whether hk.remat is applied at each autoregressive step to trade ~2× backward-pass compute for reduced activation memory during training. Default: False by default
Technology Stack
Provides JIT compilation, automatic differentiation, and vectorization for all model computation; jax.lax.scan drives the differentiable autoregressive loop
Neural network parameter management for JAX — handles weight initialization, parameter trees, and module scoping for all MLPs, GNN layers, and transformers
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)
JAX-native graph neural network library providing GraphsTuple data structure and message-passing primitives used inside typed_graph_net.py
Used for non-JAX graph topology precomputation: building KD-trees for grid-mesh connectivity, icosahedron geometry, and sparse matrix operations in the transformer
Provides frozen dataclasses (chex.dataclass) used for ModelConfig, TaskConfig, SamplerConfig — immutable configuration objects that can be used as JAX pytree leaves
3D mesh processing library used in grid_mesh_connectivity.py to find which mesh triangle each grid point falls inside (in_mesh_triangle_indices)
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.
graphcast/graphcast.py - 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.
graphcast/autoregressive.py - 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.
graphcast/deep_typed_graph_net.py - 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.
graphcast/casting.py - 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.
graphcast/dpm_solver_plus_plus_2s.py - 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.
graphcast/checkpoint.py - 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.
graphcast/icosahedral_mesh.py - 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.
graphcast/grid_mesh_connectivity.py
Explore the interactive analysis
See the full architecture map, data flow, and code patterns visualization.
Analyze on CodeSeaCompare graphcast
Related Ml Inference Repositories
Frequently Asked Questions
What is graphcast used for?
Runs and trains GraphCast/GenCast neural weather forecast models autoregressively google-deepmind/graphcast is a 8-component ml inference written in Python. Data flows through 10 distinct pipeline stages. The codebase contains 37 files.
How is graphcast architected?
graphcast is organized into 7 architecture layers: Predictor Wrapper Stack, Core Model Architectures, Graph Infrastructure, Data Preparation & Feature Engineering, and 3 more. Data flows through 10 distinct pipeline stages. This layered structure keeps concerns separated and modules independent.
How does data flow through graphcast?
Data moves 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) → .... 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. This pipeline design reflects a complex multi-stage processing system.
What technologies does graphcast use?
The core stack includes 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), chex (Provides frozen dataclasses (chex.dataclass) used for ModelConfig, TaskConfig, SamplerConfig — immutable configuration objects that can be used as JAX pytree leaves), and 1 more. A focused set of dependencies that keeps the build manageable.
What system dynamics does graphcast have?
graphcast exhibits 4 data pools (Model checkpoint (.npz file), Normalization statistics), 3 feedback loops, 6 control points, 2 delays. The feedback loops handle training-loop and convergence. These runtime behaviors shape how the system responds to load, failures, and configuration changes.
What design patterns does graphcast use?
5 design patterns detected: Decorator / Wrapper Stack (Predictor interface), Grid→Mesh→Mesh→Grid encode-process-decode, xarray-as-JAX-pytree data contract, Typed serialization via flattened key paths, Diffusion denoiser composability.
How does graphcast compare to alternatives?
CodeSea has side-by-side architecture comparisons of graphcast with earth2studio. These comparisons show tech stack differences, pipeline design, system behavior, and code patterns. See the comparison pages above for detailed analysis.
Analyzed on June 9, 2026 by CodeSea. Written by Karolina Sarna.