N
N
nirvimel2015-11-29 20:49:00
Python
nirvimel, 2015-11-29 20:49:00

How to beautifully generate a 4-dimensional array containing in the 4th dimension the indexes of the first 3rd dimensions?

The function must return a 3+1 dimensional array, in the fourth dimension in the elements 0, 1, 2 there must be values ​​equal to the indices of the 1st, 2nd and 3rd dimensions, respectively, so that the equality is fulfilled data[x, y, z] == array([x, y, z]).
I wrote like this:

def array_of_indexes(x, y, z, dtype=int):
    array = np.empty((x, y, z, 3), dtype=dtype)
    array[:, :, :, 0] = np.reshape(np.arange(x), (x, 1, 1))
    array[:, :, :, 1] = np.reshape(np.arange(y), (y, 1))
    array[:, :, :, 2] = np.arange(z)
    return array

This works, but I don't like creating 5 temporary arrays (each reshape creates another) and three passes through the resulting array.
I think there should be a simpler and prettier solution. For example, based on meshgrid. numpy.meshgrid and numpy.dstack hint at such a solution:
def array_of_indexes_alt(x, y, z):
    return np.dstack(np.mgrid[0:x, 0:y, 0:z])

But it doesn't work the way I want.

Answer the question

In order to leave comments, you need to log in

1 answer(s)
A
Alexander Movchan, 2015-11-29
@Alexander1705

def array_of_insexes(x, y, z):
    return [
               [
                   [
                       (i, j, k) for k in range(z)
                   ] for j in range(y)
               ] for i in range(x)
           ]

Didn't find what you were looking for?

Ask your question

Ask a Question

731 491 924 answers to any question