numpy - einsum 表示法:矩阵堆栈与向量堆栈的点积

numpy - einsum notation: dot product of a stack of matrices with stack of vectors

我想将 m* m 矩阵的 n 维堆栈乘以向量的 n 维堆栈(长度 m),以便生成的 m*n 数组包含矩阵点积的结果和第 n 个条目中的向量:

vec1=np.array([0,0.5,1,0.5]); vec2=np.array([2,0.5,1,0.5])
vec=np.transpose(n.stack((vec1,vec2)))
mat = np.moveaxis(n.array([[[0,1,2,3],[0,1,2,3],[0,1,2,3],[0,1,2,3]],[[-1,2.,0,1.],[0,0,-1,2.],[0,1,-1,2.],[1,0.1,1,1]]]),0,2)
outvec=np.zeros((4,2))
for i in range(2):
    outvec[:,i]=np.dot(mat[:,:,i],vec[:,i])

受此 post 的启发,我尝试了 einsum 中索引组合的所有不同扰动,并发现

np.einsum('ijk,jk->ik',mat,vec)

给出了正确的结果。

不幸的是,我真的不明白这一点——我假设我在 'ijk,jk' 部分重复条目 k 意味着我对 k 进行乘法和求和。我已经尝试阅读文档 https://docs.scipy.org/doc/numpy-1.15.1/reference/generated/numpy.einsum.html,但我还是不明白。

(包括我之前的尝试,

 np.einsum('ijk,il->ik', mat, vec)

我什至不确定这是什么意思。当我删除索引 l 时会发生什么?)

提前致谢!

In [321]: vec1=np.array([0,0.5,1,0.5]); vec2=np.array([2,0.5,1,0.5])
     ...: vec=np.transpose(np.stack((vec1,vec2)))
In [322]: vec1.shape
Out[322]: (4,)
In [323]: vec.shape
Out[323]: (4, 2)

stack 函数的一个好处是我们可以指定一个轴,跳过转置:

In [324]: np.stack((vec1,vec2), axis=1).shape
Out[324]: (4, 2)

为什么要混用 np.n.NameError: name 'n' is not defined。那种事差点把我赶走。

In [326]: mat = np.moveaxis(np.array([[[0,1,2,3],[0,1,2,3],[0,1,2,3],[0,1,2,3]],[[-1,2.,0
     ...: ,1.],[0,0,-1,2.],[0,1,-1,2.],[1,0.1,1,1]]]),0,2)
In [327]: mat.shape
Out[327]: (4, 4, 2)

In [328]: outvec=np.zeros((4,2))
     ...: for i in range(2):
     ...:     outvec[:,i]=np.dot(mat[:,:,i],vec[:,i])
     ...:     
In [329]: outvec
Out[329]: 
array([[ 4.  , -0.5 ],
       [ 4.  ,  0.  ],
       [ 4.  ,  0.5 ],
       [ 4.  ,  3.55]])

In [330]: # (4,4,2) (4,2)   'kji,ji->ki'

从您的循环中,i 轴(尺寸 2)的位置很清楚 - 在所有 3 个数组中排在最后。这为 vec 留下了一个轴,我们称之为 j。它与最后一个配对(在 mati 旁边)。 kmat 转移到 outvec

In [331]: np.einsum('kji,ji->ki', mat, vec)
Out[331]: 
array([[ 4.  , -0.5 ],
       [ 4.  ,  0.  ],
       [ 4.  ,  0.5 ],
       [ 4.  ,  3.55]])

通常 einsum 字符串会自行写入。例如,如果 mat 被描述为 (m,n,k),而 vec 被描述为 (n,k),则结果为 (m,k)

在这种情况下,只有 j 维度被求和 - 它出现在左侧,但出现在右侧。最后一个维度,i 在我的符号中,没有求和,因为 if 出现在两边,就像它在你的迭代中一样。我认为那是 'going-along-for-the-ride'。它不是 dot 产品的积极组成部分。

