1""" 2strtree 3======= 4 5Index geometry objects for efficient lookup of nearby or 6nearest neighbors. Home of the `STRtree` class which is 7an interface to the query-only GEOS R-tree packed using 8the Sort-Tile-Recursive algorithm [1]_. 9 10.. autoclass:: STRtree 11 :members: 12 13References 14---------- 15 .. [1] Leutenegger, Scott & Lopez, Mario & Edgington, Jeffrey. (1997). 16 "STR: A Simple and Efficient Algorithm for R-Tree Packing." Proc. 17 VLDB Conf. 497-506. 10.1109/ICDE.1997.582015. 18 https://www.cs.odu.edu/~mln/ltrs-pdfs/icase-1997-14.pdf 19""" 20 21import ctypes 22import logging 23from typing import Any, ItemsView, Iterable, Iterator, Optional, Sequence, Tuple, Union 24import sys 25 26from shapely.geometry.base import BaseGeometry 27from shapely.geos import lgeos 28 29log = logging.getLogger(__name__) 30 31 32class STRtree: 33 """An STR-packed R-tree spatial index. 34 35 An index is initialized from a sequence of geometry objects and 36 optionally an sequence of items. The items, if provided, are stored 37 in nodes of the tree. If items are not provided, the indices of the 38 geometry sequence will be used instead. 39 40 Stored items and corresponding geometry objects can be spatially 41 queried using another geometric object. 42 43 The tree is immutable and query-only, meaning that once created 44 nodes cannot be added or removed. 45 46 Parameters 47 ---------- 48 geoms : sequence 49 A sequence of geometry objects. 50 items : sequence, optional 51 A sequence of objects which typically serve as identifiers in an 52 application. This sequence must have the same length as geoms. 53 54 Attributes 55 ---------- 56 node_capacity : int 57 The maximum number of items per node. Default: 10. 58 59 Examples 60 -------- 61 Creating an index of polygons: 62 63 >>> from shapely.strtree import STRtree 64 >>> from shapely.geometry import Polygon 65 >>> 66 >>> polys = [Polygon(((0, 0), (1, 0), (1, 1))), 67 ... Polygon(((0, 1), (0, 0), (1, 0))), 68 ... Polygon(((100, 100), (101, 100), (101, 101)))] 69 >>> tree = STRtree(polys) 70 >>> query_geom = Polygon(((-1, -1), (2, 0), (2, 2), (-1, 2))) 71 >>> result = tree.query(query_geom) 72 >>> polys[0] in result 73 True 74 >>> polys[1] in result 75 True 76 >>> polys[2] in result 77 False 78 79 Notes 80 ----- 81 The class maintains a reverse mapping of items to geometries. These 82 items must therefore be hashable. The tree is filled using the 83 Sort-Tile-Recursive [1]_ algorithm. 84 85 References 86 ---------- 87 .. [1] Leutenegger, Scott T.; Edgington, Jeffrey M.; Lopez, Mario A. 88 (February 1997). "STR: A Simple and Efficient Algorithm for 89 R-Tree Packing". 90 https://ia600900.us.archive.org/27/items/nasa_techdoc_19970016975/19970016975.pdf 91 92 """ 93 94 def __init__( 95 self, 96 geoms: Iterable[BaseGeometry], 97 items: Iterable[Any] = None, 98 node_capacity: int = 10, 99 ): 100 self.node_capacity = node_capacity 101 102 # Keep references to geoms 103 self._geoms = list(geoms) 104 # Default enumeration index to store in the tree 105 self._idxs = list(range(len(self._geoms))) 106 107 # handle items 108 self._has_custom_items = items is not None 109 if not self._has_custom_items: 110 items = self._idxs 111 self._items = items 112 113 # initialize GEOS STRtree 114 self._tree = lgeos.GEOSSTRtree_create(self.node_capacity) 115 i = 0 116 for idx, geom in zip(self._idxs, self._geoms): 117 # filter empty geometries out of the input 118 if geom is not None and not geom.is_empty: 119 lgeos.GEOSSTRtree_insert(self._tree, geom._geom, ctypes.py_object(idx)) 120 i += 1 121 self._n_geoms = i 122 123 def __reduce__(self): 124 if self._has_custom_items: 125 return STRtree, (self._geoms, self._items) 126 else: 127 return STRtree, (self._geoms, ) 128 129 def __del__(self): 130 if self._tree is not None: 131 try: 132 lgeos.GEOSSTRtree_destroy(self._tree) 133 except AttributeError: 134 pass # lgeos might be empty on shutdown. 135 136 self._tree = None 137 138 def _query(self, geom): 139 if self._n_geoms == 0: 140 return [] 141 142 result = [] 143 144 def callback(item, userdata): 145 idx = ctypes.cast(item, ctypes.py_object).value 146 result.append(idx) 147 148 lgeos.GEOSSTRtree_query(self._tree, geom._geom, lgeos.GEOSQueryCallback(callback), None) 149 return result 150 151 def query_items(self, geom: BaseGeometry) -> Sequence[Any]: 152 """Query for nodes which intersect the geom's envelope to get 153 stored items. 154 155 Items are integers serving as identifiers for an application. 156 157 Parameters 158 ---------- 159 geom : geometry object 160 The query geometry. 161 162 Returns 163 ------- 164 An array or list of items stored in the tree. 165 166 Note 167 ---- 168 A geometry object's "envelope" is its minimum xy bounding 169 rectangle. 170 171 Examples 172 -------- 173 A buffer around a point can be used to control the extent 174 of the query. 175 176 >>> from shapely.strtree import STRtree 177 >>> from shapely.geometry import Point 178 >>> points = [Point(i, i) for i in range(10)] 179 >>> tree = STRtree(points) 180 >>> query_geom = Point(2,2).buffer(0.99) 181 >>> [o.wkt for o in tree.query(query_geom)] 182 ['POINT (2 2)'] 183 >>> query_geom = Point(2, 2).buffer(1.0) 184 >>> [o.wkt for o in tree.query(query_geom)] 185 ['POINT (1 1)', 'POINT (2 2)', 'POINT (3 3)'] 186 187 A subsequent search through the returned subset using the 188 desired binary predicate (eg. intersects, crosses, contains, 189 overlaps) may be necessary to further filter the results 190 according to their specific spatial relationships. 191 192 >>> [o.wkt for o in tree.query(query_geom) if o.intersects(query_geom)] 193 ['POINT (2 2)'] 194 195 """ 196 result = self._query(geom) 197 if self._has_custom_items: 198 return [self._items[i] for i in result] 199 else: 200 return result 201 202 def query_geoms(self, geom: BaseGeometry) -> Sequence[BaseGeometry]: 203 """Query for nodes which intersect the geom's envelope to get 204 geometries corresponding to the items stored in the nodes. 205 206 Parameters 207 ---------- 208 geom : geometry object 209 The query geometry. 210 211 Returns 212 ------- 213 An array or list of geometry objects. 214 215 """ 216 result = self._query(geom) 217 return [self._geoms[i] for i in result] 218 219 def query(self, geom: BaseGeometry) -> Sequence[BaseGeometry]: 220 """Query for nodes which intersect the geom's envelope to get 221 geometries corresponding to the items stored in the nodes. 222 223 This method is an alias for query_geoms. It may be removed in 224 version 2.0. 225 226 Parameters 227 ---------- 228 geom : geometry object 229 The query geometry. 230 231 Returns 232 ------- 233 An array or list of geometry objects. 234 235 """ 236 return self.query_geoms(geom) 237 238 def _nearest(self, geom, exclusive): 239 envelope = geom.envelope 240 241 def callback(item1, item2, distance, userdata): 242 try: 243 callback_userdata = ctypes.cast(userdata, ctypes.py_object).value 244 idx = ctypes.cast(item1, ctypes.py_object).value 245 geom2 = ctypes.cast(item2, ctypes.py_object).value 246 dist = ctypes.cast(distance, ctypes.POINTER(ctypes.c_double)) 247 if callback_userdata["exclusive"] and self._geoms[idx].equals(geom2): 248 dist[0] = sys.float_info.max 249 else: 250 lgeos.GEOSDistance(self._geoms[idx]._geom, geom2._geom, dist) 251 252 return 1 253 except Exception: 254 log.exception("Caught exception") 255 return 0 256 257 item = lgeos.GEOSSTRtree_nearest_generic( 258 self._tree, 259 ctypes.py_object(geom), 260 envelope._geom, 261 lgeos.GEOSDistanceCallback(callback), 262 ctypes.py_object({"exclusive": exclusive}), 263 ) 264 return ctypes.cast(item, ctypes.py_object).value 265 266 def nearest_item( 267 self, geom: BaseGeometry, exclusive: bool = False 268 ) -> Union[Any, None]: 269 """Query the tree for the node nearest to geom and get the item 270 stored in the node. 271 272 Items are integers serving as identifiers for an application. 273 274 Parameters 275 ---------- 276 geom : geometry object 277 The query geometry. 278 exclusive : bool, optional 279 Whether to exclude the item corresponding to the given geom 280 from results or not. Default: False. 281 282 Returns 283 ------- 284 Stored item or None. 285 286 None is returned if this index is empty. This may change in 287 version 2.0. 288 289 Examples 290 -------- 291 >>> from shapely.strtree import STRtree 292 >>> from shapely.geometry import Point 293 >>> tree = STRtree([Point(i, i) for i in range(10)]) 294 >>> tree.nearest(Point(2.2, 2.2)).wkt 295 'POINT (2 2)' 296 297 Will only return one object: 298 299 >>> tree = STRtree ([Point(0, 0), Point(0, 0)]) 300 >>> tree.nearest(Point(0, 0)).wkt 301 'POINT (0 0)' 302 303 """ 304 if self._n_geoms == 0: 305 return None 306 307 result = self._nearest(geom, exclusive) 308 if self._has_custom_items: 309 return self._items[result] 310 else: 311 return result 312 313 def nearest_geom( 314 self, geom: BaseGeometry, exclusive: bool = False 315 ) -> Union[BaseGeometry, None]: 316 """Query the tree for the node nearest to geom and get the 317 geometry corresponding to the item stored in the node. 318 319 Parameters 320 ---------- 321 geom : geometry object 322 The query geometry. 323 exclusive : bool, optional 324 Whether to exclude the given geom from results or not. 325 Default: False. 326 327 Returns 328 ------- 329 BaseGeometry or None. 330 331 None is returned if this index is empty. This may change in 332 version 2.0. 333 334 """ 335 result = self._nearest(geom, exclusive) 336 return self._geoms[result] 337 338 def nearest( 339 self, geom: BaseGeometry, exclusive: bool = False 340 ) -> Union[BaseGeometry, None]: 341 """Query the tree for the node nearest to geom and get the 342 geometry corresponding to the item stored in the node. 343 344 This method is an alias for nearest_geom. It may be removed in 345 version 2.0. 346 347 Parameters 348 ---------- 349 geom : geometry object 350 The query geometry. 351 exclusive : bool, optional 352 Whether to exclude the given geom from results or not. 353 Default: False. 354 355 Returns 356 ------- 357 BaseGeometry or None. 358 359 None is returned if this index is empty. This may change in 360 version 2.0. 361 362 """ 363 return self.nearest_geom(geom, exclusive=exclusive) 364