广播一个numpy.arange成一个多维数组

Broadcasting a numpy.arange into a multi-dimensional array

假设我有一个 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

编译

你可以先reshape 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.]]])

如果您总是从 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.]]])
>>>