遍历 ndarray 的某些（但不是全部）维度答案

【问题标题】：iterate over some (but not all) dimensions of a ndarray遍历 ndarray 的某些（但不是全部）维度
【发布时间】：2019-05-07 06:46:59
【问题描述】：

我在 python 中有一个 3D 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 ]

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

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

在我开始自己编程之前，numpy 或相关包中没有这样的迭代器吗？我检查了nditer，但它并没有达到我想要的效果。

【问题讨论】：

z[i,:,j] for i in range(z.shape[0]) for j in range(z.shape[2]) 有什么问题？
谢谢，这当然有效。我会接受这个作为答案。

标签： python numpy iterator

【解决方案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])

【讨论】：

【解决方案2】：

这回答了一个稍微不同的问题，但是，您肯定知道，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)])

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

【讨论】：