Numpy Array Broadcasting具有不同的尺寸

时间:2013-03-08 15:01:17

标签: python numpy

我对numpy的广播规则感到有些困惑。假设您想要执行更高维数组的轴方向标量积,以将数组维度减少一(基本上沿一个轴执行加权求和):

from numpy import *

A = ones((3,3,2))
v = array([1,2])

B = zeros((3,3))

# V01: this works
B[0,0] = v.dot(A[0,0])

# V02: this works
B[:,:] = v[0]*A[:,:,0] + v[1]*A[:,:,1] 

# V03: this doesn't
B[:,:] = v.dot(A[:,:]) 

为什么V03不起作用?

干杯

3 个答案:

答案 0 :(得分:4)

np.dot(a, b)运营over the last axis of a and the second-to-last of b。因此,对于您问题中的特定情况,您可以随时使用:

>>> a.dot(v)
array([[ 3.,  3.,  3.],
       [ 3.,  3.,  3.],
       [ 3.,  3.,  3.]])

如果您想保留v.dot(a)订单,则需要将轴置于适当的位置,这可以通过np.rollaxis轻松实现:

>>> v.dot(np.rollaxis(a, 2, 1))
array([[ 3.,  3.,  3.],
       [ 3.,  3.,  3.],
       [ 3.,  3.,  3.]])

我不太喜欢np.dot,除非是明显的矩阵或向量乘法,因为在使用可选的out参数时对输出dtype非常严格。 Joe Kington已经提到了它,但是如果你要做这类事情,那就习惯np.einsum:一旦你掌握了语法,它就会减少你担心的时间。把事情重塑到最低限度:

>>> a = np.ones((3, 3, 2))
>>> np.einsum('i, jki', v, a)
array([[ 3.,  3.,  3.],
       [ 3.,  3.,  3.],
       [ 3.,  3.,  3.]])

在这种情况下并不是太相关,但它的速度也非常快:

In [4]: %timeit a.dot(v)
100000 loops, best of 3: 2.43 us per loop

In [5]: %timeit v.dot(np.rollaxis(a, 2, 1))
100000 loops, best of 3: 4.49 us per loop

In [7]: %timeit np.tensordot(v, a, axes=(0, 2))
100000 loops, best of 3: 14.9 us per loop

In [8]: %timeit np.einsum('i, jki', v, a)
100000 loops, best of 3: 2.91 us per loop

答案 1 :(得分:3)

在这种特殊情况下,您也可以使用tensordot

import numpy as np

A = np.ones((3,3,2))
v = np.array([1,2])

print np.tensordot(v, A, axes=(0, 2))

这会产生:

array([[ 3.,  3.,  3.],
       [ 3.,  3.,  3.],
       [ 3.,  3.,  3.]])

axes=(0,2)表示tensordot应该在v中的第一个轴和A中的第三个轴上求和。 (另请查看einsum,它更灵活,但如果您不习惯这种表示法则更难理解。)

如果考虑速度,tensordot要比使用apply_along_axes小数组要快得多。

In [14]: A = np.ones((3,3,2))

In [15]: v = np.array([1,2])

In [16]: %timeit np.tensordot(v, A, axes=(0, 2))
10000 loops, best of 3: 21.6 us per loop

In [17]: %timeit np.apply_along_axis(v.dot, 2, A)
1000 loops, best of 3: 258 us per loop

(由于持续的开销,大型数组的差异不太明显,但tensordot始终更快。)

答案 2 :(得分:2)

您可以使用numpy.apply_along_axis()

In [35]: np.apply_along_axis(v.dot, 2, A)
Out[35]: 
array([[ 3.,  3.,  3.],
       [ 3.,  3.,  3.],
       [ 3.,  3.,  3.]])

我认为V03不起作用的原因是它与:

没有什么不同
B[:,:] = v.dot(A) 

即。它尝试沿A的最外轴计算点积。