class SimpleGraph: def __init__(self, dim, capacity=100000): self._edges = collections.defaultdict(list) self._kd = KDTree(dim, capacity) self.start_id = None self.target_id = None def __len__(self): return len(self._kd) def insert_new_node(self, point): node_id = self._kd.insert(point) return node_id def add_edge(self, node_id, neighbor_id): self._edges[node_id].append(neighbor_id) self._edges[neighbor_id].append(node_id) def get_parent(self, child_id): return self._edges[child_id] def get_point(self, node_id): return self._kd.get_node(node_id).point def get_nearest_node(self, point): return self._kd.find_nearest_point(point) def get_neighbor_within_radius(self, point, radius): """ Return a list of node_id within the radius """ return self._kd.find_points_within_radius(point, radius)
class SimpleTree: def __init__(self, dim): self._parents_map = {} self._kd = KDTree(dim) def __len__(self): return len(self._kd) def insert_new_node(self, point, parent=None): node_id = self._kd.insert(point) self._parents_map[node_id] = parent return node_id def get_parent(self, child_id): return self._parents_map[child_id] def get_point(self, node_id): return self._kd.get_node(node_id).point def get_nearest_node(self, point): return self._kd.find_nearest_point(point) def construct_path_to_root(self, leaf_node_id): path = [] node_id = leaf_node_id while node_id is not None: path.append(self.get_point(node_id)) node_id = self.get_parent(node_id) return path def get_num_nodes(self): return len(self._parents_map)
def test_size(self): tree = KDTree(1) assert tree.size == 0 tree.insert('B') assert tree.size == 1 tree.insert('A') assert tree.size == 2 tree.insert('C') assert tree.size == 3
class SimpleTree: def __init__(self, dim): self._parents_map = {} self._kd = KDTree(dim) def insert_new_node(self, point, parent=None): node_id = self._kd.insert(point) self._parents_map[node_id] = parent return node_id def get_parent(self, child_id): return self._parents_map[child_id] def get_point(self, node_id): return self._kd.get_node(node_id).point def get_nearest_node(self, point): return self._kd.find_nearest_point(point)
def test_kdtree(self): tree = KDTree() tree.insert((2, 6)) tree.insert((3, 1)) tree.insert((8, 7)) tree.insert((10, 2)) tree.insert((13, 3)) assert tree.get_nearest((9, 4)) == ((10, 2), 2.23606797749979) assert tree.get_nearest((4, 1.5))[0] == (3, 1) assert tree.get_nearest((7, 8))[0] == (8, 7) assert tree.get_nearest((11, 1))[0] == (10, 2) assert tree.get_nearest((13, 3))[0] == (13, 3)
class TestKDTree(unittest.TestCase): def setUp(self): self.dim = 3 self.count = 20 points = np.random.randint(low=0, high=20, size=(self.count, self.dim)).tolist() self.tree = KDTree(self.dim) for p in points: self.tree.insert(p) np.random.shuffle(points) self.points = points def test_contain(self): for p in self.points: self.assertTrue(self.tree.contain(p)) random_points = np.random.randint(low=30, size=(self.count, self.dim)).tolist() for p in random_points: self.assertFalse(self.tree.contain(p)) def test_find_min(self): for d in range(self.dim): sorted_at_dim = sorted(self.points, key=operator.itemgetter(d)) min_at_dim = sorted_at_dim[0][d] candidate_min_at_dim = list( filter(lambda point: point[d] == min_at_dim, sorted_at_dim)) self.assertIn(self.tree.find_min(d, 0, self.tree.root), candidate_min_at_dim) def test_delete(self): for p in self.points: self.tree.delete(p) self.assertFalse(self.tree.contain(p)) def test_nearest_neighbor(self): for p in self.points: self.assertListEqual(self.tree.find_nearest(p), p) random_points = np.random.randint(low=0, high=20, size=(20, self.dim)) for p in random_points: nearest_point = self.tree.find_nearest(p) points = np.array(self.points) self.assertEqual(min(linalg.norm(points - p, axis=1)), linalg.norm(nearest_point - p)) def test_k_nearest_neighbor(self): for p in self.points: self.assertListEqual(self.tree.find_k_nearest(p, 1), [p]) random_points = np.random.randint(low=0, high=20, size=(20, self.dim)) for p in random_points: nearest_points = self.tree.find_k_nearest(p, 5) nearest_points = np.array(nearest_points) points = np.unique(self.points, axis=0) n_dists = linalg.norm(nearest_points - p, axis=1) dists = np.sort(linalg.norm(points - p, axis=1))[:5] self.assertListEqual(n_dists.tolist(), dists.tolist())
class RRT: """Naive 2-D RRT that doesn't incorporate dynamics into the system""" def __init__(self, start, goal): self.start = start self.goal = goal # distance within which tree has 'reached' the goal self.eps = 5 # extension distance self.ext = 10 self.tree = KDTree(len(self.start), [self.start]) # edges self.connections = dict() # run the algorithm and return a path from start to goal def run(self, world): i = 0 while i < 2500: i += 1 s1 = self.sample(world) s = self.tree.kNearestNeighbors(s1, 1, [])[0][0] s1e = self.extend(s, s1, world) if not s1e: continue else: self.tree.insert(s1e) try: self.connections[s].add(s1e) except: self.connections[s] = set([s1e]) try: self.connections[s1e].add(s) except: self.connections[s1e] = set([s]) if dist(s1e, self.goal) < self.eps: return AStarSearch(self.connections, self.start, self.goal, self.eps) return [] # sample world uniformly def sample(self, world): if random.random() < 0.05: return self.goal point = (random.random() * world.displayWidth, random.random() * world.displayHeight) for obstacle in world.obstacleBeliefs: # make sure the point is at least carsize / 2 away from the nearest obstacle if obstacle.colliderect( (point[0] - world.cars[0].size[0] / 2., point[1] - world.cars[0].size[1] / 2., world.cars[0].size[0], world.cars[0].size[1])): return self.sample(world) return point def extend(self, p1, p2, world): pointDir = (p2[0] - p1[0], p2[1] - p1[1]) magDir = dist(pointDir, (0, 0)) newPoint = (p1[0] + self.ext * pointDir[0] / magDir, p1[1] + self.ext * pointDir[1] / magDir) for obstacle in world.obstacleBeliefs: if obstacle.colliderect( (newPoint[0] - world.cars[0].size[0] / 2., newPoint[1] - world.cars[0].size[1] / 2., world.cars[0].size[0], world.cars[0].size[1])): return None return newPoint
class PRM: def __init__(self, world): self.points = KDTree(2, [world.kdtreeStart]) # number of samples in the prm self.size = 2500 # number of connections per sample self.connsPerSample = 6 self.carSize = world.carSize self.getPoints(world) # edges of the prm self.connections = self.getConnections(world) # uniformly sample floating point coordinates from the world def sample(self, world): point = (random.random() * world.displayWidth, random.random() * world.displayHeight) for obstacle in world.obstacleBeliefs: if obstacle.colliderect((point[0] - self.carSize[0] / 2., point[1] - self.carSize[1] / 2., self.carSize[0], self.carSize[1])): return self.sample(world) return point # uniformly sample integer coordinates from the world def sampleInt(self, world): point = (random.randint(0, world.displayWidth), random.randint(0, world.displayHeight)) for obstacle in world.obstacleBeliefs: if obstacle.colliderect((point[0] - self.carSize[0] / 2., point[1] - self.carSize[1] / 2., self.carSize[0], self.carSize[1])): return self.sampleInt(world) return point # fill the KD tree def getPoints(self, world): for _ in range(self.size - 1): self.points.insert(self.sample(world)) # ignore this function def getPaths(self, world): paths = Paths() for p1 in self.connections: for p2 in self.connections[p1]: print p1, p2 if p1 == p2: continue try: _ = paths[p1, p2] except: paths[p1, p2] = RRT(p1, p2).run(world) return paths # generate a path (a list of points along the prm) from p1 to p2 def getPath(self, p1, p2, world): if not self.points.contains(p1): self.insertConnection(p1, world) if not self.points.contains(p2): self.insertConnection(p2, world) path = AStarSearch(self.connections, p1, p2, 5, self.size) # use RRT if no path exists along prm if not path: path = RRT(p1, p2).run(world) return path # insert a point into the prm def insertConnection(self, point, world): nns = self.points.kNearestNeighbors(point, self.connsPerSample, []) self.connections[point] = [] for p in nns: collided = False for obstacle in world.obstacleBeliefs: if obstacle.collideline((point[0], point[1], p[0][0], p[0][1]), 10): collided = True if not collided: self.connections[point].append(p[0]) self.connections[p[0]].append(point) self.points.insert(point) # form the prm from the set of sampled points def getConnections(self, world): connections = dict() pointsList = self.points.list() for point in pointsList: nns = self.points.kNearestNeighbors(point, self.connsPerSample + 1, []) connections[point] = [] for p in nns: collided = False for obstacle in world.obstacleBeliefs: if obstacle.collideline( (point[0], point[1], p[0][0], p[0][1]), 10): collided = True if not collided: connections[point].append(p[0]) connections[point].remove(point) return connections
class KDFinder(Finder): def __init__(self, p_set, measurer=SquareDistMeasurer(2)): super(KDFinder, self).__init__(None, measurer) assert p_set is not None self.count = len(p_set) assert self.count > 0 self.tree = KDTree(None, 'value', measurer.k) for i in range(len(p_set)): self.tree.insert(Element(p_set[i], i)) self.debug = False def find_closest_m(self, point, m): self.begin(m) self.search(self.tree.get_root(), point, m) if self.debug: self.do_all(point, 'F', 'Final Result', True, True) return [[ self.pq.peek().value, self.pq.peek().index, self.pq.pop().current_dis ] for i in range(m)] def begin(self, m): super(KDFinder, self).begin(m) self._b_upper = [float('inf')] * self.measurer.k self._b_lower = [float('-inf')] * self.measurer.k if self.debug: self.counter = 0 self.visited = [] def check_node(self, node, point, m): dis = self.measurer.measure(point, node.value) if self.pq.count() < m or dis < self.pq.peek().current_dis: node.obj.current_dis = dis self.add_candidate(node.obj) if self.debug: msg = 'Entering Node' self.visited.append(node.value) if node.left is not None: msg += '~HasLeft' if node.right is not None: msg += '~HasRight' else: msg += '~Discriminator: ' + CLABELS[self.tree.d_order[ node.discriminator]] self.do_all(point, 'A', msg) def search(self, node, point, m): bol = False self.check_node(node, point, m) if type(node) is Node: dim = self.tree.d_order[node.discriminator] p = node.value[dim] if point[dim] < p: if node.left is not None: # Search left subtree temp = self._b_upper[dim] self._b_upper[dim] = p if self.search(node.left, point, m): return True self._b_upper[dim] = temp # Backtracking if node.right is not None: temp = self._b_lower[dim] self._b_lower[dim] = p bol = self.bounds_overlap(self.pq.peek().current_dis, point) if self.debug: self.do_all(point, "B", "Bounds Overlap: " + str(bol)) if self.pq.count() < m or bol: if self.search(node.right, point, m): return True self._b_lower[dim] = temp else: if node.right is not None: # Search right subtree temp = self._b_lower[dim] self._b_lower[dim] = p if self.search(node.right, point, m): return True self._b_lower[dim] = temp # Backtracking if node.left is not None: temp = self._b_upper[dim] self._b_upper[dim] = p bol = self.bounds_overlap(self.pq.peek().current_dis, point) if self.debug: self.do_all(point, "B", "Bounds Overlap: " + str(bol)) if self.pq.count() < m or bol: if self.search(node.left, point, m): return True self._b_upper[dim] = temp # Instant termination condition while backtracking wb = not bol and self.within_bounds(self.pq.peek().current_dis, point) if self.debug: self.do_all(point, "C", "Within Bounds: " + str(wb)) return self.pq.count() == m and wb def bounds_overlap(self, r, point): s = 0 r_inv = self.measurer.F_inv(r) for d in range(self.measurer.k): if point[d] < self._b_lower[d]: s += self.measurer.f(point[d], self._b_lower[d]) if s > r_inv: # Same as self.measurer.F(s) > r return False elif point[d] > self._b_upper[d]: s += self.measurer.f(point[d], self._b_upper[d]) if s > r_inv: # Same as self.measurer.F(s) > r return False # If at the boundary the partial distance is zero, there's no need to alter the sum or check return True def within_bounds(self, r, point): for d in range(self.measurer.k): r_inv = self.measurer.F_inv(r) if point[d] < self._b_lower[d] or point[d] > self._b_upper[d] \ or self.measurer.f(point[d], self._b_lower[d]) < r_inv \ or self.measurer.f(point[d], self._b_upper[d]) < r_inv: return False return True # Debugging functions ################################# def setup_plot(self, xlim, ylim, save=False): ca = plt.gca() self.overlay_ax = plt.gcf().add_axes(ca.get_position(), frameon=False) self.xlim = xlim self.ylim = ylim self.overlay_ax.set_xlim(xlim) self.overlay_ax.set_ylim(ylim) if save: if os.path.exists(WORKDIR): import shutil shutil.rmtree(WORKDIR) os.makedirs(WORKDIR) self.save = save self.debug = True def do_all(self, point, label, msg, force_save=False, legend=False): if self.save or force_save: self.overlay_ax.clear() self.overlay_ax.set_xlim(self.xlim) self.overlay_ax.set_ylim(self.ylim) self.plot_bounds() self.plot_range(point) self.plot_visited() self.plot_found(point) self.overlay_ax.text(self.xlim[0], self.ylim[0], msg) if legend: plt.legend() self.save_plot(label) self.counter += 1 def plot_bounds(self): self.overlay_ax.axvline(x=self._b_lower[0], color='r', label='Lower Search Bound') self.overlay_ax.axvline(x=self._b_upper[0], color='g', label='Upper Search Bound') self.overlay_ax.axhline(y=self._b_lower[1], color='r') self.overlay_ax.axhline(y=self._b_upper[1], color='g') def plot_range(self, point): # NOTE: Using square distance if self.pq.count() > 0: r = math.sqrt(self.pq.peek().current_dis) c = Circle(point, r, fill=False, color='black', label='Candidates Region') self.overlay_ax.add_patch(c) def plot_visited(self): for p in self.visited: self.overlay_ax.plot(p[0], p[1], 'ro', markersize=3) if self.visited.count > 0: self.overlay_ax.plot(self.visited[-1][0], self.visited[-1][1], 'ro', markersize=5, label='Visited') def plot_found(self, p1): found = self.pq.h for element in found: p2 = element.value self.overlay_ax.plot([p1[0], p2[0]], [p1[1], p2[1]], color='orange', linewidth=1.2) def save_plot(self, label): path = os.path.join(WORKDIR, str(self.counter) + '_' + label + '.png') plt.savefig(path)