使用np.einsum代替for循环

Question

我想做如下计算，我用随机数组来演示：

a = np.random.randint(10, size=(100,3))
b = np.random.randint(10, size=(3,2))

result = np.zeros(100)

for i in range(100):
    result[i] = a[i] @ b @ b.T @ a[i].T

为了加快计算速度，我想过用爱因斯坦求和去掉for循环。所以我尝试了以下，使用相同的向量：

result = np.einsum('ij,jk,jk,ij->i', a, b, b, a)

我将 'i' 放在 einsum 的右侧，因为结果向量显示的大小正确。但是，结果略有不同。

我的问题可以用 einsum 解决吗？

弗朗茨

Answer 1

在一个 einsum 中，它将是 -

np.einsum('ij,jl,kl,ik->i',a,b,b,a)

引入 matrix-multiplication 和 np.dot -

np.einsum('ij,jk,ik->i',a,b.dot(b.T),a)

或更多 -

np.einsum('ij,ij->i',a.dot(b.dot(b.T)),a)

与np.matmul/@-operator in Python 3.x，它转换为-

((a@(b@b.T))[:,None,:] @ a[:,:,None])[:,0,0]

使用np.einsum代替for循环

Use np.einsum to replace for loop

python

numpy

numpy-einsum