I have a tensor which contains indices that correspond to nodes in two different graphs that were chosen based of some criteria. So let’s say indices = tensor([[0, 1], [3, 2]]
. This means that in the first graph, the node with index 0 and in the second graph the node with index 1 were chosen. Also, in the first graph the node with index 3 and in the second graph the node with index 2 were chosen. This kind of format is necessary in my case.
Now i need to compare each element in indices
with each other with a function called is_adjacent
to see which pair of nodes i still need.
My current solution:
edge_index = []
for k, i in enumerate(indices):
for l, j in enumerate(indices):
if _is_adjacent(i[0], j[0], graph_1) and _is_adjacent(i[1], j[1], graph_2):
edge_index.append([k,l])
edge_index = torch.tensor(edge_index, dtype=torch.float)
Where the following function is used:
def _is_adjacent(id_a, id_b, graph):
tens = [id_a, id_b]
adjacency_matrix = torch.reshape(graph.edge_index, (-1,2)).tolist()
return tens in adjacency_matrix
This works as it should, but it is not really fast. i would like to replace my for loops with some kind of vectorization.
I hope this makes sense.