<s> import numpy as np <EOL> import pandas as pd <EOL> from shapely import prepared <EOL> def sjoin ( left_df , right_df , how = '<STR_LIT>' , op = '<STR_LIT>' , <EOL> lsuffix = '<STR_LIT:left>' , rsuffix = '<STR_LIT:right>' , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> import rtree <EOL> allowed_hows = [ '<STR_LIT:left>' , '<STR_LIT:right>' , '<STR_LIT>' ] <EOL> if how not in allowed_hows : <EOL> raise ValueError ( "<STR_LIT>" % ( how , allowed_hows ) ) <EOL> allowed_ops = [ '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] <EOL> if op not in allowed_ops : <EOL> raise ValueError ( "<STR_LIT>" % ( op , allowed_ops ) ) <EOL> if op == "<STR_LIT>" : <EOL> left_df , right_df = right_df , left_df <EOL> if left_df . crs != right_df . crs : <EOL> print ( '<STR_LIT>' ) <EOL> tree_idx = rtree . index . Index ( ) <EOL> right_df_bounds = right_df [ '<STR_LIT>' ] . apply ( lambda x : x . bounds ) <EOL> for i in right_df_bounds . index : <EOL> tree_idx . insert ( i , right_df_bounds [ i ] ) <EOL> idxmatch = ( left_df [ '<STR_LIT>' ] . apply ( lambda x : x . bounds ) <EOL> . apply ( lambda x : list ( tree_idx . intersection ( x ) ) ) ) <EOL> idxmatch = idxmatch [ idxmatch . apply ( len ) > <NUM_LIT:0> ] <EOL> r_idx = np . concatenate ( idxmatch . values ) <EOL> l_idx = np . concatenate ( [ [ i ] * len ( v ) for i , v in idxmatch . iteritems ( ) ] ) <EOL> def find_intersects ( a1 , a2 ) : <EOL> return a1 . intersects ( a2 ) <EOL> def find_contains ( a1 , a2 ) : <EOL> return a1 . contains ( a2 ) <EOL> predicate_d = { '<STR_LIT>' : find_intersects , <EOL> '<STR_LIT>' : find_contains , <EOL> '<STR_LIT>' : find_contains } <EOL> check_predicates = np . vectorize ( predicate_d [ op ] ) <EOL> result = ( <EOL> pd . DataFrame ( <EOL> np . column_stack ( <EOL> [ l_idx , <EOL> r_idx , <EOL> check_predicates ( <EOL> left_df [ '<STR_LIT>' ] <EOL> . apply ( lambda x : prepared . prep ( x ) ) [ l_idx ] , <EOL> right_df [ '<STR_LIT>' ] [ r_idx ] ) <EOL> ] ) ) <EOL> ) <EOL> result . columns = [ '<STR_LIT>' % lsuffix , '<STR_LIT>' % rsuffix , '<STR_LIT>' ] <EOL> result = ( <EOL> pd . DataFrame ( result [ result [ '<STR_LIT>' ] == <NUM_LIT:1> ] ) <EOL> . drop ( '<STR_LIT>' , axis = <NUM_LIT:1> ) <EOL> ) <EOL> if op == "<STR_LIT>" : <EOL> left_df , right_df = right_df , left_df <EOL> result = result . rename ( columns = { <EOL> '<STR_LIT>' % ( lsuffix ) : '<STR_LIT>' % ( rsuffix ) , <EOL> '<STR_LIT>' % ( rsuffix ) : '<STR_LIT>' % ( lsuffix ) } ) <EOL> if how == '<STR_LIT>' : <EOL> result = result . set_index ( '<STR_LIT>' % lsuffix ) <EOL> return ( <EOL> left_df <EOL> . merge ( result , left_index = True , right_index = True ) <EOL> . merge ( right_df . drop ( '<STR_LIT>' , axis = <NUM_LIT:1> ) , <EOL> left_on = '<STR_LIT>' % rsuffix , right_index = True , <EOL> suffixes = ( '<STR_LIT>' % lsuffix , '<STR_LIT>' % rsuffix ) ) <EOL> ) <EOL> elif how == '<STR_LIT:left>' : <EOL> result = result . set_index ( '<STR_LIT>' % lsuffix ) <EOL> return ( <EOL> left_df <EOL> . merge ( result , left_index = True , right_index = True , how = '<STR_LIT:left>' ) <EOL> . merge ( right_df . drop ( '<STR_LIT>' , axis = <NUM_LIT:1> ) , <EOL> how = '<STR_LIT:left>' , left_on = '<STR_LIT>' % rsuffix , right_index = True , <EOL> suffixes = ( '<STR_LIT>' % lsuffix , '<STR_LIT>' % rsuffix ) ) <EOL> ) <EOL> elif how == '<STR_LIT:right>' : <EOL> return ( <EOL> left_df <EOL> . drop ( '<STR_LIT>' , axis = <NUM_LIT:1> ) <EOL> . merge ( result . merge ( right_df , <EOL> left_on = '<STR_LIT>' % rsuffix , right_index = True , <EOL> how = '<STR_LIT:right>' ) , left_index = True , <EOL> right_on = '<STR_LIT>' % lsuffix , how = '<STR_LIT:right>' ) <EOL> . set_index ( '<STR_LIT>' % rsuffix ) <EOL> ) </s>
