source: trunk/anuga_core/source/anuga/geometry/quad.py @ 9516

"""quad.py - quad tree data structure for fast indexing of regions in the plane.

This generic structure can be used to store any geometry in a quadtree.
It is naive, and does not exploit any coherency - it merely tests a point
against all bounding boxes in its heirarchy.

It returns a list of bounding boxes which intersect with the test point, which
may then be iterated over with a proper intersection test to detect actual
geometry intersections.

As of June 2010 this module has a pylint quality rating of 10/10.

"""

import anuga.utilities.log as log

            
class Cell():
    """ One cell in the plane.
        A cell is defined by an AABB, and can have smaller AABB children.
        The children can be rapidly searched for intersections in log(n) time.
    """
  
    def __init__(self, extents, parent,
         name = 'cell'):
        """ Construct a new cell.
            extents is an AABB defining a region on the plane.
            parent is the node above this one, or None if it is root.
        """
    
        self.extents = extents
        self.parent = parent
        
        # The points in this cell     
        self.leaves = []
        self.children = None
        
    
    def __repr__(self):
        """ String representation of the quadtree. """
        ret_str = '%s: leaves: %d' \
               % (self.name , len(self.leaves))    
        if self.children:
            ret_str += ', children: %d' % (len(self.children))
        return ret_str


    def insert(self, new_leaf):
        """ Insert a leaf into the quadtree.
            new_leaf is a tuple of (AABB extents, data), where data can
                     be any user data (geometry, triangle index, etc.).
        """
        if type(new_leaf)==type(list()):
            for leaf in new_leaf:
                self.insert_item(leaf)
        else:
            self.insert_item(new_leaf)


    def insert_item(self, new_leaf):   
        """ Internal recursive insert a single item.
            new_leaf is a tuple of (AABB extents, data), where data can
                     be any user data (geometry, triangle index, etc.).       
        """
        new_region, _ = new_leaf
        
        # recurse down to any children until we get an intersection
        if self.children:
            for child in self.children:
                if child.extents.is_trivial_in(new_region):
                    child.insert_item(new_leaf)
                    return
        else:            
            # try splitting this cell and see if we get a trivial in
            subregion1, subregion2 = self.extents.split()
            
            # option 1 - try splitting 4 ways
            #subregion11, subregion12 = subregion1.split()    
            #subregion21, subregion22 = subregion2.split()
            #regions = [subregion11, subregion12, subregion21, subregion22]
            #for region in regions:
                #if region.is_trivial_in(new_region):
                    #self.children = [Cell(x, parent=self) for x in regions]
                    #self.insert_item(new_leaf)
                    #return               

            # option 2 - try splitting 2 ways - no performance difference
            # noticed in practise between this and the above option.
            if subregion1.is_trivial_in(new_region):
                self.children = [Cell(subregion1, self), \
                                 Cell(subregion2, self)]    
                self.children[0].insert_item(new_leaf)
                return
            elif subregion2.is_trivial_in(new_region):
                self.children = [Cell(subregion1, self), \
                                 Cell(subregion2, self)]    
                self.children[1].insert_item(new_leaf)
                return                
    
        # recursion ended without finding a fit, so attach it as a leaf
        self.leaves.append(new_leaf)
        
     
    def retrieve(self):
        """Get all leaves from this tree.
           return a traversal of the entire tree.
        """
        
        leaves_found = list(self.leaves)
        
        if not self.children:
            return leaves_found

        for child in self.children:
            leaves_found.extend(child.retrieve())
            
        return leaves_found

    def count(self):
        """Count all leaves from this tree.
           return num of leaves in the tree.
        """
        
        leaves_found = len(self.leaves)
        
        if not self.children:
            return leaves_found

        for child in self.children:
            leaves_found += child.count()
            
        return leaves_found        

    def show(self, depth=0):
        """Traverse tree below self, dumping all information.
        """
        if depth == 0:
            log.critical() 
        print '%s%s'  % ('  '*depth, self.name), self.extents, ' [', \
            self.leaves, ']'
        if self.children:
            log.critical()
            for child in self.children:
                child.show(depth+1)

    def search(self, point):
        """
            Search the tree for intersection with leaves
            point is a test point.
            return a list of possible intersections with geometry.
        """
        intersecting_regions = self.test_leaves(point)
        
        # recurse down into nodes that the point passes through
        if self.children:
            for child in self.children:    
                if child.extents.contains(point):
                    intersecting_regions.extend(child.search(point))
             
        return intersecting_regions
 
 
    def test_leaves(self, point):
        """ Test all leaves on this node to see if they intersect x.
            Does not recurse into children.
            x is a point to test
            return a list of leaves that intersect x
        """
        intersecting_regions = []
        
        # test all leaves to see if they intersect the point
        for leaf in self.leaves:
            aabb, data = leaf
            if aabb.contains(point):
                intersecting_regions.append([data, self])
                
        return intersecting_regions                
 
 
    def get_siblings(self):
        """ return siblings of this node. If there is no parent, it
                   returns an empty list.
        """
        if not self.parent:
            return []
         
        siblings = list(self.parent.children)
        siblings.remove(self)
        return siblings