rrt.py

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()