318 lines
6.9 KiB
Python
318 lines
6.9 KiB
Python
from __future__ import absolute_import
|
|
|
|
import math
|
|
import numpy
|
|
|
|
class Vector3(object):
|
|
def __init__(self, x=0.0, y=0.0, z=0.0):
|
|
self.x = x
|
|
self.y = y
|
|
self.z = z
|
|
|
|
def __copy__(self):
|
|
return Vector3(self.x, self.y, self.z)
|
|
|
|
def copy(self):
|
|
return Vector3(self.x, self.y, self.z)
|
|
|
|
def __repr__(self):
|
|
return 'V[%s, %s, %s]' % ( self.x, self.y, self.z )
|
|
|
|
def __add__(self, v):
|
|
return Vector3( self.x + v.x, self.y + v.y, self.z + v.z )
|
|
|
|
def __sub__(self, v):
|
|
return Vector3( self.x - v.x, self.y - v.y, self.z - v.z )
|
|
|
|
def __mul__(self, v):
|
|
return Vector3( self.x * v, self.y * v, self.z * v )
|
|
|
|
def __div__(self, v):
|
|
return Vector3( self.x / v, self.y / v, self.z / v )
|
|
__truediv__ = __div__
|
|
|
|
def __neg__(self):
|
|
return Vector3( - self.x, - self.y, - self.z )
|
|
|
|
def __iadd__(self, v):
|
|
self.x += v.x
|
|
self.y += v.y
|
|
self.z += v.z
|
|
return self
|
|
|
|
def __isub__(self, v):
|
|
self.x += v.x
|
|
self.y += v.y
|
|
self.z += v.z
|
|
return self
|
|
|
|
def __imul__(self, v):
|
|
self.x *= v
|
|
self.y *= v
|
|
self.z *= v
|
|
return self
|
|
|
|
def __idiv__(self, v):
|
|
self.x /= v
|
|
self.y /= v
|
|
self.z /= v
|
|
return self
|
|
|
|
def almostEqual(self, v):
|
|
return (abs(self.x - v.x) + abs(self.y - v.y) + abs(self.z - v.z)) < 0.00001
|
|
|
|
def cross(self, v):
|
|
return Vector3(self.y * v.z - self.z * v.y, -self.x * v.z + self.z * v.x, self.x * v.y - self.y * v.x)
|
|
|
|
def vsize(self):
|
|
return math.sqrt( self.x * self.x + self.y * self.y + self.z * self.z )
|
|
|
|
def normalize(self):
|
|
f = self.vsize()
|
|
if f != 0.0:
|
|
self.x /= f
|
|
self.y /= f
|
|
self.z /= f
|
|
|
|
def min(self, v):
|
|
return Vector3(min(self.x, v.x), min(self.y, v.y), min(self.z, v.z))
|
|
|
|
def max(self, v):
|
|
return Vector3(max(self.x, v.x), max(self.y, v.y), max(self.z, v.z))
|
|
|
|
class AABB(object):
|
|
def __init__(self, vMin, vMax):
|
|
self.vMin = vMin
|
|
self.vMax = vMax
|
|
self.perimeter = numpy.sum(self.vMax - self.vMin)
|
|
|
|
def combine(self, aabb):
|
|
return AABB(numpy.minimum(self.vMin, aabb.vMin), numpy.maximum(self.vMax, aabb.vMax))
|
|
|
|
def overlap(self, aabb):
|
|
if aabb.vMin[0] - self.vMax[0] > 0.0 or aabb.vMin[1] - self.vMax[1] > 0.0 or aabb.vMin[2] - self.vMax[2] > 0.0:
|
|
return False
|
|
if self.vMin[0] - aabb.vMax[0] > 0.0 or self.vMin[1] - aabb.vMax[1] > 0.0 or self.vMin[2] - aabb.vMax[2] > 0.0:
|
|
return False
|
|
return True
|
|
|
|
def __repr__(self):
|
|
return "AABB:%s - %s" % (str(self.vMin), str(self.vMax))
|
|
|
|
class _AABBNode(object):
|
|
def __init__(self, aabb):
|
|
self.child1 = None
|
|
self.child2 = None
|
|
self.parent = None
|
|
self.height = 0
|
|
self.aabb = aabb
|
|
|
|
def isLeaf(self):
|
|
return self.child1 == None
|
|
|
|
class AABBTree(object):
|
|
def __init__(self):
|
|
self.root = None
|
|
|
|
def insert(self, aabb):
|
|
newNode = _AABBNode(aabb)
|
|
if self.root == None:
|
|
self.root = newNode
|
|
return
|
|
|
|
node = self.root
|
|
while not node.isLeaf():
|
|
child1 = node.child1
|
|
child2 = node.child2
|
|
|
|
area = node.aabb.perimeter
|
|
combinedAABB = node.aabb.combine(aabb)
|
|
combinedArea = combinedAABB.perimeter
|
|
|
|
cost = 2.0 * combinedArea
|
|
inheritanceCost = 2.0 * (combinedArea - area)
|
|
|
|
if child1.isLeaf():
|
|
cost1 = aabb.combine(child1.aabb).perimeter + inheritanceCost
|
|
else:
|
|
oldArea = child1.aabb.perimeter
|
|
newArea = aabb.combine(child1.aabb).perimeter
|
|
cost1 = (newArea - oldArea) + inheritanceCost
|
|
|
|
if child2.isLeaf():
|
|
cost2 = aabb.combine(child1.aabb).perimeter + inheritanceCost
|
|
else:
|
|
oldArea = child2.aabb.perimeter
|
|
newArea = aabb.combine(child2.aabb).perimeter
|
|
cost2 = (newArea - oldArea) + inheritanceCost
|
|
|
|
if cost < cost1 and cost < cost2:
|
|
break
|
|
|
|
if cost1 < cost2:
|
|
node = child1
|
|
else:
|
|
node = child2
|
|
|
|
sibling = node
|
|
|
|
# Create a new parent.
|
|
oldParent = sibling.parent
|
|
newParent = _AABBNode(aabb.combine(sibling.aabb))
|
|
newParent.parent = oldParent
|
|
newParent.height = sibling.height + 1
|
|
|
|
if oldParent != None:
|
|
# The sibling was not the root.
|
|
if oldParent.child1 == sibling:
|
|
oldParent.child1 = newParent
|
|
else:
|
|
oldParent.child2 = newParent
|
|
|
|
newParent.child1 = sibling
|
|
newParent.child2 = newNode
|
|
sibling.parent = newParent
|
|
newNode.parent = newParent
|
|
else:
|
|
# The sibling was the root.
|
|
newParent.child1 = sibling
|
|
newParent.child2 = newNode
|
|
sibling.parent = newParent
|
|
newNode.parent = newParent
|
|
self.root = newParent
|
|
|
|
# Walk back up the tree fixing heights and AABBs
|
|
node = newNode.parent
|
|
while node != None:
|
|
node = self._balance(node)
|
|
|
|
child1 = node.child1
|
|
child2 = node.child2
|
|
|
|
node.height = 1 + max(child1.height, child2.height)
|
|
node.aabb = child1.aabb.combine(child2.aabb)
|
|
|
|
node = node.parent
|
|
|
|
def _balance(self, A):
|
|
if A.isLeaf() or A.height < 2:
|
|
return A
|
|
|
|
B = A.child1
|
|
C = A.child2
|
|
|
|
balance = C.height - B.height
|
|
|
|
# Rotate C up
|
|
if balance > 1:
|
|
F = C.child1;
|
|
G = C.child2;
|
|
|
|
# Swap A and C
|
|
C.child1 = A;
|
|
C.parent = A.parent;
|
|
A.parent = C;
|
|
|
|
# A's old parent should point to C
|
|
if C.parent != None:
|
|
if C.parent.child1 == A:
|
|
C.parent.child1 = C
|
|
else:
|
|
C.parent.child2 = C
|
|
else:
|
|
self.root = C
|
|
|
|
# Rotate
|
|
if F.height > G.height:
|
|
C.child2 = F
|
|
A.child2 = G
|
|
G.parent = A
|
|
A.aabb = B.aabb.combine(G.aabb)
|
|
C.aabb = A.aabb.combine(F.aabb)
|
|
|
|
A.height = 1 + max(B.height, G.height)
|
|
C.height = 1 + max(A.height, F.height)
|
|
else:
|
|
C.child2 = G
|
|
A.child2 = F
|
|
F.parent = A
|
|
A.aabb = B.aabb.combine(F.aabb)
|
|
C.aabb = A.aabb.combine(G.aabb)
|
|
|
|
A.height = 1 + max(B.height, F.height)
|
|
C.height = 1 + max(A.height, G.height)
|
|
|
|
return C;
|
|
|
|
# Rotate B up
|
|
if balance < -1:
|
|
D = B.child1
|
|
E = B.child2
|
|
|
|
# Swap A and B
|
|
B.child1 = A
|
|
B.parent = A.parent
|
|
A.parent = B
|
|
|
|
# A's old parent should point to B
|
|
if B.parent != None:
|
|
if B.parent.child1 == A:
|
|
B.parent.child1 = B
|
|
else:
|
|
B.parent.child2 = B
|
|
else:
|
|
self.root = B
|
|
|
|
# Rotate
|
|
if D.height > E.height:
|
|
B.child2 = D
|
|
A.child1 = E
|
|
E.parent = A
|
|
A.aabb = C.aabb.combine(E.aabb)
|
|
B.aabb = A.aabb.combine(D.aabb)
|
|
|
|
A.height = 1 + max(C.height, E.height)
|
|
B.height = 1 + max(A.height, D.height)
|
|
else:
|
|
B.child2 = E
|
|
A.child1 = D
|
|
D.parent = A
|
|
A.aabb = C.aabb.combine(D.aabb)
|
|
B.aabb = A.aabb.combine(E.aabb)
|
|
|
|
A.height = 1 + max(C.height, D.height)
|
|
B.height = 1 + max(A.height, E.height)
|
|
|
|
return B
|
|
|
|
return A
|
|
|
|
def query(self, aabb):
|
|
resultList = []
|
|
if self.root != None:
|
|
self._query(self.root, aabb, resultList)
|
|
return resultList
|
|
|
|
def _query(self, node, aabb, resultList):
|
|
if not aabb.overlap(node.aabb):
|
|
return
|
|
if node.isLeaf():
|
|
resultList.append(node.aabb)
|
|
else:
|
|
self._query(node.child1, aabb, resultList)
|
|
self._query(node.child2, aabb, resultList)
|
|
|
|
def __repr__(self):
|
|
s = "AABBTree:\n"
|
|
s += str(self.root.aabb)
|
|
return s
|
|
|
|
if __name__ == '__main__':
|
|
tree = AABBTree()
|
|
tree.insert(AABB(Vector3(0,0,0), Vector3(0,0,0)))
|
|
tree.insert(AABB(Vector3(1,1,1), Vector3(1,1,1)))
|
|
tree.insert(AABB(Vector3(0.5,0.5,0.5), Vector3(0.5,0.5,0.5)))
|
|
print(tree)
|
|
print(tree.query(AABB(Vector3(0,0,0), Vector3(0,0,0))))
|
|
|