在Python中，使用2D掩码从（x，y）字段有效地选择元素

我有一个很大的二维位置数据字段，以两个数组形式给出

x

和

y

，其中

len（x）=len（y）

。我想返回索引数组

idx\u masked

，其中

（x[idx\u masked]，y[idx\u masked]）

被一个名为

mask

的nxn

int

数组屏蔽。也就是说，

mask[x[idx\u masked]，y[idx\u masked]==1

。

mask

数组仅由

0

s和

1

s组成

我提出了以下解决方案，但它（特别是下面的最后一行）非常慢，因为我有N x N=5000 x 5000，重复了1000次：

import numpy as np
import matplotlib.pyplot as plt

# example mask of one corner of a square
N = 100
mask = np.zeros((N, N))
mask[0:10, 0:10] = 1

# example x and y position arrays in arbitrary units
x = np.random.uniform(0, 1, 1000)
y = np.random.uniform(0, 1, 1000)

x_bins = np.linspace(np.min(x), np.max(x), N)
y_bins = np.linspace(np.min(y), np.max(y), N)

x_bin_idx = np.digitize(x, x_bins)
y_bin_idx = np.digitize(y, y_bins)

idx_masked = np.ravel(np.where(mask[y_bin_idx - 1, x_bin_idx - 1] == 1))

plt.imshow(mask[::-1, :])

---

有更有效的方法吗？

鉴于

mask

使用大小相同的存储箱覆盖您的字段，您不需要明确定义存储箱<代码>*\u bin\u idx可以通过简单的楼层划分在每个位置确定，因为您知道每个箱子的大小为

1/N

。我建议使用

1-0

作为总宽度（传入

np.random.uniform

）而不是

x.max（）-x.min（）

，当然，如果您知道范围的预期大小

x0 = 0   # or x.min()
x1 = 1   # or x.max()
x_bin = (x1 - x0) / N
x_bin_idx = ((x - x0) // x_bin).astype(int)

# ditto for y

这将比数字化更快、更简单，并避免在开始时使用额外的垃圾箱

在大多数情况下，您不需要

np.where

。90%的问题（包括这一个）不应该使用

where

。如果您想快速访问

x

和

y

的必要元素，只需使用布尔掩码。面具很简单

selction = mask[x_bin_idx, y_bin_idx].astype(bool)

如果

mask

已经是布尔值（无论如何都应该是布尔值），那么表达式

mask[x_bin_idx，y_bin_idx]

就足够了。它会生成一个与

x_bin_idx

和

y_bin_idx

大小相同的数组（与

x

和

y

大小相同），其中包含每个点的遮罩值。你可以把面具当作

x[selection]   # Elements of x in mask
y[selection]   # Elements of y in mask

如果您绝对需要整数索引，

其中

不是您的最佳选择

indices = np.flatnonzero(selection)

或

如果您的目标只是从

x

和

y

中提取值，我建议将它们叠加到一个数组中：

coords = np.stack((x, y), axis=1)

这样，就不必应用两次索引，只需使用

coords[selection, :]

或

---

根据

mask

和

x

和

y

的相对密度，布尔掩蔽或线性索引可能更快。您必须对一些相关案例进行计时，以获得更好的直觉。

您能给出一个可运行的示例吗？我没有

x

、

y

和

mask

可供使用。你能用语言解释一下你认为你的例子在做什么吗？你是想得到x和y，其中mask==1吗？mask有相同的x和y形状吗？@AlyHosny，没有一个mask和x和y的形状不同。

mask[y_bin_idx-1，x_bin_idx-1]==1

是你的mask。您可以将其应用于

x

，

y

。如果必须的话，你可以用np.flatnonzero获得索引，何必麻烦呢？您还可以将其表示为

mask[y_bin_idx-1，x_bin_idx-1].astype（bool）

。行

idx_masked=np.ravel（np.where（mask[y_bin_idx-1，x_bin_idx-1]==1））

可以简化为

idx_masked=mask[y_bin_idx-1，x_idx-1]，非零（），然后计算
非零
索引会更快，因为它们被用于索引到两个不同的数组
x
和
y
。否则会有两个循环穿过一个几乎是空的面具。@a_guest。这实际上只是一个循环，但我明白你的意思。该循环仍然需要运行以计算
非零
或等效值。接下来的问题是如何使用这些信息。如果只应用一次索引，则使用布尔值会更快。如果不止一次，线性指数可能值得计算。有一个类型问题：
x_bin_idx
是一个浮点数数组（在python 3中），不能用于索引。@Mad物理学家我刚刚用
a=np.arange（1_000_000）
和
m=np.random.random（size=len（a））进行了测试<0.5
首先计算
i=m.nonzero（）[0]
然后执行
a[i]，a[i]
要比
a[m]，a[m]
快得多。大约
m.sum（）/len（m）>=0.9时，计时变得相等。使用布尔掩码，我假设numpy首先需要对掩码求和以确定输出数组的大小，然后再次循环填充它。对于索引数组，输出大小只是
len（i）
，所以它只需要循环一次。@Haydon。谢谢你的帮助。比我添加的
astype（int）
更快的替代方法是将
x\u bin\u idx
预先分配为int数组，并调用
np.floor\u divide（x-x0，x\u bin，out=x\u bin\u idx，cast='safe'）
coords = np.stack((x, y), axis=1)

coords[selection, :]

coords[indices, :]