WeatherNext Data Format Explained: A Guide to the Structured Array Containers

WeatherNext uses a typed, tree-structured collection of arrays called data modalities that unify latitude-longitude grids, point clouds, and triangular meshes under a common JAX-compatible broadcasting convention.

The WeatherNext data format is the foundation of Google's open-source weather prediction models, including WeatherNext-1 (GraphCast) and WeatherNext-2 (FGN). Defined in the weathernext.utils.data_modalities module, this format enables seamless conversion between spatial representations while maintaining consistent shapes and metadata. In google-deepmind/weathernext, all model inputs, targets, and intermediate tensors conform to this structured convention.

Core Design: The Data Base Class

Every WeatherNext data container inherits from the abstract Data class. Arrays follow a strict shape convention: (*point_dims, *per_point_feature_dims).

• Point dimensions (point_dims): Encode spatial, batch, or temporal axes that vary across data points.
• Per-point feature dimensions (per_point_feature_dims): Feature dimensions that remain constant for each point.

The point_dims_shape property guarantees broadcast compatibility. For example, a batch of temperature fields with shape (batch, lat, lon, levels) has point_dims_shape = (batch, lat, lon) and per_point_feature_dims = (levels,).

Masks attach to leading point dimensions. The helper _initialize_mask enforces that containers use either a dense mask or per-axis masks, never both (lines 108-117).

Supported Spatial Modalities

WeatherNext provides four concrete modality implementations for different spatial structures:

LatLonGridData: Uniform Latitude-Longitude Grids

The most common representation for gridded weather data. Arrays have shape (batch, lat, lon, …) with coordinate arrays lat and lon that broadcast to the full grid shape.

import xarray as xr
from weathernext.utils import data_modalities as dm

# Load ERA5 data and wrap as LatLonGridData

ds = xr.open_zarr("gs://weatherbench2/era5/temperature.zarr")
grid_data = dm.LatLonGridData.with_xarray_dataset(ds)

print(grid_data.point_dims_shape)  # (batch, num_lat, num_lon)

print(grid_data.lat.shape)         # (1, num_lat, 1) — broadcastable

print(grid_data.lon.shape)         # (1, 1, num_lon) — broadcastable

The with_xarray_dataset constructor handles coordinate extraction and validation (lines 69-78).

LatLonPointsData: Flattened Point Lists

Useful when models require point-wise operations or when converting between grid and mesh representations. Shape is (num_points, batch, …).

# Flatten grid to points with lat-major ordering

points = dm.LatLonPointsData.with_lat_lon_grid(
    lat_lon_grid_data=grid_data,
    major_axis="lat"  # lat varies slowest

)
print(points.point_dims_shape)  # (num_points, batch)

The major_axis parameter controls memory layout. The reserved metadata key GRID_LAT_LON_POINT_MAJOR_AXIS_ATTR records this choice for later reconstruction (lines 75-92).

TriangularMeshData: Icosahedral Mesh Representations

Supports hierarchical refinement for graph-based models like GraphCast. Shape matches points data: (num_points, batch, …).

from weathernext.utils import icosahedral_mesh

mesh = dm.TriangularMeshData.with_icosahedral_mesh(
    splits_list=[2, 4, 8],  # refinement levels

    data=my_data_tree,      # arrays matching (num_points, batch, …)

)
print(mesh.face_sets[-1].shape)  # (num_faces_finest, 3)

Face sets define triangle connectivity at each refinement level (lines 814-826).

GlobalData: Non-Spatial Batch Data

Simple containers with only a batch axis, no spatial coordinates. Used for scalar targets or model parameters.

Data Conversion and Reshaping

The WeatherNext data format includes utilities for zero-copy conversions between representations:

• _batch_lat_lon_to_latlon_batch: Transposes (batch, lat, lon, …) to (lat, lon, batch, …)
• _latlon_batch_to_batch_lat_lon: Reverse operation preserving broadcast semantics

These shuffles enable efficient data loading without reallocating arrays (lines 124-138 and 168-176).

Validation and Metadata

Spatial data undergoes strict coordinate validation. The _verify_lat_lon helper enforces:

• Latitude in [-90, 90]
• Longitude in [0, 360)

# Pad latitude to multiple of 8 for efficient convolution

padded = grid_data.pad_data_to_multiple_of(
    padding_axis=1,           # latitude dimension

    multiple_of=8,
    data_padding_value=0.0,
)

Padding logic in _pad_data_impl preserves modality type and metadata (lines 290-330).

Each container carries an optional metadata dict for arbitrary descriptors. The format reserves specific keys for internal use, ensuring round-trip consistency when converting between grids and points.

Where WeatherNext Data Formats Appear in the Codebase

File	Purpose
weathernext/utils/data_modalities.py	Core Data class and all modality implementations
weathernext/weathernext2/architecture.py	WeatherNext-2 (FGN) forward pass consuming modalities
weathernext/utils/icosahedral_mesh.py	Mesh generation for TriangularMeshData
weathernext/utils/data_utils.py	Loading, normalization, batching utilities
weathernext/weathernext1_graph/graphcast.py	GraphCast implementation using LatLonGridData

Summary

• WeatherNext data format is a unified container system for gridded, point, and mesh weather data.
• Shape convention separates point dimensions (batch/spatial/temporal) from per-point features.
• Four modalities cover all use cases: LatLonGridData, LatLonPointsData, TriangularMeshData, GlobalData.
• Broadcast compatibility guaranteed by point_dims_shape property across all containers.
• Validation and metadata ensure geographic correctness and reproducible transformations.
• Conversion utilities enable zero-copy reshaping between representations.

Frequently Asked Questions

What file format does WeatherNext use for stored model data?

WeatherNext does not prescribe a storage format. The LatLonGridData.with_xarray_dataset() method accepts any xarray.Dataset, commonly loaded from NetCDF, Zarr, or GRIB sources. The modality layer abstracts storage details from model code.

How does WeatherNext handle missing or masked data?

Masks attach to point dimensions with shape matching the leading axes. The _initialize_mask helper creates fully-true default masks when none provided. Either dense masks or per-axis masks are accepted, but mixing both raises an error to prevent ambiguity.

Can WeatherNext data modalities be used with PyTorch?

The format is JAX-native but framework-agnostic in structure. Arrays are standard NumPy or JAX arrays. Convert to PyTorch tensors with torch.asarray(modality.data) while preserving the container structure for coordinate metadata.

What is the difference between LatLonGridData and LatLonPointsData?

LatLonGridData preserves 2D spatial structure (batch, lat, lon, …) with broadcastable coordinate arrays. LatLonPointsData flattens to (num_points, batch, …) for point-wise operations. The with_lat_lon_grid classmethod converts between them, recording the major axis in metadata for exact reconstruction.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}