kdtree.py

# ##### BEGIN GPL LICENSE BLOCK #####
#
#  SCA Tree Generator, a Blender addon
#  (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 #####

from copy import copy, 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()