added sparse and dense voxel cluster

Author: rusty1s

build.sh

@@ -7,4 +7,4 @@ SRC_DIR=torch_cluster/kernel
 BUILD_DIR=torch_cluster/build
 
 mkdir -p $BUILD_DIR
-nvcc -c -o $BUILD_DIR/kernel.so $SRC_DIR/kernel.cu -arch=sm_35 -Xcompiler -fPIC -shared -I$TORCH/lib/include/TH -I$TORCH/lib/include/THC -I$SRC_DIR
+$(which nvcc) -c -o $BUILD_DIR/kernel.so $SRC_DIR/kernel.cu -arch=sm_35 -Xcompiler -fPIC -shared -I$TORCH/lib/include/TH -I$TORCH/lib/include/THC -I$SRC_DIR

setup.py

@@ -2,7 +2,7 @@
 
 from setuptools import setup, find_packages
 
-__version__ = '0.2.1'
+__version__ = '0.2.2'
 url = 'https://github.com/rusty1s/pytorch_cluster'
 
 install_requires = ['cffi', 'torch-unique']

test/dense_grid.json

@@ -0,0 +1,42 @@
+[
+  {
+    "name": "One-dimensional positions without start/end",
+    "position": [2, 16],
+    "size": [5],
+    "expected": [0, 2],
+    "expected_C": 3
+  },
+  {
+    "name": "Start parameter",
+    "position": [2, 16],
+    "size": [5],
+    "start": 0,
+    "expected": [0, 3],
+    "expected_C": 4
+  },
+  {
+    "name": "End parameter",
+    "position": [2, 16],
+    "size": [5],
+    "start": 0,
+    "end": 30,
+    "expected": [0, 3],
+    "expected_C": 6
+  },
+  {
+    "name": "Two-dimensional positions",
+    "position": [[0, 0], [11, 9], [2, 8], [2, 2], [8, 3]],
+    "size": [5, 5],
+    "expected": [0, 5, 1, 0, 2],
+    "expected_C": 6
+  },
+  {
+    "name": "Batch",
+    "position": [[0, 0], [11, 9], [2, 8], [2, 2], [8, 3], [1, 1], [6, 6]],
+    "size": [5, 5],
+    "batch": [0, 0, 0, 0, 0, 1, 1],
+    "expected": [0, 5, 1, 0, 2, 6, 9],
+    "expected_C": 6
+  }
+]
+

test/sparse_grid.json

@@ -0,0 +1,35 @@
+[
+  {
+    "name": "One-dimensional positions",
+    "position": [2, 16],
+    "size": [5],
+    "expected": [0, 1]
+  },
+  {
+    "name": "Start parameter",
+    "position": [2, 6],
+    "size": [5],
+    "start": 0,
+    "expected": [0, 1]
+  },
+  {
+    "name": "Consecutive cluster",
+    "position": [0, 17, 2, 8, 3],
+    "size": [5],
+    "expected": [0, 2, 0, 1, 0]
+  },
+  {
+    "name": "Two-dimensional positions",
+    "position": [[0, 0], [11, 9], [2, 8], [2, 2], [8, 3]],
+    "size": [5, 5],
+    "expected": [0, 3, 1, 0, 2]
+  },
+  {
+    "name": "Batch",
+    "position": [[0, 0], [11, 9], [2, 8], [2, 2], [8, 3], [1, 1], [6, 6]],
+    "size": [5, 5],
+    "batch": [0, 0, 0, 0, 0, 1, 1],
+    "expected": [0, 3, 1, 0, 2, 4, 5],
+    "expected_batch": [0, 0, 0, 0, 1, 1]
+  }
+]

test/test_dense_grid.py

