update doc

README.md

@@ -90,7 +90,7 @@ tensor([0, 5, 3, 0, 1])
 
 ## FarthestPointSampling
 
-A sampling algorithm, which iteratively samples the most distant point (in metric distance) with regard to the rest points.
+A sampling algorithm, which iteratively samples the most distant point with regard to the rest points.
 
 ```python
 import torch
@@ -98,7 +98,7 @@ from torch_cluster import fps
 
 x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
 batch = torch.tensor([0, 0, 0, 0])
-sample = fps(x, batch, ratio=0.5, random_start=False)
+index = fps(x, batch, ratio=0.5, random_start=False)
 ```
 
 ```
@@ -108,7 +108,7 @@ tensor([0, 3])
 
 ## kNN-Graph
 
-Computes graph edges to the nearest *k* points in metric space.
+Computes graph edges to the nearest *k* points.
 
 ```python
 import torch
@@ -127,7 +127,7 @@ tensor([[0, 0, 1, 1, 2, 2, 3, 3],
 
 ## Radius-Graph
 
-Computes graph edges to all points within a given distance in metric space.
+Computes graph edges to all points within a given distance.
 
 ```python
 import torch
@@ -146,17 +146,17 @@ tensor([[0, 0, 1, 1, 2, 2, 3, 3],
 
 ## Nearest
 
-Clusters points which are nearest to a given query point in metric space.
+Clusters points in *x* together which are nearest to a given query point in *y*.
 
 ```python
 import torch
 from torch_cluster import nearest
 
 x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
-batch_x = torch.Tensor([0, 0, 0, 0])
-query_x = torch.Tensor([[-1, 0], [1, 0]])
-query_batch = torch.Tensor([0, 0])
-cluster = nearest(x, query_x, batch_x, query_batch)
+batch_x = torch.tensor([0, 0, 0, 0])
+y = torch.Tensor([[-1, 0], [1, 0]])
+batch_y = torch.tensor([0, 0])
+cluster = nearest(x, y, batch_x, batch_y)
 ```
 
 ```

torch_cluster/fps.py

@@ -2,31 +2,36 @@ import torch
 
 if torch.cuda.is_available():
     import fps_cuda
-    """    """
 
 
 def fps(x, batch=None, ratio=0.5, random_start=True):
-    """Iteratively samples the most distant point (in metric distance) with
-    regard to the rest points.
+    r""""A sampling algorithm from the `"PointNet++: Deep Hierarchical Feature
+    Learning on Point Sets in a Metric Space"
+    <https://arxiv.org/abs/1706.02413>`_ paper, which iteratively samples the
+    most distant point with regard to the rest points.
 
     Args:
-        x (Tensor): D-dimensional point features.
-        batch (LongTensor, optional): Vector that maps each point to its
-            example identifier. If :obj:`None`, all points belong to the same
-            example. If not :obj:`None`, points in the same example need to
-            have contiguous memory layout and :obj:`batch` needs to be
-            ascending. (default: :obj:`None`)
+        x (Tensor): Node feature matrix
+            :math:`\mathbf{X} \in \mathbb{R}^{N \times F}`.
+        batch (LongTensor, optional): Batch vector
+            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
+            node to a specific example. (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`)
+        random_start (bool, optional): If set to :obj:`False`, use the first
+            node in :math:`\mathbf{X}` as starting node. (default: obj:`True`)
 
     :rtype: :class:`LongTensor`
 
-    Examples::
+    .. testsetup::
+
+        import torch
+        from torch_cluster import fps
+
+    .. testcode::
 
         >>> x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
         >>> batch = torch.tensor([0, 0, 0, 0])
-        >>> sample = fps(x, batch, ratio=0.5)
+        >>> index = fps(x, batch, ratio=0.5)
     """
 
     if batch is None:

torch_cluster/knn.py

@@ -6,22 +6,30 @@ if torch.cuda.is_available():
 
 
 def knn(x, y, k, batch_x=None, batch_y=None):
-    """Finds for each element in `y` the `k` nearest points in `x`.
+    r"""Finds for each element in :obj:`y` the :obj:`k` nearest points in
+    :obj:`x`.
 
     Args:
-        x (Tensor): D-dimensional point features.
-        y (Tensor): D-dimensional point features.
+        x (Tensor): Node feature matrix
+            :math:`\mathbf{X} \in \mathbb{R}^{N \times F}`.
+        y (Tensor): Node feature matrix
+            :math:`\mathbf{X} \in \mathbb{R}^{M \times F}`.
         k (int): The number of neighbors.
-        batch_x (LongTensor, optional): Vector that maps each point to its
-            example identifier. If :obj:`None`, all points belong to the same
-            example. If not :obj:`None`, points in the same example need to
-            have contiguous memory layout and :obj:`batch` needs to be
-            ascending. (default: :obj:`None`)
-        batch_y (LongTensor, optional): See `batch_x` (default: :obj:`None`)
+        batch_x (LongTensor, optional): Batch vector
+            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
+            node to a specific example. (default: :obj:`None`)
+        batch_y (LongTensor, optional): Batch vector
+            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^M`, which assigns each
+            node to a specific example. (default: :obj:`None`)
 
     :rtype: :class:`LongTensor`
 
-    Examples::
+    .. testsetup::
+
+        import torch
+        from torch_cluster import knn
+
+    .. testcode::
 
         >>> x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
         >>> batch_x = torch.tensor([0, 0, 0, 0])
@@ -70,22 +78,26 @@ def knn(x, y, k, batch_x=None, batch_y=None):
 
 
 def knn_graph(x, k, batch=None, loop=False):
-    """Finds for each element in `x` the `k` nearest points.
+    r"""Computes graph edges to the nearest :obj:`k` points.
 
     Args:
-        x (Tensor): D-dimensional point features.
+        x (Tensor): Node feature matrix
+            :math:`\mathbf{X} \in \mathbb{R}^{N \times F}`.
         k (int): The number of neighbors.
-        batch (LongTensor, optional): Vector that maps each point to its
-            example identifier. If :obj:`None`, all points belong to the same
-            example. If not :obj:`None`, points in the same example need to
-            have contiguous memory layout and :obj:`batch` needs to be
-            ascending. (default: :obj:`None`)
+        batch (LongTensor, optional): Batch vector
+            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
+            node to a specific example. (default: :obj:`None`)
         loop (bool, optional): If :obj:`True`, the graph will contain
             self-loops. (default: :obj:`False`)
 
     :rtype: :class:`LongTensor`
 
-    Examples::
+    .. testsetup::
+
+        import torch
+        from torch_cluster import knn_graph
+
+    .. testcode::
 
         >>> x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
         >>> batch = torch.tensor([0, 0, 0, 0])

torch_cluster/nearest.py

@@ -6,24 +6,32 @@ if torch.cuda.is_available():
 
 
 def nearest(x, y, batch_x=None, batch_y=None):
-    """Finds for each element in `x` its nearest point in `y`.
+    """Clusters points in :obj:`x` together which are nearest to a given query
+    point in :obj:`y`.
 
     Args:
-        x (Tensor): D-dimensional point features.
-        y (Tensor): D-dimensional point features.
-        batch_x (LongTensor, optional): Vector that maps each point to its
-            example identifier. If :obj:`None`, all points belong to the same
-            example. If not :obj:`None`, points in the same example need to
-            have contiguous memory layout and :obj:`batch` needs to be
-            ascending. (default: :obj:`None`)
-        batch_y (LongTensor, optional): See `batch_x` (default: :obj:`None`)
+        x (Tensor): Node feature matrix
+            :math:`\mathbf{X} \in \mathbb{R}^{N \times F}`.
+        y (Tensor): Node feature matrix
+            :math:`\mathbf{X} \in \mathbb{R}^{M \times F}`.
+        batch_x (LongTensor, optional): Batch vector
+            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
+            node to a specific example. (default: :obj:`None`)
+        batch_y (LongTensor, optional): Batch vector
+            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^M`, which assigns each
+            node to a specific example. (default: :obj:`None`)
 
-    Examples::
+    .. testsetup::
+
+        import torch
+        from torch_cluster import nearest
+
+    .. testcode::
 
         >>> x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
-        >>> batch_x = torch.Tensor([0, 0, 0, 0])
+        >>> batch_x = torch.tensor([0, 0, 0, 0])
         >>> y = torch.Tensor([[-1, 0], [1, 0]])
-        >>> batch_x = torch.Tensor([0, 0])
+        >>> batch_x = torch.tensor([0, 0])
         >>> cluster = nearest(x, y, batch_x, batch_y)
     """
 

torch_cluster/radius.py

@@ -6,24 +6,33 @@ if torch.cuda.is_available():
 
 
 def radius(x, y, r, batch_x=None, batch_y=None, max_num_neighbors=32):
-    """Finds for each element in `y` all points in `x` within distance `r`.
+    r"""Finds for each element in :obj:`y` all points in :obj:`x` within
+    distance :obj:`r`.
 
     Args:
-        x (Tensor): D-dimensional point features.
-        y (Tensor): D-dimensional point features.
+        x (Tensor): Node feature matrix
+            :math:`\mathbf{X} \in \mathbb{R}^{N \times F}`.
+        y (Tensor): Node feature matrix
+            :math:`\mathbf{X} \in \mathbb{R}^{M \times F}`.
         r (float): The radius.
-        batch_x (LongTensor, optional): Vector that maps each point to its
-            example identifier. If :obj:`None`, all points belong to the same
-            example. If not :obj:`None`, points in the same example need to
-            have contiguous memory layout and :obj:`batch` needs to be
-            ascending. (default: :obj:`None`)
-        batch_y (LongTensor, optional): See `batch_x` (default: :obj:`None`)
+        batch_x (LongTensor, optional): Batch vector
+            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
+            node to a specific example. (default: :obj:`None`)
+        batch_y (LongTensor, optional): Batch vector
+            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^M`, which assigns each
+            node to a specific example. (default: :obj:`None`)
         max_num_neighbors (int, optional): The maximum number of neighbors to
-            return for each element in `y`. (default: :obj:`32`)
+            return for each element in :obj:`y`. (default: :obj:`32`)
 
     :rtype: :class:`LongTensor`
 
-    Examples::
+    .. testsetup::
+
+        import torch
+        from torch_cluster import radius
+
+    .. testcode::
+
 
         >>> x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
         >>> batch_x = torch.tensor([0, 0, 0, 0])
@@ -63,24 +72,28 @@ def radius(x, y, r, batch_x=None, batch_y=None, max_num_neighbors=32):
 
 
 def radius_graph(x, r, batch=None, loop=False, max_num_neighbors=32):
-    """Finds for each element in `x` all points in `x` within distance `r`.
+    r"""Computes graph edges to all points within a given distance.
 
     Args:
-        x (Tensor): D-dimensional point features.
+        x (Tensor): Node feature matrix
+            :math:`\mathbf{X} \in \mathbb{R}^{N \times F}`.
         r (float): The radius.
-        batch (LongTensor, optional): Vector that maps each point to its
-            example identifier. If :obj:`None`, all points belong to the same
-            example. If not :obj:`None`, points in the same example need to
-            have contiguous memory layout and :obj:`batch` needs to be
-            ascending. (default: :obj:`None`)
+        batch (LongTensor, optional): Batch vector
+            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
+            node to a specific example. (default: :obj:`None`)
         loop (bool, optional): If :obj:`True`, the graph will contain
             self-loops. (default: :obj:`False`)
         max_num_neighbors (int, optional): The maximum number of neighbors to
-            return for each element in `y`. (default: :obj:`32`)
+            return for each element in :obj:`y`. (default: :obj:`32`)
 
     :rtype: :class:`LongTensor`
 
-    Examples::
+    .. testsetup::
+
+        import torch
+        from torch_cluster import radius_graph
+
+    .. testcode::
 
         >>> x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
         >>> batch = torch.tensor([0, 0, 0, 0])