How to specify depth of iterator in numpy?

Question

I have a multidimensional array in numpy (e.g. 4D) and I want to specify the depth of the numpy iterator and I am not sure how to do this please?

For example, let's say I have a 4D numpy array, and I want to get elements form the iterator only for the 2D level (so then each item would be also 2D). Is there a way to specify such depth in the iterator please?

I do really want to use an iterator instead of a double loop, and I want to use numpy instead of other tools such as dictionaries and pandas.

As such, I would like this code to output [1 2]'s instead of 1,2...

x = np.array(\
(\
\
(\
(np.array([1,2]), np.array([1,2])),\
(np.array([1,2]), np.array([1,2]))\
),\
\
(\
(np.array([1,2]), np.array([1,2])),\
(np.array([1,2]), np.array([1,2]))\
)\
\
)
, dtype = np.ndarray)

for i in np.nditer(x, flags = ["refs_ok"]):
     print i

Gives me:

1
2
1
2
1
2
1
2
1
2
1
2
1
2
1
2

Instead of:

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

Answer 1

np.ndindex does a nice job of iterating over specified dimensions.

Your x is a 4d object array dtype=ndarray becomes dtype=object. Though with tuples all the same size the elements are really just scalars, not arrays.

In [385]: x
Out[385]: 
array([[[[1, 2],
         [1, 2]],

        [[1, 2],
         [1, 2]]],


       [[[1, 2],
         [1, 2]],

        [[1, 2],
         [1, 2]]]], dtype=object)

In [386]: x.shape
Out[386]: (2, 2, 2, 2)

In any case, np.ndindex generates the indices that will iterate over an array of a given shape.

In [387]: for i,j in np.ndindex(x.shape[:2]):
    print(i,j)       
    print(x[i,j])
   .....:     
0 0
[[1 2]
 [1 2]]
0 1
[[1 2]
 [1 2]]
1 0
[[1 2]
 [1 2]]
1 1
[[1 2]
 [1 2]]

The key pieces of ndindex are as_strided used to generate a dummy matrix of the right size, and nditer with multi_index mode to generate the indexes.

earlier examples of this use:

https://stackoverflow.com/a/28727290/901925

Iterating over first d axes of numpy array

More on trying to create an array of arrays (not just a higher dimensional array of numbers):

Convert a numpy array to an array of numpy arrays

To make an x that really is an array of arrays, you need to do something like:

In [397]: x=np.zeros((2,2,2),dtype=object)
In [398]: for ijk in np.ndindex(x.shape):
             x[ijk] = np.array([1,2])


In [399]: x
Out[399]: 
array([[[array([1, 2]), array([1, 2])],
        [array([1, 2]), array([1, 2])]],

       [[array([1, 2]), array([1, 2])],
        [array([1, 2]), array([1, 2])]]], dtype=object)

---

Another option is to reshape the initial dimensions, so you can do a flat iteration on those:

for i in x.reshape(-1,2):
    print(i)

nditer (and by extension ndindex) is described as being efficient, but that applies more to its C/cython use. In pure Python code the iteration mechanism doesn't matter as much. Actions in the body of the iteration usually take more time. nditer is also best when you need coordinate iteration over multiple array, as in out[...] = a[...] * b[...]. It doesn't off much special with just iterating over one array.

http://docs.scipy.org/doc/numpy-dev/reference/arrays.nditer.html

is a good nditer tutorial. The cython part at the end is best part.