Обработка углового случая одного массива Numpy для итерации

Question

У меня есть список массивов Numpy с одинаковой формой (но необязательно одинаковыми dtype), и я хочу перебирать элементы всех массивов одновременно. Например, если массивы:

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

Я хочу, чтобы итерация по [a, b] дала

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

Я обнаружил, что numpy.nditer делает почти , что мне нужно. Это работает:

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

Обратите внимание, что итератор выдает кортежи скаляров:

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

Однако в угловом случае списка, содержащего один массив, np.nditer напрямую возвращает элементы массива:

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

Мне нужно получить кортеж с одним элементом, потому что я распаковываю повторяющиеся значения в аргументах функции внутри цикла.

Как мне убедить np.nditer выдать одноэлементный кортеж при итерации по списку из одного массива?

Answer 1

Один обходной путь будет 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]

Обратите внимание, что это распаковывает массивы 0D nditer возвращает в надлежащие скаляры.Кроме того, он возвращает массивы, а не кортежи, но, поскольку вы просто хотите объединить их с функцией, это не должно иметь значения.

Answer 2

Это противоречит поведению nditer.

Может быть, вы всегда можете предоставить массив None, чтобы вы всегда получали кортеж? Вы должны обрабатывать None в функции, хотя.