Graph Storage Specification

Graph Storage Specification#

Version: 0.1.0

1. Introduction#

This document specifies the requirements for Graph disk format for neural-lam. These graphs are used by the Neural-LAM Graph Neural Network architectures for machine-learning weather prediction (MLWP) forecasting. These model architectures follow the encode-process-decode paradigm of sequential message passing, where physical variables are represented as features on so-called grid nodes, are encoded to mesh nodes, are processed on the mesh, and then decoded back to grid nodes where output tendencies or updated state are produced.

The format specified in this document was designed to support the definition of both flat (e.g. Keisler 2022, Lam et al 2022) and hierarchical (Oskarsson et al 2023) graphs for GNN-based MLWP in neural-lam.

The key words “MUST”, “MUST NOT”, “REQUIRED”, “SHALL”, “SHALL NOT”, “SHOULD”, “SHOULD NOT”, “RECOMMENDED”, “MAY”, and “OPTIONAL” in this document are to be interpreted as described in RFC 2119.

2. File and Directory Structure#

2.1 Directory Structure#

Each graph MUST be identified by a unique name and stored within the directory graph/<name>/ that in turn MUST be placed within the same directory as the datastore configuration from which the graph was derived (i.e. the spatial coordinates defining the grid coordinates are provided by the datastore).

2.2 Graph Filenames#

Each graph MUST be represented by multiple files in its graph directory. Together, these files define connectivity and static features for edges and mesh nodes.

Required files for all graphs (all of these files MUST be present):

  • m2m_edge_index.pt

  • g2m_edge_index.pt

  • m2g_edge_index.pt

  • m2m_features.pt

  • g2m_features.pt

  • m2g_features.pt

  • mesh_features.pt

Additional required files for hierarchical graphs (L > 1 mesh levels), all of which MUST be present:

  • mesh_up_edge_index.pt

  • mesh_down_edge_index.pt

  • mesh_up_features.pt

  • mesh_down_features.pt

2.2.1 Graph format versioning#

Graph directories MUST include the file metainfo.yaml. This file MUST be a YAML document containing the key spec_version with the graph storage spec version as its value. The current graph storage spec version is 0.1.0.

See 4. Differences to legacy format graphs for the legacy-specific behavior.

2.3 Graph Components and File Suffixes#

The separate files represent the different “components” of the graph, where each of the sequential message passing steps, encode, process, and decode uses a separate component, so that g2m is used in encode, m2m is used in process, and m2g is used in decode. For hierarchical graphs, the m2m component is further split into separate inter-level and intra-level message-passing steps.

Each graph component MUST be represented by two files: one for edge connectivity and one for edge features. The components are (which also define the expected file prefixes):

  • g2m: grid-to-mesh edges (sender on grid, receiver on mesh).

  • m2m: mesh-to-mesh edges (both sender and receiver on mesh).

  • m2g: mesh-to-grid edges (sender on mesh, receiver on grid).

  • mesh_up: inter-level mesh edges from lower level to upper level.

  • mesh_down: inter-level mesh edges from upper level to lower level.

Suffixes indicate content type:

  • _edge_index.pt: edge connectivity

  • _features.pt: static features associated with each edge

In addition to defining the properties of the graph edges, the mesh-node features MUST be stored separately in mesh_features.pt.

All files MUST be serialized with torch.save(...).

NOTE: Rather than the files defining inter-level edges being prefixed with m2m, they are prefixed with mesh, even though they are part of the mesh-to-mesh message passing.

3. File content requirements#

The content of the files depend on the number of mesh levels, denoted as L in the text below, so that for:

  • Non-hierarchical graphs L == 1.

  • Hierarchical graphs L > 1.

  • Entry 0 MUST always be the bottom mesh level.

Every tensor MUST be stored in a manner amenable to load with torch.load(...) (this can most easily be support by using torch.save(...) to store tensors to disk) and satisfy the requirements below.

3.1 Nodes#

The nodes in the graph are implicitly defined by the node indices appearing in the edge index tensors, and by the mesh node features stored in mesh_features.pt. Both are covered in the requirements below.

NOTE: The neural-lam graph format on disk does not explicitly store node features for grid nodes, as these are expected to be dynamic and stored separately in the dataset. However, static features for mesh nodes MUST be stored in mesh_features.pt files (as described below).

3.1.1 Node index space#

The node indices in edge index tensors MUST be defined so that for both the set of source nodes (where the edges connect from) and the set of destination nodes (where the edges connect to) the node indices being used should run from 0 to N_src-1 and 0 to N_dst-1 respectively, where N_src and N_dst are the total number of source and destination nodes being connected. This means that the node indices in the edge index tensors are numbered separately (for the source and destination node sets), and the nodes in each set MUST be labelled identically if and only if the set of nodes being connect from and to are the same set (as would be the case for mesh-to-mesh connections in non-hierarchical graphs).

NOTE: Although the integers used to label a given set of source or destination nodes must be contiguous from 0 to N-1 (where N is the number of nodes) it is not a requirement that all node indexes are actually present in the edge index tensors, i.e. it is not a requirement that all nodes in either the source or destination set of nodes has edges connecting to them (this is for example used when decoding to only a subset in the grid nodes when training in a limited-area setting where predictions are not required on the boundary)

