Repository: pyg-team/pytorch_geometric
Pull request #9756 (2024-11-01 08:34:52)

Fixing `edge_mask` handling for directed graphs in `k_hop_subgraph`
Example use case:

```python
from torch_geometric.utils import k_hop_subgraph
import torch

edge_index = torch.tensor([[1, 2, 3],
                           [0, 1, 1]])
# get the 2-hop neighbors of node 0 in the directed graph.
_, edge_index, _, edge_mask = k_hop_subgraph(0, 2, edge_index, relabel_nodes=False, directed=True)
```

This gives the following result:

Expected Outcome:
```python
>>> edge_index
tensor([[1, 2, 3],
        [0, 1, 1]])

>>> edge_mask
tensor([True,  True,  True])
```

Actual Outcome:
```python
>>> edge_index
tensor([[2, 3],
        [1, 1]])
>>> edge_mask
tensor([False,  True,  True])
```

This stems from the fact that the line `torch.index_select(node_mask, 0, row, out=edge_mask)` overwrites `edge_mask`, effectively only marking the edges used in the final hop as `True`.

To fix this, I have added `preserved_edge_mask ` that will mark the edges used in each hop as `True`.





Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9b98a9bed..7a6f07cb2 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added support for fast `Delaunay()` triangulation via the `torch_delaunay` package ([#9748](https://github.com/pyg-team/pytorch_geometric/pull/9748))
 - Added PyTorch 2.5 support ([#9779](https://github.com/pyg-team/pytorch_geometric/pull/9779))
 - Support 3D tetrahedral mesh elements of shape `[4, num_faces]` in the `FaceToEdge` transformation ([#9776](https://github.com/pyg-team/pytorch_geometric/pull/9776))
 - Added the `use_pcst` option to `WebQSPDataset` ([#9722](https://github.com/pyg-team/pytorch_geometric/pull/9722))
@@ -21,6 +22,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
+- Fixed the `k_hop_subgraph()` method for directed graphs ([#9756](https://github.com/pyg-team/pytorch_geometric/pull/9756))
 - Fixed `utils.group_cat` concatenating dimension ([#9766](https://github.com/pyg-team/pytorch_geometric/pull/9766))
 - Fixed `WebQSDataset.process` raising exceptions ([#9665](https://github.com/pyg-team/pytorch_geometric/pull/9665))
 
diff --git a/torch_geometric/transforms/delaunay.py b/torch_geometric/transforms/delaunay.py
index 8ea90333a..5c2d469ad 100644
--- a/torch_geometric/transforms/delaunay.py
+++ b/torch_geometric/transforms/delaunay.py
@@ -1,3 +1,5 @@
+from typing import List
+
 import torch
 
 from torch_geometric.data import Data
@@ -5,30 +7,78 @@ from torch_geometric.data.datapipes import functional_transform
 from torch_geometric.transforms import BaseTransform
 
 
+class _QhullTransform(BaseTransform):
+    r"""Q-hull implementation of delaunay triangulation."""
+    def forward(self, data: Data) -> Data:
+        assert data.pos is not None
+        import scipy.spatial
+
+        pos = data.pos.cpu().numpy()
+        tri = scipy.spatial.Delaunay(pos, qhull_options='QJ')
+        face = torch.from_numpy(tri.simplices)
+
+        data.face = face.t().contiguous().to(data.pos.device, torch.long)
+        return data
+
+
+class _ShullTransform(BaseTransform):
+    r"""Sweep-hull implementation of delaunay triangulation."""
+    def forward(self, data: Data) -> Data:
+        assert data.pos is not None
+        from torch_delaunay.functional import shull2d
+
+        face = shull2d(data.pos.cpu())
+        data.face = face.t().contiguous().to(data.pos.device)
+        return data
+
+
+class _SequentialTransform(BaseTransform):
+    r"""Runs the first successful transformation.
+
+    All intermediate exceptions are suppressed except the last.
+    """
+    def __init__(self, transforms: List[BaseTransform]) -> None:
+        assert len(transforms) > 0
+        self.transforms = transforms
+
+    def forward(self, data: Data) -> Data:
+        for i, transform in enumerate(self.transforms):
+            try:
+                return transform.forward(data)
+            except ImportError as e:
+                if i == len(self.transforms) - 1:
+                    raise e
+        return data
+
+
 @functional_transform('delaunay')
 class Delaunay(BaseTransform):
     r"""Computes the delaunay triangulation of a set of points
     (functional name: :obj:`delaunay`).
+
+    .. hint::
+        Consider installing the
+        `torch_delaunay <https://github.com/ybubnov/torch_delaunay>`_ package
+        to speed up computation.
     """
-    def forward(self, data: Data) -> Data:
-        import scipy.spatial
+    def __init__(self) -> None:
+        self._transform = _SequentialTransform([
+            _ShullTransform(),
+            _QhullTransform(),
+        ])
 
+    def forward(self, data: Data) -> Data:
         assert data.pos is not None
+        device = data.pos.device
 
         if data.pos.size(0) < 2:
-            data.edge_index = torch.tensor([], dtype=torch.long,
-                                           device=data.pos.device).view(2, 0)
-        if data.pos.size(0) == 2:
-            data.edge_index = torch.tensor([[0, 1], [1, 0]], dtype=torch.long,
-                                           device=data.pos.device)
+            data.edge_index = torch.empty(2, 0, dtype=torch.long,
+                                          device=device)
+        elif data.pos.size(0) == 2:
+            data.edge_index = torch.tensor([[0, 1], [1, 0]], device=device)
         elif data.pos.size(0) == 3:
-            data.face = torch.tensor([[0], [1], [2]], dtype=torch.long,
-                                     device=data.pos.device)
-        if data.pos.size(0) > 3:
-            pos = data.pos.cpu().numpy()
-            tri = scipy.spatial.Delaunay(pos, qhull_options='QJ')
-            face = torch.from_numpy(tri.simplices)
-
-            data.face = face.t().contiguous().to(data.pos.device, torch.long)
+            data.face = torch.tensor([[0], [1], [2]], device=device)
+        else:
+            data = self._transform.forward(data)
 
         return data
diff --git a/torch_geometric/utils/_subgraph.py b/torch_geometric/utils/_subgraph.py
index 5c2dcd77c..7540cc9f2 100644
--- a/torch_geometric/utils/_subgraph.py
+++ b/torch_geometric/utils/_subgraph.py
@@ -346,10 +346,12 @@ def k_hop_subgraph(
 
     subsets = [node_idx]
 
+    preserved_edge_mask = torch.zeros_like(edge_mask)
     for _ in range(num_hops):
         node_mask.fill_(False)
         node_mask[subsets[-1]] = True
         torch.index_select(node_mask, 0, row, out=edge_mask)
+        preserved_edge_mask |= edge_mask
         subsets.append(col[edge_mask])
 
     subset, inv = torch.cat(subsets).unique(return_inverse=True)
@@ -360,6 +362,8 @@ def k_hop_subgraph(
 
     if not directed:
         edge_mask = node_mask[row] & node_mask[col]
+    else:
+        edge_mask = preserved_edge_mask
 
     edge_index = edge_index[:, edge_mask]
 


Test patch:
diff --git a/test/transforms/test_delaunay.py b/test/transforms/test_delaunay.py
index fc027a2ea..4fac3917c 100644
--- a/test/transforms/test_delaunay.py
+++ b/test/transforms/test_delaunay.py
@@ -5,30 +5,60 @@ from torch_geometric.testing import withPackage
 from torch_geometric.transforms import Delaunay
 
 
-@withPackage('scipy')
-def test_delaunay():
-    assert str(Delaunay()) == 'Delaunay()'
+def assert_one_point(transform: Delaunay) -> None:
+    data = Data(pos=torch.rand(1, 2))
+    data = transform(data)
+    assert len(data) == 2
+    assert data.edge_index.tolist() == [[], []]
 
-    pos = torch.tensor([[-1, -1], [-1, 1], [1, 1], [1, -1]], dtype=torch.float)
-    data = Data(pos=pos)
-    data = Delaunay()(data)
+
+def assert_two_points(transform: Delaunay) -> None:
+    data = Data(pos=torch.rand(2, 2))
+    data = transform(data)
     assert len(data) == 2
-    assert data.face.tolist() == [[3, 1], [1, 3], [0, 2]]
+    assert data.edge_index.tolist() == [[0, 1], [1, 0]]
 
-    pos = torch.tensor([[-1, -1], [-1, 1], [1, 1]], dtype=torch.float)
-    data = Data(pos=pos)
-    data = Delaunay()(data)
+
+def assert_three_points(transform: Delaunay) -> None:
+    data = Data(pos=torch.rand(3, 2))
+    data = transform(data)
     assert len(data) == 2
     assert data.face.tolist() == [[0], [1], [2]]
 
-    pos = torch.tensor([[-1, -1], [1, 1]], dtype=torch.float)
-    data = Data(pos=pos)
-    data = Delaunay()(data)
-    assert len(data) == 2
-    assert data.edge_index.tolist() == [[0, 1], [1, 0]]
 
-    pos = torch.tensor([[-1, -1]], dtype=torch.float)
+def assert_four_points(transform: Delaunay) -> None:
+    pos = torch.tensor([
+        [-1.0, -1.0],
+        [-1.0, 1.0],
+        [1.0, 1.0],
+        [0.5, -0.5],
+    ])
     data = Data(pos=pos)
-    data = Delaunay()(data)
+    data = transform(data)
     assert len(data) == 2
-    assert data.edge_index.tolist() == [[], []]
+
+    # The order of the simplices does not matter, therefore assert sets.
+    faces = set(map(tuple, data.face.tolist()))
+    assert faces == {(3, 1), (1, 3), (0, 2)}
+
+
+@withPackage('scipy')
+def test_qhull_delaunay() -> None:
+    transform = Delaunay()
+
+    assert str(transform) == 'Delaunay()'
+    assert_one_point(transform)
+    assert_two_points(transform)
+    assert_three_points(transform)
+    assert_four_points(transform)
+
+
+@withPackage('torch_delaunay')
+def test_shull_delaunay() -> None:
+    transform = Delaunay()
+
+    assert str(transform) == 'Delaunay()'
+    assert_one_point(transform)
+    assert_two_points(transform)
+    assert_three_points(transform)
+    assert_four_points(transform)
diff --git a/test/utils/test_subgraph.py b/test/utils/test_subgraph.py
index 85161c55d..462bcb189 100644
--- a/test/utils/test_subgraph.py
+++ b/test/utils/test_subgraph.py
@@ -101,9 +101,12 @@ def test_k_hop_subgraph():
         [0, 1, 2, 3, 4, 5],
         [2, 2, 4, 4, 6, 6],
     ])
-
     subset, edge_index, mapping, edge_mask = k_hop_subgraph(
-        6, 2, edge_index, relabel_nodes=True)
+        node_idx=6,
+        num_hops=2,
+        edge_index=edge_index,
+        relabel_nodes=True,
+    )
     assert subset.tolist() == [2, 3, 4, 5, 6]
     assert edge_index.tolist() == [[0, 1, 2, 3], [2, 2, 4, 4]]
     assert mapping.tolist() == [4]
@@ -113,12 +116,29 @@ def test_k_hop_subgraph():
         [1, 2, 4, 5],
         [0, 1, 5, 6],
     ])
-
-    subset, edge_index, mapping, edge_mask = k_hop_subgraph([0, 6], 2,
-                                                            edge_index,
-                                                            relabel_nodes=True)
-
+    subset, edge_index, mapping, edge_mask = k_hop_subgraph(
+        node_idx=[0, 6],
+        num_hops=2,
+        edge_index=edge_index,
+        relabel_nodes=True,
+    )
     assert subset.tolist() == [0, 1, 2, 4, 5, 6]
     assert edge_index.tolist() == [[1, 2, 3, 4], [0, 1, 4, 5]]
     assert mapping.tolist() == [0, 5]
     assert edge_mask.tolist() == [True, True, True, True]
+
+    edge_index = torch.tensor([
+        [0, 1, 2, 3, 4, 4, 5],
+        [2, 2, 4, 4, 2, 6, 6],
+    ])
+    subset, edge_index, mapping, edge_mask = k_hop_subgraph(
+        node_idx=6,
+        num_hops=2,
+        edge_index=edge_index,
+        relabel_nodes=False,
+        directed=True,
+    )
+    assert subset.tolist() == [2, 3, 4, 5, 6]
+    assert edge_index.tolist() == [[2, 3, 4, 5], [4, 4, 6, 6]]
+    assert mapping.tolist() == [4]
+    assert edge_mask.tolist() == [False, False, True, True, False, True, True]