python - 为 3D numpy 数组中的每个元素传递多个条件的更有效方法

Question

我有一个 3 维 numpy 数组，我正在检查它的多个条件。我正在检查每个元素以查看它们是否小于某个数字。如果每个 3d 元素都由 i 索引，其中 i=[0,1,2] 在我所说的 array3 中，如果一个元素的数量大于我设置的数量，可能给出一个 bool 数组 [False, True, True] 或 [False, False, True]，然后这个索引从 array3 中删除。

对于每个小于 20 的元素，我都有一个愚蠢的方法:

import numpy as np

wx = np.where( np.abs(array3[:,0]) <= 20.0 ) # x values less than 20
xarray3x = array3[:,0][wx]
yarray3x = array3[:,1][wx]
zarray3x = array3[:,2][wx]

wy = np.where( np.abs(yarray3x) <= 20.0 ) # y values less than 20
xarray3xy = xarray3x[wy]
yarray3xy = yarray3x[wy]
zarray3xy = zarray3x[wy]

wz = np.where( np.abs(zarray3xy) <= 20.0 ) # z values less than 20
xarray3xyz = xarray3xy[wz]
yarray3xyz = yarray3xy[wz]
zarray3xyz = zarray3xy[wz]

这行得通，但跟上我命名的变量可能会很烦人。所以现在我正在尝试写一些占用更少行数的东西(希望编译时间更少)。

我正在考虑为每个索引构建一个 for 循环，如下所示:

for i in range(3):
    w = np.where( abs(array3[:,i]).all() <= 20.  )
n_array = array3[w]

但我只会构建一个值，而不是多个值。

Answer 1

我认为在这个例子中使用 4D 数组是最简单的。在这种情况下，您可以检查向量中具有最低值的元素是否低于阈值。然后，使用 np.newaxis，您可以将掩码应用于整个向量并创建一个掩码数组。

import numpy as np

n = 4 # Size of the first three dimensions.   
array3 = 100.*np.random.rand(n, n, n, 3) # Random numbers between 0 and 100.

thres = 20.
m = np.empty(array3.shape, dtype=bool)
m[:,:,:,:] = (np.min(array3, axis=-1) < thres)[:,:,:,np.newaxis]

array3_masked = np.ma.masked_array(array3, mask=m)