Function segmented_knn_graph

python/dgl/transforms/functional.py:337–487

r"""Construct multiple graphs from multiple sets of points according to k-nearest-neighbor (KNN) and return. Compared with :func:`dgl.knn_graph`, this allows multiple point sets with different capacity. The points from different sets are stored contiguously in the :attr:`x` tensor.

(
    x,
    k,
    segs,
    algorithm="bruteforce-blas",
    dist="euclidean",
    exclude_self=False,
)

# pylint: disable=invalid-name
def segmented_knn_graph(
    x,
    k,
    segs,
    algorithm="bruteforce-blas",
    dist="euclidean",
    exclude_self=False,
):
    r"""Construct multiple graphs from multiple sets of points according to
    k-nearest-neighbor (KNN) and return.

    Compared with :func:`dgl.knn_graph`, this allows multiple point sets with
    different capacity. The points from different sets are stored contiguously
    in the :attr:`x` tensor.
    :attr:`segs` specifies the number of points in each point set. The
    function constructs a KNN graph for each point set, where the predecessors
    of each point are its k-nearest neighbors measured by the Euclidean distance.
    DGL then composes all KNN graphs
    into a batched graph with multiple (:math:`len(segs)`) connected components.

    Parameters
    ----------
    x : Tensor
        Coordinates/features of points. Must be 2D. It can be either on CPU or GPU.
    k : int
        The number of nearest neighbors per node.
    segs : list[int]
        Number of points in each point set. The numbers in :attr:`segs`
        must sum up to the number of rows in :attr:`x`.
    algorithm : str, optional
        Algorithm used to compute the k-nearest neighbors.

        * 'bruteforce-blas' will first compute the distance matrix
          using BLAS matrix multiplication operation provided by
          backend frameworks. Then use topk algorithm to get
          k-nearest neighbors. This method is fast when the point
          set is small but has :math:`O(N^2)` memory complexity where
          :math:`N` is the number of points.

        * 'bruteforce' will compute distances pair by pair and
          directly select the k-nearest neighbors during distance
          computation. This method is slower than 'bruteforce-blas'
          but has less memory overhead (i.e., :math:`O(Nk)` where :math:`N`
          is the number of points, :math:`k` is the number of nearest
          neighbors per node) since we do not need to store all distances.

        * 'bruteforce-sharemem' (CUDA only) is similar to 'bruteforce'
          but use shared memory in CUDA devices for buffer. This method is
          faster than 'bruteforce' when the dimension of input points
          is not large. This method is only available on CUDA device.

        * 'kd-tree' will use the kd-tree algorithm (CPU only).
          This method is suitable for low-dimensional data (e.g. 3D
          point clouds)

        * 'nn-descent' is an approximate approach from paper
          `Efficient k-nearest neighbor graph construction for generic similarity
          measures <https://www.cs.princeton.edu/cass/papers/www11.pdf>`_. This method