Ускорить Python cKDTree

Question

В настоящее время у меня есть созданная мной функция, которая связывает синие точки с его (максимум) 3 ближайшими соседями в пределах диапазона пикселей 55. vertices_xy_list - это чрезвычайно большой список или точки (вложенный список) приблизительно 5000-10000 пар.

Пример vertices_xy_list:

[[3673.3333333333335, 2483.3333333333335],
 [3718.6666666666665, 2489.0],
 [3797.6666666666665, 2463.0],
 [3750.3333333333335, 2456.6666666666665],...]

В настоящее время я написал эту calculate_draw_vertice_lines() функцию, которая использует CKDTree внутри цикла While, чтобы найти все точки в пределах 55 пикселей и затем соединить их каждую с зеленой линией.

Видно, что это будет экспоненциально медленнее по мере увеличения списка. Есть ли способ значительно ускорить эту функцию? Такие как операции векторизации?

def calculate_draw_vertice_lines():

    global vertices_xy_list
    global cell_wall_lengths
    global list_of_lines_references

    index = 0

    while True:

        if (len(vertices_xy_list) == 1):

            break

        point_tree = spatial.cKDTree(vertices_xy_list)

        index_of_closest_points = point_tree.query_ball_point(vertices_xy_list[index], 55)

        index_of_closest_points.remove(index)

        for stuff in index_of_closest_points:

            list_of_lines_references.append(plt.plot([vertices_xy_list[index][0],vertices_xy_list[stuff][0]] , [vertices_xy_list[index][1],vertices_xy_list[stuff][1]], color = 'green'))

            wall_length = math.sqrt( (vertices_xy_list[index][0] - vertices_xy_list[stuff][0])**2 + (vertices_xy_list[index][1] - vertices_xy_list[stuff][1])**2 )

            cell_wall_lengths.append(wall_length)

        del vertices_xy_list[index]

    fig.canvas.draw()

Answer 1

Если я правильно понимаю логику правильного выбора зеленых линий, нет необходимости создавать KDTree на каждой итерации. Для каждой пары (p1, p2) синих точек линия должна быть проведена тогда и только тогда, когда выполняется следующее:

1. p1 является одним из 3 ближайших соседей p2.
2. p2 является одним из 3 ближайших соседей p1.
3. dist (p1, p2) <55. </li>

Вы можете создать KDTree один раз и эффективно создать список зеленых линий. Вот часть реализации, которая возвращает список пар индексов для точек, между которыми необходимо провести зеленые линии. Время выполнения на моей машине составляет около 0,5 секунд на 10000 баллов.

import numpy as np
from scipy import spatial


data = np.random.randint(0, 1000, size=(10_000, 2))

def get_green_lines(data):
    tree = spatial.cKDTree(data)
    # each key in g points to indices of 3 nearest blue points
    g = {i: set(tree.query(data[i,:], 4)[-1][1:]) for i in range(data.shape[0])}

    green_lines = list()
    for node, candidates in g.items():
        for node2 in candidates:
            if node2 < node:
                # avoid double-counting
                continue

            if node in g[node2] and spatial.distance.euclidean(data[node,:], data[node2,:]) < 55:
                green_lines.append((node, node2))

    return green_lines

Вы можете приступить к построению зеленых линий следующим образом:

green_lines = get_green_lines(data)
fig, ax = plt.subplots()
ax.scatter(data[:, 0], data[:, 1], s=1)
from matplotlib import collections as mc
lines = [[data[i], data[j]] for i, j in green_lines]
line_collection = mc.LineCollection(lines, color='green')
ax.add_collection(line_collection)

Пример вывода: