处理一个Numpy数组的拐角案例以进行迭代

Question

I have a list of Numpy arrays with the same shape (but not necessarily the same dtype ) and I want to iterate over the elements of all the arrays at the same time. For instance, if the arrays are:

>>> a = np.array([[1,2,3], [4,5,6]])
>>> b = np.array([['one','two','three'],['four','five','six']])

I want the iteration over [a, b] to yield

[(1, 'one'), (2, 'two'), (3, 'three'), (4, 'four'), (5, 'five'), (6, 'six')]

I have found that numpy.nditer does almost what I need. This works:

>>> for x, y in np.nditer([a, b]):
...     print('{} {}'.format(x, y))
1 one
2 two
3 three
4 four
5 five
6 six

Note that the iterator yields tuples of scalars:

>>> next(np.nditer([a, b]))
(array(1), array('one', dtype='<U5'))

However, in the corner case of a list containing one array, np.nditer directly yields the array elements:

>>> next(np.nditer([a]))
array(1)

I need it to yield a tuple with one element because I am unpacking the iterated values in the arguments of a function within the loop.

How do I convince np.nditer to yield a one-element tuple when iterating over a list of one array?

Answer 1

One workaround would be np.atleast_1D :

a = sum(np.ogrid[2:4, 3:5])
b = 2*a
for z in map(np.atleast_1d, np.nditer([a, b])):
    print(np.arange(*z))

#[5 6 7 8 9]
#[ 6  7  8  9 10 11]
#[ 6  7  8  9 10 11]
#[ 7  8  9 10 11 12 13]

for z in map(np.atleast_1d, np.nditer([a])):
    print(np.arange(*z))

#[0 1 2 3 4]
#[0 1 2 3 4 5]
#[0 1 2 3 4 5]
#[0 1 2 3 4 5 6]

Note that this unpacks the 0D arrays nditer returns into proper scalars. Also, it yields arrays, not tuples but as you just want to splat them to a function it shouldn't matter.

Answer 2

That goes against the behaviour of nditer .

Maybe you can always supply an array of None so that you always get a tuple? You have to handle the None in the function though.