fps algorithm included

setup.py

@@ -16,7 +16,7 @@ if torch.cuda.is_available():
         CUDAExtension('fps_cuda', ['cuda/fps.cpp', 'cuda/fps_kernel.cu']),
     ]
 
-__version__ = '1.1.5'
+__version__ = '1.2.0'
 url = 'https://github.com/rusty1s/pytorch_cluster'
 
 install_requires = []

torch_cluster/__init__.py

 from .graclus import graclus_cluster
 from .grid import grid_cluster
+from .fps import fps
 
-__version__ = '1.1.5'
+__version__ = '1.2.0'
 
-__all__ = ['graclus_cluster', 'grid_cluster', '__version__']
+__all__ = [
+    'graclus_cluster',
+    'grid_cluster',
+    'fps',
+    '__version__',
+]

torch_cluster/fps.py

+import torch
+
+if torch.cuda.is_available():
+    import fps_cuda
+
+
+def fps(x, batch=None, ratio=0.5, random_start=True):
+    """A clustering algorithm, which overlays a regular grid of user-defined
+    size over a point cloud and clusters all points within a voxel.
+
+    Args:
+        x (Tensor): D-dimensional node features.
+        batch (LongTensor, optional): Vector that maps each node to a graph.
+            If :obj:`None`, all node features belong to the same graph. If not
+            :obj:`None`, nodes of the same graph need to have contiguous memory
+            layout and :obj:`batch` needs to be ascending.
+            (default: :obj:`None`)
+        ratio (float, optional): Sampling ratio. (default: :obj:`0.5`)
+        random_start (bool, optional): Whether the starting node is
+            sampled randomly. (default: :obj:`True`)
+
+    Examples::
+
+        >>> x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
+        >>> batch = torch.Tensor([0, 0, 0, 0])
+        >>> sample = fps(pos, batch)
+    """
+
+    assert x.is_cuda
+
+    if batch is None:
+        batch = x.new_zeros(x.size(0), dtype=torch.long)
+
+    op = fps_cuda.fps if x.is_cuda else None
+    out = op(x, batch, ratio, random_start)
+
+    return out