nrl-tracker 0.22.5__py3-none-any.whl → 1.8.0__py3-none-any.whl

pytcl/containers/kd_tree.py

@@ -13,11 +13,17 @@ References
        Finding Best Matches in Logarithmic Expected Time," ACM TOMS, 1977.
 """
 
+import logging
 from typing import List, NamedTuple, Optional, Tuple
 
 import numpy as np
 from numpy.typing import ArrayLike, NDArray
 
+from pytcl.containers.base import BaseSpatialIndex, validate_query_input
+
+# Module logger
+_logger = logging.getLogger("pytcl.containers.kd_tree")
+
 
 class KDNode:
     """A node in the k-d tree.
@@ -66,7 +72,7 @@ class NearestNeighborResult(NamedTuple):
     distances: NDArray[np.floating]
 
 
-class KDTree:
+class KDTree(BaseSpatialIndex):
     """
     K-D Tree for efficient spatial queries.
 
@@ -97,6 +103,11 @@ class KDTree:
 
     Query complexity is O(log n) on average for nearest neighbor search,
     though worst case is O(n) for highly unbalanced queries.
+
+    See Also
+    --------
+    BaseSpatialIndex : Abstract base class defining the spatial index interface.
+    BallTree : Alternative spatial index using hyperspheres.
     """
 
     def __init__(
@@ -104,17 +115,13 @@ class KDTree:
         data: ArrayLike,
         leaf_size: int = 10,
     ):
-        self.data = np.asarray(data, dtype=np.float64)
-
-        if self.data.ndim != 2:
-            raise ValueError("Data must be 2-dimensional (n_samples, n_features)")
-
-        self.n_samples, self.n_features = self.data.shape
+        super().__init__(data)
         self.leaf_size = leaf_size
 
         # Build the tree
         indices = np.arange(self.n_samples)
         self.root = self._build_tree(indices, depth=0)
