Using KNN to find distance

[edshkim98]

From now on, we recommend using our discussion forum (https://github.com/rusty1s/pytorch_geometric/discussions) for general questions.

❓ Questions & Help

Hello, I am trying to calculate the distance between every point by using knn function in torch_geometric.
However, it only returns index values rather than distances.
Is there any way to get the values of distance rather than index values?
Thank you!

[Km3888]

You could use broadcasting to find the distances between the points and their nearest neighbors using the indices given to you by knn. The first step would be getting the values of each neighbor by using the knn output to index the original tensor.

[rusty1s]

src_index, dst_index = knn(...)
dist = (x_src[src_index] - x_dst[dst_index]).norm(dim=-1)

[chenzhekl]

Can we make knn_graph output pair-wise distances except for the edge_index to avoid duplicate computation? Probably, we can add a return_distances=False parameter following the convention of numpy.unique.

[rusty1s]

Yes, that is something which I would like to add, but hadn't have time to do so. Let me know if you want to contribute this :)

[chenzhekl]

@rusty1s Thanks for your reply. I am not quite familiar with the codebase of pytorch_cluster. I will see if I can help this weekend.

[rusty1s]

Sure, please feel free to ping me anytime.

[chenzhekl]

Hi @rusty1s, I tried rewriting the signature of knn as follows

@torch.jit.script
def knn(
    x: torch.Tensor,
    y: torch.Tensor, k: int,
    batch_x: Optional[torch.Tensor] = None,
    batch_y: Optional[torch.Tensor] = None,
    cosine: bool = False,
    num_workers: int = 1,
    return_distances: bool = False
) -> Union[Tuple[torch.Tensor, torch.Tensor], torch.Tensor]:

to include the distances in the return values. However, It seems that there is no support for Union types in TorchScript. Do you have any suggestions?

[rusty1s]

For this, you can use the concept of overloads, e.g., GATConv uses it to define different return types based on return_attention_weights:

    @torch.jit._overload_method
    def forward(self, x, edge_index, edge_attr=None, size=None, return_attention_weights=None):
        # type: (Tensor, Tensor, OptTensor, Size, NoneType) -> Tensor
        pass

    @torch.jit._overload_method
    def forward(self, x, edge_index, edge_attr=None, size=None, return_attention_weights=None):
        # type: (Tensor, Tensor, OptTensor, Size, bool) -> Tuple[Tensor, Tuple[Tensor, Tensor]]
        pass

    def forward(self, x, edge_index, edge_attr=None, size=None, return_attention_weights=None):
        # Implement function here