diff --git a/CHANGELOG.md b/CHANGELOG.md index 62dcb764..8dddca07 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/neural_lam/models/step_predictors/graph/base.py b/neural_lam/models/step_predictors/graph/base.py index 449f4bb6..23768dd0 100644 --- a/neural_lam/models/step_predictors/graph/base.py +++ b/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): diff --git a/neural_lam/models/step_predictors/graph/graph_lam.py b/neural_lam/models/step_predictors/graph/graph_lam.py index fdb12a12..7a718b64 100644 --- a/neural_lam/models/step_predictors/graph/graph_lam.py +++ b/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 ( diff --git a/neural_lam/utils/__init__.py b/neural_lam/utils/__init__.py index 7ac299d5..75ec61d0 100644 --- a/neural_lam/utils/__init__.py +++ b/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", diff --git a/neural_lam/utils/heterodata.py b/neural_lam/utils/heterodata.py new file mode 100644 index 00000000..87c162f9 --- /dev/null +++ b/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 diff --git a/tests/test_heterodata.py b/tests/test_heterodata.py new file mode 100644 index 00000000..46669dc7 --- /dev/null +++ b/tests/test_heterodata.py @@ -0,0 +1,286 @@ +"""Tests for the graph tensor-dict <-> pyg.HeteroData conversion (issue #385). + +Two layers of testing: + +- unit tests on ``graph_dict_to_heterodata`` / ``graph_tensors_from_heterodata`` + with small hand-built tensor-dicts (no datastore/model needed); +- a model-level equivalence test showing that a ``GraphLAM`` built with + ``use_heterodata=True`` is identical to one built the existing way, both in + its parameters/buffers and in its forward output. +""" + +# Standard library +from pathlib import Path + +# Third-party +import pytest +import torch +from torch_geometric.data import HeteroData + +# First-party +from neural_lam.create_graph import create_graph_from_datastore +from neural_lam.models import GraphLAM +from neural_lam.utils import ( + graph_dict_to_heterodata, + graph_tensors_from_heterodata, +) +from neural_lam.utils.buffer_list import BufferList +from neural_lam.utils.heterodata import ( + G2M_EDGE_TYPE, + GRID_NODE_TYPE, + M2G_EDGE_TYPE, + M2M_EDGE_TYPE, + MESH_NODE_TYPE, +) + +# Keys that ``load_graph`` returns; used to assert the round-trip is exact. +_GRAPH_DICT_KEYS = { + "g2m_edge_index", + "m2g_edge_index", + "m2m_edge_index", + "mesh_up_edge_index", + "mesh_down_edge_index", + "g2m_features", + "m2g_features", + "m2m_features", + "mesh_up_features", + "mesh_down_features", + "mesh_static_features", +} + + +def _flat_graph_dict(edge_feature_dim: int = 3, num_mesh: int = 4): + """Build a minimal flat graph tensor-dict for the unit tests. + + Parameters + ---------- + edge_feature_dim : int, default 3 + Number of edge feature columns (3 for 2D graphs, 4 for 3D). + num_mesh : int, default 4 + Number of mesh nodes. + + Returns + ------- + dict + A tensor-dict with the same keys and structure as the flat output of + :func:`neural_lam.utils.load_graph`. + """ + + def edges(n): + return torch.tensor([[0, 1, 2], [1, 2, 0]][:2], dtype=torch.int64)[ + :, :n + ] + + def feats(n): + return torch.arange(n * edge_feature_dim, dtype=torch.float32).reshape( + n, edge_feature_dim + ) + + return { + "g2m_edge_index": torch.tensor( + [[0, 1, 2], [0, 1, 2]], dtype=torch.int64 + ), + "m2g_edge_index": torch.tensor( + [[0, 1, 2], [3, 4, 5]], dtype=torch.int64 + ), + "m2m_edge_index": edges(3), + "mesh_up_edge_index": BufferList([], persistent=False), + "mesh_down_edge_index": BufferList([], persistent=False), + "g2m_features": feats(3), + "m2g_features": feats(3), + "m2m_features": feats(3), + "mesh_up_features": BufferList([], persistent=False), + "mesh_down_features": BufferList([], persistent=False), + "mesh_static_features": torch.arange( + num_mesh * 2, dtype=torch.float32 + ).reshape(num_mesh, 2), + } + + +def test_builder_flat_structure(): + """The builder produces the expected node/edge types and tensors.""" + graph_dict = _flat_graph_dict() + num_grid = 6 + graph = graph_dict_to_heterodata(graph_dict, num_grid_nodes=num_grid) + + assert isinstance(graph, HeteroData) + assert set(graph.node_types) == {GRID_NODE_TYPE, MESH_NODE_TYPE} + assert set(graph.edge_types) == { + G2M_EDGE_TYPE, + M2M_EDGE_TYPE, + M2G_EDGE_TYPE, + } + + # Grid node count comes from the datastore, not from edge indices. + assert graph[GRID_NODE_TYPE].num_nodes == num_grid + # Mesh node features are stored on .x, unchanged. + assert torch.equal( + graph[MESH_NODE_TYPE].x, graph_dict["mesh_static_features"] + ) + + # Each component maps to the right typed edge, with identical tensors. + for edge_type, ei_key, ef_key in ( + (G2M_EDGE_TYPE, "g2m_edge_index", "g2m_features"), + (M2M_EDGE_TYPE, "m2m_edge_index", "m2m_features"), + (M2G_EDGE_TYPE, "m2g_edge_index", "m2g_features"), + ): + assert torch.equal(graph[edge_type].edge_index, graph_dict[ei_key]) + assert torch.equal(graph[edge_type].edge_attr, graph_dict[ef_key]) + + +def test_builder_roundtrip_is_exact(): + """dict -> HeteroData -> dict reproduces the original tensor-dict.""" + graph_dict = _flat_graph_dict() + graph = graph_dict_to_heterodata(graph_dict, num_grid_nodes=6) + restored = graph_tensors_from_heterodata(graph) + + assert set(restored.keys()) == _GRAPH_DICT_KEYS + for key in _GRAPH_DICT_KEYS: + original = graph_dict[key] + if isinstance(original, torch.Tensor): + assert torch.equal(restored[key], original), key + else: + # The unused hierarchical entries stay empty BufferLists. + assert isinstance(restored[key], BufferList), key + assert len(restored[key]) == 0, key + + +def test_builder_variable_edge_feature_dim(): + """Edge feature width is not hardcoded (e.g. 4 cols for 3D graphs).""" + graph_dict = _flat_graph_dict(edge_feature_dim=4) + graph = graph_dict_to_heterodata(graph_dict, num_grid_nodes=6) + assert graph[G2M_EDGE_TYPE].edge_attr.shape[1] == 4 + assert graph[M2M_EDGE_TYPE].edge_attr.shape[1] == 4 + + +def test_builder_grid_count_not_inferred_from_edges(): + """num_grid_nodes is taken from the argument, not from edge maxima.""" + graph_dict = _flat_graph_dict() + # 100 is far larger than any index present in the edge tensors. + graph = graph_dict_to_heterodata(graph_dict, num_grid_nodes=100) + assert graph[GRID_NODE_TYPE].num_nodes == 100 + + +def test_builder_rejects_hierarchical(): + """Hierarchical graphs are explicitly not supported yet.""" + graph_dict = _flat_graph_dict() + with pytest.raises(NotImplementedError): + graph_dict_to_heterodata( + graph_dict, num_grid_nodes=6, hierarchical=True + ) + graph = graph_dict_to_heterodata(graph_dict, num_grid_nodes=6) + with pytest.raises(NotImplementedError): + graph_tensors_from_heterodata(graph, hierarchical=True) + + +def _build_graphlam(datastore, graph_name, use_heterodata): + """Construct a small GraphLAM with a fixed seed for equivalence tests.""" + torch.manual_seed(0) + return GraphLAM( + datastore=datastore, + graph_name=graph_name, + hidden_dim=4, + hidden_layers=1, + processor_layers=1, + mesh_aggr="sum", + num_past_forcing_steps=0, + num_future_forcing_steps=0, + output_std=False, + use_heterodata=use_heterodata, + ) + + +def test_graphlam_heterodata_equivalence(tmp_path): + """GraphLAM(use_heterodata=True) is identical to the dict-based model. + + Directly addresses the issue #385 requirement that training proceed + identically with the existing and the new HeteroData datastructure: + identical parameters, identical graph buffers, identical forward output + for the same inputs, and, when trained with the same optimizer steps, + identical per-step losses and identical weights afterwards. + """ + # First-party + from tests.dummy_datastore import DummyDatastore + + datastore = DummyDatastore() + graph_name = "1level" + graph_dir_path = Path(datastore.root_path) / "graph" / graph_name + if not graph_dir_path.exists(): + create_graph_from_datastore( + datastore=datastore, + output_root_path=str(graph_dir_path), + n_max_levels=1, + ) + + model_dict = _build_graphlam(datastore, graph_name, use_heterodata=False) + model_hd = _build_graphlam(datastore, graph_name, use_heterodata=True) + + # The HeteroData model exposes the graph as a HeteroData object. + assert isinstance(model_hd.graph, HeteroData) + assert not hasattr(model_dict, "graph") + + # Identical parameters. + params_dict = dict(model_dict.named_parameters()) + params_hd = dict(model_hd.named_parameters()) + assert params_dict.keys() == params_hd.keys() + for name in params_dict: + assert torch.equal(params_dict[name], params_hd[name]), name + + # Identical graph buffers (edge indices + features + mesh features). + for name in ( + "g2m_edge_index", + "m2g_edge_index", + "m2m_edge_index", + "g2m_features", + "m2g_features", + "m2m_features", + "mesh_static_features", + ): + assert torch.equal( + getattr(model_dict, name), getattr(model_hd, name) + ), name + + # Identical forward output for the same inputs. + num_grid = model_dict.num_grid_nodes + d_state = model_dict.grid_output_dim + forcing_dim = ( + model_dict.grid_input_dim + - 2 * d_state + - (model_dict.grid_static_features.shape[1]) + ) + torch.manual_seed(1) + prev_state = torch.randn(1, num_grid, d_state) + prev_prev_state = torch.randn(1, num_grid, d_state) + forcing = torch.randn(1, num_grid, max(forcing_dim, 0)) + + model_dict.eval() + model_hd.eval() + with torch.no_grad(): + out_dict, _ = model_dict(prev_state, prev_prev_state, forcing) + out_hd, _ = model_hd(prev_state, prev_prev_state, forcing) + + assert torch.equal(out_dict, out_hd) + + # Training proceeds identically: running the same optimizer steps on both + # models yields identical losses at every step and identical weights + # afterwards. + model_dict.train() + model_hd.train() + target = torch.randn(1, num_grid, d_state) + opt_dict = torch.optim.SGD(model_dict.parameters(), lr=0.1) + opt_hd = torch.optim.SGD(model_hd.parameters(), lr=0.1) + for _ in range(3): + step_losses = [] + for model, optimizer in ((model_dict, opt_dict), (model_hd, opt_hd)): + optimizer.zero_grad() + pred, _ = model(prev_state, prev_prev_state, forcing) + loss = torch.nn.functional.mse_loss(pred, target) + loss.backward() + optimizer.step() + step_losses.append(loss) + assert torch.equal(step_losses[0], step_losses[1]) + + trained_dict = dict(model_dict.named_parameters()) + trained_hd = dict(model_hd.named_parameters()) + for name in trained_dict: + assert torch.equal(trained_dict[name], trained_hd[name]), name