KD Tree visualization

animation.py

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

geometry.py

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

kdtree.py

#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

main.py

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

rect.py

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

rendering.py

#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