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

> Explore the WeatherNext data format a typed tree-structured collection of arrays that unifies grids point clouds and meshes for JAX compatibility.

- Repository: [Google DeepMind/weathernext](https://github.com/google-deepmind/weathernext)
- Tags: deep-dive
- Published: 2026-08-16

---

**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](https://github.com/google-deepmind/weathernext/blob/main/weathernext/utils/data_modalities.py#L108-L117)).

## 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.

```python
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](https://github.com/google-deepmind/weathernext/blob/main/weathernext/utils/data_modalities.py#L69-L78)).

### `LatLonPointsData`: Flattened Point Lists

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

```python

# 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](https://github.com/google-deepmind/weathernext/blob/main/weathernext/utils/data_modalities.py#L75-L92)).

### `TriangularMeshData`: Icosahedral Mesh Representations

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

```python
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](https://github.com/google-deepmind/weathernext/blob/main/weathernext/utils/data_modalities.py#L814-L826)).

### `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](https://github.com/google-deepmind/weathernext/blob/main/weathernext/utils/data_modalities.py#L124-L138) and [168-176](https://github.com/google-deepmind/weathernext/blob/main/weathernext/utils/data_modalities.py#L168-L176)).

## Validation and Metadata

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

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

```python

# 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](https://github.com/google-deepmind/weathernext/blob/main/weathernext/utils/data_modalities.py#L290-L330)).

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`](https://github.com/google-deepmind/weathernext/blob/main/weathernext/utils/data_modalities.py) | Core `Data` class and all modality implementations |
| [`weathernext/weathernext2/architecture.py`](https://github.com/google-deepmind/weathernext/blob/main/weathernext/weathernext2/architecture.py) | WeatherNext-2 (FGN) forward pass consuming modalities |
| [`weathernext/utils/icosahedral_mesh.py`](https://github.com/google-deepmind/weathernext/blob/main/weathernext/utils/icosahedral_mesh.py) | Mesh generation for `TriangularMeshData` |
| [`weathernext/utils/data_utils.py`](https://github.com/google-deepmind/weathernext/blob/main/weathernext/utils/data_utils.py) | Loading, normalization, batching utilities |
| [`weathernext/weathernext1_graph/graphcast.py`](https://github.com/google-deepmind/weathernext/blob/main/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.