+        _logger.debug("KDTree built with leaf_size=%d", leaf_size)
 
     def _build_tree(
         self,
@@ -173,12 +180,9 @@ class KDTree:
         >>> result.indices
         array([[0, 1]])
         """
-        X = np.asarray(X, dtype=np.float64)
-
-        if X.ndim == 1:
-            X = X.reshape(1, -1)
-
+        X = validate_query_input(X, self.n_features)
         n_queries = X.shape[0]
+        _logger.debug("KDTree.query: %d queries, k=%d", n_queries, k)
 
         all_indices = np.zeros((n_queries, k), dtype=np.intp)
         all_distances = np.full((n_queries, k), np.inf)
@@ -263,11 +267,7 @@ class KDTree:
         >>> tree.query_radius([[0, 0]], r=1.5)
         [[0, 1, 2]]
         """
-        X = np.asarray(X, dtype=np.float64)
-
-        if X.ndim == 1:
-            X = X.reshape(1, -1)
-
+        X = validate_query_input(X, self.n_features)
         n_queries = X.shape[0]
         results: List[List[int]] = []
 
@@ -331,7 +331,7 @@ class KDTree:
         return self.query_radius(X, r)
 
 
-class BallTree:
+class BallTree(BaseSpatialIndex):
     """
     Ball Tree for efficient spatial queries.
 
@@ -357,6 +357,11 @@ class BallTree:
     -----
     Ball trees have O(n log n) construction and O(log n) average-case
     query time. They can outperform k-d trees in high dimensions.
+
+    See Also
+    --------
+    BaseSpatialIndex : Abstract base class defining the spatial index interface.
+    KDTree : Alternative spatial index using axis-aligned splits.
     """
 
     def __init__(
@@ -364,12 +369,7 @@ class BallTree:
         data: ArrayLike,
         leaf_size: int = 10,
     ):
-        self.data = np.asarray(data, dtype=np.float64)
-
-        if self.data.ndim != 2:
-            raise ValueError("Data must be 2-dimensional")
-
-        self.n_samples, self.n_features = self.data.shape
+        super().__init__(data)
         self.leaf_size = leaf_size
 
         # Build tree using indices
@@ -382,6 +382,7 @@ class BallTree:
         self._leaf_indices: List[Optional[NDArray[np.intp]]] = []
 
         self._build_tree(self._indices)
+        _logger.debug("BallTree built with leaf_size=%d", leaf_size)
 
     def _build_tree(
         self,
@@ -459,11 +460,7 @@ class BallTree:
         result : NearestNeighborResult
             Indices and distances of k nearest neighbors.
         """
-        X = np.asarray(X, dtype=np.float64)
-
-        if X.ndim == 1:
-            X = X.reshape(1, -1)
-
+        X = validate_query_input(X, self.n_features)
         n_queries = X.shape[0]
         all_indices = np.zeros((n_queries, k), dtype=np.intp)
         all_distances = np.full((n_queries, k), np.inf)
@@ -478,6 +475,74 @@ class BallTree:
 
         return NearestNeighborResult(indices=all_indices, distances=all_distances)
 
+    def query_radius(
+        self,
+        X: ArrayLike,
+        r: float,
+    ) -> List[List[int]]:
+        """
+        Query the tree for all points within radius r.
+
+        Parameters
+        ----------
+        X : array_like
+            Query points of shape (n_queries, n_features) or (n_features,).
+        r : float
+            Query radius.
+
+        Returns
+        -------
+        indices : list of lists
+            For each query, a list of indices of points within radius r.
+        """
+        X = validate_query_input(X, self.n_features)
+        n_queries = X.shape[0]
+        results: List[List[int]] = []
+
+        for i in range(n_queries):
+            indices = self._query_radius_single(X[i], r)
+            results.append(indices)
+
+        return results
+
+    def _query_radius_single(
+        self,
+        query: NDArray[np.floating],
+        r: float,
+    ) -> List[int]:
+        """Find all points within radius r of query point."""
+        indices: List[int] = []
+
+        def _search(node_id: int) -> None:
+            if node_id < 0:
+                return
+
+            centroid = self._centroids[node_id]
+            radius = self._radii[node_id]
+
+            # Distance to ball surface
+            dist_to_center = np.sqrt(np.sum((query - centroid) ** 2))
+
+            # Prune if ball is farther than radius
+            if dist_to_center - radius > r:
+                return
+
+            if self._is_leaf[node_id]:
+                # Check all points in leaf
+                leaf_indices = self._leaf_indices[node_id]
+                if leaf_indices is not None:
+                    for idx in leaf_indices:
+                        dist = np.sqrt(np.sum((query - self.data[idx]) ** 2))
+                        if dist <= r:
+                            indices.append(idx)
+            else:
+                # Visit both children
+                _search(self._left[node_id])
+                _search(self._right[node_id])
+
+        _search(0)
+        return indices
+
     def _query_single(
         self,
         query: NDArray[np.floating],

pytcl/containers/rtree.py

@@ -13,11 +13,17 @@ References
        Method for Points and Rectangles," ACM SIGMOD, 1990.
 """
 
+import logging
 from typing import List, NamedTuple, Optional, Tuple
 
 import numpy as np
 from numpy.typing import ArrayLike, NDArray
 
+from pytcl.containers.base import SpatialQueryResult, validate_query_input
+
+# Module logger
+_logger = logging.getLogger("pytcl.containers.rtree")
+
 
 class BoundingBox(NamedTuple):
     """Axis-aligned bounding box.
@@ -160,6 +166,9 @@ class RTree:
     An R-tree groups nearby objects and represents them with their
     minimum bounding rectangle. This allows efficient spatial queries.
 
+    Unlike KDTree and BallTree which only index points, RTree can index
+    bounding boxes of arbitrary size. It also supports dynamic insertion.
+
     Parameters
     ----------
     max_entries : int, optional
@@ -167,6 +176,13 @@ class RTree:
     min_entries : int, optional
         Minimum entries per node (except root). Default max_entries // 2.
 
+    Attributes
+    ----------
+    n_entries : int
+        Number of entries in the tree.
+    n_features : int
+        Dimensionality of the data (set after first insertion).
+
     Examples
     --------
     >>> tree = RTree()
@@ -179,6 +195,11 @@ class RTree:
     -----
     This implementation uses a simplified insertion algorithm.
     For production use, consider using R*-tree or packed R-tree variants.
+
+    See Also
+    --------
+    KDTree : Point-based spatial index using axis-aligned splits.
+    BallTree : Point-based spatial index using hyperspheres.
     """
 
     def __init__(
@@ -190,11 +211,70 @@ class RTree:
         self.min_entries = min_entries or max_entries // 2
         self.root = RTreeNode(is_leaf=True)
         self.n_entries = 0
+        self.n_features: Optional[int] = None
         self._data: List[BoundingBox] = []
+        self._points: Optional[NDArray[np.floating]] = None
+        _logger.debug("RTree initialized with max_entries=%d", max_entries)
+
+    @classmethod
+    def from_points(
+        cls,
+        data: ArrayLike,
+        max_entries: int = 10,
+        min_entries: Optional[int] = None,
+    ) -> "RTree":
+        """
+        Create an RTree from point data.
+
+        This factory method provides an interface similar to KDTree and BallTree,
+        allowing RTree to be used interchangeably for point queries.
+
+        Parameters
+        ----------
+        data : array_like
+            Data points of shape (n_samples, n_features).
+        max_entries : int, optional
+            Maximum entries per node. Default 10.
+        min_entries : int, optional
+            Minimum entries per node. Default max_entries // 2.
+
+        Returns
+        -------
+        tree : RTree
+            RTree with all points inserted.
+
+        Examples
+        --------
+        >>> points = np.array([[0, 0], [1, 0], [0, 1], [1, 1]])
+        >>> tree = RTree.from_points(points)
+        >>> result = tree.query([[0.1, 0.1]], k=2)
+        """
+        data = np.asarray(data, dtype=np.float64)
+        if data.ndim != 2:
+            raise ValueError(
+                f"Data must be 2-dimensional (n_samples, n_features), "
+                f"got shape {data.shape}"
+            )
+
+        tree = cls(max_entries=max_entries, min_entries=min_entries)
+        tree._points = data
+        tree.n_features = data.shape[1]
+        tree.insert_points(data)
+        _logger.debug(
+            "RTree.from_points: indexed %d points in %d dimensions",
+            data.shape[0],
+            data.shape[1],
+        )
+        return tree
 
     def __len__(self) -> int:
         return self.n_entries
 
+    def __repr__(self) -> str:
+        if self.n_features is not None:
+            return f"RTree(n_entries={self.n_entries}, n_features={self.n_features})"
+        return f"RTree(n_entries={self.n_entries})"
+
     def insert(self, bbox: BoundingBox, data_index: Optional[int] = None) -> int:
         """
         Insert a bounding box into the tree.
@@ -214,6 +294,10 @@ class RTree:
         if data_index is None:
             data_index = self.n_entries
 
+        # Track dimensionality
+        if self.n_features is None:
+            self.n_features = len(bbox.min_coords)
+
         self._data.append(bbox)
 
         # Find leaf to insert into
@@ -330,6 +414,121 @@ class RTree:
             current.update_bbox()
             current = current.parent
 
+    def query(
+        self,
+        X: ArrayLike,
+        k: int = 1,
+    ) -> SpatialQueryResult:
+        """
+        Query the tree for k nearest neighbors.
+
+        This method provides API compatibility with KDTree and BallTree.
+
+        Parameters
+        ----------
+        X : array_like
+            Query points of shape (n_queries, n_features) or (n_features,).
+        k : int, optional
+            Number of nearest neighbors. Default 1.
+
+        Returns
+        -------
+        result : SpatialQueryResult
+            Indices and distances of k nearest neighbors for each query.
+
+        Examples
+        --------
+        >>> tree = RTree.from_points(np.array([[0, 0], [1, 1], [2, 2]]))
+        >>> result = tree.query([[0.5, 0.5]], k=2)
+        >>> result.indices
+        array([[0, 1]])
+        """
+        if self.n_features is None:
+            raise ValueError("Cannot query empty RTree")
+
+        X = validate_query_input(X, self.n_features)
+        n_queries = X.shape[0]
+        _logger.debug("RTree.query: %d queries, k=%d", n_queries, k)
+
+        all_indices = np.zeros((n_queries, k), dtype=np.intp)
+        all_distances = np.full((n_queries, k), np.inf)
+
+        for i in range(n_queries):
+            indices, distances = self.nearest(X[i], k=k)
+            n_found = len(indices)
+            if n_found > 0:
+                all_indices[i, :n_found] = indices
+                all_distances[i, :n_found] = distances
+
+        return SpatialQueryResult(indices=all_indices, distances=all_distances)
+
+    def query_radius(
+        self,
+        X: ArrayLike,
+        r: float,
+    ) -> List[List[int]]:
+        """
+        Query the tree for all points within radius r.
+
+        This method provides API compatibility with KDTree and BallTree.
+
+        Parameters
+        ----------
+        X : array_like
+            Query points of shape (n_queries, n_features) or (n_features,).
+        r : float
+            Query radius.
+
+        Returns
+        -------
+        indices : list of lists
+            For each query, a list of indices of points within radius r.
+
+        Examples
+        --------
+        >>> tree = RTree.from_points(np.array([[0, 0], [1, 0], [0, 1], [5, 5]]))
+        >>> tree.query_radius([[0, 0]], r=1.5)
+        [[0, 1, 2]]
+        """
+        if self.n_features is None:
+            raise ValueError("Cannot query empty RTree")
+
+        X = validate_query_input(X, self.n_features)
+        n_queries = X.shape[0]
+        results: List[List[int]] = []
+
+        for i in range(n_queries):
+            query = X[i]
+            indices: List[int] = []
+
+            def search(node: RTreeNode) -> None:
+                if node.bbox is None:
+                    return
+
+                # Minimum distance from query point to node's bounding box
+                clamped = np.clip(query, node.bbox.min_coords, node.bbox.max_coords)
+                min_dist = float(np.sqrt(np.sum((query - clamped) ** 2)))
+
+                # Prune if node is entirely outside radius
+                if min_dist > r:
+                    return
+
+                if node.is_leaf:
+                    for bbox, idx in node.entries:
+                        # Distance to point (center of zero-volume box)
+                        clamped_pt = np.clip(query, bbox.min_coords, bbox.max_coords)
+                        dist = float(np.sqrt(np.sum((query - clamped_pt) ** 2)))
+                        if dist <= r:
+                            indices.append(idx)
+                else:
+                    for child in node.children:
+                        search(child)
+
+            search(self.root)
+            results.append(indices)
+
+        return results
+
     def query_intersect(self, query_bbox: BoundingBox) -> RTreeResult:
         """
         Find all entries intersecting a query box.
@@ -455,7 +654,7 @@ class RTree:
         query = np.asarray(query_point, dtype=np.float64)
         neighbors: List[Tuple[float, int]] = []
 
-        def min_dist_to_box(point: NDArray, bbox: BoundingBox) -> float:
+        def min_dist_to_box(point: NDArray[np.floating], bbox: BoundingBox) -> float:
             """Minimum distance from point to bounding box."""
             clamped = np.clip(point, bbox.min_coords, bbox.max_coords)
             return float(np.sqrt(np.sum((point - clamped) ** 2)))

pytcl/containers/vptree.py

@@ -11,11 +11,17 @@ References
        neighbor search in general metric spaces," SODA 1993.
 """
 
-from typing import Callable, List, NamedTuple, Optional, Tuple
+import logging
+from typing import Any, Callable, List, NamedTuple, Optional, Tuple
 
 import numpy as np
 from numpy.typing import ArrayLike, NDArray
 
+from pytcl.containers.base import MetricSpatialIndex, validate_query_input
+
+# Module logger
+_logger = logging.getLogger("pytcl.containers.vptree")
+
 
 class VPTreeResult(NamedTuple):
     """Result of VP-tree query.
@@ -56,7 +62,7 @@ class VPNode:
         self.right: Optional["VPNode"] = None
 
 
-class VPTree:
+class VPTree(MetricSpatialIndex):
     """
     Vantage Point Tree for metric space nearest neighbor search.
 
@@ -89,32 +95,27 @@ class VPTree:
 
     Query complexity is O(log n) on average but can degrade to O(n)
     for pathological distance distributions.
+
+    See Also
+    --------
+    MetricSpatialIndex : Abstract base class for metric-based spatial indices.
+    CoverTree : Alternative metric space index with theoretical guarantees.
     """
 
     def __init__(
         self,
         data: ArrayLike,
-        metric: Optional[Callable[[NDArray, NDArray], float]] = None,
+        metric: Optional[
+            Callable[[np.ndarray[Any, Any], np.ndarray[Any, Any]], float]
+        ] = None,
     ):
-        self.data = np.asarray(data, dtype=np.float64)
-
-        if self.data.ndim != 2:
-            raise ValueError("Data must be 2-dimensional")
-
-        self.n_samples, self.n_features = self.data.shape
-
-        if metric is None:
-            self.metric = self._euclidean_distance
-        else:
-            self.metric = metric
+        super().__init__(data, metric)
 
         # Build tree
         indices = np.arange(self.n_samples)
         self.root = self._build_tree(indices)
-
-    def _euclidean_distance(self, x: NDArray, y: NDArray) -> float:
-        """Default Euclidean distance metric."""
-        return float(np.sqrt(np.sum((x - y) ** 2)))
+        metric_name = metric.__name__ if metric else "euclidean"
+        _logger.debug("VPTree built with metric=%s", metric_name)
 
     def _build_tree(self, indices: NDArray[np.intp]) -> Optional[VPNode]:
         """Recursively build the VP-tree."""
@@ -169,11 +170,7 @@ class VPTree:
         result : VPTreeResult
             Indices and distances of k nearest neighbors.
         """
-        X = np.asarray(X, dtype=np.float64)
-
-        if X.ndim == 1:
-            X = X.reshape(1, -1)
-
+        X = validate_query_input(X, self.n_features)
         n_queries = X.shape[0]
 
         all_indices = np.zeros((n_queries, k), dtype=np.intp)
@@ -259,11 +256,7 @@ class VPTree:
         indices : list of lists
             For each query, list of indices within radius.
         """
-        X = np.asarray(X, dtype=np.float64)
-
-        if X.ndim == 1:
-            X = X.reshape(1, -1)
-
+        X = validate_query_input(X, self.n_features)
         n_queries = X.shape[0]
         results: List[List[int]] = []