384 lines
13 KiB
Python
Executable file
384 lines
13 KiB
Python
Executable file
#!/usr/bin/env python
|
|
|
|
import geometry
|
|
from geometry import x,y
|
|
|
|
# import enum
|
|
|
|
def as_box( quadrant ):
|
|
""""Convert a quadrant of the form: ((x_min,y_min),width) to a box: ((x_min,y_min),(x_max,y_max))."""
|
|
width = quadrant[1]
|
|
minp = quadrant[0]
|
|
maxp = tuple(xy+width for xy in minp)
|
|
assert( x(minp) <= x(maxp) and y(minp) <= y(maxp) )
|
|
return (minp,maxp)
|
|
|
|
|
|
def as_rect( quadrant ):
|
|
""""Convert a quadrant of the form: ((x_min,y_min),width) to a rectangle: ((x0,y0),(x1,y1),(x2,y2),(x3,y3))."""
|
|
qx,qy = quadrant[0]
|
|
w = quadrant[1]
|
|
return [(qx,qy),(qx+w,qy),(qx+w,qy+w),(qx,qy+w)]
|
|
|
|
|
|
class QuadTree(object):
|
|
|
|
def __init__( self, points = [] ):
|
|
"""Build a quadtree on the given set of points.
|
|
|
|
Points must be an iterable containing 2-tuples of the form: (x,y)"""
|
|
|
|
# Initialize the root quadrant as the box around the points
|
|
self.root, self.quadrants = self.init( points = points )
|
|
|
|
# Each leaf of the quadtree may contains one resident point.
|
|
self.residents = { self.root: None }
|
|
|
|
# Each node of the quadtree may contains four children.
|
|
self.children = { self.root: [] }
|
|
|
|
# Status of quadrants
|
|
# class Status(enum.Enum):
|
|
class Status:
|
|
Leaf = 1
|
|
Node = 2
|
|
Empty = 3
|
|
Out = 4
|
|
self.Status = Status()
|
|
|
|
# Choose one of the two available functions for walking the tree:
|
|
# self.walk = self.recursive_walk
|
|
self.walk = self.iterative_walk
|
|
|
|
# Generate the complete tree.
|
|
self.build( points )
|
|
|
|
|
|
def init( self, quadrant = None, box = None, points = None ):
|
|
"""Initialize the root quadrant with the given quadrant, the given box or the given set of points."""
|
|
|
|
if len([k for k in (box,points,quadrant) if k]) > 1:
|
|
raise BaseException("ERROR: you should specify a box, a quadrant or points")
|
|
|
|
# Initialize the root quadrant as the given box
|
|
if box:
|
|
minp,maxp = box
|
|
width = max( x(maxp)-x(minp), y(maxp)-y(minp) )
|
|
|
|
# Initialize the root quadrant as the box around the points
|
|
elif points:
|
|
minp,maxp = geometry.box( points )
|
|
width = max( x(maxp)-x(minp), y(maxp)-y(minp) )
|
|
|
|
# Initialize the root quadrant as the given origin point and width
|
|
elif quadrant:
|
|
minp = quadrant[0]
|
|
width = quadrant[1]
|
|
|
|
assert( x(minp) <= x(minp)+width and y(minp) <= y(minp)+width )
|
|
|
|
# There is always the root quadrant in the list of available ones.
|
|
root = (minp,width)
|
|
quadrants = [ root ]
|
|
|
|
return root,quadrants
|
|
|
|
|
|
def status( self, point, quadrant ):
|
|
"""Return Status.Empty if the given point can be appended in the given quadrant."""
|
|
|
|
assert(point is not None)
|
|
assert(len(point) == 2)
|
|
assert(quadrant is not None)
|
|
assert(len(quadrant) == 2)
|
|
|
|
box = as_box( quadrant )
|
|
|
|
# if the point lies inside the given quadrant
|
|
if geometry.in_box( point, box):
|
|
if self.residents[quadrant]:
|
|
# external: a quadrant that already contains a point
|
|
assert( not self.children[quadrant] )
|
|
return self.Status.Leaf
|
|
elif self.children[quadrant]:
|
|
# internal: a quadrant that contains other quadrants
|
|
return self.Status.Node
|
|
else:
|
|
# empty: there is not point yet in this quadrant
|
|
return self.Status.Empty
|
|
else:
|
|
# point is out of the quadrant
|
|
return self.Status.Out
|
|
|
|
|
|
def split( self, quadrant ):
|
|
"""Split an existing quadrant in four children quadrants.
|
|
|
|
Move the existing resident to the correct child."""
|
|
|
|
# We cannot split a quadrant if it already have sub-quadrants
|
|
if quadrant != self.root:
|
|
assert( not self.children[quadrant] )
|
|
|
|
qx, qy = quadrant[0]
|
|
w = quadrant[1] / 2
|
|
|
|
# For each four children quadrant's origins
|
|
self.children[quadrant] = []
|
|
for origin in ( (qx,qy), (qx,qy+w), (qx+w,qy+w), (qx+w,qy) ):
|
|
# Create a child quadrant of half its width
|
|
q = (origin, w)
|
|
self.quadrants.append(q)
|
|
# Default resident to None, because we will test for this key later on.
|
|
self.residents[q] = None
|
|
|
|
# Add this new child to the current parent.
|
|
self.children[quadrant].append(q)
|
|
# This new quadrant has no child.
|
|
self.children[q] = []
|
|
|
|
# Move the resident to the related children node
|
|
p = self.residents[quadrant]
|
|
if p is not None:
|
|
# Find the suitable children quadrant
|
|
for child in self.children[quadrant]:
|
|
if self.status(p,child) == self.Status.Empty:
|
|
self.residents[child] = p
|
|
break
|
|
# Forget we had resident here
|
|
# Do not pop the key, because we have tests on it elsewhere
|
|
self.residents[quadrant] = None
|
|
|
|
|
|
def append( self, point, quadrant = None ):
|
|
"""Try to inset the given point in the existing quadtree, under the given quadrant.
|
|
|
|
The default quadrant is the root one.
|
|
Returns True if the point was appended, False if it is impossible to append it."""
|
|
|
|
# Default to the root quadrant
|
|
if not quadrant:
|
|
quadrant = self.root
|
|
assert(quadrant in self.quadrants)
|
|
|
|
# The point should not be out of the root quadrant
|
|
assert( self.status(point,self.root) != self.Status.Out )
|
|
|
|
# FIXME use a recursive walk and prune branches with the Out status.
|
|
for q in self.walk(quadrant):
|
|
status = self.status( point, q )
|
|
if status == self.Status.Leaf:
|
|
# Create sub-quadrants
|
|
self.split(q)
|
|
# Try to attach the point in children quadrants, recursively
|
|
for child in self.children[q]:
|
|
if self.append( point, child ):
|
|
return True
|
|
elif status == self.Status.Empty:
|
|
# add the point as an resident of the quadrant q
|
|
self.residents[q] = point
|
|
return True
|
|
return False
|
|
|
|
|
|
def build( self, points ):
|
|
"""Append all the given points in the quadtree."""
|
|
for p in points:
|
|
self.append(p)
|
|
assert( len(points) == len(self) )
|
|
|
|
|
|
def iterative_walk( self, at_quad = None ):
|
|
|
|
# Default to the root quadrant
|
|
if not at_quad:
|
|
at_quad = self.root
|
|
|
|
# First, consider the root quadrant
|
|
yield at_quad
|
|
|
|
# then children of the root quadrant
|
|
quads = list(self.children[at_quad])
|
|
for child in quads:
|
|
yield child
|
|
# then children of the current child
|
|
quads.extend( self.children[child] )
|
|
|
|
|
|
def recursive_walk( self, at_quad = None ):
|
|
|
|
# Default to the root quadrant
|
|
if not at_quad:
|
|
at_quad = self.root
|
|
|
|
yield at_quad
|
|
for child in self.children[at_quad]:
|
|
for q in self.recursive_walk(child):
|
|
yield q
|
|
|
|
|
|
def repr( self, quadrant = None, depth = 0 ):
|
|
"""Return a string representing the quadtree in a JSON-like format."""
|
|
|
|
# Default to the root quadrant
|
|
if not quadrant:
|
|
quadrant = self.root
|
|
|
|
head = " "*depth
|
|
r = head+"{"
|
|
quadrep = '"origin" : %s, "width" : %f' % quadrant
|
|
if self.residents[quadrant]: # external
|
|
r += ' "resident" : %s, \t%s },\n' % (self.residents[quadrant],quadrep)
|
|
elif self.children[quadrant]: # internal
|
|
r += ' "children_ids" : %s, \t%s, "children" : [\n' % (self.children[quadrant],quadrep)
|
|
for child in self.children[quadrant]:
|
|
r += self.repr(child, depth+1)
|
|
r+="%s]},\n" % head
|
|
else: # empty
|
|
r += ' "resident" : (), \t\t\t%s},\n' % (quadrep)
|
|
return r
|
|
|
|
|
|
def points( self ):
|
|
"""Return the set of points attached to the quadtree.
|
|
|
|
In a random order."""
|
|
return [p for p in self.residents.values() if p is not None]
|
|
|
|
|
|
def covers( self, this, that ):
|
|
"""Return true if the given quadrants does intersects each other."""
|
|
|
|
# Convert quadrants ((x,y),w) as box ((a,b),(c,d)).
|
|
this_box = as_box(this)
|
|
that_box = as_box(that)
|
|
|
|
# Convert boxes as list of edges.
|
|
this_segments = tuple(utils.tour(as_rect(this)))
|
|
that_segments = tuple(utils.tour(as_rect(that)))
|
|
|
|
# If at least one of the segment of "this" intersects with "that".
|
|
intersects = any( geometry.segment_intersection(s0,s1) for s0 in this_segments for s1 in that_segments )
|
|
|
|
# Transform nested list of segments in flat list of points without any duplicates.
|
|
this_points = as_rect(this)
|
|
that_points = as_rect(that)
|
|
|
|
# If all the points of "this" are inside "that".
|
|
# Note: what we would want to test here is if ALL the points are comprised,
|
|
# as the case where at least one is already tested by the "intersects" stage.
|
|
# But we use an "any" anyway, because it is sufficient in this case and
|
|
# that testing all the points takes more time.
|
|
this_in = any( geometry.in_box(p,this_box) for p in that_points )
|
|
that_in = any( geometry.in_box(p,that_box) for p in this_points )
|
|
|
|
return intersects or this_in or that_in
|
|
|
|
|
|
def query( self, query_quad, at_quad = None ):
|
|
"""Return all the points (currently attached to the quad tree) that are located within the query_quad quadrant."""
|
|
if not at_quad:
|
|
at_quad = self.root
|
|
|
|
query_box = as_box(query_quad)
|
|
|
|
# If we ask for a quadrant that intersects with the current one.
|
|
if self.covers( query_quad, at_quad ):
|
|
# If the current quadrant contains sub-quadrants.
|
|
if len(self.children[at_quad]) > 0:
|
|
# Then go explore them.
|
|
points = []
|
|
for quad in self.children[at_quad]:
|
|
points += self.query(query_quad,quad)
|
|
return points
|
|
else:
|
|
# Else, just return the point within the current quadrant.
|
|
resident = self.residents[at_quad]
|
|
if resident:
|
|
if geometry.in_box(resident,query_box):
|
|
# In a list, because we will concatenate.
|
|
return [resident]
|
|
# If there is no intersection, there is no points.
|
|
return []
|
|
|
|
|
|
# Pythonesque API:
|
|
|
|
def __getitem__( self, quadrant ):
|
|
"""Return all the points that are located within the given quadrant.
|
|
|
|
Ex.: points = quad[quad.root] # get all the points"""
|
|
return self.query(quadrant,self.root)
|
|
|
|
|
|
def __iter__(self):
|
|
"""Iterate over the attached points."""
|
|
return iter(self.points())
|
|
|
|
|
|
def __call__(self, points):
|
|
"""Append all the given points in the quadtree."""
|
|
self.build(points)
|
|
|
|
|
|
def __len__(self):
|
|
"""Return the number of points attached to the quad tree."""
|
|
return len(self.points())
|
|
|
|
|
|
def __repr__(self):
|
|
"""Return a string representing the quadtree in a JSON-like format."""
|
|
return self.repr()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
import sys
|
|
import math
|
|
import random
|
|
|
|
import utils
|
|
import uberplot
|
|
import matplotlib.pyplot as plot
|
|
|
|
if len(sys.argv) == 2:
|
|
seed = sys.argv[1]
|
|
else:
|
|
seed = None
|
|
|
|
random.seed(seed)
|
|
|
|
n=200
|
|
points = [ ( round(random.uniform(-n,n),2),round(random.uniform(-n,n),2) ) for i in range(n) ]
|
|
|
|
quad = QuadTree( points )
|
|
# print(quad)
|
|
# sys.stderr.write( "%i points in the quadtree / %i points\n" % (len(quad), len(points)) )
|
|
|
|
|
|
fig = plot.figure()
|
|
|
|
ax = fig.add_subplot(111)
|
|
ax.set_aspect('equal')
|
|
|
|
# Plot the whole quad tree and its points.
|
|
# Iterating over the quadtree will generate points, thus list(quad) is equivalent to quad.points()
|
|
uberplot.scatter_points( ax, list(quad), facecolor="green", edgecolor="None")
|
|
|
|
for q in quad.quadrants:
|
|
edges = list( utils.tour(as_rect(q)) )
|
|
uberplot.plot_segments( ax, edges, edgecolor = "blue", alpha = 0.1, linewidth = 2 )
|
|
|
|
# Plot a random query on the quad tree.
|
|
# Remember a quadrant is ( (orig_y,orig_y), width )
|
|
minp = ( round(random.uniform(-n,n),2), round(random.uniform(-n,n),2) )
|
|
rand_quad = ( minp, round(random.uniform(0,n),2) )
|
|
# Asking for a quadrant will query the quad tree and return the corresponding points.
|
|
uberplot.scatter_points( ax, quad[rand_quad], facecolor="None", edgecolor="red", alpha=0.5, linewidth = 2 )
|
|
edges = list( utils.tour(as_rect(rand_quad)) )
|
|
uberplot.plot_segments( ax, edges, edgecolor = "red", alpha = 0.5, linewidth = 2 )
|
|
|
|
plot.show()
|
|
|
|
assert(len(points) == len(quad))
|
|
|