Эффективно выбирать элементы из поля (x, y) с 2D-маской в ​​Python

Question

У меня есть большое поле данных 2D-позиции, представленное в виде двух массивов 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 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, :])

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

Есть ли более эффективный способ сделать это?

Answer 1

Учитывая, что mask перекрывает ваше поле с ячейками одинакового размера, вам не нужно явно определять ячейки. *_bin_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

Если вам абсолютно необходимы целочисленные индексы, 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, либо логическое маскирование, либо линейное индексация может быть быстрее. Вам понадобится время для некоторых важных случаев, чтобы получить лучшую интуицию.