在numpy中加速双倍for循环

Question

我的 Python 代码中目前有以下双循环：

for i in range(a):
    for j in range(b):
        A[:,i]*=B[j][:,C[i,j]]   

(A 是浮点矩阵。B 是浮点矩阵列表。C 是整数矩阵。矩阵是指 m x n np.arrays。

确切地说，大小是： A: mxa B: b 大小为 mxl 的矩阵（每个矩阵的 l 不同） C: axb。这里m很大，a很大，b很小，l比b还要小 )

我试图通过这样做来加快速度

for j in range(b):
    A[:,:]*=B[j][:,C[:,j]]

但令我惊讶的是，这表现得更糟。

更准确地说，这确实提高了 m 和 a 的小值（“大”数）的性能，但从 m=7000,a=700 开始，第一个方法的速度大约是原来的两倍。

还有什么我可以做的吗？

也许我可以并行化？但我真的不知道怎么做。

（我不致力于 Python 2 或 3）

Answer 1

这是一种矢量化方法，假设 B 是具有相同形状的数组列表 -

# Convert B to a 3D array
B_arr = np.asarray(B)

# Use advanced indexing to index into the last axis of B array with C
# and then do product-reduction along the second axis. 
# Finally, we perform elementwise multiplication with A
A *= B_arr[np.arange(B_arr.shape[0]),:,C].prod(1).T

---

对于较小的a，我们可以运行一个循环来迭代a 的长度。此外，为了获得更高的性能，最好将这些元素存储到一个单独的二维数组中，并在我们退出循环后只执行一次元素乘法。

因此，我们会有这样的替代实现 -

range_arr = np.arange(B_arr.shape[0])
out = np.empty_like(A)
for i in range(a):
    out[:,i] = B_arr[range_arr,:,C[i,:]].prod(0)
A *= out