diff --git a/src/weather_model_graphs/save.py b/src/weather_model_graphs/save.py index 74ff5ef..571e97f 100644 --- a/src/weather_model_graphs/save.py +++ b/src/weather_model_graphs/save.py @@ -81,11 +81,13 @@ def to_pyg( if len(set(graph.nodes)) != len(graph.nodes): raise ValueError("Node labels must be unique.") + graph_to_save = graph.copy() + # remove all node attributes but the ones we want to keep - for node in graph.nodes: - for attr in list(graph.nodes[node].keys()): + for node in graph_to_save.nodes: + for attr in list(graph_to_save.nodes[node].keys()): if attr not in node_features: - del graph.nodes[node][attr] + del graph_to_save.nodes[node][attr] def _get_edge_indecies(pyg_g): return pyg_g.edge_index @@ -114,14 +116,14 @@ def _concat_pyg_features( value for key, value in sorted( split_graph_by_edge_attribute( - graph=graph, attr=list_from_attribute + graph=graph_to_save, attr=list_from_attribute ).items() ) ] except MissingEdgeAttributeError: # neural-lam still expects a list of graphs, so if the attribute is missing # we just return the original graph as a list - sub_graphs = [graph] + sub_graphs = [graph_to_save] # Nodes must be sorted if we want to preserve the ordering in node # labels when we convert to a pyg object. This conversion does not care # about node labels inherently. @@ -129,7 +131,7 @@ def _concat_pyg_features( pyg_convert.from_networkx(sort_nodes_in_graph(g)) for g in sub_graphs ] else: - pyg_graphs = [pyg_convert.from_networkx(sort_nodes_in_graph(graph))] + pyg_graphs = [pyg_convert.from_networkx(sort_nodes_in_graph(graph_to_save))] edge_features_values = [ _concat_pyg_features(pyg_g, features=edge_features) for pyg_g in pyg_graphs diff --git a/tests/test_save.py b/tests/test_save.py index f889ab9..3b3ebf1 100644 --- a/tests/test_save.py +++ b/tests/test_save.py @@ -1,5 +1,8 @@ import tempfile +from pathlib import Path +from types import SimpleNamespace +import networkx import pytest from loguru import logger @@ -40,3 +43,70 @@ def test_save_to_pyg(list_from_attribute): name=name, list_from_attribute=list_from_attribute, ) + + +def test_to_pyg_does_not_mutate_node_attributes(monkeypatch, tmp_path): + graph = networkx.DiGraph() + graph.add_node(0, pos=[0.0, 0.0], unexported="keep me") + graph.add_node(1, pos=[1.0, 0.0], unexported="keep me too") + graph.add_edge(0, 1, len=1.0, vdiff=0.0) + + original_node_attrs = {node: attrs.copy() for node, attrs in graph.nodes(data=True)} + converted_graphs = [] + + class FakeTensor: + ndim = 1 + + def unsqueeze(self, dim): + return self + + def to(self, dtype): + return self + + class FakePygGraph: + edge_index = FakeTensor() + + def __getitem__(self, key): + return FakeTensor() + + class FakeTorch: + Tensor = FakeTensor + float32 = "float32" + + @staticmethod + def cat(values, dim): + return FakeTensor() + + @staticmethod + def save(value, path): + Path(path).write_text("saved") + + def fake_from_networkx(converted_graph): + converted_graphs.append(converted_graph) + return FakePygGraph() + + monkeypatch.setattr(wmg.save, "HAS_PYG", True) + monkeypatch.setattr(wmg.save, "torch", FakeTorch, raising=False) + monkeypatch.setattr( + wmg.save, + "pyg_convert", + SimpleNamespace(from_networkx=fake_from_networkx), + raising=False, + ) + + wmg.save.to_pyg( + graph=graph, + output_directory=tmp_path, + name="graph", + node_features=["pos"], + ) + + assert { + node: attrs for node, attrs in graph.nodes(data=True) + } == original_node_attrs + assert converted_graphs + assert all( + "unexported" not in attrs + for converted_graph in converted_graphs + for _, attrs in converted_graph.nodes(data=True) + )