перебрать некоторые (но не все) измерения ndarray

Question

У меня есть трехмерный ndarray в python, и я хотел бы поэтапно выполнить его по двум из трех полей.

Если говорить буквально, я, например, хотел бы перебратьвсе (x, y) пары, но хранят данные z вместе в виде массива.

Как псевдокод, выражение, за которым я, в конечном итоге, буду, будет примерно таким:

[ f(z) for z in all_xy_pairs(the_ndarray) if g(z) == True ]

Я подумал об использовании 'reshape' следующим образом

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 ]

, но я быпредпочитаю итератор, которому я мог бы передать поля массива для перебора (в приведенном выше примере поля = [0,1]. В псевдокоде приведенный выше пример станет

[ z for z in iterate_over_margins(ii, margins=[0,1]) if z[1]>10 ]

, прежде чем я начну программироватьэто сам, разве нет такого итератора в numpy или в связанном пакете? Я проверил nditer, но он не делает то, что мне нужно.

Answer 1

Это отвечает на несколько иной вопрос, но, как вы наверняка знаете, 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)])

Это, однако, требует еще немногопамяти, так как вы создаете новый массив со всей последовательностью элементов вместо того, чтобы обрабатывать по одному.

Answer 2

Вы можете выбрать определенные строки / столбцы массива с нумерацией путем индексации вдоль этих столбцов, например, z[i,j,k].Чтобы выбрать все элементы из определенного измерения, вы можете использовать :.Например, чтобы перебрать первое и последнее измерения трехмерного массива:

for i in range(z.shape[0]):
    for j in range(z.shape[2]):
        print(z[i,:,j])