How can I efficiently mask out certain pairs in (2, N) tensor?

Question

I have a torch tensor edge_index of shape (2, N) that represents edges in a graph. For each (x, y) there is also a (y, x), where x and y are node IDs (ints). During the forward pass of my model I need to mask out certain edges. So, for example, I have:

n1 = [0, 3, 4]   # list of node ids as x
n2 = [1, 2, 1]   # list of node ids as y
edge_index = [[1, 2, 0, 1, 3, 4, 2, 3, 1, 4, 2, 4],  # actual edges as (x, y) and (y, x)
              [2, 1, 1, 0, 4, 3, 3, 2, 4, 1, 4, 2]]

# do something that efficiently removes (x, y) and (y, x) edges as formed by n1 and n2

Final edge_index should look like:

>>> edge_index
[[1, 2, 3, 4, 2, 4],
 [2, 1, 4, 3, 4, 2]]

Preferably we need to efficiently make some kind of boolean mask that I can apply to edge index e.g. as edge_index[:, mask] or something like that.

Could also be done in numpy but I'd like to avoid converting back and forth.

Edit #1:

If that can't be done, then I can think of a way so that, instead of n1 and n2, I have access to the indices of the positions I need to exclude in one tensor e.g. _except=[2, 3, 6, 7, 8, 9] (by making a dict/index once in the beginning).

Is there a way to get the desired result by "telling" edge_index to drop the indices in except? edge_index[:, _except] gives me the ones I want to get rid of. I need its complement operation.

Edit #2:

I managed to do it like this:

mask = torch.ones(edge_index.shape[1], dtype=torch.bool)
for i in range(len(n1)):
    mask = mask &  ~(torch.tensor([n1[i], n2[i]], dtype=torch.long) == edge_index.T).all(dim=1) & ~(torch.tensor([n2[i], n1[i]], dtype=torch.long) == edge_index.T).all(dim=1)

edge_index[:, mask]

but it is too slow and I can't use it. How can I speed it up?

Edit #3: I managed to solve this Edit#1 efficiently with:

mask = torch.ones(edge_index.shape[1], dtype=torch.bool)
mask[_except] = False
edge_index[:, mask]

Still interested in solving the original problem if someone comes up with something...

user_16115033 · Answer

If you're ok with the way you suggested at Edit#1,

you get the complement result by:

edge_index[:, [i for i in range(edge_index.shape[1]) if not (i in _except)]]

hope this is fast enough for your requirement.

Edit 1:

from functools import reduce

ids = torch.stack([torch.tensor(n1), torch.tensor(n2)], dim=1)
ids = torch.cat([ids, ids[:, [1,0]]], dim=0)
res = edge_index.unsqueeze(0).repeat(6, 1, 1) == ids.unsqueeze(2).repeat(1, 1, 12)
mask = ~reduce(lambda x, y: x | (reduce(lambda p, q: p & q, y)), res, reduce(lambda p, q: p & q, res[0]))
edge_index[:, mask]

Edit 2:

ids = torch.stack([torch.tensor(n1), torch.tensor(n2)], dim=1)
ids = torch.cat([ids, ids[:, [1,0]]], dim=0)
res = edge_index.unsqueeze(0).repeat(6, 1, 1) == ids.unsqueeze(2).repeat(1, 1, 12)
mask = ~(res.sum(1) // 2).sum(0).bool()
edge_index[:, mask]

How can I efficiently mask out certain pairs in (2, N) tensor?

1 Answers