NOTE: There is no requirement that the node indices for different nodesets be non-overlapping, in fact they should be overlapping, as the node indices for each nodeset are defined to run from 0 to N-1 for that nodeset (with N nodes). The key requirement is that the node indices for each nodeset are contiguous and defined in a consistent manner across all edge index tensors.

3.1.2 Mesh node features#

mesh_features.pt files MUST be lists of length L (number of mesh levels), where each entry is a tensor containing static features for the mesh nodes at that level. Each tensor MUST satisfy the following requirements:

  • mesh_features entries MUST have shape [N_level, 2], where N_level is the number of mesh nodes at that level. The two features MUST be the horizontal coordinates of the mesh nodes at level i.

  • mesh_features[i][:, 0] and mesh_features[i][:, 1] MUST contain the horizontal coordinates of the mesh nodes at level i, with column 0 and column 1 representing the two orthogonal horizontal axes.

NOTE: These coordinates SHOULD be provided in a local equal-area projection, because the two mesh node features are normalized together using the maximum span of the grid coordinates after graph loading.

  • Mesh node features SHOULD NOT be normalized. Instead, normalization will be performed inside neural-lam after graph loading.

  • Dtype MUST be torch.float32.

3.2 Edges#

3.2.1 Edge indices#

The following edge index files MUST be defined:

  • g2m_edge_index.pt

  • m2g_edge_index.pt

  • m2m_edge_index.pt

  • mesh_up_edge_index.pt (hierarchical graphs only, L > 1)

  • mesh_down_edge_index.pt (hierarchical graphs only, L > 1)

g2m_edge_index.pt and m2g_edge_index.pt MUST each contain a single tensor with shape [2, E], where E is the number of edges in each of the g2m (grid-to-mesh) and m2g (mesh-to-grid) graph components. m2m_edge_index.pt MUST contain a list of tensors of length L, i.e. one edge-index tensor per mesh level. Each entry MUST have shape [2, E_level], where E_level is the number of edges at that level. For hierarchical graphs, mesh_up_edge_index.pt and mesh_down_edge_index.pt MUST each contain a list of length L - 1 of tensors, i.e. one per inter-level connection, so that entry i connects level i and level i+1. Each entry MUST have shape [2, E_interlevel], where E_interlevel is the number of edges going either up (mesh_up_edge_index.pt) or down (mesh_down_edge_index.pt) between that level pair. For every edge-index tensor above:

  • Row 0 MUST be sender node index, row 1 MUST be receiver node index.

  • Dtype MUST be torch.int64.

3.2.2 Edge features#

The following edge feature files MUST be defined:

  • g2m_features.pt

  • m2g_features.pt

  • m2m_features.pt

  • mesh_up_features.pt (hierarchical graphs only, L > 1)

  • mesh_down_features.pt (hierarchical graphs only, L > 1)

g2m_features.pt and m2g_features.pt MUST each contain a single tensor with shape [E, N_f], where E matches the number of edges in the corresponding *_edge_index.pt file. m2m_features.pt MUST contain a list of length L, i.e. one feature tensor per mesh level. Entry i MUST have shape [E_level, N_f], where E_level matches the edge count in entry i of m2m_edge_index.pt. For hierarchical graphs, mesh_up_features.pt and mesh_down_features.pt MUST each contain a list of length L - 1, i.e. one feature tensor per inter-level connection between level i and i+1. Entry i MUST have shape [E_interlevel, N_f], where E_interlevel matches the edge count in entry i of the corresponding mesh_*_edge_index.pt file. For every edge feature tensor above:

  • The shape MUST be [E, N_f] where E is the number of edges in a given subgraph (e.g. g2m, m2g or a inter/intra-level edgeset in m2m)

  • N_f MUST be exactly 3 (for 2D edges) or exactly 4 (for 3D edges). The value of N_f MUST be consistent across all edge feature tensors in the graph.

  • The first column (<feature_tensor>[:, 0]) MUST contain the total edge length (e.g., the Euclidean distance between the sender and receiver nodes).

NOTE: The reason for requiring that the first column be the total edge length is that neural-lam uses this column to compute the normalization factor (the longest edge length found across the m2m edge features). Since edge length and the Cartesian displacements all measure distance and share the same physical scale, all edge feature columns are normalized jointly by this single factor after loading.

  • The following columns MUST contain the Cartesian coordinate displacements (vdiff = receiver_pos - sender_pos). For 2D edges (N_f == 3), columns 1 and 2 are the x- and y-displacements respectively. For 3D edges (N_f == 4), columns 1, 2, and 3 are the x-, y-, and z-displacements respectively.

  • Edge features SHOULD NOT be normalized. Instead, normalization will be performed inside neural-lam after graph loading.

  • Dtype MUST be torch.float32.

4. Differences to legacy format graphs#

Legacy pre-spec graphs do not include metainfo.yaml. When that file is absent, neural-lam will attempt to load the graph as a legacy pre-spec graph. Their mesh_features.pt files are assumed to already be normalized.

Legacy pre-spec graphs store edge indices in one combined offset node-index layout across all node sets rather than the zero-based per-node-set index space required by this specification. When loading those graphs, neural-lam zero-offsets the edge indices so that each source and destination node set starts at 0.

This validator does not validate legacy pre-spec graphs; regenerate or migrate them to the current graph storage spec first.

Current-format graphs include metainfo.yaml with a spec_version of 0.1.0. Their mesh_features.pt files are normalized on load, and their edge-index tensors are expected to already use the zero-based per-node-set index space required by this specification.