docs

torch_cluster/fps.py

@@ -2,28 +2,31 @@ 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.
+    """Samples a specified ratio of points for each element in a batch using
+    farthest iterative point sampling.
 
     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`)
+        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`)
         ratio (float, optional): Sampling ratio. (default: :obj:`0.5`)
         random_start (bool, optional): Whether the starting node is
             sampled randomly. (default: :obj:`True`)
 
+    :rtype: :class:`LongTensor`
+
     Examples::
 
         >>> x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
         >>> batch = torch.Tensor([0, 0, 0, 0])
-        >>> sample = fps(pos, batch)
+        >>> sample = fps(x, batch)
     """
 
     if batch is None:

torch_cluster/graclus.py

@@ -15,6 +15,8 @@ def graclus_cluster(row, col, weight=None, num_nodes=None):
         weight (Tensor, optional): Edge weights. (default: :obj:`None`)
         num_nodes (int, optional): The number of nodes. (default: :obj:`None`)
 
+    :rtype: :class:`LongTensor`
+
     Examples::
 
         >>> row = torch.tensor([0, 1, 1, 2])

torch_cluster/grid.py

@@ -17,6 +17,8 @@ def grid_cluster(pos, size, start=None, end=None):
         end (Tensor, optional): End position of the grid (in each
             dimension). (default: :obj:`None`)
 
+    :rtype: :class:`LongTensor`
+
     Examples::
 
         >>> pos = torch.Tensor([[0, 0], [11, 9], [2, 8], [2, 2], [8, 3]])

torch_cluster/nearest.py

@@ -5,6 +5,27 @@ 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`.
+
+    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`)
+
+    Examples::
+
+        >>> x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
+        >>> batch_x = torch.Tensor([0, 0, 0, 0])
+        >>> y = torch.Tensor([[-1, 0], [1, 0]])
+        >>> batch_x = torch.Tensor([0, 0])
+        >>> out = fps(x, y, batch_x, batch_y)
+    """
+
     if batch_x is None:
         batch_x = x.new_zeros(x.size(0), dtype=torch.long)
 

torch_cluster/radius.py

@@ -5,6 +5,31 @@ 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`.
+
+    Args:
+        x (Tensor): D-dimensional point features.
+        y (Tensor): D-dimensional point features.
+        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`)
+        max_num_neighbors (int, optional): The maximum number of neighbors to
+            return for each element in `y`. (default: :obj:`32`)
+
+    :rtype: :class:`LongTensor`
+
+    Examples::
+
+        >>> x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
+        >>> batch_x = torch.Tensor([0, 0, 0, 0])
+        >>> y = torch.Tensor([[-1, 0], [1, 0]])
+        >>> batch_x = torch.Tensor([0, 0])
+        >>> out = radius(x, y, 1.5, batch_x, batch_y)
+    """
 
     if batch_x is None:
         batch_x = x.new_zeros(x.size(0), dtype=torch.long)
@@ -29,6 +54,30 @@ def radius(x, y, r, batch_x=None, batch_y=None, max_num_neighbors=32):
 
 
 def radius_graph(x, r, batch=None, max_num_neighbors=32):
+    """Finds for each element in `x` all points in `x` within distance `r`.
+
+    Args:
+        x (Tensor): D-dimensional point features.
+        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`)
+        max_num_neighbors (int, optional): The maximum number of neighbors to
+            return for each element in `y`. (default: :obj:`32`)
+
+    :rtype: :class:`LongTensor`
+
+    Examples::
+
+        >>> x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
+        >>> batch_x = torch.Tensor([0, 0, 0, 0])
+        >>> y = torch.Tensor([[-1, 0], [1, 0]])
+        >>> batch_x = torch.Tensor([0, 0])
+        >>> out = radius(x, y, 1.5, batch_x, batch_y)
+    """
+
     edge_index = radius(x, x, r, batch, batch, max_num_neighbors + 1)
     row, col = edge_index
     mask = row != col