graclus assert method

setup.py

@@ -5,7 +5,7 @@ from setuptools import setup, find_packages
 __version__ = '1.0.0'
 url = 'https://github.com/rusty1s/pytorch_cluster'
 
-install_requires = ['cffi', 'torch-unique']
+install_requires = ['cffi']
 setup_requires = ['pytest-runner', 'cffi']
 tests_require = ['pytest', 'pytest-cov']
 

test/test_graclus.py

 import torch
+import numpy as np
 from torch_cluster import graclus_cluster
 
 
+def assert_correct_graclus(row, col, cluster):
+    row, col, cluster = row.numpy(), col.numpy(), cluster.numpy()
+
+    # Every node was assigned a cluster.
+    assert cluster.min() >= 0
+
+    # There are no more than two nodes in each cluster.
+    _, count = np.unique(cluster, return_counts=True)
+    assert (count > 2).max() == 0
+
+    # Corresponding clusters must be adjacent.
+    for n in range(cluster.shape[0]):
+        assert (cluster[col[row == n]] == cluster[n]).max() == 1
+
+
 def test_graclus_cluster_cpu():
     row = torch.LongTensor([0, 0, 1, 1, 1, 2, 2, 2, 3, 3])
     col = torch.LongTensor([1, 2, 0, 2, 3, 0, 1, 3, 1, 2])
-    weight_ = torch.Tensor([1, 2, 1, 3, 2, 2, 3, 1, 2, 1])
+    weight = torch.Tensor([1, 2, 1, 3, 2, 2, 3, 1, 2, 1])
+
     cluster = graclus_cluster(row, col)
-    cluster = graclus_cluster(row, col, weight_)
+    assert_correct_graclus(row, col, cluster)
+
+    cluster = graclus_cluster(row, col, weight)
+    assert_correct_graclus(row, col, cluster)