radius_graph with max_num_neighbors to take the nearest neighbours

[Akhil-CM]

When I use radius_graph in my pointcloud data and try to get all neighbours(by specifying a large number like int(1e6) to max_num_neighbors) in that specified radius, say 20., I sometimes run out of memory.

So, I want to know if it is possible to use radius_graph from torch_geometric, with radius say 20., and specify the parameter max_num_neighbors, eg 100, and make sure that they are the nearest ones? Some of the nodes in a graph might have less than 100 neighbours like 35, 55, 10, and some might have more than 100 neighbours, like 200, 300. For the less than 100, I want all the neighbours and for more than 100, I want the 100 nearest.

[Akhil-CM]

I looked at the source code at torch_geometric nn pool radius_graph and the cuda kernel

It seems there's no nearest neighbor selection while choosing the max_num_neighbours and more like a first come first serve basis insertion.

I have now used knn_graph to make my own implementation of nearest_radius_graph:

    def radiusGraphNearest(self):
        def f(x, batch, r, loop, max_num_neighbors):
            edge_index = knn_graph(x, k=max_num_neighbors, batch=batch, loop=loop)
            row, col = edge_index
            distances = (x[col] - x[row]).norm(dim=1)
            mask = distances <= r
            final_edge_index = edge_index[:, mask]

            return final_edge_index

        return f

Maybe someone (@rusty1s ?) can confirm whether this function is indeed effective without any side-effects