Broadcasting a numpy.arange into a multi-dimensional array

Question

Suppose I have a 3 dimensional array of zeros

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

and I wished to broadcast a 1 dimensional array to it via the second dimension

X = np.arange(3)

such that the resulting multidimensional array was

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.]]]

What would be the most pythonic and efficient way of achieving this without loops (or even list comprehensions)? I wish to be able to scale this up with larger dimension sizes and run the code with the numba module, which does not compile with np.tile or np.repeat

akuiper · Accepted Answer · 2017-09-15 03:30:45Z

2

You can reshape X firstly so that it can be broadcasted to Y.shape and then use np.broadcast_to method:

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]]])

Another option is, create a new empty array with the same shape as Y and then assign X values to it:

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.]]])

edited Sep 15, 2017 at 3:30

answered Sep 15, 2017 at 3:19

akuiper

216k33 gold badges362 silver badges379 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

Golden_Sundown Over a year ago

Ah that's something I missed looking through the numpy documentation, works like a charm. Only issue is, its not one of the accepted numpy functions for the numba module.

akuiper Over a year ago

I am not familiar with numba. so can't give you much suggestion. You might try the numpy solution firstly, if it's still not efficient enough, then resort to numba.

wwii · Accepted Answer · 2017-09-15 04:15:42Z

If you are always starting with a zeros array, just add x so that it will broadcast

>>> 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

Or

>>> 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.]]])
>>>

I wasn't sure you could broadcast using assignment but it sems you can:

>>> 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.]]])
>>>

Collectives™ on Stack Overflow

Broadcasting a numpy.arange into a multi-dimensional array

2 Answers 2

2 Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

2 Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related