迭代ndarray的某些（但不是所有）维度

Question

我在python中有一个三维的ndarray，并且希望以元素方式沿着三个边距中的两个对其进行迭代。在

更确切地说，例如，我希望迭代所有（x，y）对，但将z数据作为一个数组保存在一起。在

作为伪代码，我最终想要的表达式应该是这样的

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

我考虑使用“重塑”如下

^{pr2}$

但是我更喜欢一个迭代器，我可以将数组边距传递给它进行迭代（在上面的例子中margins=[0,1]）。在伪代码中，上面的示例将变成

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

在我自己开始编程之前，numpy或相关的包中难道没有这样的迭代器吗？我检查了nditer，但它不能满足我的要求。在

Answer 1

这回答了一个稍有不同的问题，但是，正如您肯定知道的那样，NumPy通常从使用向量化操作中受益匪浅，因此，如果您的f和{}可以被矢量化，您还可以考虑在包含所有按顺序迭代的元素的数组上操作。你可以做一些重塑：

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是向量化的，你可以这样做

^{pr2}$

但是，这确实需要更多的内存，因为您正在使用整个元素序列生成一个新数组，而不是一次处理一个元素。在

Answer 2

您可以通过沿着这些列建立索引来选择numpy数组的某些行/列，即z[i,j,k]。为了从特定维度选择所有元素，可以使用:。例如，要迭代三维阵列的第一个和最后一个维度，请执行以下操作：

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