Python – Initialize a numpy array of a specific shape with a specific value

Hello numpy beginner:

I’m trying to create an array of shape NxWxHx2 and initialize it with the corresponding index value. In this case, W=H always.

For

example: For an array of shape Nx5x5x2, if I write it down on paper, it should be:

N times or less

(0,0)(0,1)(0,2)(0,3)(0,4)

(1,0) (1,1) (1,2) (1,3) (1,4)

(2,0) (2,1) (2,2) (2,3) (2,4)

(3,0) (3,1) (3,2) (3,3) (3,4)

(4,0) (4,1) (4,2) (4,3) (4,4)

I looked at “arange” fkt as well as expanding the array with “newaxis” and couldn’t get the desired result.

Sorry for the bad formatting.

Thanks for your help!

Edit: I thought of something like this, but it’s not good.
For arrays of shape 1x3x3x2

t = np.empty([1,3,3,2])
for n in range(1):
    for i in range(3):
        for p in range(3):
            for r in range(2):
                if r == 0:
                    t[n,i,p,r]=i
                else:
                    t[n,i,p,r]=p    

One way is to allocate an empty array

>> out = np.empty((N, 5, 5, 2), dtype=int)

Then use broadcasting, for example

>>> out[...] = np.argwhere(np.ones((5, 5), dtype=np.int8)).reshape(5, 5, 2)

or

>>> out[...] = np.moveaxis(np.indices((5, 5)), 0, 2)

or

>>> out[..., 0] = np.arange(5)[None, :, None]
>>> out[..., 1] = np.arange(5)[None, None, :]