如何矢量化这个 for 循环？

Question

我有一个长度为 n 的 numpy 数组 f 和一个长度为 n x 的 numpy 矩阵 A em>m。我想将 f 和 A 分成 r 部分 f1,...,fr 和 A1,...,Ar，然后进行计算 fi*Ai（数学意义上的向量 x 矩阵乘法）每个 fi 是一个行向量，其列数等于 Ai 的行数。结果将是一个 1 x m 的行向量。这个想法是连接所有这些行向量以形成矩阵 B = [ [f1*A1], [f2*A2],..., [fr*Ar] ] （请注意，这将是一个大小为 r x m) 的矩阵。

假设 f 和 A 已经定义。还假设片段的相应索引在列表 [0,d1,...dr] 中。例如，f1 = f[d[0]:d[1]] 和 f2 = f[d[1]:d[2]]）。我正在使用以下代码来解决我的问题：

B = numpy.zeros([r,m])
for i in range(0,r):
    lower = d[i]
    upper = d[i+1]
    B[i,:] = f[lower:upper].dot(A[lower:upper,:])

问题是这段代码将在我的程序中计算多次。我之前听说 Python for 循环很慢，实际上我的代码的瓶颈就是这部分。我不知道如何对其进行矢量化，但我觉得这是可能的。我希望这里有人能给我指路。谢谢。

Answer 1

你可以使用np.add.reduceat:

# example data
>>> f = np.arange(10)
>>> A = np.arange(50).reshape(10, 5)
>>> split = [0, 3, 5, 10]
>>> 
# reduceat
>>> np.add.reduceat(f[:, None] * A, split[:-1], axis=0)
array([[  25,   28,   31,   34,   37],
       [ 125,  132,  139,  146,  153],
       [1275, 1310, 1345, 1380, 1415]])
>>> 
# double check against list comprehension
>>> [fi @ Ai for fi, Ai in zip(*map(np.split, (f, A), 2*(split[1:-1],)))]
[array([25, 28, 31, 34, 37]), array([125, 132, 139, 146, 153]), array([1275, 1310, 1345, 1380, 1415])]

如果列表理解或 @hpaulj 的解决方案或 OP 的循环由于 blas 加速矩阵乘法而更快，我不会感到惊讶。

Answer 2

我认为这是一个有效的 MCVE：

In [139]: f = np.arange(10)
In [140]: A = np.arange(20).reshape(10,2)
In [141]: f.dot(A)
Out[141]: array([570, 615])
In [142]: d = [0,2,5,10]
In [143]: for i,j in zip(d[:-1],d[1:]):
     ...:     print(f[i:j].dot(A[i:j,:]))
     ...:     
[2 3]
[58 67]
[510 545]

570 = 2+58+510.

In [145]: np.array([f[i:j].dot(A[i:j,:]) for i,j in zip(d[:-1],d[1:])])
Out[145]: 
array([[  2,   3],
       [ 58,  67],
       [510, 545]])

鉴于i:j 切片的长度可能不同，因此可能很难真正“矢量化”它。我们可以隐藏迭代，但是以将所有迭代移动到编译代码中的方式编写它会很棘手。像cumsum 这样的累积操作通常是最好的选择。我们经常不得不退后一步，从不同的角度看待问题（而不是简单地移除循环）。

numba 和 cython 通常用于加速迭代解决方案，但我不会涉及这些。

---

如果d 将数组分成相等的部分，我们可以使用 reshaping 来计算部分：

In [228]: A.shape
Out[228]: (10, 2)
In [229]: f.shape
Out[229]: (10,)
In [230]: f2 = f.reshape(2,5)
In [231]: A2 = A.reshape(2,5,2)

In [233]: np.einsum('ij,ijk->ik',f2,A2)
Out[233]: 
array([[ 60,  70],
       [510, 545]])

matmul 运算符也可以使用，但需要对尺寸进行一些调整：

In [236]: (f2[:,None,:]@A2)[:,0,:]
Out[236]: 
array([[ 60,  70],
       [510, 545]])

---

如果d 将数组分成几个大小，我想我们可以对常见大小进行分组，并对每个组执行上述 reshape 和 einsum，但我还没有弄清楚细节：

In [238]: d = [0,2,5,7,10]
In [239]: np.array([f[i:j].dot(A[i:j,:]) for i,j in zip(d[:-1],d[1:])])
Out[239]: 
array([[  2,   3],
       [ 58,  67],
       [122, 133],
       [388, 412]])
In [240]: [f[i:j] for i,j in zip(d[:-1],d[1:])]
Out[240]: [array([0, 1]), array([2, 3, 4]), array([5, 6]), array([7, 8, 9])]

这里我们有 2 个组，一个长度为 2，另一个长度为 3。