Repository: pyg-team/pytorch_geometric
Pull request #7199 (2023-04-18 06:43:40)

Add `GraphMode` with `bidirectional` sampling support (1/2)

codecov[bot]: ## [Codecov](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7199?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7199](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7199?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (35ccf1d) into [master](https://codecov.io/gh/pyg-team/pytorch_geometric/commit/bff64052b43ecd80d2157e271292947b833c7700?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (bff6405) will **decrease** coverage by `0.43%`.
> The diff coverage is `97.95%`.

> :exclamation: Current head 35ccf1d differs from pull request most recent head 3b6b4e8. Consider uploading reports for the commit 3b6b4e8 to get more accurate results

```diff
@@            Coverage Diff             @@
##           master    #7199      +/-   ##
==========================================
- Coverage   92.10%   91.68%   -0.43%     
==========================================
  Files         436      436              
  Lines       23937    23980      +43     
==========================================
- Hits        22048    21985      -63     
- Misses       1889     1995     +106     
```


| [Impacted Files](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7199?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/testing/data.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7199?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3Rlc3RpbmcvZGF0YS5weQ==) | `40.62% <80.00%> (+4.91%)` | :arrow_up: |
| [torch\_geometric/sampler/base.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7199?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3NhbXBsZXIvYmFzZS5weQ==) | `98.33% <100.00%> (+0.30%)` | :arrow_up: |
| [torch\_geometric/sampler/utils.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7199?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3NhbXBsZXIvdXRpbHMucHk=) | `93.65% <100.00%> (+1.49%)` | :arrow_up: |
| [torch\_geometric/utils/coalesce.py](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7199?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3V0aWxzL2NvYWxlc2NlLnB5) | `100.00% <100.00%> (ø)` | |

