Перебирать некоторые (но не все) измерения ndarray
У меня есть трехмерный ndarray в python, и я хотел бы поэтапно перебирать его по двум из трех полей.
Говоря более буквально, я, например, хотел бы перебрать все пары (x, y), но сохранить данные z вместе как массив.
В качестве псевдокода выражение, которое я в конечном итоге ищу, будет примерно таким:
[ f(z) for z in all_xy_pairs(the_ndarray) if g(z) == True ]
Я подумывал об использовании "изменения формы" следующим образом
import numpy as np
# silly example
ii=np.arange(0,3*9,1).reshape(3,3,3)
[ z for z in ii.reshape(9,-1) if z[1]>10 ]
но я бы предпочел итератор, которому я мог бы передать поля массива для итерации (в приведенном выше примере margins = [0,1]. В псевдокоде приведенный выше пример станет
[ z for z in iterate_over_margins(ii, margins=[0,1]) if z[1]>10 ]
Прежде чем я начну программировать это сам, разве нет такого итератора в numpy или связанном пакете? Я проверил nditer, но он не делает то, что мне нужно.
Спасибо, конечно работает. Я бы принял это как ответ.
Ответы 2
Вы можете выбрать определенные строки / столбцы массива numpy путем индексации по этим столбцам, то есть z[i,j,k]. Чтобы выбрать элементы все из определенного измерения, вы можете использовать :. Например, чтобы перебрать первое и последнее измерения трехмерного массива:
for i in range(z.shape[0]):
for j in range(z.shape[2]):
print(z[i,:,j])
Это дает ответ на несколько другой вопрос, но, как вы наверняка знаете, NumPy обычно значительно выигрывает от использования векторизованных операций, поэтому, если ваши f и g могут быть векторизованы, вы также можете рассмотреть возможность работы с массивом, содержащим все повторяющиеся элементы в последовательности. Вы можете сделать это, просто изменив форму:
import numpy as np
# "Unrolls" an array along the given axes
def unroll_axis(a, axis):
a = np.asarray(a)
# This so it works with a single dimension or a sequence of them
axis = np.atleast_1d(axis)
# Put unrolled axes at the beginning
a = np.moveaxis(a, axis, range(len(axis)))
# Unroll
return a.reshape((-1,) + a.shape[len(axis):])
# Example
a = np.arange(27).reshape((3, 3, 3))
print(unroll_axis(a, (0, 2)))
# [[ 0 3 6]
# [ 1 4 7]
# [ 2 5 8]
# [ 9 12 15]
# [10 13 16]
# [11 14 17]
# [18 21 24]
# [19 22 25]
# [20 23 26]]
Итак, если g и f векторизованы, вы можете просто сделать
the_array_unrolled = unroll_axis(the_array, (0, 2))
result = f(the_array_unrolled[g(the_array_unrolled)])
Однако это требует немного больше памяти, поскольку вы создаете новый массив со всей последовательностью элементов, а не обрабатываете их по одному.
Что не так с
z[i,:,j] for i in range(z.shape[0]) for j in range(z.shape[2])?