使用连接运算符进行广播

Question

numpy.broadcasting 让我们可以对形状不同的数组进行基本操作（比如加法、乘法等），但这需要满足一些条件。举个例子：

>>> a = np.array([0.0, 10.0, 20.0, 30.0])
>>> b = np.array([1.0, 2.0, 3.0])
>>> a[:, np.newaxis] + b
array([[ 1.,  2.,  3.],
       [11., 12., 13.],
       [21., 22., 23.],
       [31., 32., 33.]])

输出数组的每个元素都是输入数组中对应元素的和。这种操作可以称为“外部加法操作”。

有没有简单的方法可以用连接操作符来实现呢？换句话说，是否有简单的方法可以进行“外部连接操作”？也就是说，按照之前的例子，输出数组应该是：

array([[[ 0.,  1.],
        [ 0.,  2.],
        [ 0.,  3.]],

       [[10.,  1.],
        [10.,  2.],
        [10.,  3.]],

       [[20.,  1.],
        [20.,  2.],
        [20.,  3.]],

       [[30.,  1.],
        [30.,  2.],
        [30.,  3.]]])

经过几次尝试，我用 numpy.meshgrid 找到了这个解决方案，但我觉得它的效率不高，可能还有改进的空间：

>>> a = np.array([0.0, 10.0, 20.0, 30.0])
>>> b = np.array([1.0, 2.0, 3.0])
>>> c,d = np.meshgrid(a,b)
>>> np.concatenate((c.T[..., None], d.T[..., None]), axis=-1)
array([[[ 0.,  1.],
        [ 0.,  2.],
        [ 0.,  3.]],

       [[10.,  1.],
        [10.,  2.],
        [10.,  3.]],

       [[20.,  1.],
        [20.,  2.],
        [20.,  3.]],

       [[30.,  1.],
        [30.,  2.],
        [30.,  3.]]])

有没有更“符合 Python 风格”的方法来做到这一点？

Answer 1

我在最后添加了一个快速的广播赋值解决方案。

---

什么叫做更“pythonic”（或者说更“numpy-onic”）呢？

如果代码能运行并且得到想要的结果，这样就够了吗？在perl中，一行代码很酷，但在python中就不一定了。broadcasting（广播）在使用像+这样的运算符时是个有效的概念，但concatenate（连接）就不是一个运算符。

所以我们最终要比较的就是速度。不过所有的速度测试都需要考虑数组的大小。

无论如何，对于你的测试案例：

In [33]: %%timeit 
    ...: c,d = np.meshgrid(a,b)
    ...: np.concatenate((c.T[..., None], d.T[..., None]), axis=-1)

82.7 µs ± 337 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

broadcast_arrays的方法会生成两个数组，类似于meshgrid（实际上只是它们的转置）：

In [34]: np.broadcast_arrays(a[:, None], b)
Out[34]: 
[array([[ 0.,  0.,  0.],
        [10., 10., 10.],
        [20., 20., 20.],
        [30., 30., 30.]]),
 array([[1., 2., 3.],
        [1., 2., 3.],
        [1., 2., 3.],
        [1., 2., 3.]])]

速度稍微快一点，但并没有“快很多”：

In [35]: timeit np.stack(np.broadcast_arrays(a[:, None], b),axis=2)
77.4 µs ± 2.89 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

不过我们可以用repeat构造相同的broadcasted数组（这是一个不错的编译过的numpy方法）：

In [36]: (a[:,None].repeat(3,axis=1), b[None,:].repeat(4, axis=0))
Out[36]: 
(array([[ 0.,  0.,  0.],
        [10., 10., 10.],
        [20., 20., 20.],
        [30., 30., 30.]]),
 array([[1., 2., 3.],
        [1., 2., 3.],
        [1., 2., 3.],
        [1., 2., 3.]]))

结果令人惊讶，它更快：

In [37]: timeit np.stack((a[:,None].repeat(3,axis=1), b[None,:].repeat(4, axis=0)),axis=2)
29.7 µs ± 99.6 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

那么哪个更“pythonic”呢？在这个小例子中，repeat更快，尽管它的代码更长。我猜对于更大的数组，broadcast_arrays会表现得更好。

广播代码

有人发现并给我之前关于广播问题的回答点了反对票。

广播numpy数组的底层机制是什么

评论建议查看np.lib.stride_tricks中的源代码。

broadcast_array是python代码，并使用了_broadcast_to。而_broadcast_to又使用了np.nditer。众所周知，这在python接口版本中是比较慢的。

通过meshgrid和我的repeats生成的数组是复制的。但通过broadcast_arrays（Out[34]）生成的数组是视图，是通过调整a和b的shape和strides得到的。通常创建视图是很快的，因为它不需要复制base值。但我猜stride_tricks函数的设置代码是比较耗时的。

In [44]: _34[0].base
Out[44]: array([ 0., 10., 20., 30.])

In [45]: _34[1].base
Out[45]: array([1., 2., 3.])

In [46]: _34[0].strides, _34[1].strides
Out[46]: ((8, 0), (0, 8))

正如模块名stride_tricks所暗示的，使用这个broadcast_arrays函数是一个不错的小技巧，但并不是理解或使用numpy的必要条件，尤其不是在使用运算符时的正常广播。

广播赋值

这里有一个更好的版本。它不是使用连接，而是直接创建目标数组，并用a和b的值填充它：

In [75]: timeit res=np.empty((4,3,2)); res[:,:,0]=a[:,None]; res[:,:,1]=b
4.56 µs ± 14 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

这确实利用了broadcasting，可以通过一个错误的赋值来看到：

In [76]: res=np.empty((4,3,2)); res[:,:,0]=a
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[76], line 1
----> 1 res=np.empty((4,3,2)); res[:,:,0]=a

ValueError: could not broadcast input array from shape (4,) into shape (4,3)

broadcasting不仅仅在二元运算符中使用，它也在advanced indexing（高级索引）中使用，正如这里在给数组的切片赋值时所示。

在这里，一个(4,1)的a可以适应(4,3)的槽位，而(3,)的b也可以适应这个大小的槽位。