fastquadtree 1.5.0__cp38-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl

fastquadtree/_obj_store.py

@@ -0,0 +1,194 @@
+# _bimap.py
+from __future__ import annotations
+
+from operator import itemgetter
+from typing import Any, Generic, Iterable, Iterator, Sequence, TypeVar
+
+from ._item import Item  # base class for PointItem and RectItem
+
+TItem = TypeVar("TItem", bound=Item)
+
+
+class ObjStore(Generic[TItem]):
+    """
+    High-performance id <-> object store for dense, auto-assigned ids.
+
+    Storage
+      - _arr[id]  -> Item or None
+      - _objs[id] -> Python object or None
+      - _obj_to_id: reverse identity map id(obj) -> id
+      - _free: LIFO free-list of reusable ids
+
+    Assumptions
+      - Ids are assigned by the shim and are dense [0..len) with possible holes
+        created by deletes. New inserts reuse holes via the free-list.
+    """
+
+    __slots__ = ("_arr", "_free", "_len", "_obj_to_id", "_objs")
+
+    def __init__(self, items: Iterable[TItem] | None = None) -> None:
+        self._arr: list[TItem | None] = []
+        self._objs: list[Any | None] = []
+        self._obj_to_id: dict[int, int] = {}
+        self._free: list[int] = []  # LIFO
+        self._len: int = 0  # live items
+
+        if items:
+            for it in items:
+                self.add(it, handle_out_of_order=True)
+
+    # ---- Serialization ----
+    def to_dict(self) -> dict[str, Any]:
+        """
+        Serialize to a dict suitable for JSON or other serialization.
+
+        Returns:
+            A dict with 'items' key containing list of serialized items.
+        """
+        items = [it.to_dict() for it in self._arr if it is not None]
+        return {"items": items}
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any], item_factory: Any) -> ObjStore[TItem]:
+        """
+        Deserialize from a dict.
+
+        Args:
+            data: A dict with 'items' key containing list of serialized items.
+            item_factory: A callable that takes (id, obj) and returns an Item.
+
+        Returns:
+            An ObjStore instance populated with the deserialized items.
+        """
+        items = []
+        for item_data in data.get("items", []):
+            item = Item.from_dict(item_data)
+            items.append(item_factory(item.id_, item.geom, item.obj))
+        return cls(items)
+
+    # -------- core --------
+
+    def add(self, item: TItem, handle_out_of_order: bool = False) -> None:
+        """
+        Insert or replace the mapping at item.id_. Reverse map updated so obj points to id.
+        """
+        id_ = item.id_
+        obj = item.obj
+
+        # ids must be dense and assigned by the caller
+        if id_ > len(self._arr):
+            if not handle_out_of_order:
+                raise AssertionError(
+                    "ObjStore.add received an out-of-order id, use alloc_id() to get the next available id"
+                )
+            # fill holes with None
+            while len(self._arr) < id_:
+                self._arr.append(None)
+                self._objs.append(None)
+
+        if id_ == len(self._arr):
+            # append
+            self._arr.append(item)
+            self._objs.append(obj)
+            self._len += 1
+        else:
+            # replace or fill a hole
+            old = self._arr[id_]
+            if old is None:
+                self._len += 1
+            elif old.obj is not None:
+                self._obj_to_id.pop(id(old.obj), None)
+            self._arr[id_] = item
+            self._objs[id_] = obj
+
+        if obj is not None:
+            self._obj_to_id[id(obj)] = id_
+
+    def by_id(self, id_: int) -> TItem | None:
+        return self._arr[id_] if 0 <= id_ < len(self._arr) else None
+
+    def by_obj(self, obj: Any) -> TItem | None:
+        id_ = self._obj_to_id.get(id(obj))
+        return self.by_id(id_) if id_ is not None else None
+
+    def pop_id(self, id_: int) -> TItem | None:
+        """Remove by id. Dense ids go to the free-list for reuse."""
+        if not (0 <= id_ < len(self._arr)):
+            return None
+        it = self._arr[id_]
+        if it is None:
+            return None
+        self._arr[id_] = None
+        self._objs[id_] = None
+        if it.obj is not None:
+            self._obj_to_id.pop(id(it.obj), None)
+        self._free.append(id_)
+        self._len -= 1
+        return it
+
+    # -------- allocation --------
+
+    def alloc_id(self) -> int:
+        """
+        Get a reusable dense id. Uses free-list else appends at the tail.
+        Build your Item with this id then call add(item).
+        """
+        return self._free.pop() if self._free else len(self._arr)
+
+    # -------- fast batch gathers --------
+
+    def get_many_by_ids(self, ids: Sequence[int], *, chunk: int = 2048) -> list[TItem]:
+        """
+        Batch: return Items for ids, preserving order.
+        Uses C-level itemgetter on the dense array in chunks.
+        """
+        out: list[TItem] = []
+        extend = out.extend
+        arr = self._arr
+        for i in range(0, len(ids), chunk):
+            block = ids[i : i + chunk]
+            vals = itemgetter(*block)(arr)  # tuple or single item
+            extend(vals if isinstance(vals, tuple) else (vals,))
+        return out
+
+    def get_many_objects(self, ids: Sequence[int], *, chunk: int = 2048) -> list[Any]:
+        """
+        Batch: return Python objects for ids, preserving order.
+        Mirrors get_many_by_ids but reads from _objs.
+        """
+
+        out: list[Any] = []
+        extend = out.extend
+        objs = self._objs
+        for i in range(0, len(ids), chunk):
+            block = ids[i : i + chunk]
+            vals = itemgetter(*block)(objs)  # tuple or single object
+            extend(vals if isinstance(vals, tuple) else (vals,))
+        return out
+
+    # -------- convenience and iteration --------
+
+    def __len__(self) -> int:
+        return self._len
+
+    def clear(self) -> None:
+        self._arr.clear()
+        self._objs.clear()
+        self._obj_to_id.clear()
+        self._free.clear()
+        self._len = 0
+
+    def contains_id(self, id_: int) -> bool:
+        return 0 <= id_ < len(self._arr) and self._arr[id_] is not None
+
+    def contains_obj(self, obj: Any) -> bool:
+        return id(obj) in self._obj_to_id
+
+    def items_by_id(self) -> Iterator[tuple[int, TItem]]:
+        for i, it in enumerate(self._arr):
+            if it is not None:
+                yield i, it
+
+    def items(self) -> Iterator[TItem]:
+        for _, it in self.items_by_id():
+            yield it

