python - 迭代 ndarray 的某些(但不是全部)维度

Question

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

更字面地说，例如，我想迭代所有 (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 数组的某些行/列，即 z[i,j,k]。为了从特定维度选择所有元素，您可以使用:。例如，要迭代 3d 数组的第一个维度和最后一个维度:

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