Created
October 23, 2017 16:35
-
-
Save mklingen/13dceaf4a2c44648f16b0f46499e7e23 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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