def main():
# configuration
N = 64000
D = 3
leaf_size = 4
min_extent = 0.05
k = 8
r = 0.372
# generate point cloud:
point_cloud = np.random.rand(N, D)
octree = OCTree(point_cloud=point_cloud,
leaf_size=leaf_size,
min_extent=min_extent)
# octree.traverse()
# random query test:
for _ in range(100):
# generate query point:
query = np.random.rand(D)
# 01 -- knn: brute-force as baseline:
dists = np.linalg.norm(point_cloud - query, axis=1)
sorting_idx = np.argsort(dists)
brute_force_result = {i for i in sorting_idx[:k]}
knn_result_set = KNNResultSet(capacity=k)
octree.knn_search(query, knn_result_set)
knn_result = {i.index for i in knn_result_set.dist_index_list}
assert len(brute_force_result - knn_result) == 0
# 02 -- rnn: brute-force as baseline:
dists = np.linalg.norm(point_cloud - query, axis=1)
brute_force_result = {i for i, d in enumerate(dists) if d <= r}
rnn_result_set = RadiusNNResultSet(radius=r)
octree.rnn_search(query, rnn_result_set)
rnn_result = {i.index for i in rnn_result_set.dist_index_list}
assert len(brute_force_result - rnn_result) == 0
print('[OCTree kNN & RNN Random Query Test]: Successful')
begin_t = time.time()
print("[OCTree]: RNN search normal:")
for i in range(100):
query = np.random.rand(3)
rnn_result_set = RadiusNNResultSet(radius=0.5)
octree.rnn_search(query, rnn_result_set)
# print(result_set)
print("\tSearch takes %.3fms\n" % ((time.time() - begin_t) * 1000))
begin_t = time.time()
print("[OCTree]: RNN search fast:")
for i in range(100):
query = np.random.rand(3)
rnn_result_set = RadiusNNResultSet(radius=0.5)
octree.rnn_fast_search(query, rnn_result_set)
# print(result_set)
print("\tSearch takes %.3fms\n" % ((time.time() - begin_t) * 1000))
def kdtree_radius_search(root: Node, db: np.ndarray, result_set: RadiusNNResultSet, query: np.ndarray):
if root is None:
return False
if root.is_leaf():
# compare the contents of a leaf
leaf_points = db[root.point_indices, :]
diff = np.linalg.norm(np.expand_dims(query, 0) - leaf_points, axis=1)
for i in range(diff.shape[0]):
result_set.add_point(diff[i], root.point_indices[i])
return False
# 作业3
# 提示：通过递归的方式实现搜索
# 屏蔽开始
if query[root.axis] <= root.value:
kdtree_radius_search(root.left, db, result_set, query)
if math.fabs(query[root.axis] - root.value) < result_set.worstDist():
kdtree_radius_search(root.right, db, result_set, query)
else:
kdtree_radius_search(root.right, db, result_set, query)
if math.fabs(query[root.axis] - root.value) < result_set.worstDist():
kdtree_radius_search(root.left, db, result_set, query)
# 屏蔽结束
return False
def kdtree_radius_search(root: Node, db: np.ndarray,
result_set: RadiusNNResultSet, query: np.ndarray):
if root is None:
return False
if root.is_leaf():
# compare the contents of a leaf
leaf_points = db[root.point_indices, :]
diff = np.linalg.norm(np.expand_dims(query, 0) - leaf_points, axis=1)
for i in range(diff.shape[0]):
result_set.add_point(diff[i], root.point_indices[i])
return False
# 作业3
# 提示：通过递归的方式实现搜索
# 屏蔽开始
if root.value >= query[root.axis]:
if kdtree_knn_search(root.left, db, result_set, query):
return True
elif math.fabs(root.value - query[root.axis]) < result_set.worst_dist:
return kdtree_knn_search(root.right, db, result_set, query)
return False # 3、左右子树都不满足，就需要返回上一层树，所以为False
else:
if kdtree_knn_search(root.right, db, result_set, query):
return True
elif math.fabs(root.value - query[root.axis]) < result_set.worst_dist:
return kdtree_knn_search(root.left, db, result_set, query)
return False # 3、左右子树都不满足，就需要返回上一层树，所以为False
# 屏蔽结束
return False
def octree_radius_search_fast(root: Octant, db: np.ndarray, result_set: RadiusNNResultSet, query: np.ndarray):
if root is None:
return False
if contains(query, result_set.worstDist(), root):
# compare the contents of the octant
leaf_points = db[root.point_indices, :]
diff = np.linalg.norm(np.expand_dims(query, 0) - leaf_points, axis=1)
for i in range(diff.shape[0]):
result_set.add_point(diff[i], root.point_indices[i])
# don't need to check any child
return False
if root.is_leaf and len(root.point_indices) > 0:
# compare the contents of a leaf
leaf_points = db[root.point_indices, :]
diff = np.linalg.norm(np.expand_dims(query, 0) - leaf_points, axis=1)
for i in range(diff.shape[0]):
result_set.add_point(diff[i], root.point_indices[i])
# check whether we can stop search now
return inside(query, result_set.worstDist(), root)
# no need to go to most relevant child first, because anyway we will go through all children
for c, child in enumerate(root.children):
if child is None:
continue
if False == overlaps(query, result_set.worstDist(), child):
continue
if octree_radius_search_fast(child, db, result_set, query):
return True
return inside(query, result_set.worstDist(), root)
def octree_radius_search_fast(root: Octant, db: np.ndarray, result_set: RadiusNNResultSet, query: np.ndarray):
if root is None:
return False
# 1. if query ball contains the octant, no need to check child, just compare all point in it!
# 只是不需要向下递归，这里的child是只当前节点的下一层，回溯还是有可能的！所以这里return的false是有原因的嗷！
if contains(query, result_set.worstDist(), root):
# compare all points:
leaf_points = db[root.point_indices, :]
diff = np.linalg.norm(np.expand_dims(query, 0) - leaf_points, axis=1)
for i in range(diff.shape[0]):
result_set.add_point(diff[i], root.point_indices[i])
# no need to check child
return False
# consider leaf point:
if root.is_leaf and len(root.point_indices) > 0:
# compare all points:
leaf_points = db[root.point_indices, :]
diff = np.linalg.norm(np.expand_dims(query, 0) - leaf_points, axis=1)
for i in range(diff.shape[0]):
result_set.add_point(diff[i], root.point_indices[i])
return inside(query, result_set.worstDist(), root)
# 2. check all children
for c, child in enumerate(root.children):
if child is None:
continue
if False == overlaps(query, result_set.worstDist(), child):
continue
if octree_radius_search_fast(child, db, result_set,query):
return True
return inside(query, result_set.worstDist(), root)
def main():
# configuration
db_size = 64
dim = 3
leaf_size = 4
k = 1
db_np = np.random.rand(db_size, dim)
begin_t = time.time()
root = kdtree_construction(db_np, leaf_size=leaf_size)
print("construction with sort: %.3f", (time.time() - begin_t) * 1000)
begin_t = time.time()
root2 = kdtree_construction_2(db_np, leaf_size=leaf_size)
print("construction without sort: %.3f", (time.time() - begin_t) * 1000)
# depth = [0]
# max_depth = [0]
# traverse_kdtree(root, depth, max_depth)
# print("tree max depth: %d" % max_depth[0])
# query = np.asarray([0, 0, 0])
# result_set = KNNResultSet(capacity=k)
# knn_search(root, db_np, result_set, query)
#
# print(result_set)
#
# diff = np.linalg.norm(np.expand_dims(query, 0) - db_np, axis=1)
# nn_idx = np.argsort(diff)
# nn_dist = diff[nn_idx]
# print(nn_idx[0:k])
# print(nn_dist[0:k])
#
#
print("Radius search:")
query = np.asarray([0, 0, 0])
result_set = RadiusNNResultSet(radius=0.5)
result_set2 = RadiusNNResultSet(radius=0.5)
begin_t = time.time()
kdtree_radius_search(root, db_np, result_set, query)
print("search time with sort: %.3fms", (time.time() - begin_t) * 1000)
print(result_set)
begin_t = time.time()
kdtree_radius_search(root2, db_np, result_set2, query)
print("search time without sort: %.3fms", (time.time() - begin_t) * 1000)
print(result_set2)
def main():
# configuration
db_size = 100
k = 5
radius = 2.0
#data = np.random.permutation(db_size).tolist()
data = [3, 4, 7, 0, 6, 5, 1, 2]
root = None
for i, point in enumerate(data):
root = insert(root, point, i)
query_key = 6
result_set = KNNResultSet(capacity=k)
knn_search(root, result_set, query_key)
print('kNN Search:')
print('index - distance')
print(result_set)
result_set = RadiusNNResultSet(radius=radius)
radius_search(root, result_set, query_key)
print('Radius NN Search:')
print('index - distance')
print(result_set)
def check_neighbor(self, idx, cluster_idx, layer):
#if it is visited already
radius = self.search_radius
self.data_status_[idx] = False
result_set = RadiusNNResultSet(radius)
#print(np.shape(self.data_[idx]),(self.data_[idx]))
kdtree.kdtree_radius_search(self.root_, self.data_, result_set,
self.data_[idx])
nis = result_set.index_list
ndists = result_set.dist_list
#print("idx ",idx," nis",nis," ndists",ndists)
true_nis = np.count_nonzero(self.data_status_[nis])
if (true_nis > self.min_samples):
layer += 1
#print("set idx ",idx,"as label",cluster_idx)
self.data_label_[idx] = cluster_idx
for sub_idx in nis:
#print("nis" ,nis)
if self.data_status_[sub_idx] == True:
#print("set idx inner ",sub_idx,"as label",cluster_idx)
self.data_label_[sub_idx] = cluster_idx
self.check_neighbor(sub_idx, cluster_idx, layer)
#print("iter layer ",layer)
return layer
else:
#print("layer ",layer)
return layer
def main():
# configuration
db_size = 100
k = 5
radius = 2.0
data = np.random.permutation(db_size).tolist()
root = None
for i, point in enumerate(data):
root = insert(root, point, i) # values = i, 数据中点的索引，对后面的NN搜索有用。
print("data = ", data)
query_key = 6
result_set = KNNResultSet(capacity=k) ## k等于5 查找与 query_key最近的k个数据
knn_search(root, result_set, query_key)
print('kNN Search:')
print('index - distance')
print(result_set)
result_set = RadiusNNResultSet(radius=radius)
radius_search(root, result_set, query_key)
print('Radius NN Search:')
print('index - distance')
print(result_set)
def main():
# configuration
db_size = 64
dim = 3
leaf_size = 4
k = 1
db_np = np.random.rand(db_size, dim)
root = kdtree_construction(db_np, leaf_size=leaf_size)
depth = [0]
max_depth = [0]
traverse_kdtree(root, depth, max_depth)
print("tree max depth: %d" % max_depth[0])
query = np.asarray([0, 0, 0])
result_set = KNNResultSet(capacity=k)
kdtree_knn_search(root, db_np, result_set, query)
print(result_set)
diff = np.linalg.norm(np.expand_dims(query, 0) - db_np, axis=1)
nn_idx = np.argsort(diff)
nn_dist = diff[nn_idx]
print(nn_idx[0:k])
print(nn_dist[0:k])
print("Radius search:")
query = np.asarray([0, 0, 0])
result_set = RadiusNNResultSet(radius=0.5)
kdtree_radius_search(root, db_np, result_set, query)
print(result_set)
def octree_radius_search_fast(root: Octant, db: np.ndarray, result_set: RadiusNNResultSet, query: np.ndarray):
if root is None:
return False
# 作业5
# 提示：尽量利用上面的inside、overlaps、contains等函数
# 屏蔽开始
if contains(query, result_set.worstDist(), root):
#compare the contents of the octant
leaf_points = db[root.point_indices, :]
diff = np.linalg.norm(np.expand_dims(query,0) - leaf_points, axis =1)
for i in range(diff.shape[0]):
# 屏蔽结束
return inside(query, result_set.worstDist(), root)
def main():
# configuration
leaf_size = 32
min_extent = 0.0001
k = 8
radius = 1
N = 100000 # 对N个点做搜索
# read data
filename = "../000000.bin"
db_np = read_velodyne_bin(filename)
root = kd_tree_construction(db_np, leaf_size)
depth = [0]
max_depth = [0]
traverse_kdtree(root, depth, max_depth)
knn_result_set = KNNResultSet(k)
query = db_np[0, :]
kd_tree_knn_search(root, db_np, knn_result_set, query)
print(knn_result_set)
radius_result_set = RadiusNNResultSet(radius)
kd_tree_radius_search(root, db_np, radius_result_set, query)
print(radius_result_set)
def main():
# 生成模拟数据
db_size = 64
dim = 3
leaf_size = 4
k = 8
db_np = np.random.rand(db_size, dim)
root = kdtree_construction(db_np, leaf_size=leaf_size)
# 测试Kd-Tree遍历
depth = [0]
max_depth = [0]
traverse_kdtree(root, depth, max_depth)
print("Tree max depth: %d" % max_depth[0])
# 测试KNN search
query = np.asarray([0, 0, 0])
result_set = KNNResultSet(capacity=k)
kdtree_knn_search(root, db_np, result_set, query)
print(result_set)
# 测试brute-force法
diff = np.linalg.norm(np.expand_dims(query, 0) - db_np, axis=1)
nn_idx = np.argsort(diff)
nn_dist = diff[nn_idx]
print(nn_idx[0:k])
print(nn_dist[0:k])
# 测试Radius search
result_set = RadiusNNResultSet(radius=0.5)
kdtree_radius_search(root, db_np, result_set, query)
print(result_set)
def main():
# 生成模拟数据
db_size = 3000
dim = 3
leaf_size = 4
min_extent = 0.0001
k = 8
db_np = np.random.rand(db_size, dim)
root = octree_construction(db_np, leaf_size, min_extent)
# 测试Octree遍历
depth = [0]
max_depth = [0]
traverse_octree(root, depth, max_depth)
print("Tree max depth: %d" % max_depth[0])
# 测试KNN search
query = np.asarray([0, 0, 0])
result_set = KNNResultSet(capacity=k)
octree_knn_search(root, db_np, result_set, query)
print(result_set)
# 测试brute-force法
diff = np.linalg.norm(np.expand_dims(query, 0) - db_np, axis=1)
nn_idx = np.argsort(diff)
nn_dist = diff[nn_idx]
print(nn_idx[0:k])
print(nn_dist[0:k])
# 测试Radius search (normal)
begin_t = time.time()
print("Radius search normal:")
for i in range(100):
query = np.random.rand(3)
result_set = RadiusNNResultSet(radius=0.5)
octree_radius_search(root, db_np, result_set, query)
print("Search takes %.3fms\n" % ((time.time() - begin_t) * 1000))
# 测试Radius search (fast)
begin_t = time.time()
print("Radius search fast:")
for i in range(100):
query = np.random.rand(3)
result_set = RadiusNNResultSet(radius=0.5)
octree_radius_search_fast(root, db_np, result_set, query)
print("Search takes %.3fms\n" % ((time.time() - begin_t) * 1000))
def octree_radius_search_plus(root: Octant, db: np.ndarray, result_set: RadiusNNResultSet, query: np.ndarray):
if root is None:
return False
if root.is_leaf and len(root.point_indices) > 0:
# compare the contents of a leaf
leaf_points = db[root.point_indices, :]
diff = np.linalg.norm(np.expand_dims(query, 0) - leaf_points, axis=1)
for i in range(diff.shape[0]):
result_set.add_point(diff[i], root.point_indices[i])
# check whether we can stop search now
return inside(query, result_set.worstDist(), root)
# 提前结束：核心-八叉树对3个维度有限制；
# 结束条件：其实是两个方向：向下递归，向上回溯
# 1-当一个节点返回True时，表示找到knn了，维度限制导致不会有更优的了，不再向下，立即结束；
# 2-最坏距离球如果 inside 当前节点，不需要再向上回溯检查其他节点，立即结束；
# 跳过条件：在检查
# 1. search the first relevant child:
# 找到最近的孩子，根据查询点的莫顿码
morton_code = 0
if query[0] > root.center[0]:
morton_code = morton_code | 1
if query[1] > root.center[1]:
morton_code = morton_code | 2
if query[2] > root.center[2]:
morton_code = morton_code | 4
if octree_knn_search(root.children[morton_code],
db,
result_set,
query):
return True
# 2. check other children
inv_worstDist = 1 / result_set.worstDist()
for c, child in enumerate(root.children):
if c == morton_code or child is None:
continue
if False == overlaps_plus(query, inv_worstDist, child):
continue
if octree_knn_search(child, db, result_set, query):
return True
# final check of if we can stop search
return inside(query, result_set.worstDist(), root)
def main():
# configuration
db_size = 64000
dim = 3
leaf_size = 4
min_extent = 0.0001
k = 8
db_np = np.random.rand(db_size, dim)
root = octree_construction(db_np, leaf_size, min_extent)
# depth = [0]
# max_depth = [0]
# traverse_octree(root, depth, max_depth)
# print("tree max depth: %d" % max_depth[0])
# query = np.asarray([0, 0, 0])
# result_set = KNNResultSet(capacity=k)
# octree_knn_search(root, db_np, result_set, query)
# print(result_set)
#
# diff = np.linalg.norm(np.expand_dims(query, 0) - db_np, axis=1)
# nn_idx = np.argsort(diff)
# nn_dist = diff[nn_idx]
# print(nn_idx[0:k])
# print(nn_dist[0:k])
begin_t = time.time()
print("Radius search normal:")
for i in range(100):
query = np.random.rand(3)
result_set = RadiusNNResultSet(radius=0.5)
octree_radius_search(root, db_np, result_set, query)
# print(result_set)
print("Search takes %.3fms\n" % ((time.time() - begin_t) * 1000))
begin_t = time.time()
print("Radius search fast:")
for i in range(100):
query = np.random.rand(3)
result_set = RadiusNNResultSet(radius = 0.5)
octree_radius_search_fast(root, db_np, result_set, query)
# print(result_set)
print("Search takes %.3fms\n" % ((time.time() - begin_t)*1000))
def octree_radius_search(root: Octant, db: np.ndarray, result_set: RadiusNNResultSet, query: np.ndarray):
if root is None:
return False
if root.is_leaf and len(root.point_indices) > 0:
# compare the contents of a leaf
leaf_points = db[root.point_indices, :]
diff = np.linalg.norm(np.expand_dims(query, 0) - leaf_points, axis=1)
for i in range(diff.shape[0]):
result_set.add_point(diff[i], root.point_indices[i])
# check whether we can stop search now
return inside(query, result_set.worstDist(), root)
# 作业6
# 屏蔽开始
# 屏蔽结束
# final check of if we can stop search
return inside(query, result_set.worstDist(), root)
def kdtree_radius_search(root: Node, db: np.ndarray,
result_set: RadiusNNResultSet, query: np.ndarray):
if root is None:
return False
if root.is_leaf():
# compare the contents of a leaf
leaf_points = db[root.point_indices, :]
diff = np.linalg.norm(np.expand_dims(query, 0) - leaf_points, axis=1)
for i in range(diff.shape[0]):
result_set.add_point(diff[i], root.point_indices[i])
return False
# 作业3
# 提示：通过递归的方式实现搜索
# 屏蔽开始
# 屏蔽结束
return False
def octree_radius_search_fast(root: Octant, db: np.ndarray, result_set: RadiusNNResultSet, query: np.ndarray):
if root is None:
return False
# 作业5
# 提示：尽量利用上面的inside、overlaps、contains等函数
# 屏蔽开始
# 屏蔽结束
return inside(query, result_set.worstDist(), root)
def octree_radius_search(root: Octant, db: np.ndarray,
result_set: RadiusNNResultSet, query: np.ndarray):
if root is None:
return False
if root.is_leaf and len(root.point_indices) > 0:
# compare the contents of a leaf
leaf_points = db[root.point_indices, :]
diff = np.linalg.norm(np.expand_dims(query, 0) - leaf_points, axis=1)
for i in range(diff.shape[0]):
result_set.add_point(diff[i], root.point_indices[i])
# check whether we can stop search now
return inside(query, result_set.worstDist(), root)
# 作业6
# 屏蔽开始
children_idx = 0
if query[0] > root.center[0]: # x轴
children_idx = children_idx | 1
if query[1] > root.center[1]: # y轴
children_idx = children_idx | 2
if query[2] > root.center[2]: # z轴
children_idx = children_idx | 4
# 如果在这个子octant中发现，查询点在该子octant中，同时由worst_dist构成的球也被octant包含，所以直接返回
if octree_knn_search(root.children[children_idx], db, result_set, query):
return True
# 如果不满足上边的情况则需要遍历其他子octant
for c, child in enumerate(root.children):
if c == children_idx or root.children[c]:
continue
if overlaps(query, result_set.worstDist(), child) == False:
continue
if octree_knn_search(root.children[c], db, result_set, query):
return True
# 屏蔽结束
# final check of if we can stop search
return inside(query, result_set.worstDist(), root)
def octree_radius_search(root: Octant, db: np.ndarray,
result_set: RadiusNNResultSet, query: np.ndarray):
if root is None:
return False
if root.is_leaf and len(root.point_indices) > 0:
# compare the contents of a leaf
leaf_points = db[root.point_indices, :]
diff = np.linalg.norm(np.expand_dims(query, 0) - leaf_points, axis=1)
for i in range(diff.shape[0]):
result_set.add_point(diff[i], root.point_indices[i])
# check whether we can stop search now
return inside(query, result_set.worstDist(), root)
# 作业6
# 屏蔽开始
# go to the relevant child first
morton_code = 0
if query[0] > root.center[0]:
morton_code = morton_code | 1
if query[1] > root.center[1]:
morton_code = morton_code | 2
if query[2] > root.center[2]:
morton_code = morton_code | 4
if octree_radius_search(root.children[morton_code], db, result_set, query):
return True
# check other children
for c, child in enumerate(root.children):
if c == morton_code or child is None:
continue
if False == overlaps(query, result_set.worstDist(), child):
continue
if octree_radius_search(child, db, result_set, query):
return True
# 屏蔽结束
# final check of if we can stop search
return inside(query, result_set.worstDist(), root)
def radius_search(root: Node, result_set: RadiusNNResultSet, key):
if root is None:
return False
# compare the root itself
result_set.add_point(math.fabs(root.key - key), root.value)
if root.key >= key:
# iterate left branch first
if radius_search(root.left, result_set, key):
return True
elif math.fabs(root.key-key) < result_set.worstDist():
return radius_search(root.right, result_set, key)
return False
else:
# iterate right branch first
if radius_search(root.right, result_set, key):
return True
elif math.fabs(root.key-key) < result_set.worstDist():
return radius_search(root.left, result_set, key)
return False
def KDTreeBenchmark(root_dir,
files,
k,
leaf_size,
radius,
feature=None,
feature2=None):
construction_time_sum = 0
knn_time_sum = 0
radius_time_sum = 0
brute_time_sum = 0
iteration_num = 0
for file in files:
if file.find('bin') == -1:
continue
iteration_num += 1
filename = os.path.join(root_dir, file)
db_np = read_velodyne_bin(filename)
begin_t = time.time()
root = kdtree.kdtree_construction(db_np, leaf_size, feature, feature2)
construction_time_sum += time.time() - begin_t
query = db_np[0, :]
begin_t = time.time()
result_set = KNNResultSet(capacity=k)
kdtree.kdtree_knn_search(root, db_np, result_set, query)
print("result set from KD Tree\n", result_set)
knn_time_sum += time.time() - begin_t
# print("--------")
begin_t = time.time()
result_set = RadiusNNResultSet(radius=radius)
kdtree.kdtree_radius_search(root, db_np, result_set, query)
#print(result_set)
radius_time_sum += time.time() - begin_t
#print("--------")
begin_t = time.time()
diff = np.linalg.norm(np.expand_dims(query, 0) - db_np, axis=1)
nn_idx = np.argsort(diff)
nn_dist = diff[nn_idx]
#print(nn_idx[0:k])
#print(nn_dist[0:k])
brute_time_sum += time.time() - begin_t
depth = [0]
max_depth = [0]
kdtree.traverse_kdtree(root, depth, max_depth)
print("tree depth: %d, max depth: %d" % (depth[0], max_depth[0]))
print("Kdtree: build %.3f, knn %.3f, radius %.3f, brute %.3f" %
(construction_time_sum * 1000 / iteration_num, knn_time_sum * 1000 /
iteration_num, radius_time_sum * 1000 / iteration_num,
brute_time_sum * 1000 / iteration_num))
def kdtree_radius_search(root: Node, db: np.ndarray,
result_set: RadiusNNResultSet, query: np.ndarray):
if root is None:
return
if root.is_leaf():
leaf_points = db[root.point_indices]
diff = np.linalg.norm(np.expand_dims(query, 0) - leaf_points, axis=1)
for i in range(diff.shape[0]):
result_set.add_point(diff[i], root.point_indices[i])
return
if query[root.axis] <= root.value:
kdtree_radius_search(root.left, db, result_set, query)
if math.fabs(query[root.axis] -
root.value) < result_set.get_worst_dist():
kdtree_radius_search(root.right, db, result_set, query)
else:
kdtree_radius_search(root.right, db, result_set, query)
if math.fabs(query[root.axis] -
root.value) < result_set.get_worst_dist():
kdtree_radius_search(root.left, db, result_set, query)
def main():
# configuration
db_size = 1280
dim = 3
leaf_size = 4
k = 1
db_np = np.random.rand(db_size, dim)#生成size×dim的array，每个数都是0-1之间的随机数
radius_time_sum = 0
brute_time_sum = 0
root = kdtree_construction(db_np, leaf_size)
# root = kdtree_construction_median(db_np, leaf_size)
# depth = [0]
# max_depth = [0]
# traverse_kdtree(root, depth, max_depth)
# print("tree max depth: %d" % max_depth[0])
query = np.asarray([0, 0, 0])
result_set = KNNResultSet(capacity=k)
kdtree_knn_search(root, db_np, result_set, query)
# print(result_set)
#
diff = np.linalg.norm(np.expand_dims(query, 0) - db_np, axis=1)
nn_idx = np.argsort(diff)
nn_dist = diff[nn_idx]
# print(nn_idx[0:k])
# print(nn_dist[0:k])
print("Radius search:")
query = np.asarray([7, 3, 4])
begin_t = time.time()
result_set = RadiusNNResultSet(radius = 2)
kdtree_radius_search(root, db_np, result_set, query)
radius_time_sum = time.time() - begin_t
# print(result_set)
begin_t = time.time()
diff = np.linalg.norm(np.expand_dims(query, 0) - db_np, axis=1)
nn_idx = np.argsort(diff)
nn_dist = diff[nn_idx]
brute_time_sum = time.time() - begin_t
print('KD-tree: radius_search/brute',radius_time_sum/brute_time_sum)
def fit(self,data,method="radius"):
data_num = np.shape(data)[0]
print("number of data ",data_num)
self.simi_graph=np.zeros((data_num,data_num))
if method=="fully_connect":
for i in range(np.shape(data)[0]):
self.simi_graph[i] = get_dist_array(data,data[i])
#self.simi_graph=np.max(self.simi_graph)*2- self.simi_graph
self.simi_graph= gauss(self.simi_graph)
#for i in range(np.shape(data)[0]):
# self.simi_graph[i,i]=0
#print(self.simi_graph)
elif (method=="radius"):
root = kdtree.kdtree_construction(data, 16)
for di, datum in enumerate(data):
#print(datum)
result_set = RadiusNNResultSet(radius=10)
kdtree.kdtree_radius_search(root, data, result_set, datum)
nis = result_set.index_list
ndists = result_set.dist_list
#print("nis",nis)
#print("ndists",ndists)
for ni,ndist in zip(nis,ndists):
self.simi_graph[di][ni]=self.simi_graph[ni][di]=gauss(ndist,0.2)
#print("graph",self.simi_graph)
elif (method=="knn"):
print("knn is implemented")
tree=KDTree(data)
for di, datum in enumerate(data):
ndists,nis = tree.query([datum],20)
nis=nis[0]
ndists = ndists[0]
for ni,ndist in zip(nis,ndists):
if ni==di: continue
#print("HHHHHH",ni,di,ndist)
self.simi_graph[di][ni]=self.simi_graph[ni][di]=gauss(ndist)
else:
print("not available")
def octree_radius_search_fast(root: Octant, db: np.ndarray,
result_set: RadiusNNResultSet,
query: np.ndarray):
if root is None:
return False
# 作业5
# 提示：尽量利用上面的inside、overlaps、contains等函数
# 屏蔽开始
if contains(query, result_set.worstDist(), root):
# the octant is contained by the query ball with resule_set.worstDist()
leaf_points = db[root.point_indices, :]
diff = np.linalg.norm(np.expand_dims(query, 0) - leaf_points, axis=1)
for i in range(diff.shape[0]):
result_set.add_point(diff[i], root.point_indices[i])
# check whether we can stop search now
return False
if root.is_leaf and len(root.point_indices) > 0:
# compare the contents of a leaf
leaf_points = db[root.point_indices, :]
diff = np.linalg.norm(np.expand_dims(query, 0) - leaf_points, axis=1)
for i in range(diff.shape[0]):
result_set.add_point(diff[i], root.point_indices[i])
# check whether we can stop search now
return inside(query, result_set.worstDist(), root)
# check other children
for c, child in enumerate(root.children):
if child is None:
continue
# if an octant is not overlapping with query ball, skip
if overlaps(query, result_set.worstDist(), child) == False:
continue
if octree_radius_search(child, db, result_set, query):
return True
# 屏蔽结束
return inside(query, result_set.worstDist(), root)
def main():
# configuration
N = 64000
D = 3
leaf_size = 4
k = 8
r = 0.372
point_cloud = np.random.rand(N, D)
kd_tree = KDTree(
point_cloud = point_cloud,
init_axis = 0,
leaf_size = leaf_size
)
# kd_tree.traverse()
# random query test:
for _ in range(100):
# generate query point:
query = np.random.rand(D)
# 01 -- knn: brute-force as baseline:
dists = np.linalg.norm(point_cloud - query, axis=1)
sorting_idx = np.argsort(dists)
brute_force_result = {i for i in sorting_idx[:k]}
knn_result_set = KNNResultSet(capacity=k)
kd_tree.knn_search(query, knn_result_set)
knn_result = {i.index for i in knn_result_set.dist_index_list}
assert len(brute_force_result - knn_result) == 0
# 02 -- rnn: brute-force as baseline:
dists = np.linalg.norm(point_cloud - query, axis=1)
brute_force_result = {i for i, d in enumerate(dists) if d <= r}
rnn_result_set = RadiusNNResultSet(radius = r)
kd_tree.rnn_search(query, rnn_result_set)
rnn_result = {i.index for i in rnn_result_set.dist_index_list}
assert len(brute_force_result - rnn_result) == 0
print('[KDTree kNN & RNN Random Query Test]: Successful')
def main():
# configuration
db_size = 100
k = 5
radius = 2.0
data = np.random.permutation(db_size).tolist()
root = None
for i, point in enumerate(data):
root = insert(root, point, i)
query_key = 6
result_set = KNNResultSet(capacity=k)
knn_search(root, result_set, query_key)
print('kNN Search:')
print('index - distance')
print(result_set)
result_set = RadiusNNResultSet(radius=radius)
radius_search(root, result_set, query_key)
print('Radius NN Search:')
print('index - distance')
print(result_set)
print("inorder")
inorder(root)
print("preorder")
preorder(root)
print("postorder")
postorder(root)
node = search_recursive(root, 2)
print(node)
node = search_iterative(root, 2)
print(node)
def test_search():
# Data generation
db_size = 100
k = 5 #搜寻5个点
radius = 2.0
query_key = 6 # 查找点
data = np.random.permutation(db_size).tolist() #random.permutation 随机排列一个数组
root =None
for i,point in enumerate(data): #返回序号和数据
root = insert(root,point,i)
# result_set = KNNResultSet(capacity=k)
# knn_search(root, result_set, query_key)
# print('kNN Search:')
# print('index - distance')
# print(result_set)
result_set = RadiusNNResultSet(radius=radius)
radius_search(root, result_set, query_key)
print('Radius NN Search:')
print('index - distance')
print(result_set)
