fastdev.geom.knn
================

.. py:module:: fastdev.geom.knn




Module Contents
---------------

.. py:function:: knn_points(p1: torch.Tensor, p2: torch.Tensor, lengths1: Union[torch.Tensor, None] = None, lengths2: Union[torch.Tensor, None] = None, norm: int = 2, K: int = 1, version: int = -1, return_nn: bool = False, return_sorted: bool = True) -> _KNN

   K-Nearest neighbors on point clouds.

   :param p1: Tensor of shape (N, P1, D) giving a batch of N point clouds, each
              containing up to P1 points of dimension D.
   :param p2: Tensor of shape (N, P2, D) giving a batch of N point clouds, each
              containing up to P2 points of dimension D.
   :param lengths1: LongTensor of shape (N,) of values in the range [0, P1], giving the
                    length of each pointcloud in p1. Or None to indicate that every cloud has
                    length P1.
   :param lengths2: LongTensor of shape (N,) of values in the range [0, P2], giving the
                    length of each pointcloud in p2. Or None to indicate that every cloud has
                    length P2.
   :param norm: Integer indicating the norm of the distance. Supports only 1 for L1, 2 for L2.
   :param K: Integer giving the number of nearest neighbors to return.
   :param version: Which KNN implementation to use in the backend. If version=-1,
                   the correct implementation is selected based on the shapes of the inputs.
   :param return_nn: If set to True returns the K nearest neighbors in p2 for each point in p1.
   :param return_sorted: (bool) whether to return the nearest neighbors sorted in
                         ascending order of distance.

   :returns:

             Tensor of shape (N, P1, K) giving the squared distances to
                 the nearest neighbors. This is padded with zeros both where a cloud in p2
                 has fewer than K points and where a cloud in p1 has fewer than P1 points.

             idx: LongTensor of shape (N, P1, K) giving the indices of the
                 K nearest neighbors from points in p1 to points in p2.
                 Concretely, if `p1_idx[n, i, k] = j` then `p2[n, j]` is the k-th nearest
                 neighbors to `p1[n, i]` in `p2[n]`. This is padded with zeros both where a cloud
                 in p2 has fewer than K points and where a cloud in p1 has fewer than P1
                 points.

             nn: Tensor of shape (N, P1, K, D) giving the K nearest neighbors in p2 for
                 each point in p1. Concretely, `p2_nn[n, i, k]` gives the k-th nearest neighbor
                 for `p1[n, i]`. Returned if `return_nn` is True.
                 The nearest neighbors are collected using `knn_gather`

                 .. code-block::

                     p2_nn = knn_gather(p2, p1_idx, lengths2)

                 which is a helper function that allows indexing any tensor of shape (N, P2, U) with
                 the indices `p1_idx` returned by `knn_points`. The output is a tensor
                 of shape (N, P1, K, U).
   :rtype: dists


.. py:function:: knn_gather(x: torch.Tensor, idx: torch.Tensor, lengths: Union[torch.Tensor, None] = None)

   A helper function for knn that allows indexing a tensor x with the indices `idx`
   returned by `knn_points`.

   For example, if `dists, idx = knn_points(p, x, lengths_p, lengths, K)`
   where p is a tensor of shape (N, L, D) and x a tensor of shape (N, M, D),
   then one can compute the K nearest neighbors of p with `p_nn = knn_gather(x, idx, lengths)`.
   It can also be applied for any tensor x of shape (N, M, U) where U != D.

   :param x: Tensor of shape (N, M, U) containing U-dimensional features to
             be gathered.
   :param idx: LongTensor of shape (N, L, K) giving the indices returned by `knn_points`.
   :param lengths: LongTensor of shape (N,) of values in the range [0, M], giving the
                   length of each example in the batch in x. Or None to indicate that every
                   example has length M.

   :returns:

             Tensor of shape (N, L, K, U) resulting from gathering the elements of x
                 with idx, s.t. `x_out[n, l, k] = x[n, idx[n, l, k]]`.
                 If `k > lengths[n]` then `x_out[n, l, k]` is filled with 0.0.
   :rtype: x_out