nrl-tracker 1.1.3__py3-none-any.whl → 1.2.0__py3-none-any.whl

pytcl/containers/covertree.py

@@ -11,11 +11,17 @@ References
        neighbor," ICML 2006.
 """
 
+import logging
 from typing import Callable, List, NamedTuple, Optional, Set, 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.covertree")
+
 
 class CoverTreeResult(NamedTuple):
     """Result of Cover tree query.
@@ -60,7 +66,7 @@ class CoverTreeNode:
         self.children[level].append(child)
 
 
-class CoverTree:
+class CoverTree(MetricSpatialIndex):
     """
     Cover Tree for metric space nearest neighbor search.
 
@@ -93,6 +99,11 @@ class CoverTree:
 
     The implementation uses a simplified version of the original
     algorithm for clarity.
+
+    See Also
+    --------
+    MetricSpatialIndex : Abstract base class for metric-based spatial indices.
+    VPTree : Alternative metric space index using vantage points.
     """
 
     def __init__(
@@ -101,19 +112,9 @@ class CoverTree:
         metric: Optional[Callable[[NDArray, NDArray], float]] = None,
         base: float = 2.0,
     ):
-        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, metric)
         self.base = base
 
-        if metric is None:
-            self.metric = self._euclidean_distance
-        else:
-            self.metric = metric
-
         # Compute distance cache for small datasets
         self._distance_cache: dict[Tuple[int, int], float] = {}
 
@@ -124,10 +125,12 @@ class CoverTree:
 
         if self.n_samples > 0:
             self._build_tree()
-
-    def _euclidean_distance(self, x: NDArray, y: NDArray) -> float:
-        """Default Euclidean distance metric."""
-        return float(np.sqrt(np.sum((x - y) ** 2)))
+            _logger.debug(
+                "CoverTree built with base=%.1f, levels=%d to %d",
+                base,
+                self.min_level,
+                self.max_level,
+            )
 
     def _distance(self, i: int, j: int) -> float:
         """Get distance between points i and j (with caching)."""
@@ -245,11 +248,7 @@ class CoverTree:
         result : CoverTreeResult
             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)
@@ -368,11 +367,7 @@ class CoverTree:
         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]] = []
 

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/vptree.py

@@ -11,11 +11,17 @@ References
        neighbor search in general metric spaces," SODA 1993.
 """
 
+import logging
 from typing import 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,6 +95,11 @@ 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__(
@@ -96,25 +107,13 @@ class VPTree:
         data: ArrayLike,
         metric: Optional[Callable[[NDArray, NDArray], 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 +168,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 +254,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]] = []
 

pytcl/core/__init__.py

@@ -24,10 +24,19 @@ from pytcl.core.constants import (
     PhysicalConstants,
 )
 from pytcl.core.validation import (
+    ArraySpec,
+    ScalarSpec,
+    ValidationError,
+    check_compatible_shapes,
     ensure_2d,
     ensure_column_vector,
+    ensure_positive_definite,
     ensure_row_vector,
+    ensure_square_matrix,
+    ensure_symmetric,
     validate_array,
+    validate_inputs,
+    validate_same_shape,
 )
 
 __all__ = [
@@ -40,10 +49,19 @@ __all__ = [
     "WGS84",
     "PhysicalConstants",
     # Validation
+    "ValidationError",
     "validate_array",
+    "validate_inputs",
+    "validate_same_shape",
+    "check_compatible_shapes",
+    "ArraySpec",
+    "ScalarSpec",
     "ensure_2d",
     "ensure_column_vector",
     "ensure_row_vector",
+    "ensure_square_matrix",
+    "ensure_symmetric",
+    "ensure_positive_definite",
     # Array utilities
     "wrap_to_pi",
     "wrap_to_2pi",