Simpler Way To Create A Matrix/list Of Indices?

I wonder what could be the easiest way to create a bi-dimensional array, that has for each row the indices to another multi-dimensional array. For example, let's say I have a cube

Solution 1:

You could use np.indices:

>>> a = np.random.random((4,4,4))
>>> np.indices(a.shape).reshape((a.ndim, -1)).T
array([[0, 0, 0],
       [0, 0, 1],
       [0, 0, 2],
       [0, 0, 3],
       [0, 1, 0],
       [0, 1, 1],
[...]
       [3, 3, 2],
       [3, 3, 3]])

There are also other utilities like np.ndindex, depending on your use case. (FWIW I don't think getting the coordinates in the form you're looking for is going to be as helpful as you might think it is, but YMMV.)

Solution 2:

I think

list(itertools.product(range(4),range(4),range(4)))

is more pythonic. .. (note you could use numpy.array instead of list if you were so inclined...)

Solution 3:

It would be nice if existed a function that just takes an arbitrary multi-dimensional array and returns it's index table.

If I understand your question, there is: indices:

i = np.indices(a.shape)

This doesn't give you the results in the shape you wanted:

>>> a = np.array([[1,2], [3,4], [5,6]])
>>> print(np.indices(a.shape))
[[[0 0]
  [1 1]
  [2 2]][[0 1]
  [0 1]
  [0 1]]]

… but you can flatten it and transpose it:

[[0 0][0 1][1 0][1 1][2 0][2 1]]

Solution 4:

Here's how to do what I think you actually want to do. From your comment:

I need to test some probability values that are stored in a 3-dimensional array. If they pass some test condition, then I will add them to a list of coordinates that will be visualized in a 3d scatter plot.

Let's say, for the same of simplicity, that the test is something simple, like "is positive". So, we just transform the array into a boolean array of "element is positive" for each element, which is just arr > 0, and then use nonzero to get the true indices of that boolean array:

>>> arr = np.array([[-1, 1], [2, -2], [-3, -3]])
>>> print(arr > 0)
[[False  True]
 [ True False]
 [False False]]
>>> print(np.nonzero(arr > 0))
(array([0, 1]), array([1, 0]))

Can't get much simpler than that.