python - Create an array where each element stores its indices

Question

I want to create a 2d numpy array where every element is a tuple of its indices.

Example (4x5):

array([[[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]]])

I would create an python list with the following list comprehension:

[[(y,x) for x in range(width)] for y in range(height)]

Is there a faster way to achieve the same, maybe with numpy methods?

Answer 1

Do you do this because you need it or just for sport? In the former case:

np.moveaxis(np.indices((4,5)), 0, -1)

np.indices does precisely what its name suggests. Only it arranges axes differently to you. So we move them with moveaxis

As @Eric points out one attractive feature of this method is that it works unmodified at arbitrary number of dimensions:

dims = tuple(np.multiply.reduceat(np.zeros(16,int)+2, np.r_[0, np.sort(np.random.choice(16, np.random.randint(10)))]))
# len(dims) == ?
np.moveaxis(np.indices(dims), 0, -1) # works