遍历布尔 Numpy 数组的“True”条目

Question

我希望每个索引 i 都有一个循环，其中数组 X（布尔）是 True。

有没有比在np.nditer 中包装np.nonzero 更高效/pythonic 的方法如下？

for i in np.nditer(np.nonzero(X), flags=['zerosize_ok']):
    myfunction(Y[i],Z2[Z[i]])

这里的问题是它迭代了两次而不是一次，并且占用了内存（首先，np.nonzero 遍历 X 并将其存储到一个大数组中，然后 np.nditer 遍历该数组）。

是否有一个命令（可以说有点类似于np.nditer）用于直接有效地迭代布尔数组的True 条目，而不首先用np.nonzero 明确列出它们？ （遍历所有条目并使用 if 语句检查每个条目的效率可能低于 Numpy 提供的某些迭代器（如果存在）。）

Answer 1

人们投了反对票，因为循环遍历 numpy 数组的条目是一个很大的问题。我们使用 numpy 是因为它速度快，并且单独处理每个元素而不是一次对整个数组进行操作，因此您可以获得 python 级别的性能而不是 numpy/c 性能。

想要通过给一个包含真假值的数组来排除值是很常见的，这被称为掩码。你可以通过索引到真假数组来做到这一点。要在 numpy 中执行此操作，请使用索引。例如。你做np.array([1,2,3])[np.array([True,False,True])]。它会给你np.array([1, 3])。

所以基本上尝试以你可以做的方式安排事情

myfunction(Y[mask],Z2[Z[maks]]).

有几种技术可以做到这一点。一种方法是仅使用 numpy 函数来创建 myfunction，另一种方法是使用装饰器，如 numba.vectorize、numba.guvectorize 或 numba.njit 等等。

Answer 2

将numpy.vectorise 与布尔选择器一起用作索引怎么样？

np.vectorise 接受一个函数并返回该函数的向量化版本，该函数接受数组。然后，您可以使用选择器在数组或其子集上运行该函数。

在 numpy 中，您可以使用 index of an array 使用数字列表进行子选择。

numpy.where 函数返回与某个值或函数匹配的数组的索引。

综合起来：

import numpy as np
import random

random.seed(0) # For repeatability

def myfunction(y, z, z2):
    return y*z2[z]

Y = np.array(range(100))  # Array of any length
Z = np.array(range(len(Y)))  # Must be the same length as Y
Z2 = np.array(range(len(Y))) # Must be constrained to indexes of X, Y and Z
X = [random.choice([True, False]) for toss in range(len(Y))] # Throw some coins

# X is now an array of True/False
# print( "\n".join(f"{n}: {v}" for n, v in enumerate(X)))
print(np.where(X)) # Where gets the indexes of True items

# Vectorise our function
vfunc = np.vectorize(myfunction, excluded={2}) # We exclude {2} as it's the Z2 array

# Do the work
res = vfunc(Y[np.where(X)],Z[np.where(X)],Z2)

# Check our work
for outpos, inpos in enumerate(np.where(X)[0]):
    assert myfunction(Y[inpos], Z[inpos], Z2) == res[outpos], f"Mismatch at in={inpos}, out={outpos}"

# Print the results
print(res)