Does Pytorch Geometric have a method for integrating the top K neighbors?

[michaelma01]

Hello all PyTorch Geometric Team and @rusty1s,

I'm wondering whether the Pytorch Geometric has a method for extracting a partial graph based on the edge weight. Similar to NeighborSampler, for each node in a graph, I want to only integrate the information from its partial neighbor nodes when doing the message passing. However, unlike NeighborSampler, I don't want to randomly pick the specified number of nodes. Instead, I want to pick the top K neighbor nodes based on the edges' weights.

For example, say I have a node nA, and it has five neighbor nodes nB, nC, nD, nE, nF. Their edge weights are separately 0.1, 0.2, 0.4, 0.5, 0.05. I want to only consider three of its neighbor nodes based on its weights. So when doing the message passing, I want the node nA to only integrate the information from nC, nD and nE because these three neighbor nodes have higher edge weight.

Does the Pytorch Geometric have such a method? Or do you have any suggestions how to modify the existing methods to achieve this goal? Thank you!

[rusty1s]

There does not exist such a functionality right now, sorry for that. There do exists plans though to add support for biased neighbor sampling which should cover this use-case as well.

[michaelma01]

I sse. Thank you @rusty1s.