Скудное пространственное KD-дерево Python медленнее, чем евклидовы расстояния грубой силы?

Question

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

Кто-нибудь знает, почему мой тестовый код дает такие разные результаты?Я использую это неправильно?Является ли тестовый пример непригодным для kd-деревьев?

PS: Это сокращенная версия концепции, которую я использовал для проверки концепции.Полный код, в котором я также храню и преобразую результаты, можно найти здесь , но он дает те же результаты.

Импорт

import numpy as np
from time import time
from scipy.spatial import KDTree as kd
from functools import reduce
import matplotlib.pyplot as plt

Реализации

def euclid(c, cs, r):
    return ((cs[:,0] - c[0]) ** 2 + (cs[:,1] - c[1]) ** 2 + (cs[:,2] - c[2]) ** 2) < r ** 2

def find_nn_naive(cells, radius):
    for i in range(len(cells)):
        cell = cells[i]
        cands = euclid(cell, cells, radius)

def find_nn_kd_seminaive(cells, radius):
    tree = kd(cells)
    for i in range(len(cells)):
        res = tree.query_ball_point(cells[i], radius)

def find_nn_kd_by_tree(cells, radius):
    tree = kd(cells)
    return tree.query_ball_tree(tree, radius)

Тестовая установка

min_iter = 5000
max_iter = 10000
step_iter = 1000

rng = range(min_iter, max_iter, step_iter)
elapsed_naive = np.zeros(len(rng))
elapsed_kd_sn = np.zeros(len(rng))
elapsed_kd_tr = np.zeros(len(rng))

ei = 0
for i in rng:
    random_cells = np.random.rand(i, 3) * 400.
    t = time()
    r1 = find_nn_naive(random_cells, 50.)
    elapsed_naive[ei] = time() - t
    t = time()
    r2 = find_nn_kd_seminaive(random_cells, 50.)
    elapsed_kd_sn[ei] = time() - t
    t = time()
    r3 = find_nn_kd_by_tree(random_cells, 50.)
    elapsed_kd_tr[ei] = time() - t
    ei += 1

Участок

plt.plot(rng, elapsed_naive, label='naive')
plt.plot(rng, elapsed_kd_sn, label='semi kd')
plt.plot(rng, elapsed_kd_tr, label='full kd')
plt.legend()
plt.show(block=True)

Answer 1

Как указано в scipy.spatial.KDTree():

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

(эта заметка присутствует и в scipy.spatial.cKDTree(), хотя, вероятно, это копирование-вставкаошибка документации).

Я позволил себе переписать ваш код с соответствующими функциями, чтобы я мог запустить некоторые автоматизированные тесты (на основе этого шаблона ).Я также включил реализацию Numba методом «грубой силы»:

import numpy as np
import scipy as sp
import numba as nb

import scipy.spatial

SCALE = 400.0
RADIUS = 50.0 


def find_nn_np(points, radius=RADIUS, p=2):
    n_points, n_dim = points.shape
    result = np.empty(n_points, dtype=object)
    for i in range(n_points):
        result[i] = np.where(np.sum(np.abs(points - points[i:i + 1, :]) ** p, axis=1) < radius ** p)[0].tolist()
    return result


def find_nn_kd_tree(points, radius=RADIUS):
    tree = sp.spatial.KDTree(points)
    return tree.query_ball_point(points, radius)


def find_nn_kd_tree_cy(points, radius=RADIUS):
    tree = sp.spatial.cKDTree(points)
    return tree.query_ball_point(points, radius)


@nb.jit
def neighbors_indexes_jit(radius, center, points, p=2):
    n_points, n_dim = points.shape
    k = 0
    res_arr = np.empty(n_points, dtype=nb.int64)
    for i in range(n_points):
        dist = 0.0
        for j in range(n_dim):
            dist += abs(points[i, j] - center[j]) ** p
        if dist < radius ** p:
            res_arr[k] = i
            k += 1
    return res_arr[:k]


@nb.jit(forceobj=True, parallel=True)
def find_nn_jit(points, radius=RADIUS):
    n_points, n_dim = points.shape
    result = np.empty(n_points, dtype=object)
    for i in nb.prange(n_points):
        result[i] = neighbors_indexes_jit(radius, points[i], points, 2)
    return result

Это тесты, которые я получил (я пропустил scipy.spatial.KDTree(), потому что это было далеко от графика, в соответствии с вашими выводами):

---

(для полноты ниже приведен код, необходимый для адаптации шаблона)

def gen_input(n, dim=2, scale=SCALE):
    return scale * np.random.rand(n, dim)


def equal_output(a, b):
    return all(sorted(a_i) == sorted(b_i) for a_i, b_i in zip(a, b))


funcs = find_nn_np, find_nn_jit, find_nn_kd_tree_cy


input_sizes = tuple(int(2 ** (2 + (1 * i) / 4)) for i in range(32, 32 + 16 + 1))
print('Input Sizes:\n', input_sizes, '\n')


runtimes, input_sizes, labels, results = benchmark(
    funcs, gen_input=gen_input, equal_output=equal_output,
    input_sizes=input_sizes)


plot_benchmarks(runtimes, input_sizes, labels, units='s')

Answer 2

TL; DR:

Переключиться на scipy.spatial.cKDTree или sklearn.neighbors.KDTree для характеристик, ожидаемых от алгоритмов дерева kd.