Load flat graphs into pyg.HeteroData and wire into BaseGraphModel

CHANGELOG.md

@@ -8,6 +8,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [unreleased](https://github.com/mllam/neural-lam/compare/v0.6.0...HEAD)
 
 ### Added
+- Add `neural_lam.utils.heterodata` for representing the loaded graph as a
+  `pyg.HeteroData` object, with grid nodes as the `grid` node type, mesh nodes
+  as the `mesh` node type and `g2m`/`m2m`/`m2g` as typed edges.
+  `BaseGraphModel` reads its graph tensors out of that object when constructed
+  with `use_heterodata=True` (default `False`), which leaves the model's
+  behaviour, and therefore training, unchanged. Flat (non-hierarchical) graphs
+  only. [\#711](https://github.com/mllam/neural-lam/pull/711) @prajwal-tech07
+
 - Add `--num_sanity_val_steps` CLI argument to control sanity validation steps before training (#694)
 
 - Add `--train_steps_to_log` CLI option to log training loss for individual unroll steps, and deduplicate common prediction and loss computation steps across loops [\#674](https://github.com/mllam/neural-lam/issues/674) @GiGiKoneti

neural_lam/models/step_predictors/graph/base.py

@@ -33,6 +33,7 @@ def __init__(
         output_clamping_upper: dict[str, float] | None = None,
         g2m_gnn_type: str = "InteractionNet",
         m2g_gnn_type: str = "InteractionNet",
+        use_heterodata: bool = False,
     ) -> None:
         """
         Initialize the BaseGraphModel.
@@ -62,6 +63,10 @@ def __init__(
             Lower clamping limits for state variables.
         output_clamping_upper : dict, optional
             Upper clamping limits for state variables.
+        use_heterodata : bool, default False
+            If True, represent the loaded graph as a ``pyg.HeteroData``
+            object and take the model's graph tensors from it. Only supported
+            for flat (non-hierarchical) graphs.
         """
         super().__init__(
             datastore=datastore,
@@ -111,6 +116,25 @@ def __init__(
             mesh_node_features_scaling=grid_xy_max_span,
         )
 
+        # Optionally represent the graph as a pyg.HeteroData object (issue
+        # #385). The loaded tensors are placed on the typed node/edge stores
+        # of that object, and every graph tensor the model uses below is then
+        # read back out of it, so the HeteroData is the source of the graph
+        # data. The tensors themselves are unchanged, so the model builds and
+        # trains identically either way.
+        self.use_heterodata = use_heterodata
+        if use_heterodata:
+            if self.hierarchical:
+                raise NotImplementedError(
+                    "use_heterodata is currently only supported for flat "
+                    "(non-hierarchical) graphs."
+                )
+            self.graph = utils.graph_dict_to_heterodata(
+                graph_ldict,
+                num_grid_nodes=self.num_grid_nodes,
+            )
+            graph_ldict = utils.graph_tensors_from_heterodata(self.graph)
+
         for name, attr_value in graph_ldict.items():
             # Make BufferLists module members and register tensors as buffers
             if isinstance(attr_value, torch.Tensor):

neural_lam/models/step_predictors/graph/graph_lam.py

@@ -36,6 +36,7 @@ def __init__(
         output_clamping_upper: dict[str, float] | None = None,
         g2m_gnn_type: str = "InteractionNet",
         m2g_gnn_type: str = "InteractionNet",
+        use_heterodata: bool = False,
     ) -> None:
         """
         Initialize the GraphLAM model.
@@ -64,6 +65,10 @@ def __init__(
             Lower clamping limits for state variables.
         output_clamping_upper : dict, optional
             Upper clamping limits for state variables.
+        use_heterodata : bool, default False
+            If True, represent the loaded graph as a ``pyg.HeteroData``
+            object and take the model's graph tensors from it. Only supported
+            for flat (non-hierarchical) graphs.
         """
         super().__init__(
             datastore=datastore,
@@ -79,6 +84,7 @@ def __init__(
             output_clamping_upper=output_clamping_upper,
             g2m_gnn_type=g2m_gnn_type,
             m2g_gnn_type=m2g_gnn_type,
+            use_heterodata=use_heterodata,
         )
 
         assert (

neural_lam/utils/__init__.py

@@ -8,6 +8,10 @@
     zero_index_g2m,
     zero_index_m2g,
 )
+from .heterodata import (
+    graph_dict_to_heterodata,
+    graph_tensors_from_heterodata,
+)
 from .logging import (
     init_training_logger_metrics,
     log_on_rank_zero,
@@ -22,6 +26,8 @@
     "BufferList",
     "fractional_plot_bundle",
     "get_integer_time",
+    "graph_dict_to_heterodata",
+    "graph_tensors_from_heterodata",
     "has_working_latex",
     "init_training_logger_metrics",
     "inverse_sigmoid",

neural_lam/utils/heterodata.py

@@ -0,0 +1,160 @@
+"""Convert between neural-lam's graph tensor-dict and ``pyg.HeteroData``.
+
+The tensor-dict is the representation returned by
+:func:`neural_lam.utils.load_graph` (a dictionary of edge-index and
+node/edge-feature tensors, see that function's docstring for the keys).
+
+The ``HeteroData`` representation uses the ``"grid"`` and ``"mesh"`` node
+types, matching the terminology used throughout the rest of ``neural-lam``
+(``grid_static_features``, ``mesh_static_features``, ``g2m``/``m2m``/``m2g``).
+The three graph components map onto typed edges:
+
+- ``g2m`` (grid-to-mesh)  -> ``("grid", "to", "mesh")``
+- ``m2m`` (mesh-to-mesh)  -> ``("mesh", "to", "mesh")``
+- ``m2g`` (mesh-to-grid)  -> ``("mesh", "to", "grid")``
+
+The node/edge-type names are defined as module-level constants so the naming
+convention can be changed in one place (the reference implementations in
+issue #385, ``leifdenby/weatherduck`` and ``matschreiner/equicast``, use
+``"data"``/``"hidden"`` instead).
+
+Only flat (single mesh level) graphs are supported for now; hierarchical
+graphs raise :class:`NotImplementedError` and are handled in a follow-up.
+"""
+
+from __future__ import annotations
+
+# Standard library
+from typing import Any, Dict
+
+# Third-party
+from torch_geometric.data import HeteroData
+
+# Local
+from .buffer_list import BufferList
+
+# Node/edge-type names, kept in one place so the naming convention (see the
+# issue #385 discussion) can be changed without touching call sites.
+GRID_NODE_TYPE = "grid"
+MESH_NODE_TYPE = "mesh"
+G2M_EDGE_TYPE = (GRID_NODE_TYPE, "to", MESH_NODE_TYPE)
+M2M_EDGE_TYPE = (MESH_NODE_TYPE, "to", MESH_NODE_TYPE)
+M2G_EDGE_TYPE = (MESH_NODE_TYPE, "to", GRID_NODE_TYPE)
+
+# (edge_type, edge_index dict key, edge_feature dict key) for the three
+# always-present components, iterated by both conversion directions.
+_FLAT_EDGE_SPEC = (
+    (G2M_EDGE_TYPE, "g2m_edge_index", "g2m_features"),
+    (M2M_EDGE_TYPE, "m2m_edge_index", "m2m_features"),
+    (M2G_EDGE_TYPE, "m2g_edge_index", "m2g_features"),
+)
+
+
+def graph_dict_to_heterodata(
+    graph_dict: Dict[str, Any],
+    num_grid_nodes: int,
+    hierarchical: bool = False,
+) -> HeteroData:
+    """Build a ``pyg.HeteroData`` from a neural-lam graph tensor-dict.
+
+    Parameters
+    ----------
+    graph_dict : dict
+        Graph tensors as returned by :func:`neural_lam.utils.load_graph`. For
+        a flat graph ``m2m_edge_index``, ``m2m_features`` and
+        ``mesh_static_features`` are single tensors (not lists).
+    num_grid_nodes : int
+        Number of grid nodes. Taken from the datastore rather than inferred
+        from the edge indices, because the graph spec allows grid nodes that
+        no edge connects to.
+    hierarchical : bool, default False
+        Whether the graph is hierarchical. Only ``False`` is supported for
+        now.
+
+    Returns
+    -------
+    torch_geometric.data.HeteroData
+        Typed graph with ``"grid"`` and ``"mesh"`` node types and the
+        ``g2m`` / ``m2m`` / ``m2g`` edge types. Mesh node features are stored
+        on ``graph["mesh"].x`` and edge features on ``edge_attr``. Tensors
+        are the same objects as in ``graph_dict`` (no copy).
+
+    Raises
+    ------
+    NotImplementedError
+        If ``hierarchical`` is ``True``.
+    """
+    if hierarchical:
+        raise NotImplementedError(
+            "graph_dict_to_heterodata currently supports only flat "
+            "(single-level) graphs; hierarchical support is a follow-up."
+        )
+
+    graph = HeteroData()
+
+    # Grid nodes have no stored static features here (grid features are
+    # dynamic and passed at forward time), so only the count is set.
+    graph[GRID_NODE_TYPE].num_nodes = int(num_grid_nodes)
+
+    # Mesh nodes carry their static coordinate features.
+    graph[MESH_NODE_TYPE].x = graph_dict["mesh_static_features"]
+
+    for edge_type, edge_index_key, edge_feature_key in _FLAT_EDGE_SPEC:
+        graph[edge_type].edge_index = graph_dict[edge_index_key]
+        graph[edge_type].edge_attr = graph_dict[edge_feature_key]
+
+    return graph
+
+
+def graph_tensors_from_heterodata(
+    graph: HeteroData,
+    hierarchical: bool = False,
+) -> Dict[str, Any]:
+    """Read the model's graph tensors out of a ``HeteroData`` object.
+
+    This is how the model obtains its graph tensors when it represents the
+    graph as a ``HeteroData``: every tensor is looked up on the typed
+    node/edge stores of ``graph``. They are returned under the names the
+    model refers to them by (the same names
+    :func:`neural_lam.utils.load_graph` uses), so only the *source* of the
+    tensors changes, not the rest of the model's setup. It is also the exact
+    inverse of :func:`graph_dict_to_heterodata`.
+
+    Parameters
+    ----------
+    graph : torch_geometric.data.HeteroData
+        Graph as produced by :func:`graph_dict_to_heterodata`.
+    hierarchical : bool, default False
+        Whether the graph is hierarchical. Only ``False`` is supported for
+        now.
+
+    Returns
+    -------
+    dict
+        The model's graph tensors, keyed by the names used in
+        :func:`neural_lam.utils.load_graph`.
+
+    Raises
+    ------
+    NotImplementedError
+        If ``hierarchical`` is ``True``.
+    """
+    if hierarchical:
+        raise NotImplementedError(
+            "graph_tensors_from_heterodata currently supports only flat "
+            "(single-level) graphs; hierarchical support is a follow-up."
+        )
+
+    graph_dict: Dict[str, Any] = {
+        "mesh_static_features": graph[MESH_NODE_TYPE].x,
+        # No inter-level edges for a flat graph.
+        "mesh_up_edge_index": BufferList([], persistent=False),
+        "mesh_down_edge_index": BufferList([], persistent=False),
+        "mesh_up_features": BufferList([], persistent=False),
+        "mesh_down_features": BufferList([], persistent=False),
+    }
+    for edge_type, edge_index_key, edge_feature_key in _FLAT_EDGE_SPEC:
+        graph_dict[edge_index_key] = graph[edge_type].edge_index
+        graph_dict[edge_feature_key] = graph[edge_type].edge_attr
+
+    return graph_dict