geocell.py

import math

def between(a, b, c):
  return a <= b and b <= c


class Region(object):
  def intersects(self, other):
    raise NotImplemented(
        "%s does not know how to check for intersection with %s"
        % (self.__class__.__name__, other.__class__.__name__))

  def getCovering(self):
    """Returns the minimum range that covers this region."""
    raise NotImplemented()


class BoundingBox(Region):
  def __init__(self, ul, br):
    self.ul = ul
    self.br = br
  
  def __str__(self):
    return '(%f, %f) - (%f, %f)' % (self.ul + self.br)
  
  def intersects(self, other):
    if isinstance(other, BoundingBox):
      return ((between(self.ul[0], other.ul[0], self.br[0]) or
               between(other.ul[0], self.ul[0], other.br[0])) and
              (between(self.ul[1], other.ul[1], self.br[1]) or
               between(other.ul[1], self.ul[1], other.br[1])))
    else:
      super(BoundingBox, self).insertsects(other)

  def covers(self, other):
    return (self.ul[0] <= other.ul[0] and
            self.ul[1] <= other.ul[1] and
            self.br[0] >= other.br[0] and
            self.br[1] >= other.br[1])

  def getCovering(self):
    """Returns the smallest box that covers this bounding box.
    
    This is achieved by starting with an arbitrary point and repeatedly getting
    a prefix that covers the prefix-so-far and the next point.
    """
    points = [self.ul, self.br, (self.ul[0], self.br[1]),
              (self.br[0], self.ul[1])]
    covering = GeoCell.fromLatLon(*points[0])
    for point in points[1:]:
      covering = covering.getCoveringRange(GeoCell.fromLatLon(*point))
    return covering
  
  def getCenter(self):
    return ((self.ul[0] + self.br[0]) / 2, (self.ul[1] + self.br[1]) / 2)


class GeoCell(long):
  """Represents points and regions on a 2D surface in a query-friendly manner.
  
  A geocell is a 63-bit bit string, consisting of consecutive pairs of bits.
  Each pair indicates which quarter of the region currently being considered the
  final region or point lies in. For example, a geocell starting with the bit
  pattern 01 11 indicates that the region being described lies in the upper left
  quarter of the original range, and in the upper right quarter of that. In
  lat/long terms, this would be the region (-90, 90) - (0, 45). This recursive
  subdivision approach is similar to that used in a quadtree. With a little more
  work this would form a space-filling curve, but currently the ordering of
  child cells is not specified in such a way as to achieve this.
  
  The 'level' of a geocell is the number of bit-pairs _not_ specified by the
  cell. This is indicated by terminating the sequence of bit pairs with a single
  1 bit, followed by 0 bits. Thus, a 'leaf' cell (a cell at the highest possible
  level of detail) will always end with a 1, a level 1 cell will end with 100,
  a level 2 cell with 10000, and so forth.
  
  Querying is accomplished by finding a set of cells that covers the requested
  region and as little extra as possible, without using more than the requested
  number of cells. These cells are then translated into integer ranges that can
  be used in a query.
  """
  @staticmethod
  def fromLatLon(lat, lon):
    node = 0L
    lonrange = (-180.0, 180.0)
    latrange = (-90.0, 90.0)
    for i in range(31):
      node <<= 2
      
      midlon = (lonrange[0] + lonrange[1]) / 2
      if lon >= midlon:
        node |= 2
        lonrange = (midlon, lonrange[1])
      else:
        lonrange = (lonrange[0], midlon)
      
      midlat = (latrange[0] + latrange[1]) / 2
      if lat >= midlat:
        node |= 1
        latrange = (midlat, latrange[1])
      else:
        latrange = (latrange[0], midlat)
    
    node = (node << 1) | 1
    return GeoCell(node)
  
  @staticmethod
  def getRanges(cells):
    """Utility function to return an iterator of ranges from a cell iterator.
    
    This function automatically coalesces consecutive ranges.
    """
    if len(cells) == 0:
      return
    cells = sorted(cells)
    current = cells[0].getRange()
    for cell in cells[1:]:
      range = cell.getRange()
      if range[0] == current[1] + 1:
        current = (current[0], range[1])
      else:
        yield current
        current = range
    yield current
  
  def __repr__(self):
    return '0x%X' % self
  
  def getLevel(self):
    """Returns the level of this GeoCell, where 0 is a 'leaf' cell."""
    # The level is determined by the number of trailing 0 bits.
    # If we and the number with its negation, the only bit left is the last one.
    x = self & -self
    return int(math.log(x, 2)) / 2
  
  def toBoundingBox(self):
    lonrange = (-180.0, 180.0)
    latrange = (-90.0, 90.0)
    for i in range(61, self.getLevel() * 2, -2):
      midlon = (lonrange[0] + lonrange[1]) / 2
      if self & (2 << i):
        lonrange = (midlon, lonrange[1])
      else:
        lonrange = (lonrange[0], midlon)
      
      midlat = (latrange[0] + latrange[1]) / 2
      if self & (1 << i):
        latrange = (midlat, latrange[1])
      else:
        latrange = (latrange[0], midlat)
    return BoundingBox((latrange[0], lonrange[0]), (latrange[1], lonrange[1]))
  
  def toLatLon(self):
    return self.toBoundingBox().getCenter()
  
  def getCoveringRange(self, other):
    """Returns the smallest range that encloses both points.
    
    This is represented as the longest common prefix of the two (binary) numbers
    followed by a 1 bit.
    """
    # Finds the least significant set bit and constructs a mask covering it 
    # and everything after it
    mask_self = ((self & -self) << 1) - 1
    mask_other = ((other & -other) << 1) - 1
    
    # The first set bit in mask is the first differing bit, or the first
    # trailing 1, whichever comes first
    mask = (self ^ other) | (self & -self) | (other & -other)
    # Set the first bit of every pair to 1 if the pair differs
    mask = (mask | (mask << 1)) & 0x5555555555555555
    # Expand the mask to cover all bits after the first set bit
    mask |= mask >> 1
    mask |= mask >> 2
    mask |= mask >> 4
    mask |= mask >> 8
    mask |= mask >> 16
    mask |= mask >> 32
    return GeoCell((self & ~mask) | ((mask + 1) >> 1))
  
  def isLeaf(self):
    return bool(self&1)

  def getChildren(self):
    if self.isLeaf():
      return []
    shift = self.getLevel() * 2 - 2
    return [GeoCell(self ^ (x<<shift)) for x in range(1, 8, 2)]
  
  def getRange(self):
    """Returns a (start, end) tuple indicating the range this GeoCell covers"""
    level = self.getLevel()
    mask = (1 << 2*level+1) - 1
    start = self & ~mask
    end = start + mask
    return (start, end)


