假设我有两个三维矩阵,就像这样(取自这个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中的等效操作,最好是矢量化选项,以避免必须在整个数组中编写一个double for循环。我似乎无法理解np.tensordot
或np.inner
应该做什么,但它们可能是选项。
答案 0 :(得分:4)
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
答案 1 :(得分:2)
这是一个解决方案:
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)
答案 2 :(得分:2)
使用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]])