将 numpy.arange 广播到多维数组中答案

【问题标题】：Broadcasting a numpy.arange into a multi-dimensional array将 numpy.arange 广播到多维数组中
【发布时间】：2017-09-15 03:14:00
【问题描述】：

假设我有一个 3 维零数组

Y = np.zeros((2,3,4))

我希望通过第二维向它广播一维数组

X = np.arange(3)

这样得到的多维数组是

Y = [[[ 0. 0. 0. 0.]
      [ 1. 1. 1. 1.]
      [ 2. 2. 2. 2.]]

     [[ 0. 0. 0. 0.]
      [ 1. 1. 1. 1.]
      [ 2. 2. 2. 2.]]]

在没有循环（甚至列表推导）的情况下实现这一目标的最 Pythonic 和最有效的方法是什么？我希望能够使用更大的尺寸来扩展它并使用 numba 模块运行代码，该模块不能使用 np.tile 或 np.repeat 编译

【问题讨论】：

标签： python arrays numpy multidimensional-array numba

【解决方案1】：

您可以先重塑X，使其可以广播到Y.shape，然后使用np.broadcast_to方法：

Y = np.zeros((2,3,4))  
np.broadcast_to(X[:,None], Y.shape)

#array([[[0, 0, 0, 0],
#        [1, 1, 1, 1],
#        [2, 2, 2, 2]],

#       [[0, 0, 0, 0],
#        [1, 1, 1, 1],
#        [2, 2, 2, 2]]])

另一种选择是，创建一个与Y 形状相同的新空数组，然后为其分配 X 值：

new_X = np.empty_like(Y)    
new_X[:] = X[:,None]

new_X
#array([[[ 0.,  0.,  0.,  0.],
#        [ 1.,  1.,  1.,  1.],
#        [ 2.,  2.,  2.,  2.]],

#       [[ 0.,  0.,  0.,  0.],
#        [ 1.,  1.,  1.,  1.],
#        [ 2.,  2.,  2.,  2.]]])

【讨论】：

啊，这是我在浏览 numpy 文档时错过的东西，就像一个魅力。唯一的问题是，它不是 numba 模块可接受的 numpy 函数之一。
我不熟悉numba。所以不能给你太多建议。您可以先尝试numpy 解决方案，如果仍然不够高效，请使用numba。

【解决方案2】：

如果您总是从 zeros 数组开始，只需添加 x 以便广播

>>> y + x[None, :, None]
array([[[ 0.,  0.,  0.,  0.],
        [ 1.,  1.,  1.,  1.],
        [ 2.,  2.,  2.,  2.]],

       [[ 0.,  0.,  0.,  0.],
        [ 1.,  1.,  1.,  1.],
        [ 2.,  2.,  2.,  2.]]])

General Broadcasting Rules

或者

>>> np.add(y, x[None, :, None], out = y)
array([[[ 0.,  0.,  0.,  0.],
        [ 1.,  1.,  1.,  1.],
        [ 2.,  2.,  2.,  2.]],

       [[ 0.,  0.,  0.,  0.],
        [ 1.,  1.,  1.,  1.],
        [ 2.,  2.,  2.,  2.]]])
>>> y
array([[[ 0.,  0.,  0.,  0.],
        [ 1.,  1.,  1.,  1.],
        [ 2.,  2.,  2.,  2.]],

       [[ 0.,  0.,  0.,  0.],
        [ 1.,  1.,  1.,  1.],
        [ 2.,  2.,  2.,  2.]]])
>>>

我不确定您是否可以使用作业进行广播，但似乎可以：

>>> y[:] = x[None, :, None]
>>> y
array([[[ 0.,  0.,  0.,  0.],
        [ 1.,  1.,  1.,  1.],
        [ 2.,  2.,  2.,  2.]],

       [[ 0.,  0.,  0.,  0.],
        [ 1.,  1.,  1.,  1.],
        [ 2.,  2.,  2.,  2.]]])
>>> y[...] = x[None, :, None]
>>> y
array([[[ 0.,  0.,  0.,  0.],
        [ 1.,  1.,  1.,  1.],
        [ 2.,  2.,  2.,  2.]],

       [[ 0.,  0.,  0.,  0.],
        [ 1.,  1.,  1.,  1.],
        [ 2.,  2.,  2.,  2.]]])
>>>

【讨论】：