fastquadtree/point_quadtree.py

@@ -0,0 +1,206 @@
+# point_quadtree.py
+from __future__ import annotations
+
+from typing import Any, Literal, SupportsFloat, Tuple, overload
+
+from ._base_quadtree import Bounds, _BaseQuadTree
+from ._item import Point, PointItem
+from ._native import QuadTree as QuadTreeF32, QuadTreeF64, QuadTreeI32, QuadTreeI64
+
+_IdCoord = Tuple[int, SupportsFloat, SupportsFloat]
+
+DTYPE_MAP = {
+    "f32": QuadTreeF32,
+    "f64": QuadTreeF64,
+    "i32": QuadTreeI32,
+    "i64": QuadTreeI64,
+}
+
+
+class QuadTree(_BaseQuadTree[Point, _IdCoord, PointItem]):
+    """
+    Point version of the quadtree. All geometries are 2D points (x, y).
+    High-level Python wrapper over the Rust quadtree engine.
+
+    Performance characteristics:
+        Inserts: average O(log n) <br>
+        Rect queries: average O(log n + k) where k is matches returned <br>
+        Nearest neighbor: average O(log n) <br>
+
+    Thread-safety:
+        Instances are not thread-safe. Use external synchronization if you
+        mutate the same tree from multiple threads.
+
+    Args:
+        bounds: World bounds as (min_x, min_y, max_x, max_y).
+        capacity: Max number of points per node before splitting.
+        max_depth: Optional max tree depth. If omitted, engine decides.
+        track_objects: Enable id <-> object mapping inside Python.
+        dtype: Data type for coordinates and ids in the native engine. Default is 'f32'. Options are 'f32', 'f64', 'i32', 'i64'.
+
+    Raises:
+        ValueError: If parameters are invalid or inserts are out of bounds.
+    """
+
+    def __init__(
+        self,
+        bounds: Bounds,
+        capacity: int,
+        *,
+        max_depth: int | None = None,
+        track_objects: bool = False,
+        dtype: str = "f32",
+    ):
+        super().__init__(
+            bounds,
+            capacity,
+            max_depth=max_depth,
+            track_objects=track_objects,
+            dtype=dtype,
+        )
+
+    @overload
+    def query(
+        self, rect: Bounds, *, as_items: Literal[False] = ...
+    ) -> list[_IdCoord]: ...
+    @overload
+    def query(self, rect: Bounds, *, as_items: Literal[True]) -> list[PointItem]: ...
+    def query(
+        self, rect: Bounds, *, as_items: bool = False
+    ) -> list[PointItem] | list[_IdCoord]:
+        """
+        Return all points inside an axis-aligned rectangle.
+
+        Args:
+            rect: Query rectangle as (min_x, min_y, max_x, max_y).
+            as_items: If True, return Item wrappers. If False, return raw tuples.
+
+        Returns:
+            If as_items is False: list of (id, x, y) tuples.
+            If as_items is True: list of Item objects.
+
+        Example:
+            ```python
+            results = qt.query((10.0, 10.0, 20.0, 20.0), as_items=True)
+            for item in results:
+                print(f"Found point id={item.id_} at {item.geom} with obj={item.obj}")
+            ```
+        """
+        if not as_items:
+            return self._native.query(rect)
+        if self._store is None:
+            raise ValueError("Cannot return results as items with track_objects=False")
+        return self._store.get_many_by_ids(self._native.query_ids(rect))
+
+    def query_np(self, rect: Bounds) -> tuple[Any, Any]:
+        """
+        Return all points inside an axis-aligned rectangle as NumPy arrays.
+        The first array is an array of IDs, and the second is a corresponding array of point coordinates.
+
+        Requirements:
+            NumPy must be installed to use this method.
+
+        Args:
+            rect: Query rectangle as (min_x, min_y, max_x, max_y).
+
+        Returns:
+            Tuple of (ids, locations) where:
+                ids:       NDArray[np.uint64] with shape (N,)
+                locations: NDArray[np.floating] with shape (N, 2)
+
+        Example:
+            ```python
+            ids, locations = qt.query_np((10.0, 10.0, 20.0, 20.0))
+            for id_, (x, y) in zip(ids, locations):
+                print(f"Found point id={id_} at ({x}, {y})")
+            ```
+        """
+
+        return self._native.query_np(rect)
+
+    @overload
+    def nearest_neighbor(
+        self, xy: Point, *, as_item: Literal[False] = ...
+    ) -> _IdCoord | None: ...
+    @overload
+    def nearest_neighbor(
+        self, xy: Point, *, as_item: Literal[True]
+    ) -> PointItem | None: ...
+    def nearest_neighbor(
+        self, xy: Point, *, as_item: bool = False
+    ) -> PointItem | _IdCoord | None:
+        """
+        Return the single nearest neighbor to the query point.
+
+        Args:
+            xy: Query point (x, y).
+            as_item: If True, return Item. If False, return (id, x, y).
+
+        Returns:
+            The nearest neighbor or None if the tree is empty.
+        """
+        t = self._native.nearest_neighbor(xy)
+        if t is None or not as_item:
+            return t
+        if self._store is None:
+            raise ValueError("Cannot return result as item with track_objects=False")
+        id_, _x, _y = t
+        it = self._store.by_id(id_)
+        if it is None:
+            raise RuntimeError("Internal error: missing tracked item")
+        return it
+
+    @overload
+    def nearest_neighbors(
+        self, xy: Point, k: int, *, as_items: Literal[False] = ...
+    ) -> list[_IdCoord]: ...
+    @overload
+    def nearest_neighbors(
+        self, xy: Point, k: int, *, as_items: Literal[True]
+    ) -> list[PointItem]: ...
+    def nearest_neighbors(
+        self, xy: Point, k: int, *, as_items: bool = False
+    ) -> list[PointItem] | list[_IdCoord]:
+        """
+        Return the k nearest neighbors to the query point in order of increasing distance.
+
+        Args:
+            xy: Query point (x, y).
+            k: Number of neighbors to return.
+            as_items: If True, return Item wrappers. If False, return raw tuples.
+
+        Returns:
+            If as_items is False: list of (id, x, y) tuples. <br>
+            If as_items is True: list of Item objects. <br>
+        """
+        raw = self._native.nearest_neighbors(xy, k)
+        if not as_items:
+            return raw
+        if self._store is None:
+            raise ValueError("Cannot return results as items with track_objects=False")
+        out: list[PointItem] = []
+        for id_, _x, _y in raw:
+            it = self._store.by_id(id_)
+            if it is None:
+                raise RuntimeError("Internal error: missing tracked item")
+            out.append(it)
+        return out
+
+    def _new_native(self, bounds: Bounds, capacity: int, max_depth: int | None) -> Any:
+        """Create the native engine instance."""
+        rust_cls = DTYPE_MAP.get(self._dtype)
+        if rust_cls is None:
+            raise TypeError(f"Unsupported dtype: {self._dtype}")
+        return rust_cls(bounds, capacity, max_depth)
+
+    @classmethod
+    def _new_native_from_bytes(cls, data: bytes, dtype: str = "f32") -> Any:
+        """Create a new native engine instance from serialized bytes."""
+        rust_cls = DTYPE_MAP.get(dtype)
+        if rust_cls is None:
+            raise TypeError(f"Unsupported dtype: {dtype}")
+        return rust_cls.from_bytes(data)
+
+    @staticmethod
+    def _make_item(id_: int, geom: Point, obj: Any | None) -> PointItem:
+        return PointItem(id_, geom, obj)

