Coo_matrix 切片问题

import torch
import torch_sparse
from torch_sparse import index_select
from torch_sparse.tensor import SparseTensor

#创建一个COO矩阵
##创建coo矩阵有使用Scipy库和torch库的两种方法

    #scipy库
import numpy as np
from scipy.sparse import coo_matrix
    # 创建一个3x3的COO矩阵
data = np.array([1,1, 2, 3])
row  = np.array([0, 0,1, 2])
col  = np.array([0, 1,1, 2])
scipy_coo_matrtix = coo_matrix((data, (row, col)), shape=(3, 3))
    #查看coo矩阵内容
print(scipy_coo_matrtix.todense())

    #pytorch库
    # 创建一个3x3的COO矩阵
values = torch.tensor([1,1,2, 3])
indices = torch.tensor([[0,0, 1, 2], [0,1, 1, 2]])
coo = torch.sparse_coo_tensor(indices, values,size=(3, 3))

    # 查看矩阵内容
print(coo.to_dense())

#对比矩阵
ori_matirx =torch.Tensor([
    [0,1,2],
    [1,2,3],
    [4,5,6]
])


#在分batch的训练中 常常只需要邻接矩阵的一部分，即子图，则需要从原来的边中进行分割


# 转化了可以切片操作的Sparse_tensor
sparse_coo_1 = SparseTensor.from_torch_sparse_coo_tensor(coo)
sparse_coo_2 = SparseTensor.from_scipy(scipy_coo_matrtix)

#切片
coo_new_1 = sparse_coo_1[[0,1],:]
coo_new_2 = sparse_coo_2[[0,1],:]
print("coo_1:",coo_new_1.to_dense())
print("coo_2:",coo_new_2.to_dense())

#将新的邻接矩阵转为原来coo矩阵
coo_ori_1 = coo_new_1.to_torch_sparse_coo_tensor()
coo_ori_2 = coo_new_2.to_torch_sparse_coo_tensor()


print("finished")

#上面的代码是用于将Batch中的点的相应的邻接矩阵中的边拿出来，我们还需要将这些Batch中的点的邻接节点也拿出来

    #重新写一个矩阵(6,6)
value_new = torch.tensor([1]*14)
indice_new = torch.tensor([[0,0,1,1,2,2,3,3,3,4,4,5,5,5], [0,2,1,3,0,2,1,3,5,4,5,3,4,5]])
coo_new = torch.sparse_coo_tensor(indice_new,value_new,size=(6,6))

print("new_coo",coo_new.to_dense())

#这里我们选择[2,5]的点
    #转换矩阵
new_coo = SparseTensor.from_torch_sparse_coo_tensor(coo_new)
    #取处指定的点的边列表,可将Sparsez转为标量,避免相加的时候出错
specific_coo = new_coo[[2,5],:]
    #找到这两个点所包含的所有邻接节点
neighbor = torch.sparse.sum(specific_coo.to_torch_sparse_coo_tensor(),dim=0)._indices()[0]
    #在原来的边列表中找到对应的点的边（Batch中的点加上他们的邻接节点）
new_indice = indice_new[:,(indice_new[0].view(-1,1)==neighbor).any(-1)]
    #得到的最终的子图上的边
final = torch.sparse_coo_tensor(new_indice,[1]*new_indice.shape[1],size=[6,6])