added position translate

setup.py

@@ -2,7 +2,7 @@ from os import path as osp
 
 from setuptools import setup, find_packages
 
-__version__ = '0.1.1'
+__version__ = '0.2.0'
 url = 'https://github.com/rusty1s/pytorch_cluster'
 
 install_requires = ['cffi', 'torch-unique']

test/test_grid.py

@@ -7,62 +7,98 @@ from .utils import tensors, Tensor
 
 @pytest.mark.parametrize('tensor', tensors)
 def test_grid_cluster_cpu(tensor):
-    position = Tensor(tensor, [0, 9, 2, 8, 3])
+    position = Tensor(tensor, [2, 6])
     size = torch.LongTensor([5])
-    expected = torch.LongTensor([0, 1, 0, 1, 0])
-    output = grid_cluster(position, size)
+    expected = torch.LongTensor([0, 0])
+    output, _ = grid_cluster(position, size)
+    assert output.tolist() == expected.tolist()
+
+    expected = torch.LongTensor([0, 1])
+    output, _ = grid_cluster(position, size, offset=0)
+    assert output.tolist() == expected.tolist()
+
+    position = Tensor(tensor, [0, 17, 2, 8, 3])
+    expected = torch.LongTensor([0, 2, 0, 1, 0])
+    output, _ = grid_cluster(position, size)
+    assert output.tolist() == expected.tolist()
+
+    output, _ = grid_cluster(position, size, fake_nodes=True)
+    expected = torch.LongTensor([0, 3, 0, 1, 0])
     assert output.tolist() == expected.tolist()
 
     position = Tensor(tensor, [[0, 0], [9, 9], [2, 8], [2, 2], [8, 3]])
     size = torch.LongTensor([5, 5])
     expected = torch.LongTensor([0, 3, 1, 0, 2])
-    output = grid_cluster(position, size)
+    output, _ = grid_cluster(position, size)
     assert output.tolist() == expected.tolist()
 
-    position = Tensor(tensor, [[0, 9, 2, 2, 8], [0, 9, 8, 2, 3]]).t()
-    output = grid_cluster(position, size)
+    position = Tensor(tensor, [[0, 11, 2, 2, 8], [0, 9, 8, 2, 3]]).t()
+    output, _ = grid_cluster(position, size)
     assert output.tolist() == expected.tolist()
 
-    output = grid_cluster(position.expand(2, 5, 2), size)
+    output, _ = grid_cluster(position.expand(2, 5, 2), size)
     assert output.tolist() == expected.expand(2, 5).tolist()
 
     position = position.repeat(2, 1)
     batch = torch.LongTensor([0, 0, 0, 0, 0, 1, 1, 1, 1, 1])
     expected = torch.LongTensor([0, 3, 1, 0, 2, 4, 7, 5, 4, 6])
-    expected_batch = torch.LongTensor([0, 0, 0, 0, 1, 1, 1, 1])
-    output, reduced_batch = grid_cluster(position, size, batch)
+    expected_batch2 = torch.LongTensor([0, 0, 0, 0, 1, 1, 1, 1])
+    output, batch2 = grid_cluster(position, size, batch)
+    assert output.tolist() == expected.tolist()
+    assert batch2.tolist() == expected_batch2.tolist()
+
+    output, C = grid_cluster(position, size, batch, fake_nodes=True)
+    expected = torch.LongTensor([0, 5, 1, 0, 2, 6, 11, 7, 6, 8])
     assert output.tolist() == expected.tolist()
-    assert reduced_batch.tolist() == expected_batch.tolist()
+    assert C == 6
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason='no CUDA')
 @pytest.mark.parametrize('tensor', tensors)
 def test_grid_cluster_gpu(tensor):  # pragma: no cover
-    position = Tensor(tensor, [0, 9, 2, 8, 3]).cuda()
+    position = Tensor(tensor, [2, 6]).cuda()
     size = torch.cuda.LongTensor([5])
