Repository: pyg-team/pytorch_geometric
Pull request #8736 (2024-01-07 17:37:05)

Include all synthetic data referenced in `GNNExplainer`
A smaller PR for https://github.com/pyg-team/pytorch_geometric/pull/8704
### Include grid motifs and tree-based datasets in GNNExplainer.
- add torch_geometric/datasets/graph_generator/tree_graph.py
- add torch_geometric/datasets/motif_generator/grid.py
- add node label in torch_geometric/datasets/motif_generator/cycle.py (according to bug report https://github.com/pyg-team/pytorch_geometric/issues/8509, it is better to directly add node label since we know what the label should be, rather than be captured at `else` sentense in line 127 of `./torch_geometric/datasets/explainer_dataset.py`)

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8736?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`3fb6ff1`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/3fb6ff13dc890e2aa4763c9e4e687af6a24c1124?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 89.56% compared to head [(`2b21de3`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8736?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 89.57%.


<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##           master    #8736   +/-   ##
=======================================
  Coverage   89.56%   89.57%           
=======================================
  Files         481      481           
  Lines       31048    31048           
=======================================
+ Hits        27809    27810    +1     
+ Misses       3239     3238    -1     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8736?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?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 a3f4fe46b..4e71ac0e5 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 `TreeGraph` and `GridMotif` generators ([#8736](https://github.com/pyg-team/pytorch_geometric/pull/8736))
 - Added an example for edge-level temporal sampling on a heterogenous graph ([#8383](https://github.com/pyg-team/pytorch_geometric/pull/8383))
 - Added the `num_graphs` option to the `StochasticBlockModelDataset` ([#8648](https://github.com/pyg-team/pytorch_geometric/pull/8648))
 - Added noise scheduler utility for diffusion based graph generative models ([#8347](https://github.com/pyg-team/pytorch_geometric/pull/8347))
diff --git a/torch_geometric/datasets/graph_generator/__init__.py b/torch_geometric/datasets/graph_generator/__init__.py
index a6e64e8ac..65298bb92 100644
--- a/torch_geometric/datasets/graph_generator/__init__.py
+++ b/torch_geometric/datasets/graph_generator/__init__.py
@@ -2,10 +2,12 @@ from .base import GraphGenerator
 from .ba_graph import BAGraph
 from .er_graph import ERGraph
 from .grid_graph import GridGraph
+from .tree_graph import TreeGraph
 
 __all__ = classes = [
     'GraphGenerator',
     'BAGraph',
     'ERGraph',
     'GridGraph',
+    'TreeGraph',
 ]
diff --git a/torch_geometric/datasets/graph_generator/tree_graph.py b/torch_geometric/datasets/graph_generator/tree_graph.py
new file mode 100644
index 000000000..6af40858e
--- /dev/null
+++ b/torch_geometric/datasets/graph_generator/tree_graph.py
@@ -0,0 +1,80 @@
+from typing import List, Optional, Tuple
+
+import torch
+from torch import Tensor
+
+from torch_geometric.data import Data
+from torch_geometric.datasets.graph_generator import GraphGenerator
+from torch_geometric.utils import to_undirected
+
+
+def tree(
+    depth: int,
+    branch: int = 2,
+    undirected: bool = False,
+    device: Optional[torch.device] = None,
+) -> Tuple[Tensor, Tensor]:
+    """Generates a tree graph with the given depth and branch size, along with
+    node-level depth indicators.
+
+    Args:
+        depth (int): The depth of the tree.
+        branch (int, optional): The branch size of the tree.
+            (default: :obj:`2`)
+        undirected (bool, optional): If set to :obj:`True`, the tree graph will
+            be undirected. (default: :obj:`False`)
+        device (torch.device, optional): The desired device of the returned
+            tensors. (default: :obj:`None`)
+    """
+    edges: List[Tuple[int, int]] = []
+    depths: List[int] = [0]
+
+    def add_edges(node: int, current_depth: int) -> None:
+        node_count = len(depths)
+
+        if current_depth < depth:
+            for i in range(branch):
+                edges.append((node, node_count + i))
+                depths.append(current_depth + 1)
+
+            for i in range(branch):
+                add_edges(node=node_count + i, current_depth=current_depth + 1)
+
+    add_edges(node=0, current_depth=0)
+
+    edge_index = torch.tensor(edges, device=device).t().contiguous()
+    if undirected:
+        edge_index = to_undirected(edge_index, num_nodes=len(depths))
+
+    return edge_index, torch.tensor(depths, device=device)
+
+
+class TreeGraph(GraphGenerator):
+    r"""Generates tree graphs.
+
+    Args:
+        depth (int): The depth of the tree.
+        branch (int, optional): The branch size of the tree.
+            (default: :obj:`2`)
+        undirected (bool, optional): If set to :obj:`True`, the tree graph will
+            be undirected. (default: :obj:`False`)
+    """
+    def __init__(
+        self,
+        depth: int,
+        branch: int = 2,
+        undirected: bool = False,
+    ) -> None:
+        super().__init__()
+        self.depth = depth
+        self.branch = branch
+        self.undirected = undirected
+
+    def __call__(self) -> Data:
+        edge_index, depth = tree(self.depth, self.branch, self.undirected)
+        num_nodes = depth.numel()
+        return Data(edge_index=edge_index, depth=depth, num_nodes=num_nodes)
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}(depth={self.depth}, '
+                f'branch={self.branch}, undirected={self.undirected})')
diff --git a/torch_geometric/datasets/motif_generator/__init__.py b/torch_geometric/datasets/motif_generator/__init__.py
index b97b63946..ca42aaa54 100644
--- a/torch_geometric/datasets/motif_generator/__init__.py
+++ b/torch_geometric/datasets/motif_generator/__init__.py
@@ -2,10 +2,12 @@ from .base import MotifGenerator
 from .custom import CustomMotif
 from .house import HouseMotif
 from .cycle import CycleMotif
+from .grid import GridMotif
 
 __all__ = classes = [
     'MotifGenerator',
     'CustomMotif',
     'HouseMotif',
     'CycleMotif',
+    'GridMotif',
 ]
diff --git a/torch_geometric/datasets/motif_generator/cycle.py b/torch_geometric/datasets/motif_generator/cycle.py
index 1bb3c6f17..78eaf5542 100644
--- a/torch_geometric/datasets/motif_generator/cycle.py
+++ b/torch_geometric/datasets/motif_generator/cycle.py
@@ -24,7 +24,6 @@ class CycleMotif(CustomMotif):
             num_nodes=num_nodes,
             edge_index=torch.stack([row, col], dim=0),
         )
-
         super().__init__(structure)
 
     def __repr__(self) -> str:
diff --git a/torch_geometric/datasets/motif_generator/grid.py b/torch_geometric/datasets/motif_generator/grid.py
new file mode 100644
index 000000000..997627e30
--- /dev/null
+++ b/torch_geometric/datasets/motif_generator/grid.py
@@ -0,0 +1,44 @@
+import torch
+
+from torch_geometric.data import Data
+from torch_geometric.datasets.motif_generator import CustomMotif
+
+
+class GridMotif(CustomMotif):
+    r"""Generates the grid-structured motif from the
+    `"GNNExplainer: Generating Explanations for Graph Neural Networks"
+    <https://arxiv.org/abs/1903.03894>`__ paper.
+    """
+    def __init__(self) -> None:
+        edge_indices = [
+            [0, 1],
+            [0, 3],
+            [1, 4],
+            [3, 4],
+            [1, 2],
+            [2, 5],
+            [4, 5],
+            [3, 6],
+            [6, 7],
+            [4, 7],
+            [5, 8],
+            [7, 8],
+            [1, 0],
+            [3, 0],
+            [4, 1],
+            [4, 3],
+            [2, 1],
+            [5, 2],
+            [5, 4],
+            [6, 3],
+            [7, 6],
+            [7, 4],
+            [8, 5],
+            [8, 7],
+        ]
+        structure = Data(
+            num_nodes=9,
+            edge_index=torch.tensor(edge_indices).t().contiguous(),
+            y=torch.tensor([0, 1, 0, 1, 2, 1, 0, 1, 0]),
+        )
+        super().__init__(structure)


Test patch:
diff --git a/test/datasets/graph_generator/test_tree_graph.py b/test/datasets/graph_generator/test_tree_graph.py
new file mode 100644
index 000000000..d32ee0adb
--- /dev/null
+++ b/test/datasets/graph_generator/test_tree_graph.py
@@ -0,0 +1,25 @@
+import pytest
+
+from torch_geometric.datasets.graph_generator import TreeGraph
+
+
+@pytest.mark.parametrize('undirected', [False, True])
+def test_tree_graph(undirected):
+    graph_generator = TreeGraph(depth=2, branch=2, undirected=undirected)
+    assert str(graph_generator) == (f'TreeGraph(depth=2, branch=2, '
+                                    f'undirected={undirected})')
+
+    data = graph_generator()
+    assert len(data) == 3
+    assert data.num_nodes == 7
+    assert data.depth.tolist() == [0, 1, 1, 2, 2, 2, 2]
+    if not undirected:
+        assert data.edge_index.tolist() == [
+            [0, 0, 1, 1, 2, 2],
+            [1, 2, 3, 4, 5, 6],
+        ]
+    else:
+        assert data.edge_index.tolist() == [
+            [0, 0, 1, 1, 1, 2, 2, 2, 3, 4, 5, 6],
+            [1, 2, 0, 3, 4, 0, 5, 6, 1, 1, 2, 2],
+        ]
diff --git a/test/datasets/motif_generator/test_grid_motif.py b/test/datasets/motif_generator/test_grid_motif.py
new file mode 100644
index 000000000..84271a094
--- /dev/null
+++ b/test/datasets/motif_generator/test_grid_motif.py
@@ -0,0 +1,17 @@
+from torch_geometric.datasets.motif_generator import GridMotif
+
+
+def test_grid_motif():
+    motif_generator = GridMotif()
+    assert str(motif_generator) == 'GridMotif()'
+
+    motif = motif_generator()
+    assert len(motif) == 3
+    assert motif.num_nodes == 9
+    assert motif.num_edges == 24
+    assert motif.edge_index.size() == (2, 24)
+    assert motif.edge_index.min() == 0
+    assert motif.edge_index.max() == 8
+    assert motif.y.size() == (9, )
+    assert motif.y.min() == 0
+    assert motif.y.max() == 2