Repository: pyg-team/pytorch_geometric
Pull request #8502 (2023-12-01 11:50:36)

Add `EdgeIndex(is_undirected=True)` option


Patch:
diff --git a/torch_geometric/data/edge_index.py b/torch_geometric/data/edge_index.py
index 49082fde4..de48c24d8 100644
--- a/torch_geometric/data/edge_index.py
+++ b/torch_geometric/data/edge_index.py
@@ -69,6 +69,12 @@ def assert_contiguous(tensor: Tensor):
                          "`edge_index.contiguous()` before proceeding.")
 
 
+def assert_symmetric(size: Tuple[Optional[int], Optional[int]]):
+    if size[0] is not None and size[1] is not None and size[0] != size[1]:
+        raise ValueError("'EdgeIndex' is undirected but received a "
+                         "non-symmetric size")
+
+
 class EdgeIndex(Tensor):
     r"""An advanced :obj:`edge_index` representation with additional (meta)data
     attached.
@@ -83,14 +89,14 @@ class EdgeIndex(Tensor):
 
     * :obj:`sparse_size`: The underlying sparse matrix size
     * :obj:`sort_order`: The sort order (if present), either by row or column.
+    * :obj:`is_undirected`: Whether edges are bidirectional.
 
     Additionally, :class:`EdgeIndex` caches data for fast CSR or CSC conversion
     in case its representation is sorted, such as its :obj:`rowptr` or
-    :obj:`colptr`, or the permutation vectors for fast conversion from CSR to
-    CSC and vice versa.
+    :obj:`colptr`, or the permutation vectors from CSR to CSC and vice versa.
     Caches are filled based on demand (*e.g.*, when calling
     :meth:`EdgeIndex.sort_by`), or when explicitly requested via
-    :meth:`EdgeIndex.fill_cache`, and are maintained and adjusted over its
+    :meth:`EdgeIndex.fill_cache_`, and are maintained and adjusted over its
     lifespan (*e.g.*, when calling :meth:`EdgeIndex.flip`).
 
     This representation ensures for optimal computation in GNN message passing
@@ -107,6 +113,9 @@ class EdgeIndex(Tensor):
     # based on row or column values.
     _sort_order: Optional[SortOrder] = None
 
+    # Whether the `edge_index` is undirected:
+    _is_undirected: bool = False
+
     # An additional data cache:
     _rowptr: Optional[Tensor] = None  # The CSR `rowptr` in case sorted by row.
     _colptr: Optional[Tensor] = None  # The CSC `colptr` in case sorted by col.
@@ -121,6 +130,7 @@ class EdgeIndex(Tensor):
         *args,
         sparse_size: Tuple[Optional[int], Optional[int]] = (None, None),
         sort_order: Optional[Union[str, SortOrder]] = None,
+        is_undirected: bool = False,
         **kwargs,
     ):
         if not isinstance(data, Tensor):
@@ -138,11 +148,19 @@ class EdgeIndex(Tensor):
         assert_two_dimensional(data)
         assert_contiguous(data)
 
+        if is_undirected:
+            assert_symmetric(sparse_size)
+            if sparse_size[0] is not None and sparse_size[1] is None:
+                sparse_size = (sparse_size[0], sparse_size[0])
+            elif sparse_size[0] is None and sparse_size[1] is not None:
+                sparse_size = (sparse_size[1], sparse_size[1])
+
         out = super().__new__(cls, data)
 
         # Attach metadata:
         out._sparse_size = sparse_size
         out._sort_order = None if sort_order is None else SortOrder(sort_order)
+        out._is_undirected = is_undirected
 
         return out
 
@@ -154,6 +172,8 @@ class EdgeIndex(Tensor):
         assert_valid_dtype(self)
         assert_two_dimensional(self)
         assert_contiguous(self)
+        if self.is_undirected:
+            assert_symmetric(self.sparse_size)
 
         if self.numel() > 0 and self.min() < 0:
             raise ValueError(f"'{self.__class__.__name__}' contains negative "
@@ -181,6 +201,13 @@ class EdgeIndex(Tensor):
             raise ValueError(f"'{self.__class__.__name__}' is not sorted by "
                              f"column indices")
 
+        if self.is_undirected:
+            flat_index1 = (self[0] * self.get_num_rows() + self[1]).sort()[0]
+            flat_index2 = (self[1] * self.get_num_cols() + self[0]).sort()[0]
+            if not torch.equal(flat_index1, flat_index2):
+                raise ValueError(f"'{self.__class__.__name__}' is not "
+                                 f"undirected")
+
         return self
 
     @property
@@ -204,10 +231,12 @@ class EdgeIndex(Tensor):
         Automatically computed and cached when not explicitly set.
         """
         if self.num_rows is None:
