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

pytcl/containers/base.py

@@ -0,0 +1,219 @@
+"""
+Base classes for spatial data structures.
+
+This module provides abstract base classes that define the common interface
+for spatial indexing data structures like KD-trees, VP-trees, R-trees, and
+Cover trees.
+"""
+
+import logging
+from abc import ABC, abstractmethod
+from typing import Callable, List, NamedTuple, Optional
+
+import numpy as np
+from numpy.typing import ArrayLike, NDArray
+
+# Module logger
+_logger = logging.getLogger("pytcl.containers")
+
+
+class SpatialQueryResult(NamedTuple):
+    """Result of a spatial query.
+
+    Attributes
+    ----------
+    indices : ndarray
+        Indices of matching points in the original data.
+    distances : ndarray
+        Distances to matching points.
+    """
+
+    indices: NDArray[np.intp]
+    distances: NDArray[np.floating]
+
+
+class BaseSpatialIndex(ABC):
+    """
+    Abstract base class for spatial indexing data structures.
+
+    All spatial index implementations (KDTree, VPTree, RTree, CoverTree)
+    should inherit from this class and implement the required methods.
+
+    This provides a consistent interface for:
+    - Building the index from point data
+    - k-nearest neighbor queries
+    - Range/radius queries
+    - Dimension and size introspection
+
+    Parameters
+    ----------
+    data : array_like
+        Data points of shape (n_samples, n_features).
+
+    Attributes
+    ----------
+    data : ndarray
+        The indexed data points.
+    n_samples : int
+        Number of data points.
+    n_features : int
+        Dimensionality of data points.
+    """
+
+    def __init__(self, data: ArrayLike):
+        self.data = np.asarray(data, dtype=np.float64)
+
+        if self.data.ndim != 2:
+            raise ValueError(
+                f"Data must be 2-dimensional (n_samples, n_features), "
+                f"got shape {self.data.shape}"
+            )
+
+        self.n_samples, self.n_features = self.data.shape
+        _logger.debug(
+            "%s initialized with %d points in %d dimensions",
+            self.__class__.__name__,
+            self.n_samples,
+            self.n_features,
+        )
+
+    @abstractmethod
+    def query(
+        self,
+        X: ArrayLike,
+        k: int = 1,
+    ) -> SpatialQueryResult:
+        """
+        Query the index for k nearest neighbors.
+
+        Parameters
+        ----------
+        X : array_like
+            Query points of shape (n_queries, n_features) or (n_features,).
+        k : int, optional
+            Number of nearest neighbors to return. Default is 1.
+
+        Returns
+        -------
+        result : SpatialQueryResult
+            Named tuple with indices and distances of k nearest neighbors
+            for each query point.
+        """
+        pass
+
+    @abstractmethod
+    def query_radius(
+        self,
+        X: ArrayLike,
+        r: float,
+    ) -> List[List[int]]:
+        """
+        Query the index for all points within radius r.
+
+        Parameters
+        ----------
+        X : array_like
+            Query points of shape (n_queries, n_features) or (n_features,).
+        r : float
+            Search radius.
+
+        Returns
+        -------
+        indices : list of list of int
+            For each query point, a list of indices of data points
+            within distance r.
+        """
+        pass
+
+    def __len__(self) -> int:
+        """Return number of indexed points."""
+        return self.n_samples
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}("
+            f"n_samples={self.n_samples}, n_features={self.n_features})"
+        )
+
+
+class MetricSpatialIndex(BaseSpatialIndex):
+    """
+    Base class for metric space spatial indices.
+
+    Extends BaseSpatialIndex with support for custom distance metrics.
+    Used by VP-trees and Cover trees which can work with any metric.
+
+    Parameters
+    ----------
+    data : array_like
+        Data points of shape (n_samples, n_features).
+    metric : callable, optional
+        Distance function with signature metric(x, y) -> float.
+        Default is Euclidean distance.
+    """
+
+    def __init__(
+        self,
+        data: ArrayLike,
+        metric: Optional[Callable[[NDArray, NDArray], float]] = None,
+    ):
+        super().__init__(data)
+
+        if metric is None:
+            self.metric = self._euclidean_distance
+        else:
+            self.metric = metric
+
+    @staticmethod
+    def _euclidean_distance(x: NDArray, y: NDArray) -> float:
+        """Default Euclidean distance metric."""
+        return float(np.sqrt(np.sum((x - y) ** 2)))
+
+
+def validate_query_input(
+    X: ArrayLike,
+    n_features: int,
+) -> NDArray[np.floating]:
+    """
+    Validate and reshape query input.
+
+    Parameters
+    ----------
+    X : array_like
+        Query points.
+    n_features : int
+        Expected number of features.
+
+    Returns
+    -------
+    X : ndarray
+        Validated query array of shape (n_queries, n_features).
+
+    Raises
+    ------
+    ValueError
+        If query has wrong number of features.
+    """
+    X = np.asarray(X, dtype=np.float64)
+
+    if X.ndim == 1:
+        X = X.reshape(1, -1)
+
+    if X.shape[1] != n_features:
+        _logger.warning(
+            "Query feature mismatch: got %d, expected %d", X.shape[1], n_features
+        )
+        raise ValueError(f"Query has {X.shape[1]} features, expected {n_features}")
+
+    _logger.debug(
+        "Validated query input: %d queries, %d features", X.shape[0], X.shape[1]
+    )
+    return X
+
+
+__all__ = [
+    "SpatialQueryResult",
+    "BaseSpatialIndex",
+    "MetricSpatialIndex",
+    "validate_query_input",
+]

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],