fastquadtree/pyqtree.py

@@ -0,0 +1,176 @@
+"""
+Python Shim to mimic the interface of pyqtree and allow for a
+drop-in replacement to fastquadtree.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from operator import itemgetter
+from typing import Any, SupportsFloat, Tuple
+
+from ._native import RectQuadTree
+
+Point = Tuple[SupportsFloat, SupportsFloat]  # only for type hints in docstrings
+
+# Default parameters from pyqtree
+MAX_ITEMS = 10
+MAX_DEPTH = 20
+
+
+# Helper to gather objects by ids in chunks
+# Performance improvement over list comprehension for large result sets
+# 2.945 median query time --> 2.030 median query time (500k items, 500 queries)
+def gather_objs(objs, ids, chunk=2048):
+    out = []
+    for i in range(0, len(ids), chunk):
+        getter = itemgetter(*ids[i : i + chunk])
+        vals = getter(objs)  # tuple or single object
+        if isinstance(vals, tuple):
+            out.extend(vals)
+        else:
+            out.append(vals)
+    return out
+
+
+class Index:
+    """
+    The interface of the class below is taken from the pyqtree package, but the implementation
+    has been modified to use the fastquadtree package as a backend instead of
+    the original pure-python implementation.
+    Based on the benchmarks, this gives a overall performance boost of 6.514x.
+    See the benchmark section of the docs for more details and the latest numbers.
+
+    Index is  the top-level class for creating and using a quadtree spatial index
+    with the original pyqtree interface. If you are not migrating from pyqtree,
+    consider using the RectQuadTree class for detailed control and better performance.
+
+    This class wraps a RectQuadTree instance and provides methods to insert items with bounding boxes,
+    remove items, and query for items intersecting a given bounding box.
+
+    Example usage:
+    ```python
+    from fastquadtree.pyqtree import Index
+
+
+    spindex = Index(bbox=(0, 0, 100, 100))
+    spindex.insert('duck', (50, 30, 53, 60))
+    spindex.insert('cookie', (10, 20, 15, 25))
+    spindex.insert('python', (40, 50, 95, 90))
+    results = spindex.intersect((51, 51, 86, 86))
+    sorted(results) # ['duck', 'python']
+    ```
+    """
+
+    __slots__ = ("_free", "_item_to_id", "_objects", "_qt")
+
+    def __init__(
+        self,
+        bbox: Iterable[SupportsFloat] | None = None,
+        x: float | int | None = None,
+        y: float | int | None = None,
+        width: float | int | None = None,
+        height: float | int | None = None,
+        max_items: int = MAX_ITEMS,
+        max_depth: int = MAX_DEPTH,
+    ):
+        """
+        Initiate by specifying either 1) a bbox to keep track of, or 2) with an xy centerpoint and a width and height.
+
+        Args:
+          bbox: The coordinate system bounding box of the area that the quadtree should
+            keep track of, as a 4-length sequence (xmin,ymin,xmax,ymax)
+          x:
+            The x center coordinate of the area that the quadtree should keep track of.
+          y:
+            The y center coordinate of the area that the quadtree should keep track of.
+          width:
+            How far from the xcenter that the quadtree should look when keeping track.
+          height:
+            How far from the ycenter that the quadtree should look when keeping track
+          max_items (optional): The maximum number of items allowed per quad before splitting
+              up into four new subquads. Default is 10.
+          max_depth (optional): The maximum levels of nested subquads, after which no more splitting
+            occurs and the bottommost quad nodes may grow indefinately. Default is 20.
+
+        Note:
+            Either the bbox argument must be set, or the x, y, width, and height
+            arguments must be set.
+        """
+        if bbox is not None:
+            x1, y1, x2, y2 = bbox
+            self._qt = RectQuadTree((x1, y1, x2, y2), max_items, max_depth=max_depth)
+
+        elif (
+            x is not None and y is not None and width is not None and height is not None
+        ):
+            self._qt = RectQuadTree(
+                (x - width / 2, y - height / 2, x + width / 2, y + height / 2),
+                max_items,
+                max_depth=max_depth,
+            )
+
+        else:
+            raise ValueError(
+                "Either the bbox argument must be set, or the x, y, width, and height arguments must be set"
+            )
+
+        self._objects = []
+        self._free = []
+        self._item_to_id = {}
+
+    def insert(self, item: Any, bbox: Iterable[SupportsFloat]):
+        """
+        Inserts an item into the quadtree along with its bounding box.
+
+        Args:
+          item: The item to insert into the index, which will be returned by the intersection method
+          bbox: The spatial bounding box tuple of the item, with four members (xmin,ymin,xmax,ymax)
+        """
+        if type(bbox) is not tuple:  # Handle non-tuple input
+            bbox = tuple(bbox)
+
+        if self._free:
+            rid = self._free.pop()
+            self._objects[rid] = item
+        else:
+            rid = len(self._objects)
+            self._objects.append(item)
+        self._qt.insert(rid, bbox)
+        self._item_to_id[id(item)] = rid
+
+    def remove(self, item: Any, bbox: Iterable[SupportsFloat]):
+        """
+        Removes an item from the quadtree.
+
+        Args:
+          item: The item to remove from the index
+          bbox: The spatial bounding box tuple of the item, with four members (xmin,ymin,xmax,ymax)
+
+        Note:
+            Both parameters need to exactly match the parameters provided to the insert method.
+        """
+        if type(bbox) is not tuple:  # Handle non-tuple input
+            bbox = tuple(bbox)
+
+        rid = self._item_to_id.pop(id(item))
+        self._qt.delete(rid, bbox)
+        self._objects[rid] = None
+        self._free.append(rid)
+
+    def intersect(self, bbox: Iterable[SupportsFloat]) -> list:
+        """
+        Intersects an input bounding box rectangle with all of the items
+        contained in the quadtree.
+
+        Args:
+          bbox: A spatial bounding box tuple with four members (xmin,ymin,xmax,ymax)
+
+        Returns:
+          A list of inserted items whose bounding boxes intersect with the input bbox.
+        """
+        if type(bbox) is not tuple:  # Handle non-tuple input
+            bbox = tuple(bbox)
+        result = self._qt.query_ids(bbox)
+        # result = [id1, id2, ...]
+        return gather_objs(self._objects, result)