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)