-    expected = torch.cuda.LongTensor([0, 1, 0, 1, 0])
-    output = grid_cluster(position, size)
+    expected = torch.LongTensor([0, 0])
+    output, _ = grid_cluster(position, size)
+    assert output.cpu().tolist() == expected.tolist()
+
+    expected = torch.LongTensor([0, 1])
+    output, _ = grid_cluster(position, size, offset=0)
+    assert output.cpu().tolist() == expected.tolist()
+
+    position = Tensor(tensor, [0, 17, 2, 8, 3]).cuda()
+    expected = torch.LongTensor([0, 2, 0, 1, 0])
+    output, _ = grid_cluster(position, size)
+    assert output.cpu().tolist() == expected.tolist()
+
+    output, _ = grid_cluster(position, size, fake_nodes=True)
+    expected = torch.LongTensor([0, 3, 0, 1, 0])
     assert output.cpu().tolist() == expected.tolist()
 
     position = Tensor(tensor, [[0, 0], [9, 9], [2, 8], [2, 2], [8, 3]])
     position = position.cuda()
     size = torch.cuda.LongTensor([5, 5])
-    expected = torch.cuda.LongTensor([0, 3, 1, 0, 2])
-    output = grid_cluster(position, size)
+    expected = torch.LongTensor([0, 3, 1, 0, 2])
+    output, _ = grid_cluster(position, size)
     assert output.cpu().tolist() == expected.tolist()
 
-    position = Tensor(tensor, [[0, 9, 2, 2, 8], [0, 9, 8, 2, 3]])
+    position = Tensor(tensor, [[0, 11, 2, 2, 8], [0, 9, 8, 2, 3]])
     position = position.cuda().t()
-    output = grid_cluster(position, size)
+    output, _ = grid_cluster(position, size)
     assert output.cpu().tolist() == expected.tolist()
 
-    output = grid_cluster(position.expand(2, 5, 2), size)
+    output, _ = grid_cluster(position.expand(2, 5, 2), size)
     assert output.tolist() == expected.expand(2, 5).tolist()
 
     position = position.repeat(2, 1)
     batch = torch.cuda.LongTensor([0, 0, 0, 0, 0, 1, 1, 1, 1, 1])
     expected = torch.LongTensor([0, 3, 1, 0, 2, 4, 7, 5, 4, 6])
-    expected_batch = torch.LongTensor([0, 0, 0, 0, 1, 1, 1, 1])
-    output, reduced_batch = grid_cluster(position, size, batch)
+    expected_batch2 = torch.LongTensor([0, 0, 0, 0, 1, 1, 1, 1])
+    output, batch2 = grid_cluster(position, size, batch)
+    assert output.cpu().tolist() == expected.tolist()
+    assert batch2.cpu().tolist() == expected_batch2.tolist()
+
+    output, C = grid_cluster(position, size, batch, fake_nodes=True)
+    expected = torch.LongTensor([0, 5, 1, 0, 2, 6, 11, 7, 6, 8])
     assert output.cpu().tolist() == expected.tolist()
-    assert reduced_batch.cpu().tolist() == expected_batch.tolist()
+    assert C == 6

torch_cluster/__init__.py

 from .functions.grid import grid_cluster
 
-__version__ = '0.1.1'
+__version__ = '0.2.0'
 
 __all__ = ['grid_cluster', '__version__']

torch_cluster/functions/grid.py

@@ -3,7 +3,7 @@ import torch
 from .utils import get_func, consecutive
 
 
-def grid_cluster(position, size, batch=None):
+def grid_cluster(position, size, batch=None, fake_nodes=False, offset=None):
     # Allow one-dimensional positions.
     if position.dim() == 1:
         position = position.unsqueeze(-1)
@@ -22,8 +22,12 @@ def grid_cluster(position, size, batch=None):
         size = torch.cat([size.new(1).fill_(1), size], dim=-1)
 
     # Translate to minimal positive positions.
-    min = position.min(dim=-2, keepdim=True)[0]
-    position = position - min
+    if offset is None:
+        min = position.min(dim=-2, keepdim=True)[0]
+        position = position - min
+    else:
+        position = position + offset
+        assert position.min() >= 0, 'Offset resulting in negative positions'
 
     # Compute cluster count for each dimension.
     max = position.max(dim=0)[0]
@@ -43,10 +47,9 @@ def grid_cluster(position, size, batch=None):
     func = get_func('grid', position)
     func(C, cluster, position, size, c_max)
     cluster = cluster.squeeze(dim=-1)
-    cluster, u = consecutive(cluster)
 
-    if batch is None:
-        return cluster
-    else:
-        batch = (u / c_max[1:].prod()).long()
-        return cluster, batch
+    if fake_nodes:
+        return cluster, C // c_max[0]
+
+    cluster, u = consecutive(cluster)
+    return cluster, None if batch is None else (u / (C // c_max[0])).long()