How to randomly sample edges while weighting according to Euclidean distance

[amorehead]

Hello.

Having recently read the Chroma paper by Generate Biomedicines (https://www.biorxiv.org/content/10.1101/2022.12.01.518682v1?rss=1), I was curious to see how difficult it would be to implement the random edge sampling algorithm described in this preprint. Specifically, for each node, they determine the 20-nearest neighbors to define edges between these neighbors (already doable using knn_graph). However, where the implementation gets a bit tricky is where they describe also including 40 random, long-range (far Euclidean distance) neighbors in their edge construction for each node.

Would anyone here happen to know whether such an implementation is currently tractable using PyG?

[wsad1]

I can't think of any efficient implementation in pyg for this. But some things you could try out are

1. Check out fps. You might be able to modify it for your use case.
2. knn_graph set num_neighbors to say 100 and use the first 20 and last 40 points by distance as neighbors.
3. If you don't have too many points, you could compute pairwisedistance and then use weightedsampler. Probably something you already thought of.