r"""For each element in each segment in :attr:`y`, find :attr:`k` nearest points in the same segment in :attr:`x`. If :attr:`y` is None, perform a self-query over :attr:`x`. This function allows multiple point sets with different capacity. The points from different sets are stored c
(
k, x, x_segs, y=None, y_segs=None, algorithm="bruteforce", dist="euclidean"
)
| 639 | |
| 640 | |
| 641 | def knn( |
| 642 | k, x, x_segs, y=None, y_segs=None, algorithm="bruteforce", dist="euclidean" |
| 643 | ): |
| 644 | r"""For each element in each segment in :attr:`y`, find :attr:`k` nearest |
| 645 | points in the same segment in :attr:`x`. If :attr:`y` is None, perform a self-query |
| 646 | over :attr:`x`. |
| 647 | |
| 648 | This function allows multiple point sets with different capacity. The points |
| 649 | from different sets are stored contiguously in the :attr:`x` and :attr:`y` tensor. |
| 650 | :attr:`x_segs` and :attr:`y_segs` specifies the number of points in each point set. |
| 651 | |
| 652 | Parameters |
| 653 | ---------- |
| 654 | k : int |
| 655 | The number of nearest neighbors per node. |
| 656 | x : Tensor |
| 657 | The point coordinates in x. It can be either on CPU or GPU (must be the |
| 658 | same as :attr:`y`). Must be 2D. |
| 659 | x_segs : Union[List[int], Tensor] |
| 660 | Number of points in each point set in :attr:`x`. The numbers in :attr:`x_segs` |
| 661 | must sum up to the number of rows in :attr:`x`. |
| 662 | y : Tensor, optional |
| 663 | The point coordinates in y. It can be either on CPU or GPU (must be the |
| 664 | same as :attr:`x`). Must be 2D. |
| 665 | (default: None) |
| 666 | y_segs : Union[List[int], Tensor], optional |
| 667 | Number of points in each point set in :attr:`y`. The numbers in :attr:`y_segs` |
| 668 | must sum up to the number of rows in :attr:`y`. |
| 669 | (default: None) |
| 670 | algorithm : str, optional |
| 671 | Algorithm used to compute the k-nearest neighbors. |
| 672 | |
| 673 | * 'bruteforce' will compute distances pair by pair and |
| 674 | directly select the k-nearest neighbors during distance |
| 675 | computation. This method is slower than 'bruteforce-blas' |
| 676 | but has less memory overhead (i.e., :math:`O(Nk)` where :math:`N` |
| 677 | is the number of points, :math:`k` is the number of nearest |
| 678 | neighbors per node) since we do not need to store all distances. |
| 679 | |
| 680 | * 'bruteforce-sharemem' (CUDA only) is similar to 'bruteforce' |
| 681 | but use shared memory in CUDA devices for buffer. This method is |
| 682 | faster than 'bruteforce' when the dimension of input points |
| 683 | is not large. This method is only available on CUDA device. |
| 684 | |
| 685 | * 'kd-tree' will use the kd-tree algorithm (CPU only). |
| 686 | This method is suitable for low-dimensional data (e.g. 3D |
| 687 | point clouds) |
| 688 | |
| 689 | * 'nn-descent' is an approximate approach from paper |
| 690 | `Efficient k-nearest neighbor graph construction for generic similarity |
| 691 | measures <https://www.cs.princeton.edu/cass/papers/www11.pdf>`_. This method |
| 692 | will search for nearest neighbor candidates in "neighbors' neighbors". |
| 693 | |
| 694 | Note: Currently, 'nn-descent' only supports self-query cases, i.e. :attr:`y` is None. |
| 695 | (default: 'bruteforce') |
| 696 | dist : str, optional |
| 697 | The distance metric used to compute distance between points. It can be the following |
| 698 | metrics: |
no test coverage detected