迭代 ndarray 的一些（但不是全部）維度

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 ]

但我更喜歡一個迭代器，我可以將數組邊距傳遞給它進行迭代（在上面的例子中邊距=[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)])

然而，這確實需要更多的內存，因為您正在使用整個元素序列創建一個新數組，而不是一次處理一個。