Numpy中dot(A,B,3)的等价实现

Question

假设我有两个三维矩阵，像这样（这个例子来自于这个matlab的链接 http://www.mathworks.com/help/matlab/ref/dot.html）：

A = cat(3,[1 1;1 1],[2 3;4 5],[6 7;8 9])
B = cat(3,[2 2;2 2],[10 11;12 13],[14 15; 16 17])

如果我想在第三个维度上进行成对的点积，我可以在matlab中这样做：

C = dot(A,B,3)

这样做会得到以下结果：

C =
  106   140
  178   220

那么在numpy中，有什么等效的操作，最好是使用向量化的方法，这样就不需要通过整个数组写双重循环。我似乎搞不懂 np.tensordot 或 np.inner 是干什么的，但它们可能是可行的选项。

Answer 1

这里有一个解决方案：

A = dstack([[[1, 1],[1 ,1]],[[2 ,3],[4, 5]],[[6, 7],[8, 9]]])
B = dstack([[[2, 2],[2, 2]],[[10, 11],[12, 13]],[[14, 15], [16, 17]]])

C = einsum('...k,...k',A,B)

简单来说，dstack 是在第三个维度上把数组拼接在一起，具体可以参考这个文档。然后，你可以使用 numpy 提供的强大工具 einsum 来进行计算，详细信息可以查看这个文档。

Answer 2

In [169]:

A = np.dstack([[[1, 1],[1 ,1]],[[2 ,3],[4, 5]],[[6, 7],[8, 9]]])
B = np.dstack([[[2, 2],[2, 2]],[[10, 11],[12, 13]],[[14, 15], [16, 17]]])
c=np.tensordot(A, B.T,1)
np.vstack([np.diag(c[:,i,i]) for i in range(A.shape[0])]).T
Out[169]:
array([[106, 140],
       [178, 220]])

In [170]:

%%timeit
c=np.tensordot(A, B.T,1)
np.vstack([np.diag(c[:,i,i]) for i in range(A.shape[0])]).T
10000 loops, best of 3: 95.2 µs per loop
In [171]:

%timeit np.einsum('i...,i...',a,b)
100000 loops, best of 3: 6.93 µs per loop
In [172]:

%timeit inner1d(A,B)
100000 loops, best of 3: 4.51 µs per loop

但令人惊讶的是，它是最慢的：

Answer 3

使用 np.einsum：

In [9]: B = np.array([[[2, 2],[2, 2]],[[10, 11],[12, 13]],[[14, 15],[16, 17]]])

In [10]: A = np.array([[[1, 1],[1, 1]],[[2, 3],[4, 5]],[[6, 7],[8, 9]]]) 

In [11]: np.einsum('i...,i...',A,B)
Out[11]:
array([[106, 140],
       [178, 220]])

或者这里还有一个有趣的例子：

In [37]: from numpy.core.umath_tests import inner1d

In [38]: inner1d(A,B)
Out[38]:
array([[106, 140],
       [178, 220]])

编辑：根据 @flebool 的评论，inner1d 对于形状为 (2,2,3) 和 (3,2,2) 的数组都适用：

In [41]: A = dstack([[[1, 1],[1 ,1]],[[2 ,3],[4, 5]],[[6, 7],[8, 9]]])

In [42]: B = dstack([[[2, 2],[2, 2]],[[10, 11],[12, 13]],[[14, 15], [16, 17]]])

In [43]: inner1d(A,B)
Out[43]:
array([[106, 140],
       [178, 220]])