# 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//` 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 (`[:, 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.