self-loops flag

cuda/knn_kernel.cu

@@ -64,7 +64,7 @@ at::Tensor knn_cuda(at::Tensor x, at::Tensor y, size_t k, at::Tensor batch_x,
 
   auto dist = at::full(y.size(0) * k, 1e38, y.options());
   auto row = at::empty(y.size(0) * k, batch_y.options());
-  auto col = at::empty(y.size(0) * k, batch_y.options());
+  auto col = at::full(y.size(0) * k, -1, batch_y.options());
 
   AT_DISPATCH_FLOATING_TYPES(x.type(), "knn_kernel", [&] {
     knn_kernel<scalar_t><<<batch_size, THREADS>>>(
@@ -73,5 +73,6 @@ at::Tensor knn_cuda(at::Tensor x, at::Tensor y, size_t k, at::Tensor batch_x,
         col.data<int64_t>(), k, x.size(1));
   });
 
-  return at::stack({row, col}, 0);
+  auto mask = col != -1;
+  return at::stack({row.masked_select(mask), col.masked_select(mask)}, 0);
 }

torch_cluster/knn.py

@@ -51,7 +51,7 @@ def knn(x, y, k, batch_x=None, batch_y=None):
     return assign_index
 
 
-def knn_graph(x, k, batch=None):
+def knn_graph(x, k, batch=None, loop=False):
     """Finds for each element in `x` the `k` nearest points.
 
     Args:
@@ -62,6 +62,8 @@ def knn_graph(x, k, batch=None):
             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`)
+        loop (bool, optional): If :obj:`True`, the graph will contain
+            self-loops. (default: :obj:`False`)
 
     :rtype: :class:`LongTensor`
 
@@ -72,8 +74,10 @@ def knn_graph(x, k, batch=None):
         >>> out = knn_graph(x, 2, batch)
     """
 
-    edge_index = knn(x, x, k + 1, batch, batch)
-    row, col = edge_index
-    mask = row != col
-    row, col = row[mask], col[mask]
-    return torch.stack([row, col], dim=0)
+    edge_index = knn(x, x, k if loop else k + 1, batch, batch)
+    if not loop:
+        row, col = edge_index
+        mask = row != col
+        row, col = row[mask], col[mask]
+        edge_index = torch.stack([row, col], dim=0)
+    return edge_index

torch_cluster/radius.py

@@ -53,7 +53,7 @@ def radius(x, y, r, batch_x=None, batch_y=None, max_num_neighbors=32):
     return assign_index
 
 
-def radius_graph(x, r, batch=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`.
 
     Args:
@@ -64,6 +64,8 @@ def radius_graph(x, r, batch=None, max_num_neighbors=32):
             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`)
+        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`)
 
@@ -78,6 +80,9 @@ def radius_graph(x, r, batch=None, max_num_neighbors=32):
 
     edge_index = radius(x, x, r, batch, batch, max_num_neighbors + 1)
     row, col = edge_index
-    mask = row != col
-    row, col = row[mask], col[mask]
-    return torch.stack([row, col], dim=0)
+    if not loop:
+        row, col = edge_index
+        mask = row != col
+        row, col = row[mask], col[mask]
+        edge_index = torch.stack([row, col], dim=0)
+    return edge_index