实际上,您是在最后一个维度上堆叠,大小为 2。通常我们堆叠在第一个,但你将两者调换放在最后一个。


您的 'failed' 尝试运行,并且可以重现为:

In [332]: np.einsum('ijk,il->ik', mat, vec)
Out[332]: 
array([[12. ,  4. ],
       [ 6. ,  1. ],
       [12. ,  4. ],
       [ 6. ,  3.1]])
In [333]: mat.sum(axis=1)*vec.sum(axis=1)[:,None]
Out[333]: 
array([[12. ,  4. ],
       [ 6. ,  1. ],
       [12. ,  4. ],
       [ 6. ,  3.1]])

jl 维度没有出现在右边,所以它们被求和了。它们可以在乘法之前求和,因为它们每个只出现在一个项中。我添加了 None 以启用广播(将 iki 相乘)。

np.einsum('ik,i->ik', mat.sum(axis=1), vec.sum(axis=1))

如果您在第一个上堆叠,并为 vec (2,4,1) 添加维度,它将 matmul 与 (2,4,4) 垫。 mat @ vec[...,None].

In [337]: m1 = mat.transpose(2,0,1)
In [338]: m1@v1[...,None]
Out[338]: 
array([[[ 4.  ],
        [ 4.  ],
        [ 4.  ],
        [ 4.  ]],

       [[-0.5 ],
        [ 0.  ],
        [ 0.5 ],
        [ 3.55]]])
In [339]: _.shape
Out[339]: (2, 4, 1)

继续阅读 Einstein summation notation

基本上,规则是:

没有 ->

  • 输入中重复的任何字母表示要相乘和求和的轴
  • 输入中未重复的任何字母都包含在输出中

有了->

  • 输入中重复的任何字母表示要乘以的轴
  • 任何不在输出中的字母表示要求和的轴

因此,例如,矩阵 AB 具有相同的形状:

np.einsum('ij, ij',       A, B)  # is A ddot B,                returns 0d scalar
np.einsum('ij, jk',       A, B)  # is A dot  B,                returns 2d tensor
np.einsum('ij, kl',       A, B)  # is outer(A, B),             returns 4d tensor
np.einsum('ji, jk, kl',   A, B)  # is A.T @ B @ A,             returns 2d tensor
np.einsum('ij, ij -> ij', A, B)  # is A * B,                   returns 2d tensor
np.einsum('ij, ij -> i' , A, A)  # is norm(A, axis = 1),       returns 1d tensor
np.einsum('ii'             , A)  # is tr(A),                   returns 0d scalar

einsum 很简单(当您玩过一段时间的索引排列时,就是...)。

让我们来处理一些简单的事情,2×2 矩阵的三重堆栈和 2×, 数组的三重堆栈

import numpy as np

a = np.arange(3*2*2).reshape((3,2,2))
b = np.arange(3*2).reshape((3,2))

我们需要知道我们将使用 einsum

计算什么
In [101]: for i in range(3): 
     ...:     print(a[i]@b[i])                                                                            
[1 3]
[23 33]
[77 95]

我们做了什么?我们有一个索引 i,当我们在一个堆叠矩阵和一个堆叠向量(均由 i 索引)之间执行点积时,它是固定的,并且单独的输出行意味着对堆叠矩阵的最后一个索引和堆叠向量的唯一索引。

这很容易在 einsum 指令中编码

  • 我们想要相同的 i 索引来指定矩阵、向量以及输出,
  • 我们想减少最后一个矩阵索引和剩余的向量索引,比如说k
  • 我们希望输出中的列数与每个堆叠矩阵中的行数一样多,比如 j

因此

In [102]: np.einsum('ijk,ik->ij', a, b)                                                                   
Out[102]: 
array([[ 1,  3],
       [23, 33],
       [77, 95]])

我希望我关于如何正确使用指令的讨论是清楚、正确和有用的。