sorting bug fix

src/pathpyG/core/graph.py

@@ -116,8 +116,10 @@ def __init__(self, data: Data, mapping: Optional[IndexMap] = None):
 
         # sort EdgeIndex and validate
         data.edge_index, sorted_idx = data.edge_index.sort_by("row")
-        for edge_attr in self.edge_attrs():
-            data[edge_attr] = self.data[edge_attr][sorted_idx]
+        # an edge index that is already sorted is returned without a permutation
+        if sorted_idx is not None:
+            for edge_attr in self.edge_attrs():
+                data[edge_attr] = self.data[edge_attr][sorted_idx]
 
         data.edge_index.validate()
 

tests/core/test_graph.py

@@ -35,6 +35,16 @@ def test_init_with_edge_index():
     assert isinstance(g.edge_to_index, dict)
 
 
+def test_init_with_presorted_edge_index_keeps_edge_attrs():
+    # An EdgeIndex that already carries sort_order="row" is returned by sort_by without a
+    # permutation, in which case the edge attributes must be left as they are.
+    edge_index = EdgeIndex([[0, 0, 1], [1, 2, 2]], sparse_size=(3, 3), sort_order="row")
+    edge_weight = torch.tensor([1.0, 2.0, 3.0])
+    g = Graph(Data(edge_index=edge_index, num_nodes=3, edge_weight=edge_weight))
+    assert g.data.edge_weight.shape == edge_weight.shape
+    assert torch.equal(g.data.edge_weight, edge_weight)
+
+
 def test_init_with_mapping():
     edge_index = get_random_edge_index(100, 100, 1000)
     data = Data(edge_index=edge_index, num_nodes=100)