@@ -0,0 +1,46 @@
+from os import path as osp
+from itertools import product
+
+import pytest
+import json
+import torch
+from torch_cluster import dense_grid_cluster
+
+from .utils import tensors, Tensor
+
+f = open(osp.join(osp.dirname(__file__), 'dense_grid.json'), 'r')
+data = json.load(f)
+f.close()
+
+
+@pytest.mark.parametrize('tensor,i', product(tensors, range(len(data))))
+def test_dense_grid_cluster_cpu(tensor, i):
+    position = Tensor(tensor, data[i]['position'])
+    size = torch.LongTensor(data[i]['size'])
+    batch = data[i].get('batch')
+    batch = None if batch is None else torch.LongTensor(batch)
+    start = data[i].get('start')
+    end = data[i].get('end')
+    expected = torch.LongTensor(data[i]['expected'])
+    expected_C = data[i]['expected_C']
+
+    output = dense_grid_cluster(position, size, batch, start, end)
+    assert output[0].tolist() == expected.tolist()
+    assert output[1] == expected_C
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason='no CUDA')
+@pytest.mark.parametrize('tensor,i', product(tensors, range(len(data))))
+def test_dense_grid_cluster_gpu(tensor, i):  # pragma: no cover
+    position = Tensor(tensor, data[i]['position']).cuda()
+    size = torch.cuda.LongTensor(data[i]['size'])
+    batch = data[i].get('batch')
+    batch = None if batch is None else torch.cuda.LongTensor(batch)
+    start = data[i].get('start')
+    end = data[i].get('end')
+    expected = torch.LongTensor(data[i]['expected'])
+    expected_C = data[i]['expected_C']
+
+    output = dense_grid_cluster(position, size, batch, start, end)
+    assert output[0].cpu().tolist() == expected.tolist()
+    assert output[1] == expected_C

test/test_grid.py

@@ -0,0 +1,52 @@
+from os import path as osp
+from itertools import product
+
+import pytest
+import json
+import torch
+from torch_cluster import sparse_grid_cluster
+
+from .utils import tensors, Tensor
+
+f = open(osp.join(osp.dirname(__file__), 'sparse_grid.json'), 'r')
+data = json.load(f)
+f.close()
+
+
+@pytest.mark.parametrize('tensor,i', product(tensors, range(len(data))))
+def test_sparse_grid_cluster_cpu(tensor, i):
+    position = Tensor(tensor, data[i]['position'])
+    size = torch.LongTensor(data[i]['size'])
+    batch = data[i].get('batch')
+    start = data[i].get('start')
+    expected = torch.LongTensor(data[i]['expected'])
+
+    if batch is None:
+        output = sparse_grid_cluster(position, size, batch, start)
+        assert output.tolist() == expected.tolist()
+    else:
+        batch = torch.LongTensor(batch)
+        expected_batch = torch.LongTensor(data[i]['expected_batch'])
+        output = sparse_grid_cluster(position, size, batch, start)
+        assert output[0].tolist() == expected.tolist()
+        assert output[1].tolist() == expected_batch.tolist()
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason='no CUDA')
+@pytest.mark.parametrize('tensor,i', product(tensors, range(len(data))))
+def test_sparse_grid_cluster_gpu(tensor, i):  # pragma: no cover
+    position = Tensor(tensor, data[i]['position']).cuda()
+    size = torch.cuda.LongTensor(data[i]['size'])
+    batch = data[i].get('batch')
+    start = data[i].get('start')
+    expected = torch.LongTensor(data[i]['expected'])
+
+    if batch is None:
+        output = sparse_grid_cluster(position, size, batch, start)
+        assert output.cpu().tolist() == expected.tolist()
+    else:
+        batch = torch.cuda.LongTensor(batch)
+        expected_batch = torch.LongTensor(data[i]['expected_batch'])
+        output = sparse_grid_cluster(position, size, batch, start)
+        assert output[0].cpu().tolist() == expected.tolist()
+        assert output[1].cpu().tolist() == expected_batch.tolist()

test/test_sparse_grid.py

@@ -1,5 +1,5 @@
-from .functions.grid import grid_cluster
+from .functions.grid import sparse_grid_cluster, dense_grid_cluster
 
-__version__ = '0.2.1'
+__version__ = '0.2.2'
 
-__all__ = ['grid_cluster', '__version__']
+__all__ = ['sparse_grid_cluster', 'dense_grid_cluster', '__version__']