pytorch dense to coo, edge_index to coo

pytorch 的dense matrix如何转成coo的sparse tensor？

给定边的下标，如何转成coo格式的稀疏矩阵？

稀疏矩阵如何转成dense的矩阵？

本文全部是pytorch的操作，如果想要numpy和pytorch结合的操作，请参考我另外一篇:

pytorch 邻接矩阵转稀疏矩阵 (dense matrix to coo matrix)

一. dense矩阵转coo稀疏矩阵

    def dense_to_coo(a:torch.Tensor):
        idx = torch.nonzero(a).T
        data = a[idx[0],idx[1]]
        coo = torch.sparse_coo_tensor(idx, data, a.shape)
        return coo

二. edge_index转coo稀疏矩阵

这里的edge index也就是边的下标，indices。

    def edge_index_to_coo(edge_index:torch.Tensor, N:int):
        v = torch.ones(edge_index.size(1))
        coo = torch.sparse_coo_tensor(edge_index, v, (N, N))
        return coo

三. coo转dense

dense = coo.to_dense()即可。