numpy - Python numpy 矩阵乘法与一个对角矩阵

Question

我有两个数组 A (4000,4000)，其中只有对角线填充了数据，而 B (4000,5) 填充了数据。有没有比 numpy.dot(a,b) 函数更快的方法来乘(点)这些数组？

到目前为止，我发现 (A * B.T).T 应该更快(其中 A 是一维的 (4000,)，填充了对角线元素)，但结果是大致慢了一倍。

在 A 是对角线数组的情况下，有没有更快的方法来计算 B.dot(A)？

Answer 1

您可以简单地提取对角线元素，然后执行广播元素乘法。

因此，B*A 的替换将是 -

np.multiply(np.diag(B)[:,None], A)

对于 A.T*B -

np.multiply(A.T,np.diag(B))

运行时测试-

In [273]: # Setup
     ...: M,N = 4000,5
     ...: A = np.random.randint(0,9,(M,N)).astype(float)
     ...: B = np.zeros((M,M),dtype=float)
     ...: np.fill_diagonal(B, np.random.randint(11,99,(M)))
     ...: A = np.matrix(A)
     ...: B = np.matrix(B)
     ...: 

In [274]: np.allclose(B*A, np.multiply(np.diag(B)[:,None], A))
Out[274]: True

In [275]: %timeit B*A
10 loops, best of 3: 32.1 ms per loop

In [276]: %timeit np.multiply(np.diag(B)[:,None], A)
10000 loops, best of 3: 33 µs per loop

In [282]: np.allclose(A.T*B, np.multiply(A.T,np.diag(B)))
Out[282]: True

In [283]: %timeit A.T*B
10 loops, best of 3: 24.1 ms per loop

In [284]: %timeit np.multiply(A.T,np.diag(B))
10000 loops, best of 3: 36.2 µs per loop