Python 的 scipy 空间 KD 树比蛮力欧氏距离慢?
Python's scipy spatial KD-tree is slower than brute force euclidean distances?
我已经快速检查了构建树并对其进行查询与仅计算所有欧氏距离的性能。如果我向这棵树查询半径内的所有其他点,它不应该大大优于蛮力方法吗?
有谁知道为什么我的测试代码会产生这些不同的结果?我用错了吗?测试用例不适合kd-trees吗?
PS:这是我使用的代码的简化概念验证版本。可以找到我存储和转换结果的完整代码 here,但它产生相同的结果。
进口
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)
如 scipy.spatial.KDTree()
中所述:
For large dimensions (20 is already large) do not expect this to run significantly faster than brute force. High-dimensional nearest-neighbor queries are a substantial open problem in computer science.
(此说明也出现在 scipy.spatial.cKDTree()
中,尽管这可能是一个复制粘贴文档错误)。
我冒昧地用适当的函数重写了你的代码,这样我就可以 运行 一些自动化基准测试(基于 this template)。我还包括了一个暴力 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')
长话短说:
切换到 scipy.spatial.cKDTree
或 sklearn.neighbors.KDTree
以获得 kd-tree 算法预期的性能。
我已经快速检查了构建树并对其进行查询与仅计算所有欧氏距离的性能。如果我向这棵树查询半径内的所有其他点,它不应该大大优于蛮力方法吗?
有谁知道为什么我的测试代码会产生这些不同的结果?我用错了吗?测试用例不适合kd-trees吗?
PS:这是我使用的代码的简化概念验证版本。可以找到我存储和转换结果的完整代码 here,但它产生相同的结果。
进口
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)
如 scipy.spatial.KDTree()
中所述:
For large dimensions (20 is already large) do not expect this to run significantly faster than brute force. High-dimensional nearest-neighbor queries are a substantial open problem in computer science.
(此说明也出现在 scipy.spatial.cKDTree()
中,尽管这可能是一个复制粘贴文档错误)。
我冒昧地用适当的函数重写了你的代码,这样我就可以 运行 一些自动化基准测试(基于 this template)。我还包括了一个暴力 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')
长话短说:
切换到 scipy.spatial.cKDTree
或 sklearn.neighbors.KDTree
以获得 kd-tree 算法预期的性能。