def generateRegionCovering(region, max_cells=8):
  # Partially covered cells
  q = [region.getCovering()]
  # Entirely covered cells
  ret = []
  
  while len(q) + len(ret) < max_cells:
    cell = q.pop(0)
    if cell.isLeaf():
      ret.append(cell)
    else:
      for child in cell.getChildren():
        bb = child.toBoundingBox()
        if region.covers(bb):
          ret.append(child)
        elif region.intersects(bb):
          q.append(child)
  ret.extend(q)
  return ret

geocell_test.py

#!/usr/bin/python
import geocell
import unittest


class GeoCellTest(unittest.TestCase):
  def testGeoCell(self):
    self.failUnlessEqual(geocell.GeoCell.fromLatLon(0, 0), 0x6000000000000001)
    self.failUnlessEqual(geocell.GeoCell.fromLatLon(-90, -180), 0x1)
    self.failUnlessEqual(geocell.GeoCell.fromLatLon(90, 180), 0x7FFFFFFFFFFFFFFF)
    self.failUnlessEqual(geocell.GeoCell.fromLatLon(0, 180), 0x7555555555555555)
    self.failUnlessEqual(geocell.GeoCell.fromLatLon(90, 0), 0x6AAAAAAAAAAAAAAB)
    for x in range(-180, 181, 2):
      for y in range(-90, 91, 2):
        lat, lon = geocell.GeoCell.fromLatLon(y, x).toLatLon()
        self.failUnlessAlmostEqual(lat, y, 6)
        self.failUnlessAlmostEqual(lon, x, 6)

  def testBoundingBox(self):
    bb = geocell.BoundingBox((5, 5), (10, 10))
    self.failUnless(bb.intersects(geocell.BoundingBox((0, 0), (6, 6))))
    self.failUnless(bb.intersects(geocell.BoundingBox((6, 6), (9, 9))))
    self.failUnless(bb.intersects(geocell.BoundingBox((0, 0), (15, 15))))
    self.failUnless(bb.intersects(bb))
    self.failIf(bb.intersects(geocell.BoundingBox((0, 0), (6, 4))))

    self.failUnless(bb.covers(geocell.BoundingBox((6, 7), (8, 9))))
    self.failUnless(bb.covers(bb))
    self.failIf(bb.covers(geocell.BoundingBox((0, 0), (6, 6))))

  def testCoveringRange(self):
    # The min covering range of a cell on itself is the original cell.
    self.failUnlessEqual(geocell.GeoCell(0x12345).getCoveringRange(
        geocell.GeoCell(0x12345)), 0x12345)
    # If only the second bit in a pair differs, both should be masked out
    # 1F = 0 00 11 11 1, 1D = 0 00 11 10 1, 1C = 0 00 11 1 00
    self.failUnlessEqual(geocell.GeoCell(0x1F).getCoveringRange(0x1D), 0x1C)
    
    # Covering of multiple sets should be the largest of the individual
    # coverings
    c1 = geocell.GeoCell(0x603F03F03F03F03F) # 60 = 0 11 00 00 0
    c2 = geocell.GeoCell(0x6FC0FC0FC0FC0FC1) # 6F = 0 11 01 11 1
    c3 = geocell.GeoCell(0x6A95A95A95A95A95)
    self.failUnlessEqual(c1.getCoveringRange(c2), 0x7000000000000000)
    self.failUnlessEqual(c1.getCoveringRange(c2).getCoveringRange(c3),
                         0x7000000000000000)
    cell = geocell.GeoCell.fromLatLon(90, 0)
    self.failUnlessEqual(cell.getCoveringRange(cell), cell)
    self.failUnlessEqual(cell.getCoveringRange(0x1), 0x4000000000000000)
    self.failUnlessEqual(cell.getCoveringRange(geocell.GeoCell.fromLatLon(90, 180)),
                         0x7000000000000000)

  def testMinCovering(self):
    bb = geocell.BoundingBox((10, 10), (80, 80))
    self.failUnlessEqual(bb.getCovering(), 0x7000000000000000)
    self.failUnless(bb.getCovering().toBoundingBox().covers(bb))
    bb2 = geocell.BoundingBox((10, 10), (20, 80))
    self.failUnless(bb2.getCovering().toBoundingBox().covers(bb2))
    self.failUnlessEqual(geocell.BoundingBox((-43.53273, 172.6321429),
                         (-43.52912, 172.6408769)).getCovering(),
                         0x5D46509000000000)

  def testGetLevel(self):
    self.failUnlessEqual(geocell.GeoCell(0x1).getLevel(), 0)
    self.failUnlessEqual(geocell.GeoCell(0x4).getLevel(), 1)
    self.failUnlessEqual(geocell.GeoCell(0xF0).getLevel(), 2)
    self.failUnlessEqual(geocell.GeoCell(0x8000000000000000).getLevel(), 31)

  def testIsLeaf(self):
    self.failUnless(geocell.GeoCell(0x1).isLeaf())
    self.failIf(geocell.GeoCell(0xC).isLeaf())
  
  def testGetChildren(self):
    self.failUnlessEqual(geocell.GeoCell(0x1).getChildren(), [])
    self.failUnlessEqual(sorted(geocell.GeoCell(0x300).getChildren()),
                         [0x240, 0x2C0, 0x340, 0x3C0])
    self.failUnlessEqual(sorted(geocell.GeoCell(0xC).getChildren()),
                         [0x9, 0xB, 0xD, 0xF])

  def testRange(self):
    # Leaf node
    self.failUnlessEqual(geocell.GeoCell(0x1).getRange(), (0x0, 0x1))
    # Level 1
    self.failUnlessEqual(geocell.GeoCell(0xC).getRange(), (0x8, 0xF))
    # Level 2
    self.failUnlessEqual(geocell.GeoCell(0x30).getRange(), (0x20, 0x3F))
    self.failUnlessEqual(geocell.GeoCell(0x1230).getRange(), (0x1220, 0x123F))

  def testRanges(self):
    # Single range
    self.failUnlessEqual(list(geocell.GeoCell.getRanges([geocell.GeoCell(0x1)])),
                         [(0x0, 0x1)])
    self.failUnlessEqual(list(geocell.GeoCell.getRanges([geocell.GeoCell(0x30)])),
                         [(0x20, 0x3F)])
    
    # All 4 children
    cell = geocell.GeoCell(0x1230)
    range = cell.getRange()
    children = cell.getChildren()
    self.failUnlessEqual(list(geocell.GeoCell.getRanges(children)), [range])

    # Multiple levels
    ranges = children[:3] + [geocell.GeoCell(0xC), geocell.GeoCell(0x1229)]
    self.failUnlessEqual(list(geocell.GeoCell.getRanges(ranges)),
                         [(0x8, 0xF), (0x1220, 0x1229), (0x1230, 0x123F)])

  def testRegionCovering(self):
    covering = geocell.generateRegionCovering(geocell.BoundingBox((10, 10), (80, 80)))
    ranges = list(geocell.GeoCell.getRanges(covering))
    self.failUnlessEqual(ranges, [(0x6000000000000000, 0x6FFFFFFFFFFFFFFF)])
    covering = geocell.generateRegionCovering(geocell.BoundingBox((10, 10), (20, 80)))
    ranges = list(geocell.GeoCell.getRanges(covering))
    self.failUnlessEqual(ranges, [(0x6000000000000000, 0x61FFFFFFFFFFFFFF),
                                  (0x6400000000000000, 0x65FFFFFFFFFFFFFF)])


if __name__ == '__main__':
  unittest.main()