nojhan/ubergeekism
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Add queries to quadtree.

Browse source 
Cleaner quadtree code and comments
This commit is contained in:
Johann Dreo 2014-07-29 09:12:58 +02:00
commit a1838de026
2 changed files with 142 additions and 51 deletions
Download patch file Download diff file Expand all files Collapse all files

3
geometry.py
View file

@ -142,8 +142,7 @@ def in_box( point, box, exclude_edges = False ):
def segment_intersection( seg0, seg1 ):
"""Return the coordinates of the intersection point of two segments, or None.
If segments are colinear, returns colinear_value."""
"""Return the coordinates of the intersection point of two segments, or None."""
assert( len(seg0) == 2 )
assert( len(seg1) == 2 )

190
quadtree.py Normal file → Executable file
View file

@ -5,18 +5,36 @@ 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."""
"""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.init( points = points )
self.root, self.quadrants = self.init( points = points )
# Data structures to handle the quadtree
# Each leaf of the quadtree may contains one resident point.
self.residents = { self.root: None }
# Quadrants may have four children
# Each node of the quadtree may contains four children.
self.children = { self.root: [] }
# Status of quadrants
@ -28,16 +46,19 @@ class QuadTree(object):
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
self.build(points)
# Generate the complete tree.
self.build( points )
def init( self, quadrant = None, box = None, points = None ):
"""Initialize the root quadrant with the given quadrant ((x,y),width), the given box or the given set of points."""
"""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 only one of the options")
raise BaseException("ERROR: you should specify a box, a quadrant or points")
# Initialize the root quadrant as the given box
if box:
@ -54,18 +75,13 @@ class QuadTree(object):
minp = quadrant[0]
width = quadrant[1]
else:
raise BaseException("ERROR: you should specify a box, a quadrant or points")
assert( x(minp) <= x(minp)+width and y(minp) <= y(minp)+width )
# There is always the root quadrant in the list of available ones.
self.root = (minp,width)
self.quadrants = [ self.root ]
root = (minp,width)
quadrants = [ root ]
def as_box( self, quadrant );
width = quadrant[1]
maxp = tuple(xy+width for xy in quadrant[0])
return (quadrant[0],maxp)
return root,quadrants
def status( self, point, quadrant ):
@ -76,14 +92,13 @@ class QuadTree(object):
assert(quadrant is not None)
assert(len(quadrant) == 2)
box = self.as_box( quadrant )
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] )
# print("is external leaf")
return self.Status.Leaf
elif self.children[quadrant]:
# internal: a quadrant that contains other quadrants
@ -96,10 +111,10 @@ class QuadTree(object):
return self.Status.Out
def split(self, quadrant ):
def split( self, quadrant ):
"""Split an existing quadrant in four children quadrants.
Spread existing residents to the children."""
Move the existing resident to the correct child."""
# We cannot split a quadrant if it already have sub-quadrants
if quadrant != self.root:
@ -110,19 +125,18 @@ class QuadTree(object):
# For each four children quadrant's origins
self.children[quadrant] = []
for orig in ((qx,qy), (qx,qy+w), (qx+w,qy+w), (qx+w,qy)):
q = (orig,w)
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 a new child to the current parent.
# Add this new child to the current parent.
self.children[quadrant].append(q)
# The new quadrant has no child.
# This new quadrant has no child.
self.children[q] = []
assert( len(self.children[quadrant]) == 4 )
# Move the resident to the related children node
p = self.residents[quadrant]
if p is not None:
@ -136,7 +150,6 @@ class QuadTree(object):
self.residents[quadrant] = None
def append( self, point, quadrant = None ):
"""Try to inset the given point in the existing quadtree, under the given quadrant.
@ -151,6 +164,7 @@ class QuadTree(object):
# 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:
@ -167,14 +181,14 @@ class QuadTree(object):
return False
def build(self, points):
"""append all the given points in the quadtree."""
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 ):
def iterative_walk( self, at_quad = None ):
# Default to the root quadrant
if not at_quad:
@ -191,7 +205,7 @@ class QuadTree(object):
quads.extend( self.children[child] )
def recursive_walk(self, at_quad = None ):
def recursive_walk( self, at_quad = None ):
# Default to the root quadrant
if not at_quad:
@ -203,20 +217,21 @@ class QuadTree(object):
yield q
def repr(self, quad=None, depth=0):
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 quad:
quad = self.root
if not quadrant:
quadrant = self.root
head = " "*depth
r = head+"{"
quadrep = '"origin" : %s, "width" : %f' % quad
if self.residents[quad]: # external
r += ' "resident" : %s, \t%s },\n' % (self.residents[quad],quadrep)
elif self.children[quad]: # internal
r += ' "children_ids" : %s, \t%s, "children" : [\n' % (self.children[quad],quadrep)
for child in self.children[quad]:
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
@ -225,14 +240,84 @@ class QuadTree(object):
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)
@ -242,6 +327,7 @@ class QuadTree(object):
def __repr__(self):
"""Return a string representing the quadtree in a JSON-like format."""
return self.repr()
@ -262,29 +348,35 @@ if __name__ == "__main__":
random.seed(seed)
n=20
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)) )
# 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')
uberplot.scatter_points( ax, points, facecolor="red", edgecolor="red")
# uberplot.scatter_points( ax, quad.points(), facecolor="green", edgecolor="None")
# 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:
qx, qy = q[0]
w = q[1]
box = [(qx,qy), (qx,qy+w), (qx+w,qy+w), (qx+w,qy)]
edges = list( utils.tour(box) )
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()