cluster

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

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_geometric.nn.pool.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_interpolate –

k-NN interpolation with inverse-distance weighting.
dense_neighbors –

Pad a [source, target] edge list into a dense \((N_t, k)\) neighbor table.
gather_neighbors –

Gather source features along a dense_neighbors table, zero where the table is padded.

knn [source] ¶

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\).

A wrapper around torch_geometric.nn.pool.knn that checks its inputs.

Important

If provided, the batch_x and batch_y tensors must be sorted in non-decreasing order; 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_geometric.nn.pool.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 [source] ¶

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\).

A wrapper around torch_geometric.nn.pool.knn_graph that checks its inputs.

Important

If provided, the batch tensor must be sorted in non-decreasing order; 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 ) –

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 –

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

fps [source] ¶

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 wraps torch_geometric.nn.pool.fps and supports sampling a fixed number of nodes per sample 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.ops.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 [source] ¶

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 Point Cloud Engine by Tang et al., which quantizes the source tensor to a local grid.

Note

If a batch tensor is provided, the quantization is applied to each sample separately, so the grid of every sample starts at its own minimum cell.

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.ops.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 [source] ¶

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_geometric.nn.pool.radius wrapper with an optional sort-by-source-index tie-breaker.

With sort=False (default) this just delegates to torch_geometric.nn.pool.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_geometric.nn.pool.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 [source] ¶

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

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

group(
    pos: Tensor,
    batch: Tensor,
    num_groups: 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_groups 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_groups 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_groups (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.ops.cluster import group

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

knn_interpolate [source] ¶

knn_interpolate(
    x: Tensor,
    pos_x: Tensor,
    pos_y: Tensor,
    batch_x: OptTensor = None,
    batch_y: OptTensor = None,
    k: int = 3,
    num_workers: int = 1,
    weighting: Literal["squared", "inverse"] = "squared",
    eps: float = 1e-16,
) -> Tensor

k-NN interpolation with inverse-distance weighting.

From PointNet++: Deep Hierarchical Feature Learning on Point Sets in a Metric Space.

For each point \(y\) with position \(\mathbf{p}(y)\), its interpolated features \(\mathbf{f}(y)\) are given by

\[ \mathbf{f}(y) = \frac{\sum_{i=1}^k w(x_i) \mathbf{f}(x_i)}{\sum_{i=1}^k w(x_i)} \]

where \(\{ x_1, \ldots, x_k \}\) are the \(k\) nearest points to \(y\) and the weights \(w(x_i)\) depend on the chosen weighting scheme:

"squared" (default, torch_geometric convention): \(w(x_i) = 1 / d(\mathbf{p}(y), \mathbf{p}(x_i))^2\)
"inverse" (PointNet++ three_interpolation convention): \(w(x_i) = 1 / d(\mathbf{p}(y), \mathbf{p}(x_i))\)

Note

Adapted from the torch_geometric package.

Parameters:

x (Tensor) –

Node feature matrix \(\mathbf{X} \in \mathbb{R}^{N \times F}\).
pos_x (Tensor) –

Node position matrix \(\in \mathbb{R}^{N \times d}\).
pos_y (Tensor) –

Upsampled node position matrix \(\in \mathbb{R}^{M \times d}\).
batch_x (OptTensor, default: None ) –

Batch vector \(\mathbf{b_x} \in \{ 0, \ldots, B-1 \}^N\), assigning each node from \(\mathbf{X}\) to a specific example.
batch_y (OptTensor, default: None ) –

Batch vector \(\mathbf{b_y} \in \{ 0, \ldots, B-1 \}^M\), assigning each node from \(\mathbf{Y}\) to a specific example.
k (int, default: 3 ) –

Number of neighbors.
num_workers (int, default: 1 ) –

Number of workers for computation. Has no effect when batch_x or batch_y is not None, or the input lies on GPU.
weighting (Literal['squared', 'inverse'], default: 'squared' ) –

Weighting scheme for neighbors. "squared" for \(1/d^2\) weights (torch_geometric default) or "inverse" for \(1/d\) weights (PointNet++ convention).
eps (float, default: 1e-16 ) –

Small value to avoid division by zero.

Returns:

Tensor –

Interpolated features \(\in \mathbb{R}^{M \times F}\).

dense_neighbors [source] ¶

dense_neighbors(
    edge_index: Tensor,
    num_sources: int,
    num_targets: int,
    max_num_neighbors: Optional[int] = None,
) -> Tuple[Tensor, Tensor]

Pad a [source, target] edge list into a dense \((N_t, k)\) neighbor table.

\(k\) is the largest target degree unless max_num_neighbors is given (the cap of the search that produced the edges, which skips the device read). Rows are filled in edge order and padded with the shadow index num_sources, so gather_neighbors needs no mask. The second output is the slot of every edge in its target's row, so per-edge values scatter into the table as values_table[target, slot] = values. Edges grouped by target, which knn and radius return, are consumed as they are; any other order is sorted first.

Parameters:

edge_index (Tensor) –

Edges as [source, target], shape \((2, E)\).
num_sources (int) –

Number of source points \(N_s\), the shadow index of the padding.
num_targets (int) –

Number of target points \(N_t\).
max_num_neighbors (Optional[int], default: None ) –

Row width \(k\); must be at least the largest target degree.

Returns:

Tuple[Tensor, Tensor] –

The \((N_t, k)\) table of source indices and the \((E,)\) slot of every edge.

Shape

Input: \((2, E)\) Output: \((N_t, k)\) and \((E,)\)

Example

>>> import torch
>>> from torch_pointcloud.ops.cluster import dense_neighbors
>>> edge_index = torch.tensor([[3, 5, 0, 1, 2], [0, 0, 1, 1, 1]])
>>> table, slot = dense_neighbors(edge_index, num_sources=6, num_targets=2)
>>> table.tolist(), slot.tolist()
([[3, 5, 6], [0, 1, 2]], [0, 1, 0, 1, 2])

gather_neighbors [source] ¶

gather_neighbors(x: Tensor, table: Tensor) -> Tensor

Gather source features along a dense_neighbors table, zero where the table is padded.

Parameters:

x (Tensor) –

Source features, shape \((N_s, C)\).
table (Tensor) –

Neighbor table, shape \((N_t, k)\), padded with the shadow index \(N_s\).

Returns:

Tensor –

The neighbor features, shape \((N_t, k, C)\).

Shape

Input: \((N_s, C)\) and \((N_t, k)\) Output: \((N_t, k, C)\)

Example

>>> import torch
>>> from torch_pointcloud.ops.cluster import gather_neighbors
>>> x = torch.arange(6, dtype=torch.float).view(3, 2)
>>> gather_neighbors(x, torch.tensor([[2, 3], [0, 1]])).tolist()
[[[4.0, 5.0], [0.0, 0.0]], [[0.0, 1.0], [2.0, 3.0]]]