import array
import geohash
import os
import math
import operator
import six
import ujson as json
from collections import defaultdict, OrderedDict
from leveldb import LevelDB
from geodata.distance.haversine import haversine_distance
class PointIndex(object):
include_only_properties = None
persistent_index = False
cache_size = 0
POINTS_DB_DIR = 'points'
GEOHASH_PRECISION = 7
PROPS_FILENAME = 'properties.json'
POINTS_FILENAME = 'points.json'
INDEX_FILENAME = 'index.json'
def __init__(self, index=None, save_dir=None,
points=None,
points_path=None,
points_db=None,
points_db_path=None,
index_path=None,
include_only_properties=None,
precision=GEOHASH_PRECISION):
if save_dir:
self.save_dir = save_dir
else:
self.save_dir = None
if include_only_properties and hasattr(include_only_properties, '__contains__'):
self.include_only_properties = include_only_properties
if not index_path:
index_path = os.path.join(save_dir or '.', self.INDEX_FILENAME)
self.index_path = index_path
if not index:
self.index = defaultdict(list)
else:
self.index = index
if not points_path:
points_path = os.path.join(save_dir or '.', self.POINTS_FILENAME)
self.points_path = points_path
if not points:
self.points = array.array('d')
else:
self.points = points
if not points_db_path:
points_db_path = os.path.join(save_dir or '.', self.POINTS_DB_DIR)
if not points_db:
self.points_db = LevelDB(points_db_path)
else:
self.points_db = points_db
self.precision = precision
self.i = 0
def index_point(self, lat, lon):
code = geohash.encode(lat, lon)[:self.precision]
for key in [code] + geohash.neighbors(code):
self.index[key].append(self.i)
self.points.extend([lat, lon])
def add_point(self, lat, lon, properties, cache=False, include_only_properties=None):
if include_only_properties is None and self.include_only_properties:
include_only_properties = self.include_only_properties
if include_only_properties is not None:
properties = {k: v for k, v in properties.iteritems() if k in include_only_properties}
self.index_point(lat, lon)
self.points_db.Put(self.properties_key(self.i), json.dumps(properties))
self.i += 1
def load_properties(self, filename):
properties = json.load(open(filename))
self.i = int(properties.get('num_points', self.i))
self.precision = int(properties.get('precision', self.precision))
def save_properties(self, out_filename):
out = open(out_filename, 'w')
json.dump({'num_points': str(self.i),
'precision': self.precision}, out)
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)))
def save_points(self):
json.dump(self.points, open(self.points_path, 'w'))
@classmethod
def load_points(cls, d):
return array.array('d', json.load(open(os.path.join(d, cls.POINTS_FILENAME))))
def properties_key(self, i):
return 'props:{}'.format(i)
def get_properties(self, i):
return json.loads(self.points_db.Get(self.properties_key(i)))
def compact_points_db(self):
self.points_db.CompactRange('\x00', '\xff')
def save(self):
self.save_index()
self.save_points()
self.compact_points_db()
self.save_properties(os.path.join(self.save_dir, self.PROPS_FILENAME))
@classmethod
def load(cls, d):
index = cls.load_index(d)
points = cls.load_points(d)
points_db = LevelDB(os.path.join(d, cls.POINTS_DB_DIR))
point_index = cls(index=index, points=points, points_db=points_db)
point_index.load_properties(os.path.join(d, cls.PROPS_FILENAME))
return point_index
def __iter__(self):
for i in xrange(self.i):
lat, lon = self.points[i * 2], self.points[i * 2 + 1]
yield self.get_properties(i), lat, lon
def __len__(self):
return self.i
def get_candidate_points(self, latitude, longitude):
code = geohash.encode(latitude, longitude)[:self.precision]
candidates = OrderedDict()
candidates.update([(k, None) for k in self.index.get(code, [])])
for neighbor in geohash.neighbors(code):
candidates.update([(k, None) for k in self.index.get(neighbor, [])])
return candidates.keys()
def point_distances(self, latitude, longitude):
candidates = self.get_candidate_points(latitude, longitude)
return [(i, self.points[i * 2], self.points[i * 2 + 1],
haversine_distance(latitude, longitude,
self.points[i * 2],
self.points[i * 2 + 1]))
for i in candidates]
def all_nearby_points(self, latitude, longitude):
distances = self.point_distances(latitude, longitude)
if not distances:
return []
return sorted(distances, key=operator.itemgetter(-1))
def points_with_properties(self, results):
return [(self.get_properties(i), lat, lon, distance)
for i, lat, lon, distance in results]
def nearest_points(self, latitude, longitude):
return self.points_with_properties(self.all_nearby_points(latitude, longitude))
def nearest_n_points(self, latitude, longitude, n=2):
return self.points_with_properties(self.all_nearby_points(latitude, longitude)[:n])
def nearest_point(self, latitude, longitude):
distances = self.all_nearby_points(latitude, longitude)
if not distances:
return None
return self.points_with_properties(distances[:1])[0]