Convert dense adjacency matrix to sparse matrix

[ChenYizhu97]

Hi, i want to convert a batched dense edge adjacency matrix of size (B,N,N) to a batched sparse edge adjacency matrix of size (2, M), in which B denotes the batch size, N denotes the maximum number of nodes each graph and M denotes the number of edges in one batch.

I could only find one function for this purpose in the package torch_geometric.utils named dense_to_sparse. However, the source code shows that this function does not work for batched data. My solution is to iterate over batch, apply this function on each graph, and aggregate the results to get the batched result. Is there any better way to achieve the same purpose?

The other question is that when i convert a batched sparse edge adjacency matrix of size (2,M) with edge attributes of size (M, F), in which F denotes the dimension of features for each edge, to a batched dense edge adjacency matrix using the function to_dense_adj provided by package torch_geometric.utils, It results a tensor of size (B, N, N, F), but I couldn’t find the function for converting such a tensor back. Is there any convenient way to convert such a tensor back to sparse edge adjacency matrix of size (2,M) and edge attributes of size (M, F)?

[rusty1s]

I modified the source-code of dense_to_sparse so that it can now convert batched adjacency matrices as well, see here.

There does not yet exist a function that can handle multi-dimensional edge features though :(

[ChenYizhu97]

Will you consider adding such function to torch_geometric? It could be helpful when we want to combine different kinds of layers in one GNN model. Especially when we want to insert the hierarchical pooling layer into several convolution layers, because such pooling layer sometimes takes dense adjacency matrixs as input and output, while the succeeding convolution layers may take the sparse adjacency matrixs as input.

For example, if we create a GNN model which contains 3 GINEConv layers, 1 diffpool layer and 2 GINEConv layer. The data should be converted to dense form first, then coverted back to sparse form after pooling with its edge attributes preserved.

[rusty1s]

I agree. I think as a workaround, the following should work just fine:

adj2 = adj.abs().sum(dim=-1)  # Find non-zero edges in `adj` with multi-dimensional edge_features
index = adj2.nonzero(as_tuple=True)
edge_attr = adj[index]
batch = index[0] * adj.size(-1)
index = (batch + index[1], batch + index[2])
edge_index = torch.stack(index, dim=0)

[ChenYizhu97]

Hi, rusty1s.
I think there might be a small problem in the code. It does not return the right edge_index.

I test it with the following code, and the function results edge_index that contains node index which is out of range.

def dense_to_sparse_with_attr(adj):
    adj2 = adj.abs().sum(dim=-1)  
    index = adj2.nonzero(as_tuple=True)
    edge_attr = adj[index]
    batch = index[0] * adj.size(-1)
    index = (batch + index[1], batch + index[2])
    edge_index = torch.stack(index, dim=0)
    return edge_index, edge_attr

if __name__ == "__main__":
    # adj is a batched dense adjacency matrix of size (2,2,2,3), 
    # which means batch size is 2, there are two nodes in each graph and the dimension of edge attributes is 3 
    adj = torch.tensor([
                [[[1,1,1], [0,1,0]],
                 [[1,0,1], [0,1,0]]],
                [[[1,1,1], [0,1,0]],
                 [[1,0,1], [0,1,0]]],
            ])
    edge_index, edge_attr = dense_to_sparse_with_attr(adj)
    print(edge_index)
    print(edge_index.max())

the output of the code is:

tensor([[0, 0, 1, 1, 3, 3, 4, 4],
[0, 1, 0, 1, 3, 4, 3, 4]])
tensor(4)

We only have 4 nodes here, which means the max value of edge_index is lower or equal to 3, but we get 4.

[rusty1s]

Yes, it needs to be:

batch = index[0] * adj.size(-2)