Question

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

更确切地说，例如，我想遍历所有（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，但它不执行我想要的操作。

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])

Answer 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)])

但是，这确实要占用更多的内存，因为您要创建一个具有整个元素序列的新数组，而不是一次处理一个数组。

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

2 个答案: