def mm(indexA, valueA, indexB, valueB, m, k, n):
assert valueA.dtype == valueB.dtype
if indexA.is_cuda:
return torch_sparse.spspmm_cuda.spspmm(indexA, valueA, indexB, valueB,
m, k, n)
A = to_scipy(indexA, valueA, m, k)
B = to_scipy(indexB, valueB, k, n)
C = A.dot(B).tocoo().tocsr().tocoo() # Force coalesce.
indexC, valueC = from_scipy(C)
return indexC, valueC
def test_convert_scipy():
index = torch.tensor([[0, 0, 1, 2, 2], [0, 2, 1, 0, 1]])
value = torch.Tensor([1, 2, 4, 1, 3])
N = 3
out = from_scipy(to_scipy(index, value, N, N))
assert out[0].tolist() == index.tolist()
assert out[1].tolist() == value.tolist()
def transpose(index, value, m, n):
"""Transposes dimensions 0 and 1 of a sparse tensor.
Args:
index (:class:`LongTensor`): The index tensor of sparse matrix.
value (:class:`Tensor`): The value tensor of sparse matrix.
m (int): The first dimension of sparse matrix.
n (int): The second dimension of sparse matrix.
:rtype: (:class:`LongTensor`, :class:`Tensor`)
"""
if value.dim() == 1 and not value.is_cuda:
mat = to_scipy(index, value, m, n).tocsc()
(col, row), value = from_scipy(mat)
index = torch.stack([row, col], dim=0)
return index, value
row, col = index
index = torch.stack([col, row], dim=0)
index, value = coalesce(index, value, n, m)
return index, value
def transpose(index, value, m, n, coalesced=True):
"""Transposes dimensions 0 and 1 of a sparse tensor.
Args:
index (:class:`LongTensor`): The index tensor of sparse matrix.
value (:class:`Tensor`): The value tensor of sparse matrix.
m (int): The first dimension of corresponding dense matrix.
n (int): The second dimension of corresponding dense matrix.
coalesced (bool, optional): If set to :obj:`False`, will not coalesce
the output. (default: :obj:`True`)
:rtype: (:class:`LongTensor`, :class:`Tensor`)
"""
if value.dim() == 1 and not value.is_cuda:
mat = to_scipy(index, value, m, n).tocsc()
(col, row), value = from_scipy(mat)
index = torch.stack([row, col], dim=0)
return index, value
row, col = index
index = torch.stack([col, row], dim=0)
if coalesced:
index, value = coalesce(index, value, n, m)
return index, value
def masked_distance(graph):
sp = torch_sparse.to_scipy(graph.edge_index, graph.edge_attr[:, 0], graph.num_nodes, graph.num_nodes)
dense = sp.todense()
masked = np.ma.masked_equal(dense, value=0)
return masked
}, f"CASP{casp_ed}/processed/{target_id}.pth")
# %%
# ! ls CASP*/processed/*.pth | wc -l
# ! du -hsc CASP*/processed
# %%
sample = torch.load('CASP13/processed/T0980s2.pth')
fig, axes = plt.subplots(1, 3, figsize=(5*3, 4))
def masked_distance(graph):
sp = torch_sparse.to_scipy(graph.edge_index, graph.edge_attr[:, 0], graph.num_nodes, graph.num_nodes)
dense = sp.todense()
masked = np.ma.masked_equal(dense, value=0)
return masked
for ax, graph in zip(axes.flat, sample['graphs']):
print(graph)
dist = torch_sparse.to_scipy(graph.edge_index, graph.edge_attr[:, 0], graph.num_nodes, graph.num_nodes).todense()
img = ax.pcolormesh(
np.where(dist>0, dist, np.nan),
vmin=0,
vmax=dist.max(),
cmap="BuGn_r",
)
cbar = fig.colorbar(img, ax=ax)
ax.set_aspect('equal')
fig.tight_layout()
display(fig)
plt.close(fig)
import torch
from torch_sparse import to_scipy, from_scipy, coalesce
def transpose(index, value, m, nm coalesce=True):
"""Transposes dimensions 0 and 1 of a sparse tensor.
Args:
index (:class:`LongTensor`): The index tensor of sparse matrix.
value (:class:`Tensor`): The value tensor of sparse matrix.
m (int): The first dimension of corresponding dense matrix.
n (int): The second dimension of corresponding dense matrix.
coalesce (bool, optional): To return coalesced index and value or not (default: :obj:`True`)
:rtype: (:class:`LongTensor`, :class:`Tensor`)
"""
if value.dim() == 1 and not value.is_cuda:
mat = to_scipy(index, value, m, n).tocsc()
(col, row), value = from_scipy(mat)
index = torch.stack([row, col], dim=0)
return index, value
row, col = index
index = torch.stack([col, row], dim=0)
if coalesce:
index, value = coalesce(index, value, n, m)
return index, value
