KNNGraph#

class KNNGraph(pos, neighbours, **kwargs)[source]#

Bases: GraphBuilder

Extends the GraphBuilder class to compute edge_index based on a K-nearest neighbors algorithm.

Instantiate the Graph class by computing the edge_index based on the K-nearest neighbors algorithm.

Parameters:

  • pos (torch.Tensor | LabelTensor) – A tensor of shape (N, D) representing the positions of N points in D-dimensional space.

  • neighbours (int) – The number of nearest neighbors to consider when building the graph.

  • kwargs (dict) – The additional keyword arguments to be passed to GraphBuilder and Graph classes.

Returns:

A Graph instance with the computed edge_index.

Return type:

Graph

static compute_knn_graph(points, neighbours)[source]#

Computes the edge_index based on the K-nearest neighbors algorithm.

Parameters:

  • points (torch.Tensor | LabelTensor) – A tensor of shape (N, D) representing the positions of N points in D-dimensional space.

  • neighbours (int) – The number of nearest neighbors to consider when building the graph.

Returns:

A tensor of shape (2, E), with E number of edges, representing the edge indices of the graph.

Return type:

torch.Tensor