python - 堆叠稀疏矩阵

Question

我有具有以下功能的密集数组的现有代码，并且也希望支持稀疏矩阵。

特别是，我有一个四维tensor的形状(k, k, n, n)，我部分地压平成这样matrix的形状(k * n, k * n)。

# Roll the third dimension to the second position to get a tensor
# with shape (k, n, k, n)
rolled = np.rollaxis(tensor, 2, 1)
# Extract the dimensions
k, n = rolled.shape[:2]
# Reshape to partially flatten the array
matrix = rolled.reshape((k * n, k * n))

在稀疏设置中，我有一个k由k列表组成的列表，其中每个元素都是一个n稀疏n矩阵。使用np.rollaxis不是一种选择，因为有些维度是稀疏的，有些是密集的。关于如何实现相同行为的任何想法，以便我最终得到一个带有 shape 的稀疏矩阵(k * n, k * n)？

一种选择是使用坐标格式并单独设置每个元素matrix，但这似乎很麻烦且容易出错。

标签： pythonnumpymatrixscipysparse-matrix