Skip to content

Instantly share code, notes, and snippets.

@mklingen mklingen/rrt.py
Created Oct 23, 2017

Embed
What would you like to do?
import numpy as np
import matplotlib.pyplot as plt
import random
world_size = (512, 512)
num_nodes = 100
max_length = 20
root = None
colors = ['red', 'green', 'blue', 'maroon', 'magenta', 'cyan', 'orange', 'gold', 'darkblue', 'darkred', 'darkgreen']
# A node is just a position and a label with children.
class Node:
def __init__(self, pos):
# Position contains (x, y) coordinates and biome (a number from 0 to 9)
self.pos = pos
self.children = []
# Squared distance from a node to a sample.
def dist(self, sample):
return pow(self.pos[0] - sample[0], 2) + pow(self.pos[1] - sample[1], 2)
# recursively find the closest node in the tree to the sample.
def get_closest(self, sample):
closest_dist = self.dist(sample)
closest = self
for child in self.children:
closest_child, dist = child.get_closest(sample)
if (dist < closest_dist):
closest = closest_child
closest_dist = dist
return (closest, closest_dist)
# grow from the node toward a random sample.
def grow(self, sample):
d = np.sqrt(self.dist(sample))
length = min(d, max_length)
diff = [(sample[0] - self.pos[0]) / d, (sample[1] - self.pos[1]) / d]
step = [self.pos[0] + diff[0] * length, self.pos[1] + diff[1] * length, sample[2]]
child = Node(step)
child_dist = child.dist(sample)
if (child_dist * 0.1 > d):
child.pos[2] = self.pos[2]
self.children.append(child)
# Plot this node with a color based on its biome with branches to all children.
def plot(self):
plt.plot(self.pos[0], self.pos[1], 'o', color=colors[self.pos[2]])
for child in self.children:
child.plot()
plt.plot([self.pos[0], child.pos[0]], [self.pos[1], child.pos[1]], color=colors[child.pos[2]])
# Print the node for debug purposes
def printme(self):
print self.pos
for child in self.children:
print str(self.pos) + "->" + str(child.pos)
# Randomly sample a node/biome within the play area.
def random_sample():
return [random.randint(0, world_size[0]), random.randint(0, world_size[1]), random.randint(0, 10)]
curr_nodes = 0
root = Node(random_sample())
# Keep sampling new nodes until the tree is of sufficient size
while (curr_nodes < num_nodes):
new = random_sample()
closest, closest_dist = root.get_closest(new)
closest.grow(new)
curr_nodes+=1
print '----'
# Plot the tree.
plt.figure()
root.plot()
plt.show()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.