cluster

Neighbor search and grouping: kNN, FPS, radius queries, and local grids.

Functions:

knn –

Find the \(k\) nearest neighbors in \(x\) for each point in \(y\).
knn_graph –

Compute the kNN graph of \(x\).
fps –

A sampling algorithm from the paper [PointNet++: Deep Hierarchical Feature
local_grid –

Applies local grid quantization to the source tensor as explained in the paper
radius –

torch_cluster.radius wrapper with an optional sort-by-source-index tie-breaker.
group –

Partition a packed point cloud into local groups with FPS centers and a \(k\)-NN neighborhood.

knn ¶

knn(
    x: Tensor,
    y: Tensor,
    k: int,
    batch_x: OptTensor = None,
    batch_y: OptTensor = None,
    cosine: bool = False,
    num_workers: int = 1,
    batch_size: Optional[int] = None,
) -> Tensor

Find the \(k\) nearest neighbors in \(x\) for each point in \(y\). This function is a wrapper around the torch_cluster.knn function, and supports the same arguments. However, in case the batch_x and batch_y tensors are provided, and the samples have the same number of nodes, this function uses a more efficient implementation that is significantly faster on GPU using torch.cdist + topk.

Important

If provided, the batch_x and batch_y tensors must be sorted in non-decreasing order (both the dense fast path and torch_cluster require it); unsorted batches raise a ValueError.

Note

When a point of \(y\) coincides with a point of \(x\) (e.g. knn(pos, pos, k)), the query point itself counts among the \(k\) neighbors, matching torch_cluster.knn.

Parameters:

x (Tensor) –

The source tensor to find the nearest neighbors of shape \((N, *)\).
y (Tensor) –

The target tensor to find the nearest neighbors of shape \((M, *)\).
k (int) –

The number of nearest neighbors to find.
batch_x (OptTensor, default: None ) –

The batch tensor of the source tensor of shape \((N,)\).
batch_y (OptTensor, default: None ) –

The batch tensor of the target tensor of shape \((M,)\).
cosine (bool, default: False ) –

Whether to use cosine distance.
num_workers (int, default: 1 ) –

The number of workers to use for the computation.
batch_size (Optional[int], default: None ) –

The batch size to use for the computation.

Returns:

Tensor –

The nearest neighbors of shape \((2, M \cdot k)\).

knn_graph ¶

knn_graph(
    x: Tensor,
    k: int,
    batch: OptTensor = None,
    loop: bool = False,
    flow: str = "source_to_target",
    cosine: bool = False,
    num_workers: int = 1,
    batch_size: Optional[int] = None,
) -> Tensor

Compute the kNN graph of \(x\).

This function is a drop-in for torch_cluster.knn_graph, except that when the batch tensor partitions the points into uniformly-sized samples this function uses a torch.cdist + topk implementation that is significantly faster on GPU than the underlying torch_cluster.knn_graph.

Important

If provided, the batch tensor must be sorted in non-decreasing order (both the dense fast path and torch_cluster require it); an unsorted batch raises a ValueError.

Parameters:

x (Tensor) –

The input tensor of shape \((N, *)\).
k (int) –

The number of nearest neighbors to find. When loop=False, the self-edge is excluded from the result.
batch (OptTensor, default: None ) –

The batch tensor of shape \((N,)\).
loop (bool, default: False ) –

Whether to include self-edges.
flow (str, default: 'source_to_target' ) –

Either "source_to_target" (PyG default, edge_index = (src, dst) where src is the neighbor and dst is the central point) or "target_to_source".
cosine (bool, default: False ) –

Whether to use cosine distance.
num_workers (int, default: 1 ) –

Forwarded to the torch_cluster fallback.
batch_size (Optional[int], default: None ) –

Forwarded to the torch_cluster fallback.

Returns:

Tensor –

Edge index of shape \((2, k \cdot N)\).

fps ¶

fps(
    src: Tensor,
    batch: Optional[Tensor] = None,
    ratio: Optional[Union[Tensor, float]] = None,
    num_nodes: Optional[Union[Tensor, float]] = None,
    random_start: bool = True,
    batch_size: Optional[int] = None,
    ptr: Optional[Union[Tensor, List[int]]] = None,
) -> Tensor

A sampling algorithm from the paper PointNet++: Deep Hierarchical Feature Learning on Point Sets in a Metric Space by Qi et al., which iteratively samples the most distant point with regard to the rest points.

This function is adapted from the torch_cluster.fps function and supports a sampling a fixed number of nodes with the num_nodes argument.

Important

If provided, the batch tensor is expected to be sorted.

Parameters:

src (Tensor) –

The source tensor to sample from of shape \((N, *)\).
batch (Optional[Tensor], default: None ) –

The batch tensor to sample from of shape \((N,)\).
ratio (Optional[Union[Tensor, float]], default: None ) –

The sampling ratio.
num_nodes (Optional[Union[Tensor, float]], default: None ) –

The number of nodes to sample. When a sample holds fewer than num_nodes points, indices repeat (sampling with replacement, matching the reference CUDA FPS), so the output always holds num_nodes indices per sample and downstream shapes stay stable.
random_start (bool, default: True ) –

Whether to start the sampling randomly.
batch_size (Optional[int], default: None ) –

The batch size.
ptr (Optional[Union[Tensor, List[int]]], default: None ) –

The pointer tensor to sample from.

Returns:

Tensor –

The sampled indices of shape \((M,)\).

Examples:

>>> import torch
>>> from torch_pointcloud.utils.cluster import fps
>>> src = torch.randn(100, 3)
>>> batch = torch.cat([torch.zeros(50), torch.ones(50)]).long()
>>> idx = fps(src, batch, num_nodes=10)  # doctest: +SKIP
>>> print(idx.shape)  # doctest: +SKIP
torch.Size([10])

local_grid ¶

local_grid(
    src: Tensor, size: float, batch: Tensor | None = None
) -> Tensor

Applies local grid quantization to the source tensor as explained in the paper TorchSparse++: Efficient node Cloud Engine by Tang et al., which quantizes the source tensor to a local grid.

Note

If a batch tensor is provided, the function will apply the quantization to each batch separately, ensuring the

Parameters:

src (Tensor) –

The source tensor to quantize of shape \((N, *)\).
size (float) –

The quantization size.
batch (Tensor | None, default: None ) –

The associated batch tensor of shape \((N,)\).

Returns:

Tensor –

The quantized source tensor of shape \((N, *)\).

Examples:

>>> import torch
>>> from torch_pointcloud.utils.cluster import local_grid
>>> src = torch.randn(100, 3)
>>> batch = torch.cat([torch.zeros(50), torch.ones(50)]).long()
>>> src_grid = local_grid(src, size=1.0, batch=batch)  # doctest: +SKIP

radius ¶

radius(
    x: Tensor,
    y: Tensor,
    r: float,
    batch_x: OptTensor = None,
    batch_y: OptTensor = None,
    max_num_neighbors: int = 32,
    sort: bool = False,
) -> tuple[Tensor, Tensor]

torch_cluster.radius wrapper with an optional sort-by-source-index tie-breaker.

With sort=False (default) this just delegates to torch_cluster.radius and returns edges in kernel-traversal order. With sort=True, when more than max_num_neighbors source points lie inside a ball, the \(k\) smallest source indices are kept (PointNet++'s reference query_ball_point behavior). Pretrained PointNet++ checkpoints from yanx27 / charlesq34 overfit to this selection rule, so reproducing their accuracy requires sort=True.

Important

If provided, the batch_x and batch_y tensors must be sorted in non-decreasing order (torch_cluster.radius requires it); unsorted batches raise a ValueError.

Parameters:

x (Tensor) –

Source positions, shape \((N_x, d)\).
y (Tensor) –

Query positions, shape \((N_y, d)\).
r (float) –

Ball radius.
batch_x (OptTensor, default: None ) –

Batch index for \(x\), shape \((N_x,)\). None for a single batch.
batch_y (OptTensor, default: None ) –

Batch index for \(y\), shape \((N_y,)\). None for a single batch.
max_num_neighbors (int, default: 32 ) –

Max neighbors \(k\) kept per query.
sort (bool, default: False ) –

Sort in-ball source indices ascending and keep the first \(k\).

Returns:

Tensor –

(row, col) edges. row is the query index, col is the source index.
Tensor –

Centroids with no in-ball neighbors emit zero edges; pooling leaves those
tuple[Tensor, Tensor] –

rows at the reduction identity.

group ¶

group(
    pos: Tensor,
    batch: Tensor,
    num_group: int,
    group_size: int,
    random_start: bool = ...,
    *,
    return_indices: Literal[False] = ...,
) -> Tuple[Tensor, Tensor]

group(
    pos: Tensor,
    batch: Tensor,
    num_group: int,
    group_size: int,
    random_start: bool = ...,
    *,
    return_indices: Literal[True],
) -> Tuple[Tensor, Tensor, Tensor]

group(
    pos: Tensor,
    batch: Tensor,
    num_group: int,
    group_size: int,
    random_start: bool = False,
    *,
    return_indices: bool = False,
) -> Union[
    Tuple[Tensor, Tensor], Tuple[Tensor, Tensor, Tensor]
]

Partition a packed point cloud into local groups with FPS centers and a \(k\)-NN neighborhood.

Farthest point sampling selects num_group centers per sample, then \(k\)-NN gathers the group_size nearest neighbors of each center, and each neighborhood is recentered on its center. Because num_group is fixed per sample, the packed result densifies to a regular \((B, G, k, 3)\) batch without padding.

Parameters:

pos (Tensor) –

Packed point coordinates of shape \((N, 3)\).
batch (Tensor) –

Per-point batch index of shape \((N,)\).
num_group (int) –

Number of groups (FPS centers) \(G\) per sample.
group_size (int) –

Number of neighbors \(k\) per group.
random_start (bool, default: False ) –

Whether to start farthest point sampling from a random point.
return_indices (bool, default: False ) –

If True, also return the flat neighbor index into the packed input.

Returns:

Union[Tuple[Tensor, Tensor], Tuple[Tensor, Tensor, Tensor]] –

(neighborhood, center), or (neighborhood, center, idx) when return_indices is True.
Union[Tuple[Tensor, Tensor], Tuple[Tensor, Tensor, Tensor]] –

neighborhood has shape \((B, G, k, 3)\) recentered on each center, center has shape
Union[Tuple[Tensor, Tensor], Tuple[Tensor, Tensor, Tensor]] –

\((B, G, 3)\), and idx has shape \((B \cdot G \cdot k,)\) indexing the packed input.

Shape

Input: \((N, 3)\) and \((N,)\).
Output: \((B, G, k, 3)\) and \((B, G, 3)\) (plus \((B \cdot G \cdot k,)\) when return_indices).

Example

import torch
from torch_pointcloud.utils.cluster import group

pos = torch.randn(2048, 3)
batch = torch.cat([torch.zeros(1024), torch.ones(1024)]).long()
neighborhood, center = group(pos, batch, num_group=64, group_size=32)
print(neighborhood.shape, center.shape)