N = 1024
D = 3
np.random.seed(123)
data = np.random.uniform(size=(N, D)).astype(bt.FLOAT_TYPE)
leaf_size = 16
@benchmark("sklearn")
def sklearn_impl():
return sklearn.neighbors.kd_tree.KDTree(data, leaf_size=leaf_size)
@benchmark("base")
def numba_impl():
return bt.get_tree_data(data, leaf_size=leaf_size)
@benchmark("BinaryTree")
def binary_tree():
return bt.BinaryTree(data, leaf_size=leaf_size)
@benchmark("KDTree")
def kd_tree():
return kd.KDTree(data, leaf_size=leaf_size)
run_benchmarks(10, 100)
max_sample_size,
max_neighbors,
)
@benchmark("numba3")
def numba3_impl():
tree = kd.KDTree3(data, leaf_size)
return tree.rejection_sample_query(
rejection_r**2,
query_r**2,
tree.get_node_indices(),
max_sample_size,
max_neighbors,
)
# @benchmark()
# def separate():
# ragged_in_place_and_down_sample_query_np(data, N, rejection_r,
# max_neighbors, max_sample_size,
# query_r, max_neighbors)
run_benchmarks(20, 100)
sample_result, query_result = numba_impl()
count = sample_result.count
print(count)
counts = query_result.counts[:count]
print(np.min(counts), np.max(counts), np.mean(counts))
data = get_data()
sklearn.neighbors._kd_tree.simultaneous_sort(*data)
return data
def numba_impl():
data = get_data()
bt.simultaneous_sort(*data)
return data
def numpy_sort(dst, idx):
i = np.argsort(dst)
expected_dst = np.take_along_axis(dst, i, axis=1)
expected_idx = np.take_along_axis(idx, i, axis=1)
return expected_dst, expected_idx
def numpy_impl():
dst, idx = get_data()
return numpy_sort(dst, idx)
run_benchmarks(
10,
100,
("sklearn", sklearn_impl),
("numba", numba_impl),
("numpy", numpy_impl),
)
