From 33e9aa5dc7ddff50c87506ad5e7fe00ca57f7626 Mon Sep 17 00:00:00 2001 From: mdbartos Date: Sat, 4 Apr 2015 03:38:51 -0700 Subject: [PATCH] Optimize sjoin --- geopandas/tools/sjoin.py | 92 +++++++++++++++++++++++++--------------- 1 file changed, 58 insertions(+), 34 deletions(-) diff --git a/geopandas/tools/sjoin.py b/geopandas/tools/sjoin.py index 896d823..634e465 100644 --- a/geopandas/tools/sjoin.py +++ b/geopandas/tools/sjoin.py @@ -1,8 +1,10 @@ +import geopandas as gpd +import numpy as np import pandas as pd -from geopandas import GeoDataFrame -from .overlay import _uniquify +import rtree +from shapely import prepared -def sjoin(left_df, right_df, how="left", op="intersects", use_sindex=True, **kwargs): +def sjoin(left_df, right_df, how='left', op='intersects', crs_convert=True, lsuffix='left', rsuffix='right', **kwargs): """Spatial join of two GeoDataFrames. left_df, right_df are GeoDataFrames @@ -15,48 +17,70 @@ def sjoin(left_df, right_df, how="left", op="intersects", use_sindex=True, **kwa use_sindex : Use the spatial index to speed up operation? Default is True kwargs: passed to op method """ + + # CHECK VALIDITY OF JOIN TYPE allowed_hows = ['left', 'right', 'inner'] if how not in allowed_hows: raise ValueError("`how` was \"%s\" but is expected to be in %s" % \ (how, allowed_hows)) + + # CHECK VALIDITY OF PREDICATE OPERATION + allowed_ops = ['contains', 'within', 'intersects'] - if how == "right": - # right outer join just implemented as the inverse of left; swap names + if op not in allowed_ops: + raise ValueError("`op` was \"%s\" but is expected to be in %s" % \ + (op, allowed_ops)) + + # IF WITHIN, SWAP NAMES + if op == "within": + # within implemented as the inverse of contains; swap names left_df, right_df = right_df, left_df - collection = [] - for i, feat in left_df.iterrows(): - geom = feat.geometry + # CONVERT CRS IF NOT EQUAL + if left_df.crs != right_df.crs: + print 'Warning: CRS does not match!' + if crs_convert == True: + print 'Converting CRS...' + if left_df.values.nbytes >= right_df.values.nbytes: + right_df = right_df.to_crs(left_df.crs) + elif left_df.values.nbytes < right_df.values.nbytes: + left_df = left_df.to_crs(right_df.crs) - if use_sindex and right_df.sindex: - candidates = [x.object for x in - right_df.sindex.intersection(geom.bounds, objects=True)] - else: - candidates = [i for i, x in right_df.iterrows()] + # CONSTRUCT SPATIAL INDEX FOR RIGHT DATAFRAME + tree_idx = rtree.index.Index() + right_df_bounds = right_df['geometry'].apply(lambda x: x.bounds) + for i in right_df_bounds.index: + tree_idx.insert(i, right_df_bounds[i]) - feature_hits = 0 - for cand_id in candidates: - candidate = right_df.ix[cand_id] - if getattr(geom, op)(candidate.geometry, **kwargs): - newseries = candidate.drop(right_df._geometry_column_name) - newfeat = pd.concat([feat, newseries]) - newfeat.index = _uniquify(newfeat.index) - collection.append(newfeat) - feature_hits += 1 + # FIND INTERSECTION OF SPATIAL INDEX + idxmatch = left_df['geometry'].apply(lambda x: x.bounds).apply(lambda x: list(tree_idx.intersection(x))) + idxmatch = idxmatch[idxmatch.str.len() > 0] - # TODO Should we perform aggregation if feature_hit > 1? - # Advantage: single step and possible performance improvement - # Disadvantage: Pandas already has groupby so user can do this later + r_idx = np.concatenate(idxmatch.values) + l_idx = np.concatenate((idxmatch.str.len()*pd.Series([[i] for i in idxmatch.index], index=idxmatch.index)).values) - # If left does not spatially join with any right features, - # Fill in the right columns with NA - if how != 'inner' and feature_hits == 0: - empty = pd.Series(dict.fromkeys(right_df.columns, None)) - empty.drop(right_df._geometry_column_name, inplace=True) + # VECTORIZE PREDICATE OPERATIONS + def find_intersects(a1, a2): + return a1.intersects(a2) - newfeat = pd.concat([feat, empty]) - newfeat.index = _uniquify(newfeat.index) - collection.append(newfeat) + def find_contains(a1, a2): + return a1.contains(a2) - return GeoDataFrame(collection, index=range(len(collection))) + predicate_d = {'intersects': find_intersects, 'contains': find_contains, 'within': find_contains} + + check_predicates = np.vectorize(predicate_d[op]) + + # CHECK PREDICATES + result = pd.DataFrame(np.column_stack([l_idx, r_idx, check_predicates(left_df['geometry'].apply(lambda x: prepared.prep(x)).values[l_idx], right_df['geometry'].values[r_idx])])) + result.columns = ['index_%s' % lsuffix, 'index_%s' % rsuffix, 'match_bool'] + result = pd.DataFrame(result[result['match_bool']==1].set_index('index_%s' % lsuffix)['index_%s' % rsuffix]) + + # IF 'WITHIN', SWAP NAMES AGAIN + if op == "within": + # within implemented as the inverse of contains; swap names + left_df, right_df = right_df, left_df + result = result.reset_index().rename(columns={'index_%s' % (lsuffix): 'index_%s' % (rsuffix), 'index_%s' % (rsuffix): 'index_%s' % (lsuffix)}).set_index('index_left').sort_index() + + # APPLY JOIN + return left_df.merge(result, left_index=True, right_index=True).merge(right_df, left_on='index_%s' % rsuffix, right_index=True, how=how, suffixes=('_%s' % lsuffix, '_%s' % rsuffix))