Neighbor sampling to have degree-bounded graphs

[sisaman]

I wonder what would be the most efficient way to apply neighbor sampling on a graph data object so that for every node, we have (inbound + outbound degree) <= K for some constant K?

Currently, I am using torch_sparse.neighbor_sample as follows:

from torch_geometric.loader.utils import to_csc, filter_data

colptr, row, perm = to_csc(data, device='cpu')
index = torch.range(0, data.num_nodes-1, dtype=int)
sample_fn = torch.ops.torch_sparse.neighbor_sample
node, row, col, edge = sample_fn(colptr, row, index, [K], False, True)
data = filter_data(data, node, row, col, edge, perm)

However, I can only bound either inbound or outbound degree, not inbound+outbound.
Is there a fast way to do this without iterating manually over all the edges?

[rusty1s]

Interesting. Besides adding this directly into the low-level C++ op, I guess one alternative is to apply a filtering of edges after sampling has been done:

1. compute the colptr representation of the col output
2. compute some edge mask based on colptr offsets (currently unsure about how to do that efficiently)
3. Apply the masking both on row and col

[sisaman]

Thanks for your suggestion. However, having spent quite a lot of time on this, I couldn't figure out any algorithm working better than O(E), with E being the number of edges. But I am still in doubt if there would exist an O(NK) algorithm 🤔

[rusty1s]

Don't you get O(NK) for free if you apply that in a post-processing step? Since in that case E == NK.

[sisaman]

Yes, that's right. Thank you!