def numba_impl():
tree = kd.KDTree(data, leaf_size)
return tree.rejection_sample_query(
rejection_r**2,
query_r**2,
tree.get_node_indices(),
max_sample_size,
max_neighbors,
)
def numba_impl(data, leaf_size, n1, r0, r1, max_neigh0, max_neigh1):
tree = kd.KDTree(data, leaf_size)
start_nodes = tree.get_node_indices()
query0 = tree.query_radius_bottom_up(data, r0 ** 2, start_nodes, max_neigh0)
sample = bt.rejection_sample_precomputed(
query0.indices, query0.counts, n1, tree.int_type, tree.bool_type
)
sample_indices = sample.indices[: sample.count]
query1 = tree.query_radius_bottom_up(
data[sample_indices], r1 ** 2, start_nodes[sample_indices], max_neigh1
)
return query0, sample, query1
from numba_neighbors import kd_tree as kd
N = 1024
n = N
D = 2
rejection_r = 0.1
query_r = 0.2
max_neighbors = 256
leaf_size = 16
np.random.seed(124)
data = np.random.uniform(size=(N, D)).astype(kd.FLOAT_TYPE)
# i = np.argsort(data[:, 0])
# data = data[i]
tree = kd.KDTree(data, leaf_size=leaf_size)
sample_result, query_result = tree.rejection_sample_query(
rejection_r**2, query_r**2, tree.get_node_indices(), n, max_neighbors)
print(np.max(query_result.counts))
print(sample_result.count)
def vis(x0,
sample_indices,
small_balls=True,
big_balls=True,
labels=False,
aspect=1):
x1 = x0[sample_indices]
x0 = x0.T
x1 = x1.T
def kd_tree():
return kd.KDTree(data, leaf_size=leaf_size)
def tree(self, data, leaf_size):
return kd.KDTree(data, leaf_size=leaf_size)