Search code examples
pythonpytorchadjacency-matrixpytorch-geometric

How to convert sparse to dense adjacency matrix?


I am trying to convert a sparse adjacency matrix/list that only contains the indices of the non-zero elements ([[rows], [columns]]) to a dense matrix that contains 1s at the indices and otherwise 0s. I found a solution using to_dense_adj from Pytorch geometric (Documentation). But this does not exactly what I want, since the shape of the dense matrix is not as expected. Here is an example:

sparse_adj = torch.tensor([[0, 1, 2, 1, 0], [0, 1, 2, 3, 4]])

So the dense matrix should be of size 5x3 (the second array "stores" the columns; with non-zero elements at (0,0), (1,1), (2,2),(1,3) and (0,4)) because the elements in the first array are lower or equal than 2.

However,

dense_adj = to_dense(sparse_adj)[0]

outputs a dense matrix, but of shape (5,5). Is it possible to define the output shape or is there a different solution to get what I want?

Edit: I have a solution to convert it back to the sparse representation now that works

dense_adj = torch.sparse.FloatTensor(sparse_adj, torch.ones(5), torch.Size([3,5])).to_dense()
ind = dense_adj.nonzero(as_tuple=False).t().contiguous()
sparse_adj = torch.stack((ind[1], ind[0]), dim=0)

Or is there any alternative way that is better?


Solution

  • You can acheive this by first constructing a sparse matrix with torch.sparse then converting it to a dense matrix. For this you will need to provide torch.sparse.FloatTensor a 2D tensor of indices, a tensor of values as well as a output size:

    sparse_adj = torch.tensor([[0, 1, 2, 1, 0], [0, 1, 2, 3, 4]])
    torch.sparse.FloatTensor(sparse_adj, torch.ones(5), torch.Size([3,5])).to_dense()
    

    You can get the size of the output matrix dynamically with

    sparse_adj.max(axis=1).values + 1
    

    So it becomes:

    torch.sparse.FloatTensor(
        sparse_adj, 
        torch.ones(sparse_adj.shape[1]), 
        (sparse_adj.max(axis=1).values + 1).tolist())