# ##### BEGIN GPL LICENSE BLOCK #####
#
# SCA Tree Generator, a Blender add-on
# (c) 2013 Michel J. Anders (varkenvarken)
#
# This module is: kdtree.py
# a pure python implementation of a kdtree
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software Foundation,
# Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
#
# ##### END GPL LICENSE BLOCK #####
# <pep8 compliant>
from copy import deepcopy
class Hyperrectangle:
'''an axis aligned bounding box of arbitrary dimension'''
def __init__(self, dim, min, max):
self.dim = dim
self.min = deepcopy(min) # min and max should never point to the same instance
self.max = deepcopy(max)
def extend(self, pos):
'''adapt the hyperectangle if necessary so it will contain pos.'''
for i in range(self.dim):
if pos[i] < self.min[i]:
self.min[i] = pos[i]
elif pos[i] > self.max[i]:
self.max[i] = pos[i]
def distance_squared(self, pos):
'''return the distance squared to the nearest edge, or zero if pos lies within the hyperrectangle'''
result = 0.0
for i in range(self.dim):
if pos[i] < self.min[i]:
result += (pos[i] - self.min[i]) ** 2
elif pos[i] > self.max[i]:
result += (pos[i] - self.max[i]) ** 2
return result
def __str__(self):
return "[(%d) %s:%s]" % (int(self.dim), str(self.min), str(self.max))
class Node:
"""implements a node in a kd-tree"""
def __init__(self, pos, data=None):
self.pos = deepcopy(pos)
self.data = data
self.left = None
self.right = None
self.dim = len(pos)
self.dir = 0
self.count = 0
self.level = 0
self.rect = Hyperrectangle(self.dim, pos, pos)
def addleft(self, node):
self.left = node
self.rect.extend(node.pos)
node.level = self.level + 1
node.dir = (self.dir + 1) % self.dim
def addright(self, node):
self.right = node
self.rect.extend(node.pos)
node.level = self.level + 1
node.dir = (self.dir + 1) % self.dim
def distance_squared(self, pos):
d = self.pos - pos
return d.dot(d)
def _str(self, level):
s = ' ' * level + str(self.dir) + ' ' + str(self.pos) + ' ' + str(self.rect) + '\n'
return s + ('' if self.left is None else 'L:' + self.left._str(level + 1)) + \
('' if self.right is None else 'R:' + self.right._str(level + 1))
def __str__(self):
return self._str(0)
class Tree:
"""implements a kd-tree"""
def __init__(self, dim):
self.root = None
self.nnearest = 0 # number of nearest neighbor queries
self.count = 0 # number of nodes visited
self.level = 0 # deepest node level
def resetcounters(self):
self.nnearest = 0 # number of nearest neighbor queries
self.count = 0 # number of nodes visited
def _insert(self, node, pos, data):
if pos[node.dir] < node.pos[node.dir]:
if node.left is None:
node.addleft(Node(pos, data))
return node.left
else:
node.rect.extend(pos)
return self._insert(node.left, pos, data)
else:
if node.right is None:
node.addright(Node(pos, data))
return node.right
else:
node.rect.extend(pos)
return self._insert(node.right, pos, data)
def insert(self, pos, data):
if self.root is None:
self.root = Node(pos, data)
self.level = self.root.level
return self.root
else:
node = self._insert(self.root, pos, data)
if node.level > self.level:
self.level = node.level
return node
def _nearest(self, node, pos, checkempty, level=0):
self.count += 1
dir = node.dir
d = pos[dir] - node.pos[dir]
result = node
distsq = None
if checkempty and (node.data is None):
result = None
else:
distsq = node.distance_squared(pos)
if d <= 0:
neartree = node.left
fartree = node.right
else:
neartree = node.right
fartree = node.left
if neartree is not None:
nearnode, neardistsq = self._nearest(neartree, pos, checkempty, level + 1)
if (result is None) or (neardistsq is not None and neardistsq < distsq):
result, distsq = nearnode, neardistsq
if fartree is not None:
if (result is None) or (fartree.rect.distance_squared(pos) < distsq):
farnode, fardistsq = self._nearest(fartree, pos, checkempty, level + 1)
if (result is None) or (fardistsq is not None and fardistsq < distsq):
result, distsq = farnode, fardistsq
return result, distsq
def nearest(self, pos, checkempty=False):
self.nnearest += 1
if self.root is None:
return None, None
self.root.count = 0
node, distsq = self._nearest(self.root, pos, checkempty)
self.count += self.root.count
return node, distsq
def __str__(self):
return str(self.root)
if __name__ == "__main__":
class vector(list):
def __init__(self, *args):
super().__init__([float(a) for a in args])
def __str__(self):
return "<%.1f %.1f %.1f>" % tuple(self[0:3])
def __sub__(self, other):
return vector(self[0] - other[0], self[1] - other[1], self[2] - other[2])
def __add__(self, other):
return vector(self[0] + other[0], self[1] + other[1], self[2] + other[2])
def __mul__(self, other):
s = sum(self[i] * other[i] for i in (0, 1, 2))
# print("ds",s,self,other,[self[i]*other[i] for i in (0,1,2)])
return s
def dot(self, other):
return sum(self[k] * other[k] for k in (0, 1, 2))
from random import random, seed, shuffle
from time import time
import unittest
class TestVector(unittest.TestCase):
def test_ops(self):
v1 = vector(1, 0, 0)
v2 = vector(2, 1, 0)
self.assertAlmostEqual(v1 * v2, 2.0)
self.assertAlmostEqual(v1.dot(v2), 2.0)
v3 = vector(-1, -1, 0)
self.assertListEqual(v1 - v2, v3)
v4 = vector(3, 1, 0)
self.assertListEqual(v1 + v2, v4)
class TestHyperRectangle(unittest.TestCase):
def setUp(self):
self.left = vector(0, 0, 0)
self.right = vector(1, 1, 1)
self.left1 = vector(-1, 0, 0)
self.left2 = vector(0, -1, 0)
self.left3 = vector(0, 0, -1)
self.right1 = vector(2, 0, 0)
self.right2 = vector(0, 2, 0)
self.right3 = vector(0, 0, 2)
def test_constructor(self):
hr = Hyperrectangle(3, self.left, self.right)
self.assertListEqual(hr.min, self.left)
self.assertListEqual(hr.max, self.right)
def test_extend(self):
hr = Hyperrectangle(3, self.left, self.right)
hr.extend(self.left1)
self.assertListEqual(hr.min, [-1, 0, 0])
self.assertListEqual(hr.max, [1, 1, 1])
hr.extend(self.left2)
self.assertListEqual(hr.min, [-1, -1, 0])
self.assertListEqual(hr.max, [1, 1, 1])
hr.extend(self.left3)
self.assertListEqual(hr.min, [-1, -1, -1])
self.assertListEqual(hr.max, [1, 1, 1])
hr.extend(self.right1)
self.assertListEqual(hr.min, [-1, -1, -1])
self.assertListEqual(hr.max, [2, 1, 1])
hr.extend(self.right2)
self.assertListEqual(hr.min, [-1, -1, -1])
self.assertListEqual(hr.max, [2, 2, 1])
hr.extend(self.right3)
self.assertListEqual(hr.min, [-1, -1, -1])
self.assertListEqual(hr.max, [2, 2, 2])
def test_distance_squared(self):
hr = Hyperrectangle(3, self.left, self.right)
self.assertAlmostEqual(hr.distance_squared(vector(0.5, 0.5, 0.5)), 0.0)
self.assertAlmostEqual(hr.distance_squared(vector(0, 0, 0)), 0.0)
self.assertAlmostEqual(hr.distance_squared(vector(-1, 0, 0)), 1.0)
self.assertAlmostEqual(hr.distance_squared(vector(2, 0, 0)), 1.0)
self.assertAlmostEqual(hr.distance_squared(vector(2, 2, 2)), 3.0)
self.assertAlmostEqual(hr.distance_squared(vector(0.5, 2, 2)), 2.0)
self.assertAlmostEqual(hr.distance_squared(vector(0.5, -1, -1)), 2.0)
self.assertAlmostEqual(hr.distance_squared(vector(0.5, 0.5, 2)), 1.0)
class TestTree(unittest.TestCase):
def setUp(self):
seed(42)
r = (-1, 0, 1)
self.points = [vector(k, l, m) for k in r for l in r for m in r]
d = (-0.1, 0, 0.1)
self.d = [vector(k, l, m) for k in d for l in d for m in d if (k * l * m) != 0]
self.repeats = 4
def test_simple(self):
tree = Tree(3)
p1 = vector(0, 0, 0)
p2 = vector(-1, 0, 0)
p3 = vector(-1, 1, 0)
d = vector(0.1, 0.1, 0.1)
tree.insert(p1, p1)
node, distsq = tree.nearest(p1)
self.assertListEqual(node.pos, p1)
self.assertAlmostEqual(distsq, 0.0)
node, distsq = tree.nearest(p1 + d)
self.assertListEqual(node.pos, p1)
self.assertAlmostEqual(distsq, 0.03)
tree.insert(p2, p2)
node, distsq = tree.nearest(p1)
self.assertListEqual(node.pos, p1)
self.assertAlmostEqual(distsq, 0.0)
node, distsq = tree.nearest(p1 + d)
self.assertListEqual(node.pos, p1)
self.assertAlmostEqual(distsq, 0.03)
node, distsq = tree.nearest(p2)
self.assertListEqual(node.pos, p2)
self.assertAlmostEqual(distsq, 0.0)
node, distsq = tree.nearest(p2 + d)
self.assertListEqual(node.pos, p2)
self.assertAlmostEqual(distsq, 0.03)
tree.insert(p3, p3)
node, distsq = tree.nearest(p1)
self.assertListEqual(node.pos, p1)
self.assertAlmostEqual(distsq, 0.0)
node, distsq = tree.nearest(p1 + d)
self.assertListEqual(node.pos, p1)
self.assertAlmostEqual(distsq, 0.03)
node, distsq = tree.nearest(p2)
self.assertListEqual(node.pos, p2)
self.assertAlmostEqual(distsq, 0.0)
node, distsq = tree.nearest(p2 + d)
self.assertListEqual(node.pos, p2)
self.assertAlmostEqual(distsq, 0.03)
node, distsq = tree.nearest(p3)
self.assertListEqual(node.pos, p3)
self.assertAlmostEqual(distsq, 0.0)
node, distsq = tree.nearest(p3 + d)
self.assertListEqual(node.pos, p3)
self.assertAlmostEqual(distsq, 0.03)
def test_nearest(self):
for n in range(self.repeats):
tree = Tree(3)
shuffle(self.points)
for p in self.points:
tree.insert(p, p) # data equal to position
for p in self.points:
for d in self.d:
node, distsq = tree.nearest(p + d)
s = "%s %s %s %s\n%s" % (str(p + d), str(p), str(d), str(node.pos), str(tree.root))
self.assertListEqual(node.pos, p, msg=s)
self.assertListEqual(node.data, p)
self.assertAlmostEqual(distsq, d.dot(d), msg=s)
for p in self.points:
node, distsq = tree.nearest(p)
self.assertListEqual(node.pos, p)
self.assertListEqual(node.data, p)
self.assertAlmostEqual(distsq, 0.0)
def test_nearest_empty(self):
for n in range(self.repeats):
tree = Tree(3)
shuffle(self.points)
for p in self.points:
tree.insert(p, p) # data equal to position
for p in self.points:
for d in self.d:
node, distsq = tree.nearest(p + d)
s = "%s %s %s %s\n%s" % (str(p + d), str(p), str(d), str(node.pos), str(tree.root))
self.assertListEqual(node.pos, p, msg=s)
self.assertListEqual(node.data, p)
self.assertAlmostEqual(distsq, d.dot(d), msg=s)
for p in self.points:
node, distsq = tree.nearest(p)
self.assertListEqual(node.pos, p)
self.assertListEqual(node.data, p)
self.assertAlmostEqual(distsq, 0.0)
# zeroing out a node should not affect retrieving any other node ...
node, _ = tree.nearest(self.points[-1]) # last point
node.data = None
for p in self.points[:-1]: # all but last
for d in self.d:
node, distsq = tree.nearest(p + d)
s = "%s %s %s %s\n%s" % (str(p + d), str(p), str(d), str(node.pos), str(tree.root))
self.assertListEqual(node.pos, p, msg=s)
self.assertListEqual(node.data, p)
self.assertAlmostEqual(distsq, d.dot(d), msg=s)
for p in self.points[:-1]: # all but last
node, distsq = tree.nearest(p)
self.assertListEqual(node.pos, p)
self.assertListEqual(node.data, p)
self.assertAlmostEqual(distsq, 0.0)
# ... also, we should be able to retrieve the node anyway ...
node, distsq = tree.nearest(self.points[-1])
self.assertListEqual(node.pos, self.points[-1])
self.assertIsNone(node.data)
self.assertAlmostEqual(distsq, 0.0)
# ... even for points nearby ...
for d in self.d:
node, distsq = tree.nearest(self.points[-1] + d)
self.assertEqual(node.pos, self.points[-1])
self.assertIsNone(node.data)
self.assertAlmostEqual(distsq, d.dot(d))
# ... unless we set checkempty
node, distsq = tree.nearest(self.points[-1], checkempty=True)
self.assertNotEqual(node.pos, self.points[-1])
self.assertIsNotNone(node.data)
self.assertAlmostEqual(distsq, 1.0) # on a perpendicular grid nearest neighbor is at most 1 away
def test_performance(self):
tree = Tree(3)
tsize = 1000
qsize = 1000
emptyq = 3
print("<performance test, may take several seconds>")
qpos = [vector(random(), random(), random()) for p in range(qsize)]
for p in range(tsize):
pos = vector(random(), random(), random())
tree.insert(pos, pos)
s = time()
for p in qpos:
node, distsq = tree.nearest(p)
e = time() - s
print("queries|tree size|tree height|empties|query load|query time")
print("{0:7d}|{2:9d}|{1.level:11d}| 0|{3:10.2f}|{4:10.1f}".format(
qsize, tree, tsize, float(tree.count) / qsize, e)
)
tree.resetcounters()
empty = []
for p in range(tsize * 9):
pos = vector(random(), random(), random())
tree.insert(pos, pos)
if (p % emptyq) == 1:
empty.append(pos)
s = time()
for p in qpos:
node, distsq = tree.nearest(p)
e2 = time() - s
print("{0:7d}|{2:9d}|{1.level:11d}| 0|{3:10.2f}|{4:10.1f}".format(
qsize, tree, tsize * 10, float(tree.count) / qsize, e2)
)
self.assertLess(e2, 3 * e, msg="a 10x bigger tree shouldn't take more than 3x the time to query")
for p in empty:
node, distsq = tree.nearest(p)
node.data = None
tree.resetcounters()
s = time()
for p in qpos:
node, distsq = tree.nearest(p, checkempty=True)
e3 = time() - s
print("{0:7d}|{2:9d}|{1.level:11d}|{5:7d}|{3:10.2f}|{4:10.1f}".format(
qsize, tree, tsize * 10, float(tree.count) / qsize, e3, tsize * 10 // emptyq)
)
unittest.main()