-            self._sparse_size = (
-                int(self[0].max()) + 1 if self.numel() > 0 else 0,
-                self.num_cols,
-            )
+            num_rows = int(self[0].max()) + 1 if self.numel() > 0 else 0
+
+            if self.is_undirected:
+                self._sparse_size = (num_rows, num_rows)
+            else:
+                self._sparse_size = (num_rows, self.num_cols)
 
         return self.num_rows
 
@@ -221,10 +250,12 @@ class EdgeIndex(Tensor):
         Automatically computed and cached when not explicitly set.
         """
         if self.num_cols is None:
-            self._sparse_size = (
-                self.num_rows,
-                int(self[1].max()) + 1 if self.numel() > 0 else 0,
-            )
+            num_cols = int(self[1].max()) + 1 if self.numel() > 0 else 0
+
+            if self.is_undirected:
+                self._sparse_size = (num_cols, num_cols)
+            else:
+                self._sparse_size = (self.num_rows, num_cols)
 
         return self.num_cols
 
@@ -250,6 +281,11 @@ class EdgeIndex(Tensor):
         r"""Returns whether indices are sorted by columns."""
         return self._sort_order == SortOrder.COL
 
+    @property
+    def is_undirected(self) -> bool:
+        r"""Returns whether indices are bidirectional."""
+        return self._is_undirected
+
     def get_rowptr(self) -> Tensor:
         r"""Returns the :obj:`rowptr` vector of :class:`EdgeIndex`, a
         compressed representation of row indices in case :class:`EdgeIndex` is
@@ -377,8 +413,7 @@ class EdgeIndex(Tensor):
         sort_order = SortOrder(sort_order)
 
         if self._sort_order == sort_order:  # Nothing to do.
-            perm = torch.arange(self.size(1), device=self.device)
-            return torch.return_types.sort([self, perm])
+            return torch.return_types.sort([self, slice(None, None, None)])
 
         # If conversion from CSR->CSC or CSC->CSR is known, make use of it:
         if self.is_sorted_by_row:


Test patch:
diff --git a/test/data/test_edge_index.py b/test/data/test_edge_index.py
index 7113d7041..32f19bef1 100644
--- a/test/data/test_edge_index.py
+++ b/test/data/test_edge_index.py
@@ -22,6 +22,7 @@ from torch_geometric.utils import scatter
 
 def test_basic():
     adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], sparse_size=(3, 3))
+    adj.validate()
     assert isinstance(adj, EdgeIndex)
     assert str(adj) == ('EdgeIndex([[0, 1, 1, 2],\n'
                         '           [1, 0, 2, 1]])')
@@ -32,11 +33,28 @@ def test_basic():
     assert not adj.is_sorted_by_row
     assert not adj.is_sorted_by_col
 
+    assert not adj.is_undirected
+
     assert not isinstance(adj.as_tensor(), EdgeIndex)
 
     assert not isinstance(adj + 1, EdgeIndex)
 
 
+def test_undirected():
+    adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], is_undirected=True)
+    assert isinstance(adj, EdgeIndex)
+    assert adj.is_undirected
+
+    assert adj.sparse_size == (None, None)
+    adj.get_num_rows()
+    assert adj.sparse_size == (3, 3)
+
+    adj.validate()
+
+    with pytest.raises(ValueError, match="'EdgeIndex' is not undirected"):
+        EdgeIndex([[0, 1, 1, 2], [0, 0, 1, 1]], is_undirected=True).validate()
+
+
 def test_fill_cache_():
     adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], sort_order='row')
     adj.validate().fill_cache_()
@@ -129,7 +147,7 @@ def test_sort_by():
     assert isinstance(out.values, EdgeIndex)
     assert not isinstance(out.indices, EdgeIndex)
     assert torch.equal(out.values, adj)
-    assert torch.equal(out.indices, torch.arange(4))
+    assert out.indices == slice(None, None, None)
 
     adj = EdgeIndex([[0, 1, 2, 1], [1, 0, 1, 2]])
     out = adj.sort_by('row')