python - 在 Python 中使用 2D 掩码从 (x,y) 字段中有效地选择元素

Question

我有一个很大的二维位置数据字段，以两个数组形式给出 x和 y ，其中 len(x) == len(y) .我想返回索引数组 idx_masked其中(x[idx_masked], y[idx_masked])被 N x N int 掩盖名为 mask 的数组.即，mask[x[idx_masked], y[idx_masked]] == 1 . mask数组由 0 组成s 和 1只。

我想出了以下解决方案，但它(特别是下面的最后一行)非常慢，因为我有 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, :])

plt.scatter(x, y, color='red')
plt.scatter(x[idx_masked], y[idx_masked], color='blue')

有没有更有效的方法来做到这一点？

Answer 1

鉴于 mask用相同大小的 bin 覆盖您的字段，您不需要明确定义 bin。 *_bin_idx可以通过简单的楼层划分在每个位置确定，因为您知道每个 bin 是 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

这将比数字化更快、更简单，并避免在开始时出现额外的 bin。

大多数情况下，您不需要 np.where . 90% 的问题(包括这个)不应该使用 where .如果您想快速访问 x 的必要元素和 y ，只需使用 bool 掩码。面膜简直了

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

如 mask已经是一个 bool 值(无论如何都应该是)，表达式 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

如果你绝对需要整数索引， where仍然不是您的最佳选择。

indices = np.flatnonzero(selection)

或者

indices = selection.nonzero()[0]

如果您的目标只是从 x 中提取值和 y ，我建议将它们堆叠成一个数组:

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

这样，您不必两次应用索引，只需使用以下命令即可提取值

coords[selection, :]

或者

coords[indices, :]

取决于 mask 的相对密度和 x和 y ， bool 掩码或线性索引可能更快。您将不得不计时一些相关案例以获得更好的直觉。