Warning

This page was created from a pull request (#250).

pypose.knn

class pypose.knn(ref, nbr, k=1, ord=2, dim=- 1, largest=False, sorted=True)

Select the k nearest neighbor points of reference from neighbors in each batch.

Parameters:

• ref (torch.Tensor) – the coordinates of the reference point sets. The shape has to be (…, N1, :).

• nbr (torch.Tensor) – the coordinates of the neighbors point sets. The shape has to be (…, N2, :).

• k (int, optional) – the number of the nearest neighbors to be selected. k has to be k \(\leq\) N2. Default: 1.

• ord (int, optional) – the order of norm to use for distance calculation. Default: 2 (Euclidean distance).

• dim (int, optional) – the dimension encompassing the point cloud coordinates, utilized for calculating distance and sorting. Default: -1 (The last dimension).

• largest (bool, optional) – controls whether to return largest (furthest) or smallest (nearest) neighbors. Default: False.

• sorted (bool, optional) – controls whether to return the neighbors in sorted order. Default: True.

Returns:

The named tuple of (values, indices).

values: The ord-norm distance between each point in ref and its sorted k nearest neighbors in nbr. The shape is (…, N1, k).

indices: The index of the k nearest neighbor points in neighbors point sets (nbr). The shape is (…, N1, k).

Return type:

torch.return_types.topk(values: torch.Tensor, indices: torch.LongTensor)

Note

If sorted is set to False, the output will be unspecified and not necessarily sorted along the index of the input point cloud.

Example

>>> import torch, pypose as pp
>>> ref = torch.tensor([[9., 2., 2.],
...                     [1., 0., 2.],
...                     [0., 1., 1.],
...                     [5., 0., 1.],
...                     [1., 0., 1.],
...                     [5., 5., 3.]])
>>> nbr = torch.tensor([[1., 0., 1.],
...                     [1., 6., 2.],
...                     [5., 1., 0.],
...                     [9., 0., 2.]])
>>> pp.knn(ref, nbr)
torch.return_types.topk(
values=tensor([[2.0000],
        [1.0000],
        [1.4142],
        [1.4142],
        [0.0000],
        [4.2426]]),
indices=tensor([[3],
        [0],
        [0],
        [2],
        [0],
        [1]]))
>>> pp.knn(ref, nbr, k=2, ord=2)
torch.return_types.topk(
values=tensor([[2.0000, 4.5826],
        [1.0000, 4.5826],
        [1.4142, 5.0990],
        [1.4142, 4.0000],
        [0.0000, 4.2426],
        [4.2426, 5.0000]]),
indices=tensor([[3, 2],
        [0, 2],
        [0, 2],
        [2, 0],
        [0, 2],
        [1, 2]]))
>>> pp.knn(ref, nbr, k=2, ord=2).values
tensor([[2.0000, 4.5826],
        [1.0000, 4.5826],
        [1.4142, 5.0990],
        [1.4142, 4.0000],
        [0.0000, 4.2426],
        [4.2426, 5.0000]])