Эффективно находить соседей с разрезами и возвращать индекс

Подход №1

Мы могли бы просто индексировать первую маску с ее собственной маской для выбора замаскированных значений True Places из второго этапа, например так:

idx[idx] = idx1

Таким образом, idx будет иметь окончательные допустимые замаскированные значения/места с хорошими значениями, соответствующие исходному массиву x и y, т.е. -

x_good = x[idx]
y_good = y[idx]

Затем эту маску можно использовать для индексации других массивов, как указано в вопросе.

Подход № 2

В качестве другого подхода мы могли бы использовать два условных оператора, таким образом создавая с ними две маски. Наконец, объедините их с AND-ing, чтобы получить комбинированную маску, которую можно проиндексировать в массивы x и y для окончательных выходных данных. Нам не нужно будет получать фактические индексы таким образом, так что это еще одно преимущество.

Следовательно, реализация -

X = x-x1
Y = y-y1
mask1 = (np.abs(X) < 1.) & (np.abs(Y) < 1.)
mask2 = X**2 + Y*2 < r**2
comb_mask = mask1 & mask2

x_good = x[comb_mask]
y_good = y[comb_mask]

Если по какой-то причине вам все еще нужны соответствующие индексы, просто сделайте -

comb_idx = np.flatnonzero(comb_mask)

Если вы выполняете эти операции для разных пар x1 и y1 для одного и того же набора данных x и y, я бы предложил использовать broadcasting для векторизации его через все эти парные наборы данных x1, y1, как показано в this post.

numpy.whereпохоже сделано для поиска индексов

векторизованная норма calc + np.where() может быть быстрее, чем цикл

sq_norm = (x - x1)**2 + (y - y1)**2  # no need to take 10000 sqrt
idcs = np.where(sq_norm < 1.)

len(idcs[0])
Out[193]: 69

np.stack((idcs[0], x[idcs], y[idcs]), axis=1)[:5]
Out[194]: 
array([[  38.        ,    9.47165956,   11.94250173],
       [  39.        ,    9.6966941 ,   11.67505453],
       [ 276.        ,   10.68835317,   12.11589316],
       [ 288.        ,    9.93632584,   11.07624915],
       [ 344.        ,    9.48644057,   12.04911857]])

норма calc тоже может включать массив r, 2-й шаг?

r_sq_norm = (x[idcs] - x1)**2 + (y[idcs] - y1)**2 - r[idcs]**2
r_idcs = np.where(r_sq_norm < 0.)

idcs[0][r_idcs]
Out[11]: array([1575, 3476, 3709], dtype=int64)

возможно, вы захотите рассчитать время двухэтапного теста по сравнению с включением r в 1-й векторизованный расчет нормы?

sq_norm = (x - x1)**2 + (y - y1)**2 - r**2
idcs = np.where(sq_norm < 0.)

idcs[0]
Out[13]: array([1575, 3476, 3709], dtype=int64)

Вы можете взять маску своих индексов, например:

def index_finder(x,y,r,x1,y1):
    idx = np.nonzero((abs(x - x1) < 1.) & (abs(y - y1) < 1.))  #numerical, not boolean
    mask = (x[idx] - x1)*(x[idx] - x1) + (y[idx] - y1)*(y[idx] - y1) < r*r
    idx1 = [i[mask] for i in idx]
    x_good = x_temp[idx1]
    y_good = y_temp[idx1]

теперь idx1 - это индексы, которые вы хотите извлечь.

В общем, более быстрый способ сделать это - использовать scipy.spatial.KDTree

from scipy.spatial import KDTree

xy = np.stack((x,y))
kdt = KDTree(xy)
kdt.query_ball_point([x1, y1], r)

Если у вас есть много точек для запроса к одному и тому же набору данных, это будет намного быстрее, чем последовательный вызов вашего index_finder приложения.

x1y1 = np.stack((x1, y1)) #`x1` and `y1` are arrays of coordinates.
kdt.query_ball_point(x1y1, r)

ТАКЖЕ НЕПРАВИЛЬНО: если у вас разные расстояния для каждой точки, вы можете сделать следующее:

def query_variable_ball(kdtree, x, y, r):
    out = []
    for x_, y_, r_ in zip(x, y, r):
        out.append(kdt.query_ball_point([x_, y_], r_)
    return out

xy = np.stack((x,y))
kdt = KDTree(xy)
query_variable_ball(kdt, x1, y1, r)

EDIT 2: Это должно работать с разными значениями r для каждой точки.

from scipy.spatial import KDTree

def index_finder_kd(x, y, r, x1, y1):  # all arrays
    xy = np.stack((x,y), axis = -1)
    x1y1 = np.stack((x1, y1), axis = -1)
    xytree = KDTree(xy)
    d, i = xytree.query(x1y1, k = None, distance_upper_bound = 1.)
    good_idx = np.zeros(x.size, dtype = bool)
    for idx, dist in zip(i, d):
        good_idx[idx] |= r[idx] > dist
    x_good = x[good_idx]
    y_good = y[good_idx]
    return x_good, y_good, np.flatnonzero(good_idx)

Это очень медленно только для одной пары (x1, y1), поскольку для заполнения KDTree требуется некоторое время. Но если у вас миллионы пар, это будет гораздо быстрее.

(Я предположил, что вы хотите объединить все хорошие точки в данных (x, y) для всех (x1, y1), если вы хотите их по отдельности, также возможно использовать аналогичный метод, удаляя элементы i[j] в зависимости от того, является ли d[j] < r[i[j]])