... and [21 files with indirect coverage changes](https://codecov.io/gh/pyg-team/pytorch_geometric/pull/7199/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0c318546d..a0614968e 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 graph mode sampling option with bidirectional sampling support ([#7199](https://github.com/pyg-team/pytorch_geometric/pull/7199))
 - Added support for `"any"`-reductions in `scatter` ([#7198](https://github.com/pyg-team/pytorch_geometric/pull/7198))
 - Added manual sampling interface to `NodeLoader` and `LinkLoader` ([#7197](https://github.com/pyg-team/pytorch_geometric/pull/7197))
 - Extending `torch.sparse` support ([#7155](https://github.com/pyg-team/pytorch_geometric/pull/7155))
diff --git a/torch_geometric/sampler/base.py b/torch_geometric/sampler/base.py
index 9b34cce97..d53530c53 100644
--- a/torch_geometric/sampler/base.py
+++ b/torch_geometric/sampler/base.py
@@ -1,5 +1,6 @@
 import copy
 import math
+import warnings
 from abc import ABC
 from dataclasses import dataclass
 from enum import Enum
@@ -9,6 +10,7 @@ import torch
 from torch import Tensor
 
 from torch_geometric.data import Data, FeatureStore, GraphStore, HeteroData
+from torch_geometric.sampler.utils import to_bidirectional
 from torch_geometric.typing import EdgeType, EdgeTypeStr, NodeType, OptTensor
 from torch_geometric.utils.mixin import CastMixin
 
@@ -35,6 +37,13 @@ class DataType(Enum):
                          f"(got '{type(data)}')")
 
 
+class GraphMode(Enum):
+    r"""The graph mode of the returned subgraph."""
+    directional = 'directional'
+    bidirectional = 'bidirectional'
+    induced_subgraph = 'induced_subgraph'
+
+
 @dataclass
 class NodeSamplerInput(CastMixin):
     r"""The sampling input of
@@ -144,6 +153,23 @@ class SamplerOutput(CastMixin):
     # API for the expected output of a sampler.
     metadata: Optional[Any] = None
 
+    def to_bidirectional(self) -> 'SamplerOutput':
+        r"""Converts the sampled subgraph into a bidirectional variant, in
+        which all sampled edges are guaranteed to be bidirectional."""
+        out = copy.copy(self)
+
+        out.row, out.col, out.edge = to_bidirectional(
+            row=self.row,
+            col=self.col,
+            rev_row=self.row,
+            rev_col=self.col,
+            edge_id=self.edge,
+            rev_edge_id=self.edge,
+        )
+        out.num_sampled_edges = None
+
+        return out
+
 
 @dataclass
 class HeteroSamplerOutput(CastMixin):
@@ -191,6 +217,45 @@ class HeteroSamplerOutput(CastMixin):
     # API for the expected output of a sampler.
     metadata: Optional[Any] = None
 
+    def to_bidirectional(self) -> 'SamplerOutput':
+        r"""Converts the sampled subgraph into a bidirectional variant, in
+        which all sampled edges are guaranteed to be bidirectional."""
+        out = copy.copy(self)
+        out.row = copy.copy(self.row)
+        out.col = copy.copy(self.col)
+
+        edge_types = self.row.keys()
+        edge_types = [k for k in edge_types if not k[1].startswith('rev_')]
+        for edge_type in edge_types:
+            src, rel, dst = edge_type
+            rev_edge_type = (dst, f'rev_{rel}', src)
+
+            if src == dst and rev_edge_type not in self.row:
+                out.row[edge_type], out.col[edge_type], _ = to_bidirectional(
+                    row=self.row[edge_type],
+                    col=self.col[edge_type],
+                    rev_row=self.row[edge_type],
+                    rev_col=self.col[edge_type],
+                )
+            elif rev_edge_type in self.row:
+                out.row[edge_type], out.col[edge_type], _ = to_bidirectional(
+                    row=self.row[edge_type],
+                    col=self.col[edge_type],
+                    rev_row=self.row[rev_edge_type],
+                    rev_col=self.col[rev_edge_type],
+                )
+                out.row[rev_edge_type] = out.col[edge_type]
+                out.col[rev_edge_type] = out.row[edge_type]
+            else:
+                warnings.warn(f"Cannot convert to bidirectional graph since "
+                              f"the edge type {edge_type} does not seem to "
+                              f"have a reverse edge type")
+
+        out.edge = None
+        out.num_sampled_edges = None
+
+        return out
+
 
 @dataclass(frozen=True)
 class NumNeighbors:
diff --git a/torch_geometric/sampler/utils.py b/torch_geometric/sampler/utils.py
index 15ee49a78..35cf9392e 100644
--- a/torch_geometric/sampler/utils.py
+++ b/torch_geometric/sampler/utils.py
@@ -7,7 +7,7 @@ from torch import Tensor
 from torch_geometric.data import Data, HeteroData
 from torch_geometric.data.storage import EdgeStorage
 from torch_geometric.typing import NodeType, OptTensor
-from torch_geometric.utils import index_sort
+from torch_geometric.utils import coalesce, index_sort
 from torch_geometric.utils.sparse import index2ptr
 
 # Edge Layout Conversion ######################################################
@@ -112,6 +112,32 @@ def to_hetero_csc(
     return colptr_dict, row_dict, perm_dict
 
 
+def to_bidirectional(
+    row: Tensor,
+    col: Tensor,
+    rev_row: Tensor,
+    rev_col: Tensor,
+    edge_id: OptTensor = None,
+    rev_edge_id: OptTensor = None,
+) -> Tuple[Tensor, Tensor, OptTensor]:
+
+    assert row.numel() == col.numel()
+    assert rev_row.numel() == rev_col.numel()
+
+    edge_index = row.new_empty(2, row.numel() + rev_row.numel())
+    edge_index[0, :row.numel()] = row
+    edge_index[1, :row.numel()] = col
+    edge_index[0, row.numel():] = rev_col
+    edge_index[1, row.numel():] = rev_row
+
+    if edge_id is not None:
+        edge_id = torch.cat([edge_id, rev_edge_id], dim=0)
+
+    (row, col), edge_id = coalesce(edge_index, edge_id, reduce='any')
+
+    return row, col, edge_id
+
+
 ###############################################################################
 
 X, Y = TypeVar('X'), TypeVar('Y')
diff --git a/torch_geometric/utils/coalesce.py b/torch_geometric/utils/coalesce.py
index 92b54643e..20e1b899d 100644
--- a/torch_geometric/utils/coalesce.py
+++ b/torch_geometric/utils/coalesce.py
@@ -117,7 +117,7 @@ def coalesce(
     edge_index = edge_index[:, mask]
 
     dim_size: Optional[int] = None
-    if isinstance(edge_attr, (Tensor, list, tuple)):
+    if isinstance(edge_attr, (Tensor, list, tuple)) and len(edge_attr) > 0:
         dim_size = edge_index.size(1)
         idx = torch.arange(0, nnz, device=edge_index.device)
         idx.sub_(mask.logical_not_().cumsum(dim=0))
@@ -128,6 +128,8 @@ def coalesce(
         edge_attr = scatter(edge_attr, idx, 0, dim_size, reduce)
         return edge_index, edge_attr
     if isinstance(edge_attr, (list, tuple)):
+        if len(edge_attr) == 0:
+            return edge_index, edge_attr
         edge_attr = [scatter(e, idx, 0, dim_size, reduce) for e in edge_attr]
         return edge_index, edge_attr
 


Test patch:
diff --git a/test/sampler/test_sampler_base.py b/test/sampler/test_sampler_base.py
index a1edf5783..9d57d1250 100644
--- a/test/sampler/test_sampler_base.py
+++ b/test/sampler/test_sampler_base.py
@@ -1,6 +1,13 @@
 import pytest
+import torch
 
-from torch_geometric.sampler.base import NumNeighbors
+from torch_geometric.sampler.base import (
+    HeteroSamplerOutput,
+    NumNeighbors,
+    SamplerOutput,
+)
+from torch_geometric.testing import get_random_edge_index
+from torch_geometric.utils import is_undirected
 
 
 def test_homogeneous_num_neighbors():
@@ -61,3 +68,51 @@ def test_heterogeneous_num_neighbors_empty_dict():
     assert values == {'A__to__B': [25, 10], 'B__to__A': [25, 10]}
 
     assert num_neighbors.num_hops == 2
+
+
+def test_homogeneous_to_bidirectional():
+    edge_index = get_random_edge_index(10, 10, num_edges=20)
+
+    obj = SamplerOutput(
+        node=torch.arange(10),
+        row=edge_index[0],
+        col=edge_index[0],
+        edge=torch.arange(edge_index.size(1)),
+    ).to_bidirectional()
+
+    assert is_undirected(torch.stack([obj.row, obj.col], dim=0))
+
+
+def test_heterogeneous_to_bidirectional():
+    edge_index1 = get_random_edge_index(10, 5, num_edges=20)
+    edge_index2 = get_random_edge_index(5, 10, num_edges=20)
+    edge_index3 = get_random_edge_index(10, 10, num_edges=20)
+
+    obj = HeteroSamplerOutput(
+        node={
+            'v1': torch.arange(10),
+            'v2': torch.arange(5)
+        },
+        row={
+            ('v1', 'to', 'v2'): edge_index1[0],
+            ('v2', 'rev_to', 'v1'): edge_index2[0],
+            ('v1', 'to', 'v1'): edge_index3[0],
+        },
+        col={
+            ('v1', 'to', 'v2'): edge_index1[1],
+            ('v2', 'rev_to', 'v1'): edge_index2[1],
+            ('v1', 'to', 'v1'): edge_index3[1],
+        },
+        edge=None,
+    ).to_bidirectional()
+
+    assert torch.equal(
+        obj.row['v1', 'to', 'v2'].sort().values,
+        obj.col['v2', 'rev_to', 'v1'].sort().values,
+    )
+    assert torch.equal(
+        obj.col['v1', 'to', 'v2'].sort().values,
+        obj.row['v2', 'rev_to', 'v1'].sort().values,
+    )
+    assert is_undirected(
+        torch.stack([obj.row['v1', 'to', 'v1'], obj.col['v1', 'to', 'v1']], 0))
diff --git a/torch_geometric/testing/data.py b/torch_geometric/testing/data.py
index 9e8c776d5..fbb74349b 100644
--- a/torch_geometric/testing/data.py
+++ b/torch_geometric/testing/data.py
@@ -4,6 +4,7 @@ import torch
 from torch import Tensor
 
 from torch_geometric.data import HeteroData, InMemoryDataset
+from torch_geometric.utils import coalesce as coalesce_fn
 
 
 def get_random_edge_index(
@@ -12,12 +13,18 @@ def get_random_edge_index(
     num_edges: int,
     dtype: Optional[torch.dtype] = None,
     device: Optional[torch.device] = None,
+    coalesce: bool = False,
 ) -> Tensor:
     row = torch.randint(num_src_nodes, (num_edges, ), dtype=dtype,
                         device=device)
     col = torch.randint(num_dst_nodes, (num_edges, ), dtype=dtype,
                         device=device)
-    return torch.stack([row, col], dim=0)
+    edge_index = torch.stack([row, col], dim=0)
+
+    if coalesce:
+        edge_index = coalesce_fn(edge_index)
+
+    return edge_index
 
 
 class FakeHeteroDataset(InMemoryDataset):