python - numpy.repeat() 创建 block 对角线索引？

Question

我正在尝试找出如何加速以下 Python 代码。
基本上，该代码构建了矩阵 C 的外部乘积矩阵，并将其存储为 block 对角稀疏矩阵。
我使用numpy.repeat()在 block 对角线中构建索引。
分析代码显示对numpy.repeat()的调用> 占用大约 50% 的执行时间。

import numpy as np
import scipy.sparse as spspar

L = 1000
K = 100
C = np.random.randn(L,K)

# From the matrix of outter products of C and store in block_diagonal 
# sparse matrix
CCt = np.einsum('ij...,i...->ij...',C,C)

# create indices into the block diagonal sparse coo matrix
i = np.tile(np.tile(np.arange(K),K),L) + K*np.repeat(np.arange(L),K*K)
j = np.tile(np.repeat(np.arange(K),K),L) + K*np.repeat(np.arange(L),K*K)

# store as block diagonal sparse coo matrix
BlckCCt = spspar.coo_matrix((CCt.flatten(),(j,i)),shape=(K*K*L,K*K*L))

最初，我按如下方式构建稀疏矩阵

BlckCCt = spspar.block_diag(CCt,"coo")

这太慢并且占用内存。

感谢您的任何意见。

编辑:我使用 ipython timeit 比较了 @hjpaul 的建议。这是我可以报告的内容

timeit K*np.repeat(np.arange(L),K*K)
10 loops, best of 3: 82.1 ms per loop

timeit (np.zeros((K*K,),int)+np.arange(L)[:,None]).flatten()*K
10 loops, best of 3: 89.9 ms per loop

timeit np.tile(np.arange(L)*K,K*K).reshape(K*K,L).T.flatten()
10 loops, best of 3: 85.5 ms per loop

所以看起来它们都花费了大约相同的数量(我是 ipython 分析的新手，所以也许我没有以正确的方式比较它们)。

Answer 1

仅供引用，您的

CCt = np.einsum('ij...,i...->ij...',C,C)

与

相同

CCt1=C[:,None,:]*C[:,:,None]

生成一个(L,K,K)数组。对于我的较小测试用例，np.einsum 速度提高了 2 倍。

sparse.block_diag 将每个子矩阵转换为 coo，并将它们传递给 sparse.bmat。 bmat 将所有子矩阵的 rows、cols、data 收集到一个类似于您的 的大数组中>j, i，并用它们调用 coo_matrix。

在不同的部分上执行ipython timeit，我同意K*np.repeat(np.arange(L),K*K) 是最慢的代码块。例如，比 tile 片段慢得多。

既然你对i和j做同样的repeat，你不能只做一次，然后使用该变量两次吗？

kk= K*np.repeat(np.arange(L),K*K)
ii=np.tile(np.tile(np.arange(K),K),L) + kk
jj=np.tile(np.repeat(np.arange(K),K),L) + kk

我会再看一下那篇文章，但这只是一个开始。

这是对重复的轻微改进(20%):

(np.zeros((K*K,),int)+np.arange(L)[:,None]).flatten()*K

更好 (>2x)

np.tile(np.arange(L)*K,K*K).reshape(L,K*K).T.flatten()

我将 *K 移至较小的 arange(L)，并使用更快的 tile。 .T.flatten 负责更改顺序。

根据评论， reshape 应该是(K*K,L)。我正在用无关紧要的值进行测试。这些替代方案的相对速度随 K 和 L 的相对大小而变化。

如果 kk(第二部分)形状为 (L,K)，则 i 和 j 的第一部分的平铺是可选的,K)(如CCt)。是否能节省时间尚不清楚。跨步比完全平铺版本 (0,4,0) v. (4,).)

更复杂

i = (np.arange(K)[None,None,:] + kk.reshape(L,K,K)).flatten()
j = (np.arange(K)[None,:,None] + kk.reshape(L,K,K)).flatten()

我们可以对 kk 做同样的事情

k1 = K*np.arange(L)[:,None,None]

np.arange(K)[None,None,:] + k1 是 (L,1,K)，所以我们需要平铺它

i = np.tile( np.arange(K)[None,None,:] + k1, (1,K,1)).flatten()
j = np.tile( np.arange(K)[None,:,None] + k1, (1,1,K)).flatten()

<小时/>

生成这些数组的另一种方法是使用 np.ix_ reshape 范围，然后对值进行求和。

i = np.sum(np.ix_(K*np.arange(L), np.arange(K), np.zeros(K)))
j = np.sum(np.ix_(K*np.arange(L), np.zeros(K), np.arange(K)))

(根据需要添加.flatten)。我已经在小尺寸上对此进行了测试，看起来不错。我不知道速度。