Build KNN graph over some subset of Node features
I have a point cloud that I want to use a graph neural network on. Each point in the point cloud is characterised by its positional coordinates as well as it's color. So a single node is (X, Y, Z, C).
Now I want to apply an Edge Convolution on this (as described in the DGL Edge-Conv example, and to do it I should build a Nearest Neighbors graph on (X, Y, Z) (And not on C), then use all the 4 properties as features for my neural network.
What would be a clean and efficient way to do this? (I have a lot of data so I want to batch and collate well)

Supposing you have a tensor pc of shape (NUM_POINTS, 4) where each row is (X, Y, Z, C), then you could use sklearn as follows:
from sklearn.neighbors import NearestNeighbors
import dgl

k = 3 # number of neighbours you want

neigh = NearestNeighbors(n_neighbors=k)
neigh.fit(pc[:, :3].numpy()) # selects only (X, Y, Z)

knn = neigh.kneighbors_graph()
graph = dgl.from_scipy(knn)
graph.ndata['x'] = pc

I would recommend saving these graphs to disk, so they are not computed each time you train etc.
