From f490606527df3bee94ec5e857cec022938938b40 Mon Sep 17 00:00:00 2001
From: rusty1s <matthias.fey@tu-dortmund.de>
Date: Sat, 17 Dec 2022 12:10:17 +0000
Subject: [PATCH] update

---
 .../motif_generator/test_cycle_motif.py       | 13 ++++++----
 .../datasets/motif_generator/__init__.py      |  5 +++-
 .../datasets/motif_generator/cycle.py         | 25 ++++++++++++-------
 3 files changed, 28 insertions(+), 15 deletions(-)

diff --git a/test/datasets/motif_generator/test_cycle_motif.py b/test/datasets/motif_generator/test_cycle_motif.py
index 6389c53f3acb..35cafb58f1af 100644
--- a/test/datasets/motif_generator/test_cycle_motif.py
+++ b/test/datasets/motif_generator/test_cycle_motif.py
@@ -2,11 +2,14 @@
 
 
 def test_cycle_motif():
-    n = 5
-    motif_generator = CycleMotif(n)
-    assert str(motif_generator) == 'CycleMotif()'
+    motif_generator = CycleMotif(5)
+    assert str(motif_generator) == 'CycleMotif(5)'
 
     motif = motif_generator()
     assert len(motif) == 2
-    assert motif.num_nodes == n
-    assert motif.num_edges == n
+    assert motif.num_nodes == 5
+    assert motif.num_edges == 10
+    assert motif.edge_index.tolist() == [
+        [0, 0, 1, 1, 2, 2, 3, 3, 4, 4],
+        [1, 4, 0, 2, 1, 3, 2, 4, 0, 3],
+    ]
diff --git a/torch_geometric/datasets/motif_generator/__init__.py b/torch_geometric/datasets/motif_generator/__init__.py
index 4d9d7c4b9083..b97b639468a1 100644
--- a/torch_geometric/datasets/motif_generator/__init__.py
+++ b/torch_geometric/datasets/motif_generator/__init__.py
@@ -4,5 +4,8 @@
 from .cycle import CycleMotif
 
 __all__ = classes = [
-    'MotifGenerator', 'CustomMotif', 'HouseMotif', 'CycleMotif'
+    'MotifGenerator',
+    'CustomMotif',
+    'HouseMotif',
+    'CycleMotif',
 ]
diff --git a/torch_geometric/datasets/motif_generator/cycle.py b/torch_geometric/datasets/motif_generator/cycle.py
index a2f56cd64724..4ffdfd8999fa 100644
--- a/torch_geometric/datasets/motif_generator/cycle.py
+++ b/torch_geometric/datasets/motif_generator/cycle.py
@@ -7,18 +7,25 @@
 class CycleMotif(CustomMotif):
     r"""Generates the cycle motif from the `"GNNExplainer:
     Generating Explanations for Graph Neural Networks"
-    <https://arxiv.org/pdf/1903.03894.pdf>`_ paper, containing n nodes and n
-    undirected edges.
+    <https://arxiv.org/pdf/1903.03894.pdf>`_ paper.
 
     Args:
-        n (int): Number of nodes (or edges) in the cycle.
+        num_nodes (int): The number of nodes in the cycle.
     """
-    def __init__(self, n: int = 6):
-        # construct edge_index based on n
+    def __init__(self, num_nodes: int):
+        self.num_nodes = num_nodes
+
+        row = torch.arange(num_nodes).view(-1, 1).repeat(1, 2).view(-1)
+        col1 = torch.arange(-1, num_nodes - 1) % num_nodes
+        col2 = torch.arange(1, num_nodes + 1) % num_nodes
+        col = torch.stack([col1, col2], dim=1).sort(dim=-1)[0].view(-1)
+
         structure = Data(
-            num_nodes=n,
-            edge_index=torch.Tensor([[x for x in range(n)],
-                                     [y for y in range(1, n)] + [0]
-                                     ]).type(torch.int32),
+            num_nodes=num_nodes,
+            edge_index=torch.stack([row, col], dim=0),
         )
+
         super().__init__(structure)
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}({self.num_nodes})'