Last active
August 29, 2015 14:24
-
-
Save kms70847/1e3ba6af6437813eb8c0 to your computer and use it in GitHub Desktop.
KD Tree visualization
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
#animation.py | |
#prerequisites: ImageMagick (http://www.imagemagick.org/script/index.php) | |
import itertools | |
import os | |
import os.path | |
import subprocess | |
import shutil | |
import math | |
def generate_unused_folder_name(): | |
base = "temp_{}" | |
for i in itertools.count(): | |
name = base.format(i) | |
if not os.path.exists(name): | |
return name | |
def make_gif(imgs, **kwargs): | |
"""creates a gif in the current directory, composed of the given images. | |
parameters: | |
- imgs: a list of PIL image objects. | |
optional parameters: | |
- name: the name of the gif file. | |
- default: "output.gif" | |
- delay: the number of 'ticks' between frames. 1 tick == 10 ms == 0.01 seconds. anything smaller than 2 will cause display problems. | |
- default: 2 | |
- delete_temp_files: True if the temporary directory containing each individual frame should be deleted, False otherwise. | |
- default: True | |
""" | |
name = kwargs.get("name", "output.gif") | |
delay = kwargs.get("delay", 2) | |
dir_name = generate_unused_folder_name() | |
#create directory and move into it | |
os.mkdir(dir_name) | |
os.chdir(dir_name) | |
#create images. Use leading zeroes to ensure lexicographic order. | |
num_digits = max(1, int(math.log(len(imgs)))) | |
for i, img in enumerate(imgs): | |
img.save("img_{:0{padding}}.png".format(i, padding=num_digits)) | |
#create gif | |
#cmd = "imgconvert -delay {} img_*.png -layers optimize output.gif".format(delay) | |
cmd = ["imgconvert", "-delay", str(delay), "img_*.png", "-layers", "optimize", "output.gif"] | |
subprocess.call(cmd) | |
#move gif out of temp directory | |
shutil.copyfile("output.gif", "../{}".format(name)) | |
#return to original directory, and delete temp | |
os.chdir("..") | |
if kwargs.get("delete_temp_files", True): | |
shutil.rmtree(dir_name) |
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 math | |
class Point(object): | |
def __init__(self, *args, **kargs): | |
self.num_dimensions = kargs.get("num_dimensions", len(args)) | |
self.coords = [0 for i in range(self.num_dimensions)] | |
for i in range(len(args)): | |
self.coords[i] = args[i] | |
"""Gives the distance from this point to the origin.""" | |
def magnitude(self): | |
return math.sqrt(sum(c*c for c in self.coords)) | |
""" | |
Gives the angle of this point above the x axis. | |
Measured in radians. | |
Ranges between -pi and pi. | |
""" | |
def angle(self): | |
assert self.num_dimensions == 2 | |
assert self.x != 0 or self.y != 0 | |
return math.atan2(self.y,self.x) | |
def tuple(self): | |
return tuple(self.coords) | |
def map(self, func): | |
new_coords = [func(a) for a in self.coords] | |
return Point(*new_coords) | |
def _applyVectorFunc(self, other, f): | |
assert self.num_dimensions == other.num_dimensions | |
new_coords = [f(a,b) for a,b in zip(self.coords, other.coords)] | |
return Point(*new_coords) | |
def _applyScalarFunc(self, val, f): | |
return self.map(lambda a: f(a,val)) | |
""" | |
new_coords = [f(a, val) for a in self.coords] | |
return Point(*new_coords) | |
""" | |
def normalized(self): | |
return self.__div__(self.magnitude()) | |
def __add__(self, other): | |
return self._applyVectorFunc(other, lambda a,b: a+b) | |
def __sub__(self, other): | |
return self._applyVectorFunc(other, lambda a,b: a-b) | |
def __mul__(self, a): | |
return self._applyScalarFunc(a, lambda b,c: b*c) | |
def __div__(self, a): | |
return self._applyScalarFunc(a, lambda b,c: b/c) | |
def __eq__(self, other): | |
try: | |
return self.num_dimensions == other.num_dimensions and self.coords == other.coords | |
except: | |
return False | |
def __ne__(self, other): | |
return not self.__eq__(other) | |
#simple comparator for sorting purposes | |
def __lt__(self, other): | |
if not isinstance(other, Point): raise Exception("can only compare points to points") | |
return self.coords < other.coords | |
def __getattr__(self, name): | |
if name == "x": return self.coords[0] | |
if name == "y": return self.coords[1] | |
if name == "z": return self.coords[2] | |
raise AttributeError(name) | |
def __setattr__(self, name, value): | |
if name == "x": self.coords[0] = value | |
elif name == "y": self.coords[1] = value | |
elif name == "z": self.coords[2] = value | |
else: object.__setattr__(self, name, value) | |
def copy(self): | |
return Point(*self.coords[:]) | |
def __hash__(self): | |
return hash(self.tuple()) | |
def __repr__(self): | |
use_nice_floats = False | |
if use_nice_floats: | |
return "Point(" + ", ".join("%.1f" % c for c in self.coords) + ")" | |
else: | |
return "Point(" + ", ".join(str(c) for c in self.coords) + ")" | |
#old version that is always three dimensions | |
""" | |
class Point: | |
def __init__(self, x=0, y=0, z=0): | |
self.x = x | |
self.y = y | |
self.z = z | |
def magnitude(self): | |
return math.sqrt(self.x*self.x + self.y*self.y + self.z*self.z) | |
def tuple(self): | |
return (self.x, self.y, self.z) | |
def _applyVectorFunc(self, other, f): | |
return Point(f(self.x, other.x), f(self.y, other.y), f(self.z, other.z)) | |
def _applyScalarFunc(self, a, f): | |
return self._applyVectorFunc(Point(a,a,a), f) | |
def __add__(self, other): | |
return self._applyVectorFunc(other, lambda a,b: a+b) | |
def __sub__(self, other): | |
return self._applyVectorFunc(other, lambda a,b: a-b) | |
def __mul__(self, a): | |
return self._applyScalarFunc(a, lambda b,c: b*c) | |
def __div__(self, a): | |
return self._applyScalarFunc(a, lambda b,c: b/c) | |
def __eq__(self, other): | |
try: | |
return self.x == other.x and self.y == other.y and self.z == other.z | |
except: | |
return False | |
def copy(self): | |
return Point(self.x, self.y, self.z) | |
def __hash__(self): | |
return hash(self.tuple()) | |
def __repr__(self): | |
#return "Point({}, {}, {})".format(self.x, self.y, self.z) | |
return "Point({}, {})".format(self.x, self.y) | |
""" | |
def distance(a,b): | |
return (a-b).magnitude() | |
def dot_product(a,b): | |
return sum(a._applyVectorFunc(b, lambda x,y: x*y).coords) | |
def cross_product(u,v): | |
#todo: support more dimensions than three, if it is possible to do so | |
x = u.y*v.z - u.z*v.y | |
y = u.z*v.x - u.x*v.z | |
z = u.x*v.y - u.y*v.x | |
return Point(x,y,z) | |
def midpoint(a,b, frac=0.5): | |
return a*(1-frac) + b*frac |
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
#kdtree.py | |
def bisect(rect, p): | |
if rect.height > rect.width: | |
return [rect.copy(bottom=p.y), rect.copy(top=p.y)] | |
else: | |
return [rect.copy(left=p.x), rect.copy(right=p.x)] | |
class KD_Tree: | |
def __init__(self, rect): | |
self.rect = rect | |
self.p = None | |
self.children = [] | |
def add(self, p): | |
assert self.rect.contains(p) | |
if self.p is None: | |
self.p = p | |
else: | |
if not self.children: | |
self.children = [KD_Tree(rect) for rect in bisect(self.rect, self.p)] | |
for child in self.children: | |
if child.rect.contains(p): | |
child.add(p) | |
return | |
raise Exception("could not find branch for {}".format(p)) | |
@staticmethod | |
def populated(rect, points): | |
ret = KD_Tree(rect) | |
for p in points: | |
ret.add(p) | |
return ret | |
def iter(self): | |
if self.p is not None: | |
yield self.p | |
for child in self.children: | |
for item in child.iter(): | |
yield item | |
def iter_in_range(self, rect): | |
if self.p is None: | |
return | |
if rect.contains(self.p): | |
yield self.p | |
for child in self.children: | |
if rect.overlaps(child.rect): | |
for item in child.iter_in_range(rect): | |
yield item |
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
#main.py | |
import math | |
import random | |
from kdtree import KD_Tree | |
from rect import Rect | |
from rendering import render | |
from animation import make_gif | |
from geometry import distance | |
def linear_interpolate(start, end, t): | |
return start * (1-t) + end * t | |
def sinusoidal_interpolate(start, end, t): | |
t = math.sin(t * math.pi / 2) | |
return linear_interpolate(start, end, t) | |
def create_path(rand_p_func, frames): | |
def first(cond, iterable): | |
for item in iterable: | |
if cond(item): | |
return item | |
raise Exception("no items in iterable satisfied condition") | |
num_waypoints = random.randint(2, 5) | |
waypoints = [rand_p_func() for i in range(num_waypoints)] | |
#close the loop by making the last waypoint the same as the first | |
waypoints.append(waypoints[0]) | |
path_lengths = [distance(a,b) for a,b in zip(waypoints, waypoints[1:])] | |
total_length = sum(path_lengths) | |
waypoint_times = [sum(path_lengths[:i])/total_length for i in range(len(path_lengths))] + [1] | |
#waypoint_times = [0] + sorted([random.random() for i in range(num_waypoints-2)]) + [1] | |
assert len(waypoints) == len(waypoint_times) | |
for i in range(frames): | |
t = float(i) / (frames-1) | |
waypoint_idx = first(lambda idx: waypoint_times[idx+1] >= t, range(len(waypoint_times)-1)) | |
start_point, end_point = waypoints[waypoint_idx], waypoints[waypoint_idx+1] | |
start_t = waypoint_times[waypoint_idx] | |
end_t = waypoint_times[waypoint_idx+1] | |
assert start_t <= t <= end_t | |
leg_length = end_t - start_t | |
leg_t = float(t - start_t) / leg_length | |
yield sinusoidal_interpolate(start_point, end_point, leg_t) | |
def transpose(matrix): | |
return zip(*matrix) | |
def main(): | |
field = Rect(0,0,10,10) | |
frames = 200 | |
size = 200 | |
points = 50 | |
print "generating paths..." | |
paths = [create_path(field.random_contained_point, frames) for i in range(points)] | |
print "done." | |
point_frames = transpose(paths) | |
print "rendering..." | |
images = [render(KD_Tree.populated(field, points), size) for points in point_frames] | |
print "done. Composing..." | |
make_gif(images, delay=4, delete_temp_files=False) | |
print "done." | |
main() |
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
#rect.py | |
from geometry import Point | |
import random | |
class Rect: | |
def __init__(self, *args): | |
if len(args) == 4: #scalars | |
self.left, self.top, self.right, self.bottom = args | |
elif len(args) == 2: #Points | |
self.__init__(*(args[0].tuple() + args[1].tuple())) | |
else: | |
raise Exception("Need 4 scalars or 2 points, got {}".format(len(args))) | |
@property | |
def height(self): | |
return self.bottom - self.top | |
@property | |
def width(self): | |
return self.right - self.left | |
def contains(self, p): | |
return self.left <= p.x < self.right and self.top <= p.y <= self.bottom | |
def overlaps(self, other): | |
#returns true if line segment AB shares a common segment with line segment CD. | |
def intersects(a,b,c,d): | |
#returns true if x lies between a and b. | |
def lies_between(x,a,b): | |
return a <= x < b | |
return lies_between(a, c, d) or lies_between(b, c, d) or lies_between(c, a, b) or lies_between(d, a, b) | |
return intersects(self.left, self.right, other.left, other.right) and intersects(self.top, self.bottom, other.top, other.bottom) | |
def copy(self, **d): | |
return Rect(*[d.get(attr, getattr(self, attr)) for attr in "left top right bottom".split()]) | |
def corners(self): | |
return (Point(self.left, self.top), Point(self.right, self.bottom)) | |
def tuple(self): | |
return (self.left, self.top, self.right, self.bottom) | |
def map(self, f): | |
new_corners = [p.map(f) for p in self.corners()] | |
return Rect(*new_corners) | |
def pmap(self, f): | |
new_corners = [f(p) for p in self.corners()] | |
return Rect(*new_corners) | |
def random_contained_point(self): | |
return Point(random.uniform(self.left, self.right), random.uniform(self.top, self.bottom)) |
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
#rendering.py | |
from geometry import Point | |
from rect import Rect | |
from kdtree import KD_Tree | |
from PIL import Image, ImageDraw | |
def render(tree, target_size=500, margin=10): | |
scale = target_size / float(max(tree.rect.width, tree.rect.height)) | |
width = int(tree.rect.width * scale) + margin*2 | |
height = int(tree.rect.height * scale) + margin*2 | |
dx = tree.rect.left | |
dy = tree.rect.top | |
delta = Point(dx,dy) | |
img = Image.new("RGB", (width, height), "white") | |
draw = ImageDraw.Draw(img) | |
def logical_to_screen(p): | |
return (((p - delta) * scale) + Point(margin, margin)).map(int) | |
def dot(p, radius, **options): | |
a = p - Point(radius, radius) | |
b = p + Point(radius, radius) | |
draw.chord(a.tuple() + b.tuple(), 0, 359, **options) | |
def box(rect, **options): | |
draw.rectangle(rect.tuple(), **options) | |
def render_node(node): | |
if node.p is not None: | |
dot(logical_to_screen(node.p), 2, fill="black") | |
for child in node.children: | |
box(child.rect.pmap(logical_to_screen), outline="black") | |
render_node(child) | |
render_node(tree) | |
radius = 1 | |
search_vector = Point(radius, radius) | |
for p in tree.iter(): | |
area = Rect(p - search_vector, p + search_vector) | |
for neighbor in tree.iter_in_range(area): | |
draw.line(logical_to_screen(p).tuple() + logical_to_screen(neighbor).tuple(), fill="red") | |
return img |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment