import fiona import geohash import os import rtree import ujson as json from leveldb import LevelDB from collections import OrderedDict, defaultdict from shapely.geometry import Point, Polygon, MultiPolygon from shapely.prepared import prep from shapely.geometry.geo import mapping from geodata.polygons.area import polygon_bounding_box_area class PolygonIndex(object): include_only_properties = None simplify_tolerance = 0.00001 preserve_topology = True large_polygon_threshold = 10000 INDEX_FILENAME = None POLYGONS_DB_DIR = 'polygons' DEFAULT_POLYS_FILENAME = 'polygons.geojson' def __init__(self, index=None, polygons=None, polygons_db=None, save_dir=None, index_filename=None, polygons_db_path=None, include_only_properties=None): if save_dir: self.save_dir = save_dir else: self.save_dir = None if not index_filename: index_filename = self.INDEX_FILENAME self.index_path = os.path.join(save_dir or '.', index_filename) if not index: self.create_index(overwrite=True) else: self.index = index if include_only_properties and hasattr(include_only_properties, '__contains__'): self.include_only_properties = include_only_properties if not polygons: self.polygons = [] else: self.polygons = polygons if not polygons_db_path: polygons_db_path = os.path.join(save_dir or '.', self.POLYGONS_DB_DIR) if not polygons_db: self.polygons_db = LevelDB(polygons_db_path) else: self.polygons_db = polygons_db self.setup() self.i = 0 def create_index(self, overwrite=False): raise NotImplementedError('Children must implement') def index_polygon(self, polygon): raise NotImplementedError('Children must implement') def setup(self): pass def simplify_polygon(self, poly, simplify_tolerance=None, preserve_topology=None): if simplify_tolerance is None: simplify_tolerance = self.simplify_tolerance if preserve_topology is None: preserve_topology = self.preserve_topology return poly.simplify(simplify_tolerance, preserve_topology=preserve_topology) def index_polygon_properties(self, properties): pass def add_polygon(self, poly, properties, include_only_properties=None): if include_only_properties is not None: properties = {k: v for k, v in properties.iteritems() if k in include_only_properties} self.polygons.append(prep(poly)) self.polygons_db.Put(str(self.i), json.dumps(properties)) self.index_polygon_properties(properties) self.i += 1 @classmethod def create_from_shapefiles(cls, inputs, output_dir, index_filename=None, include_only_properties=None): index = cls(save_dir=output_dir, index_filename=index_filename or cls.INDEX_FILENAME) for input_file in inputs: if include_only_properties is not None: include_props = include_only_properties.get(input_file, cls.include_only_properties) else: include_props = cls.include_only_properties f = fiona.open(input_file) index.add_geojson_like_file(f) return index @classmethod def fix_polygon(cls, poly): ''' Coerce to valid polygon ''' if not poly.is_valid: poly = poly.buffer(0) if not poly.is_valid: return None return poly @classmethod def to_polygon(cls, coords): ''' Create shapely polygon from list of coordinate tuples if valid ''' if not coords or len(coords) < 3: return None poly = Polygon(coords) return poly def add_geojson_like_record(self, rec, include_only_properties=None): if not rec or not rec.get('geometry') or 'type' not in rec['geometry']: return poly_type = rec['geometry']['type'] if poly_type == 'Polygon': coords = rec['geometry']['coordinates'][0] poly = self.to_polygon(coords) if poly is None or not poly.bounds or len(poly.bounds) != 4: return self.index_polygon(poly) self.add_polygon(poly, rec['properties'], include_only_properties=include_only_properties) elif poly_type == 'MultiPolygon': polys = [] poly_coords = rec['geometry']['coordinates'] for coords in poly_coords: poly = self.to_polygon(coords[0]) if poly is None or not poly.bounds or len(poly.bounds) != 4: continue polys.append(poly) self.index_polygon(poly) self.add_polygon(MultiPolygon(polys), rec['properties'], include_only_properties=include_only_properties) else: return def add_geojson_like_file(self, f, include_only_properties=None): ''' Add either GeoJSON or a shapefile record to the index ''' for rec in f: self.add_geojson_like_record(rec, include_only_properties=include_only_properties) @classmethod def create_from_geojson_files(cls, inputs, output_dir, index_filename=None, polys_filename=DEFAULT_POLYS_FILENAME, include_only_properties=None): index = cls(save_dir=output_dir, index_filename=index_filename or cls.INDEX_FILENAME) for input_file in inputs: if include_only_properties is not None: include_props = include_only_properties.get(input_file, cls.include_only_properties) else: include_props = cls.include_only_properties f = json.load(open(input_file)) index.add_geojson_like_file(f['features'], include_only_properties=include_props) return index def save(self, polys_filename=DEFAULT_POLYS_FILENAME): self.save_polygons(os.path.join(self.save_dir, polys_filename)) self.save_index() self.save_polygon_properties(self.save_dir) def save_polygons(self, out_filename): out = open(out_filename, 'w') for poly in self.polygons: feature = { 'type': 'Feature', 'geometry': mapping(poly.context), } out.write(json.dumps(feature) + u'\n') self.polygons_db.CompactRange('\x00', '\xff') def save_index(self): raise NotImplementedError('Children must implement') def load_polygon_properties(self, d): pass def save_polygon_properties(self, d): pass @classmethod def load_polygons(cls, filename): f = open(filename) polygons = [] for line in f: feature = json.loads(line.rstrip()) poly_type = feature['geometry']['type'] if poly_type == 'Polygon': poly = Polygon(feature['geometry']['coordinates'][0]) polygons.append(prep(poly)) elif poly_type == 'MultiPolygon': polys = [] for coords in feature['geometry']['coordinates']: poly = Polygon(coords[0]) polys.append(poly) polygons.append(prep(MultiPolygon(polys))) return polygons @classmethod def load_index(cls, d, index_name=None): raise NotImplementedError('Children must implement') @classmethod def load(cls, d, index_name=None, polys_filename=DEFAULT_POLYS_FILENAME, polys_db_dir=POLYGONS_DB_DIR): index = cls.load_index(d, index_name=index_name or cls.INDEX_FILENAME) polys = cls.load_polygons(os.path.join(d, polys_filename)) polygons_db = LevelDB(os.path.join(d, polys_db_dir)) polygon_index = cls(index=index, polygons=polys, polygons_db=polygons_db, save_dir=d) polygon_index.load_polygon_properties(d) return polygon_index def get_candidate_polygons(self, lat, lon): raise NotImplementedError('Children must implement') def point_in_poly(self, lat, lon, return_all=False): polys = self.get_candidate_polygons(lat, lon) pt = Point(lon, lat) containing = None if return_all: containing = [] for i in polys: poly = self.polygons[i] contains = poly.contains(pt) if contains and not return_all: try: return json.loads(self.polygons_db.Get(str(i))) except KeyError: return None elif contains: containing.append(json.loads(self.polygons_db.Get(str(i)))) return containing class RTreePolygonIndex(PolygonIndex): INDEX_FILENAME = 'rtree' def create_index(self, overwrite=False): self.index = rtree.index.Index(self.index_path, overwrite=overwrite) def index_polygon(self, polygon): self.index.insert(self.i, polygon.bounds) def get_candidate_polygons(self, lat, lon): return OrderedDict.fromkeys(self.index.intersection((lon, lat, lon, lat))).keys() def save_index(self): # need to close index before loading it self.index.close() @classmethod def load_index(cls, d, index_name=None): return rtree.index.Index(os.path.join(d, index_name or cls.INDEX_FILENAME)) class GeohashPolygonIndex(PolygonIndex): # Reference: https://www.elastic.co/guide/en/elasticsearch/guide/current/geohashes.html GEOHASH_PRECISIONS = [ (7, 152.8 ** 2), (6, 1200.0 * 610.0), (5, 4900.0 * 4900.0), (4, 39000.0 * 19500.0), (3, 156000.0 * 156000.0), (2, 1251000.0 * 625000.0), ] INDEX_FILENAME = 'index.json' def create_index(self, overwrite=False): self.index = defaultdict(list) def index_point(self, lat, lon, geohash_level): code = geohash.encode(lat, lon)[:geohash_level] for key in [code] + geohash.neighbors(code): self.index[key].append(self.i) def index_polygon(self, polygon): poly_area = polygon_bounding_box_area(polygon) for geohash_level, area in self.GEOHASH_PRECISIONS: if area * 9.0 >= poly_area: break bbox = polygon.bounds lat, lon = (bbox[1] + bbox[3]) / 2.0, (bbox[0] + bbox[2]) / 2.0 self.index_point(lat, lon, geohash_level) def get_candidate_polygons(self, lat, lon, return_all=False): candidates = [] code = geohash.encode(lat, lon) for level, area in self.GEOHASH_PRECISIONS: indices = self.index.get(code[:level]) if not indices: continue candidates.extend(indices) if not return_all: break return candidates def save_index(self): if not self.index_path: self.index_path = os.path.join(self.save_dir or '.', self.INDEX_FILENAME) json.dump(self.index, open(self.index_path, 'w')) @classmethod def load_index(cls, d, index_name=None): return json.load(open(os.path.join(d, index_name or cls.INDEX_FILENAME)))