Fix to_pyg graph attribute mutation

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,22 +116,22 @@ 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.
         pyg_graphs = [
             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

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)
+    )