Skip to content

Instantly share code, notes, and snippets.

@mklingen mklingen/
Created Oct 23, 2017

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]
# 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:
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)
print '----'
# Plot the tree.
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.