python - 如何在 numpy 中推广逐元素矩阵运算

Question

假设我有两组由 3xn 数组 V 和 W 表示的 n 个向量和一组由 3x3xn 数组 Q 表示的 n 个矩阵。我如何向量化操作来给我

a) n 个向量的集合 Q[:,:,k]*V[:,k] for k in range(n)

b) n 个标量的集合 W[:,k].T*Q[:,:,k]*V[:,k] for k in range(n)

这里的乘法理解为矩阵乘法。无法理解我认为应该使用的 einsum 文档。任何澄清/解决方案将不胜感激。

Answer 1

a) the set of n vectors Q[:,:,k]*V[:,k] for k in range(n)

np.einsum('ijk,jk->ik', Q, v)

应该生成一个 (3,n) 数组。矩阵乘积求和超过j。

b) the set of n scalars W[:,k].T*Q[:,:,k]*V[:,k] for k in range(n)

np.einsum('ik,ijk,jk->k', W, Q, V)

我凭内存工作，并且我对您需要什么的最佳猜测。所以我的“ij”表达式可能需要调整。但请尝试一下，让我知道它是如何工作的。

---

测试

In [180]: V=W=np.arange(3*n).reshape(3,n)

In [181]: Q=np.arange(3*3*n).reshape(3,3,n)

In [182]: np.einsum('ijk,jk->ik',Q,V)
Out[182]: 
array([[ 80, 107, 140, 179],
       [224, 287, 356, 431],
       [368, 467, 572, 683]])

In [183]: np.einsum('ik,ijk,jk',W,Q,V)
Out[183]: 28788    # summation over k

In [184]: np.einsum('ik,ijk,jk->k',W,Q,V)
Out[184]: array([ 3840,  5745,  8136, 11067])

有时将 einsum 分成几个步骤会更快，因为它可以防止迭代空间变得太大。我不认为这里是这种情况，但这就是它的样子。

In [185]: np.einsum('jk,jk->k',np.einsum('ik,ijk->jk',W,Q),V)
Out[185]: array([ 3840,  5745,  8136, 11067])

并使用 Jaime 的评论:

In [186]: np.einsum('i...,ij...,j...',W,Q,V)
Out[186]: array([ 3840,  5745,  8136, 11067])

In [187]: np.einsum('ij...,j...->i...',Q,V)
Out[187]: 
array([[ 80, 107, 140, 179],
       [224, 287, 356, 431],
       [368, 467, 572, 683]])