def closest(self, hpoint): """ Return point closest to HRect (from Moore's eqn. 6.6) @param hpoint: HPoint """ p = HPoint(len(hpoint.coord)) for i, val in enumerate(hpoint.coord): if val <= self.min.coord[i]: p.coord[i] = self.min.coord[i] elif val >= self.max.coord[i]: p.coord[i] = self.max.coord[i] else: p.coord[i] = val return p
def nearest_k(self, key, k, alpha): """ Find KD-tree nodes whose keys are k nearest neighbors to key. Uses algorithm above. Neighbors are returned in ascending order of distance to key. @param key: [float] - key for KD-tree node @param k: int - how many neighbors to find @param alpha: float - alpha for approximate k-nn """ if k < 0 or k > self.m_count: raise Exception( "Number of neighbors (" + str(k) + ") cannot be negative or greater than number of nodes (" + str(self.m_count) + ").") if len(key) != self.m_K: raise Exception("KDTree: wrong key size!") nbrs = [None for _ in range(k)] nnl = NearestNeighborList(k) # initial call is with infinite hyper-rectangle and max distance hrect = HRect.infinite_hrect(len(key)) max_dist_sqd = float("inf") keypoint = HPoint(key) KDNode.nnbr(self.m_root, keypoint, hrect, max_dist_sqd, 0, self.m_K, nnl, alpha) for i in range(k): kdnode = nnl.remove_highest() nbrs[k - i - 1] = kdnode.value return nbrs
def search(self, key): """ Find KD-tree node whose key is identical to key. Uses algorithm translated from 352.srch.c of Gonnet & Baeza-Yates. @param key: [float] - key for KD-tree node @return: Object - object at key, or None if not found """ if len(key) != self.m_K: raise Exception("KDTree: wrong key size!") kdnode = KDNode.srch(HPoint(key), self.m_root, self.m_K) if kdnode is None: return None else: return kdnode.value
def infinite_hrect(cls, dim): """ Create an infinite rectangle. Used in the initial conditions of KDTree.nearest() """ vmin = HPoint(dim) vmax = HPoint(dim) vmin.coord = [-float("inf") for _ in range(dim)] vmax.coord = [float("inf") for _ in range(dim)] return HRect(vmin, vmax)
def delete(self, key): """ Delete a node from a KD-tree. Instead of actually deleting node and rebuilding tree, marks node as deleted. Hence, it is up to the caller to rebuild the tree as needed for efficiency. @param key: [float] - key for KD-tree node """ if len(key) != self.m_K: raise Exception("KDTree: wrong key size!") kdnode = KDNode.srch(HPoint(key), self.m_root, self.m_K) if kdnode is None: raise Exception("KDTree: key missing!") else: kdnode.deleted = True self.m_count -= 1
def insert(self, key, value): """ Insert a node in a KD-tree. Uses algorithm translated from 352.ins.c of Book{GonnetBaezaYates1991, author = {G.H. Gonnet and R. Baeza-Yates}, title = {Handbook of Algorithms and Data Structures}, publisher = {Addison-Wesley}, year = {1991} } @param key: [float] - key for KD-tree node @param value: Object - value at the key """ if len(key) != self.m_K: raise Exception("KDTree: wrong key size!") else: self.m_root = KDNode.ins(HPoint(key), value, self.m_root, 0, self.m_K) self.m_count += 1
def nnbr(cls, kdnode, target, hrect, max_dist_sqd, level, K, nnl, alpha): """ Method Nearest Neighbor from Andrew Moore's thesis. Numbered comments are direct quotes from there. Step "SDL" is added to make the algorithm work correctly. NearestNeighborList solution courtesy of Bjoern Heckel. @param kdnode: KDNode @param target: HPoint @param hrect: HRect @param max_dist_sqd: float @param level: int @param K: int @param nnl: NearestNeighborList @param alpha: float """ # 1. if kd is empty then set dist-sqd to infinity and exit. if kdnode is None: return # 2. s := split field of kd split = level % K # 3. pivot := dom-elt field of kd pivot = kdnode.key pivot_to_target = HPoint.sqrdist(pivot, target) # 4. Cut hr into to sub-hyperrectangles left-hr and right-hr. # The cut plane is through pivot and perpendicular to the s # dimension. left_hrect = hrect # optimize by not cloning right_hrect = hrect.clone() left_hrect.max.coord[split] = pivot.coord[split] right_hrect.min.coord[split] = pivot.coord[split] # 5. target-in-left := target_s <= pivot_s target_in_left = target.coord[split] < pivot.coord[split] # 6. if target-in-left then # 6.1. nearer-kd := left field of kd and nearer-hr := left-hr # 6.2. further-kd := right field of kd and further-hr := right-hr if target_in_left: nearer_kdnode = kdnode.left nearer_hrect = left_hrect further_kdnode = kdnode.right further_hrect = right_hrect # 7. if not target-in-left then # 7.1. nearer-kd := right field of kd and nearer-hr := right-hr # 7.2. further-kd := left field of kd and further-hr := left-hr else: nearer_kdnode = kdnode.right nearer_hrect = right_hrect further_kdnode = kdnode.left further_hrect = left_hrect # 8. Recursively call Nearest Neighbor with parameters # (nearer-kd, target, nearer-hr, max-dist-sqd), storing the # results in nearest and dist-sqd cls.nnbr(nearer_kdnode, target, nearer_hrect, max_dist_sqd, level + 1, K, nnl, alpha) # furthest node in acceptable set if not nnl.is_capacity_reached(): dist_sqd = float("inf") else: # furthest distance in accepted set dist_sqd = nnl.get_max_priority() # 9. max-dist-sqd := minimum of max-dist-sqd and dist-sqd if dist_sqd < max_dist_sqd: max_dist_sqd = dist_sqd # 10. A nearer point could only lie in further-kd if there were some # part of further-hr within distance sqrt(max-dist-sqd) of # target. If this is the case then # CHAD: Scale this comparison by alpha? closest = further_hrect.closest(target) if HPoint.eucdist(closest, target) < math.sqrt(max_dist_sqd) / alpha: # 10.1 if (pivot-target)^2 < dist-sqd then if pivot_to_target < dist_sqd: # 10.1.1 nearest := (pivot, range-elt field of kd) # 10.1.2 dist-sqd = (pivot-target)^2 dist_sqd = pivot_to_target # add to nnl if not kdnode.deleted: nnl.insert(kdnode, dist_sqd) # 10.1.3 max-dist-sqd = dist-sqd # max_dist_sqd = dist_sqd; if nnl.is_capacity_reached(): max_dist_sqd = nnl.get_max_priority() else: max_dist_sqd = float("inf") # 10.2 Recursively call Nearest Neighbor with parameters # (further-kd, target, further-hr, max-dist_sqd), # storing results in temp-nearest and temp-dist-sqd cls.nnbr(further_kdnode, target, further_hrect, max_dist_sqd, level + 1, K, nnl, alpha) temp_dist_sqd = nnl.get_max_priority() # 10.3 If tmp-dist-sqd < dist-sqd then if temp_dist_sqd < dist_sqd: # 10.3.1 nearest := temp_nearest and dist_sqd := temp_dist_sqd dist_sqd = temp_dist_sqd # SDL: otherwise, current point is nearest elif pivot_to_target < max_dist_sqd: #nearest = kdnode dist_sqd = pivot_to_target
def nnbr(cls, kdnode, target, hrect, max_dist_sqd, level, K, nnl, alpha): """ Method Nearest Neighbor from Andrew Moore's thesis. Numbered comments are direct quotes from there. Step "SDL" is added to make the algorithm work correctly. NearestNeighborList solution courtesy of Bjoern Heckel. @param kdnode: KDNode @param target: HPoint @param hrect: HRect @param max_dist_sqd: float @param level: int @param K: int @param nnl: NearestNeighborList @param alpha: float """ # 1. if kd is empty then set dist-sqd to infinity and exit. if kdnode is None or kdnode.key is None or target is None: return # 2. s := split field of kd split = level % K # 3. pivot := dom-elt field of kd pivot = kdnode.key; pivot_to_target = HPoint.sqrdist(pivot, target); # 4. Cut hr into to sub-hyperrectangles left-hr and right-hr. # The cut plane is through pivot and perpendicular to the s # dimension. left_hrect = hrect # optimize by not cloning right_hrect = hrect.clone() left_hrect.max.coord[split] = pivot.coord[split] right_hrect.min.coord[split] = pivot.coord[split] # 5. target-in-left := target_s <= pivot_s target_in_left = target.coord[split] < pivot.coord[split] # 6. if target-in-left then # 6.1. nearer-kd := left field of kd and nearer-hr := left-hr # 6.2. further-kd := right field of kd and further-hr := right-hr if target_in_left: nearer_kdnode = kdnode.left nearer_hrect = left_hrect further_kdnode = kdnode.right further_hrect = right_hrect # 7. if not target-in-left then # 7.1. nearer-kd := right field of kd and nearer-hr := right-hr # 7.2. further-kd := left field of kd and further-hr := left-hr else: nearer_kdnode = kdnode.right; nearer_hrect = right_hrect; further_kdnode = kdnode.left; further_hrect = left_hrect; # 8. Recursively call Nearest Neighbor with parameters # (nearer-kd, target, nearer-hr, max-dist-sqd), storing the # results in nearest and dist-sqd cls.nnbr(nearer_kdnode, target, nearer_hrect, max_dist_sqd, level + 1, K, nnl, alpha) # furthest node in acceptable set if not nnl.is_capacity_reached(): dist_sqd = float("inf") else: # furthest distance in accepted set dist_sqd = nnl.get_max_priority() # 9. max-dist-sqd := minimum of max-dist-sqd and dist-sqd if dist_sqd < max_dist_sqd: max_dist_sqd = dist_sqd # 10. A nearer point could only lie in further-kd if there were some # part of further-hr within distance sqrt(max-dist-sqd) of # target. If this is the case then # CHAD: Scale this comparison by alpha? closest = further_hrect.closest(target) if HPoint.eucdist(closest, target) < math.sqrt(max_dist_sqd) / alpha: # 10.1 if (pivot-target)^2 < dist-sqd then if pivot_to_target < dist_sqd: # 10.1.1 nearest := (pivot, range-elt field of kd) # 10.1.2 dist-sqd = (pivot-target)^2 dist_sqd = pivot_to_target # add to nnl if not kdnode.deleted: nnl.insert(kdnode, dist_sqd) # 10.1.3 max-dist-sqd = dist-sqd # max_dist_sqd = dist_sqd; if nnl.is_capacity_reached(): max_dist_sqd = nnl.get_max_priority() else: max_dist_sqd = float("inf") # 10.2 Recursively call Nearest Neighbor with parameters # (further-kd, target, further-hr, max-dist_sqd), # storing results in temp-nearest and temp-dist-sqd cls.nnbr(further_kdnode, target, further_hrect, max_dist_sqd, level + 1, K, nnl, alpha) temp_dist_sqd = nnl.get_max_priority() # 10.3 If tmp-dist-sqd < dist-sqd then if temp_dist_sqd < dist_sqd: # 10.3.1 nearest := temp_nearest and dist_sqd := temp_dist_sqd dist_sqd = temp_dist_sqd # SDL: otherwise, current point is nearest elif pivot_to_target < max_dist_sqd: #nearest = kdnode dist_sqd = pivot_to_target