Add queries to quadtree.

Cleaner quadtree code and comments

geometry.py

@@ -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.
+    """Return the coordinates of the intersection point of two segments, or None."""
-    If segments are colinear, returns colinear_value."""
     assert( len(seg0) == 2 )
     assert( len(seg1) == 2 )
 

quadtree.py

@@ -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:
+        assert( x(minp) <= x(minp)+width and y(minp) <= y(minp)+width )
-            raise BaseException("ERROR: you should specify a box, a quadrant or points")
 
         # There is always the root quadrant in the list of available ones.
-        self.root = (minp,width)
+        root = (minp,width)
-        self.quadrants = [ self.root ]
+        quadrants = [ root ]
 
-
+        return root,quadrants
-    def as_box( self, quadrant );
-        width = quadrant[1]
-        maxp = tuple(xy+width for xy in quadrant[0])
-        return (quadrant[0],maxp)
 
 
     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)):
+        for origin in ( (qx,qy), (qx,qy+w), (qx+w,qy+w), (qx+w,qy) ):
-            q = (orig,w)
             # 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):
+    def build( self, points ):
-        """append all the given points in the quadtree."""
+        """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:
+        if not quadrant:
-            quad = self.root
+            quadrant = self.root
 
         head = "  "*depth
         r = head+"{"
-        quadrep = '"origin" : %s, "width" : %f' % quad
+        quadrep = '"origin" : %s, "width" : %f' % quadrant
-        if self.residents[quad]: # external
+        if self.residents[quadrant]: # external
-            r += ' "resident" : %s, \t%s },\n' % (self.residents[quad],quadrep)
+            r += ' "resident" : %s, \t%s },\n' % (self.residents[quadrant],quadrep)
-        elif self.children[quad]: # internal
+        elif self.children[quadrant]: # internal
-            r += ' "children_ids" : %s, \t%s, "children" : [\n' % (self.children[quad],quadrep)
+            r += ' "children_ids" : %s, \t%s, "children" : [\n' % (self.children[quadrant],quadrep)
-            for child in self.children[quad]:
+            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)
+    # print(quad)
-    sys.stderr.write( "%i points in the quadtree / %i points\n" % (len(quad), len(points)) )
+    # 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]
+        edges = list( utils.tour(as_rect(q)) )
-        w = q[1]
-        box = [(qx,qy), (qx,qy+w), (qx+w,qy+w), (qx+w,qy)]
-        edges = list( utils.tour(box) )
         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()