Access all elements at given x, y position in 3-dimensional numpy array

Question

mat_a = np.random.random((5, 5))
mat_b = np.random.random((5, 5))
mat_c = np.random.random((5, 5))
bigmat = np.stack((mat_a, mat_b, mat_c)) # this is a 3, 5, 5 array

for (x, y, z), value in np.ndenumerate(bigmat):
    print (x, y, z)

In the example above, how can I loop so that I iterate only across the 5 x 5 array and at each position I get 3 values i.e. loop should run 25 times and each time, I get an array with 3 values (one from each of mat_a, mat_b and mat_c)

• EDIT: Please note that I need to be able to access elements by position later i.e. if bigmat is reshaped, there should be a way to access element based on specific y, z

Answer 1

If you don't actually need to stack the arrays, and only want to iterate over all three arrays, element-wise, at once, numpy.nditer works - I'm still fuzzy on all its parameters I don't know if it is any faster, test it on a subset.

a1 = np.arange(9).reshape(3,3) + 10
a2 = np.arange(9).reshape(3,3) + 20
a3 = np.arange(9).reshape(3,3) + 30

c = np.nditer((a1, a2, a3))
for thing in c:
    print(np.array(thing))

>>> 
[10 20 30]
[11 21 31]
[12 22 32]
[13 23 33]
[14 24 34]
[15 25 35]
[16 26 36]
[17 27 37]
[18 28 38]
>>>

Answer 2

There is a function that generates all indices for a given shape, ndindex.

for y,z in np.ndindex(bigmat.shape[1:]):
    print(y,z,bigmat[:,y,z])

0 0 [ 0 25 50]
0 1 [ 1 26 51]
0 2 [ 2 27 52]
0 3 [ 3 28 53]
0 4 [ 4 29 54]
1 0 [ 5 30 55]
1 1 [ 6 31 56]
...

For a simple case like this it isn't much easier than the double for range loop. Nor will it be faster; but you asked for an iteration.

Another iterator is itertools.product(range(5),range(5))

Timewise, product is pretty good:

In [181]: timeit [bigmat[:,y,z] for y,z in itertools.product(range(5),range(5
     ...: ))]
10000 loops, best of 3: 26.5 µs per loop

In [191]: timeit [bigmat[:,y,z] for (y,z),v in np.ndenumerate(bigmat[0,...])]
     ...: 
10000 loops, best of 3: 61.9 µs per loop

transposing and reshaping is the fastest way to get a list (or array) of the triplets - but it does not give the indices as well:

In [198]: timeit list(bigmat.transpose(1,2,0).reshape(-1,3))
100000 loops, best of 3: 15.1 µs per loop

But the same operation gets the indices from np.mgrid (or np.meshgrid):

np.mgrid[0:5,0:5].transpose(1,2,0).reshape(-1,2)

(though this is surprisingly slow)

Answer 3

Simon's answer is fine. If you reshape things properly you can get them all in a nice array without any looping.

In [33]: bigmat
Out[33]: 
array([[[ 0.51701737,  0.90723012,  0.42534365,  0.3087416 ,  0.44315561],
        [ 0.3902181 ,  0.59261932,  0.21231607,  0.61440961,  0.24910501],
        [ 0.63911556,  0.16333704,  0.62123781,  0.6298554 ,  0.29012245],
        [ 0.95260313,  0.86813746,  0.26722519,  0.14738102,  0.60523372],
        [ 0.33189713,  0.6494197 ,  0.30269686,  0.47312059,  0.84690451]],

       [[ 0.95974972,  0.09659425,  0.06765838,  0.36025411,  0.91492751],
        [ 0.92421874,  0.31670119,  0.99623178,  0.30394588,  0.30970197],
        [ 0.53590091,  0.04273708,  0.97876218,  0.09686119,  0.78394054],
        [ 0.5463358 ,  0.29239676,  0.6284822 ,  0.96649507,  0.05261606],
        [ 0.91733464,  0.77312656,  0.45962704,  0.06446105,  0.58643379]],

       [[ 0.75161903,  0.43286354,  0.09633492,  0.52275049,  0.40827006],
        [ 0.51816158,  0.05330978,  0.49134325,  0.73652136,  0.14437844],
        [ 0.83833791,  0.2072704 ,  0.18345275,  0.57282927,  0.7218022 ],
        [ 0.56180415,  0.85591746,  0.35482315,  0.94562085,  0.92706479],
        [ 0.2994697 ,  0.99724253,  0.66386017,  0.0121033 ,  0.43448805]]])

Reshaping things...

new_bigmat =  bigmat.T.reshape([25,3])

In [36]: new_bigmat
Out[36]: 
array([[ 0.51701737,  0.95974972,  0.75161903],
       [ 0.3902181 ,  0.92421874,  0.51816158],
       [ 0.63911556,  0.53590091,  0.83833791],
       [ 0.95260313,  0.5463358 ,  0.56180415],
       [ 0.33189713,  0.91733464,  0.2994697 ],
       [ 0.90723012,  0.09659425,  0.43286354],
       [ 0.59261932,  0.31670119,  0.05330978],
       [ 0.16333704,  0.04273708,  0.2072704 ],
       [ 0.86813746,  0.29239676,  0.85591746],
       [ 0.6494197 ,  0.77312656,  0.99724253],
       [ 0.42534365,  0.06765838,  0.09633492],
       [ 0.21231607,  0.99623178,  0.49134325],
       [ 0.62123781,  0.97876218,  0.18345275],
       [ 0.26722519,  0.6284822 ,  0.35482315],
       [ 0.30269686,  0.45962704,  0.66386017],
       [ 0.3087416 ,  0.36025411,  0.52275049],
       [ 0.61440961,  0.30394588,  0.73652136],
       [ 0.6298554 ,  0.09686119,  0.57282927],
       [ 0.14738102,  0.96649507,  0.94562085],
       [ 0.47312059,  0.06446105,  0.0121033 ],
       [ 0.44315561,  0.91492751,  0.40827006],
       [ 0.24910501,  0.30970197,  0.14437844],
       [ 0.29012245,  0.78394054,  0.7218022 ],
       [ 0.60523372,  0.05261606,  0.92706479],
       [ 0.84690451,  0.58643379,  0.43448805]])

Edit: To keep track of indices, you might try the following (open to other ideas here). Each row in xy_index gives your x,y values respectively for the corresponding row in the new_bigmat array. This answer doesn't require any loops. If looping is acceptable you can borrow Simon's suggestion in the comments or np.ndindex as suggested in hpaulj's answer.

row_index, col_index = np.meshgrid(range(5),range(5))
xy_index = np.array([row_index.flatten(), col_index.flatten()]).T

In [48]: xy_index
Out[48]: 
array([[0, 0],
       [1, 0],
       [2, 0],
       [3, 0],
       [4, 0],
       [0, 1],
       [1, 1],
       [2, 1],
       [3, 1],
       [4, 1],
       [0, 2],
       [1, 2],
       [2, 2],
       [3, 2],
       [4, 2],
       [0, 3],
       [1, 3],
       [2, 3],
       [3, 3],
       [4, 3],
       [0, 4],
       [1, 4],
       [2, 4],
       [3, 4],
       [4, 4]])

Answer 4

The required result can be obtained by slicing, e.g.:

for x in range(5):
    for y in range(5):
        print (bigmat[:,x,y])