numpy.repeat() 来创建块对角线索引？

Question

我想弄清楚如何加速以下 Python 代码。
基本上，该代码构建矩阵 C 的外积矩阵并存储它作为块对角线稀疏矩阵。
我使用 numpy.repeat() 将索引构建到块对角线中。
分析代码显示对 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

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

repeat

略有改进 (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 负责更改顺序。

根据评论，整形应该是 (K*K,L)。我正在测试无关紧要的价值观。这些替代方案的相对速度随 K 和 L.

的相对大小而变化

i 和 j 的第一部分的平铺是可选的，如果 kk（第二部分）是形状 (L,K,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_ 重塑范围，然后对值求和。

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）。我已经在小尺寸上对此进行了测试，它看起来不错。我不知道速度。

numpy.repeat() 来创建块对角线索引？

numpy.repeat() to create block-diagonal indices?

python

numpy

scipy