nrl-tracker 0.21.4__py3-none-any.whl → 1.7.5__py3-none-any.whl

pytcl/clustering/dbscan.py

@@ -12,7 +12,7 @@ References
        with Noise," KDD 1996.
 """
 
-from typing import List, NamedTuple, Set
+from typing import Any, List, NamedTuple, Set
 
 import numpy as np
 from numba import njit
@@ -42,7 +42,7 @@ class DBSCANResult(NamedTuple):
 
 
 @njit(cache=True)
-def _compute_distance_matrix(X: np.ndarray) -> np.ndarray:
+def _compute_distance_matrix(X: np.ndarray[Any, Any]) -> np.ndarray[Any, Any]:
     """Compute pairwise Euclidean distance matrix (JIT-compiled)."""
     n = X.shape[0]
     dist = np.zeros((n, n), dtype=np.float64)

pytcl/clustering/gaussian_mixture.py

@@ -674,9 +674,9 @@ class GaussianMixture:
 
     def _gaussian_pdf(
         self,
-        x: NDArray,
-        mean: NDArray,
-        cov: NDArray,
+        x: NDArray[np.floating],
+        mean: NDArray[np.floating],
+        cov: NDArray[np.floating],
     ) -> float:
         """Evaluate single Gaussian PDF."""
         n = len(x)

pytcl/clustering/hierarchical.py

@@ -12,7 +12,7 @@ References
 """
 
 from enum import Enum
-from typing import List, Literal, NamedTuple, Optional
+from typing import Any, List, Literal, NamedTuple, Optional
 
 import numpy as np
 from numba import njit
@@ -72,7 +72,7 @@ class HierarchicalResult(NamedTuple):
 
 
 @njit(cache=True)
-def _compute_distance_matrix_jit(X: np.ndarray) -> np.ndarray:
+def _compute_distance_matrix_jit(X: np.ndarray[Any, Any]) -> np.ndarray[Any, Any]:
     """JIT-compiled pairwise Euclidean distance computation."""
     n = X.shape[0]
     n_features = X.shape[1]
@@ -112,43 +112,43 @@ def compute_distance_matrix(
 
 
 def _single_linkage(
-    dist_i: NDArray,
-    dist_j: NDArray,
+    dist_i: NDArray[Any],
+    dist_j: NDArray[Any],
     size_i: int,
     size_j: int,
-) -> NDArray:
+) -> NDArray[Any]:
     """Single linkage: minimum of distances."""
     return np.minimum(dist_i, dist_j)
 
 
 def _complete_linkage(
-    dist_i: NDArray,
-    dist_j: NDArray,
+    dist_i: NDArray[Any],
+    dist_j: NDArray[Any],
     size_i: int,
     size_j: int,
-) -> NDArray:
+) -> NDArray[Any]:
     """Complete linkage: maximum of distances."""
     return np.maximum(dist_i, dist_j)
 
 
 def _average_linkage(
-    dist_i: NDArray,
-    dist_j: NDArray,
+    dist_i: NDArray[Any],
+    dist_j: NDArray[Any],
     size_i: int,
     size_j: int,
-) -> NDArray:
+) -> NDArray[Any]:
     """Average linkage: weighted average of distances."""
     return (size_i * dist_i + size_j * dist_j) / (size_i + size_j)
 
 
 def _ward_linkage(
-    dist_i: NDArray,
-    dist_j: NDArray,
+    dist_i: NDArray[Any],
+    dist_j: NDArray[Any],
     size_i: int,
     size_j: int,
-    size_k: NDArray,
+    size_k: NDArray[Any],
     dist_ij: float,
-) -> NDArray:
+) -> NDArray[Any]:
     """Ward's linkage: minimum variance merge."""
     total = size_i + size_j + size_k
     return np.sqrt(

pytcl/clustering/kmeans.py

@@ -10,7 +10,7 @@ References
        Careful Seeding," SODA 2007.
 """
 
-from typing import Literal, NamedTuple, Optional, Union
+from typing import Any, Literal, NamedTuple, Optional, Union
 
 import numpy as np
 from numpy.typing import ArrayLike, NDArray
@@ -305,7 +305,7 @@ def _kmeans_single(
 
         # Handle empty clusters: keep old center
         for k in range(n_clusters):
-            if np.all(new_centers[k] == 0) and np.any(labels == k) is False:
+            if np.all(new_centers[k] == 0) and not np.any(labels == k):
                 new_centers[k] = centers[k]
 
         # Check convergence
@@ -336,8 +336,8 @@ def _kmeans_single(
 def kmeans_elbow(
     X: ArrayLike,
     k_range: Optional[range] = None,
-    **kwargs,
-) -> dict:
+    **kwargs: Any,
+) -> dict[str, Any]:
     """
     Compute K-means for a range of k values for elbow method.
 

pytcl/containers/__init__.py

@@ -3,8 +3,27 @@ Containers module.
 
 This module provides spatial data structures for efficient
 nearest neighbor queries, spatial indexing, and tracking containers.
+
+Spatial Index Hierarchy
+-----------------------
+All spatial index structures inherit from BaseSpatialIndex which defines
+a common interface for k-nearest neighbor and radius queries:
+
+    BaseSpatialIndex (abstract)
+    ├── KDTree - K-dimensional tree (Euclidean space)
+    ├── BallTree - Ball tree variant of KD-tree
+    ├── RTree - Rectangle tree for bounding boxes
+    └── MetricSpatialIndex (abstract)
+        ├── VPTree - Vantage point tree (any metric)
+        └── CoverTree - Cover tree (any metric)
 """
 
+from pytcl.containers.base import (
+    BaseSpatialIndex,
+    MetricSpatialIndex,
+    SpatialQueryResult,
+    validate_query_input,
+)
 from pytcl.containers.cluster_set import (
     ClusterSet,
     ClusterStats,
@@ -13,17 +32,8 @@ from pytcl.containers.cluster_set import (
     cluster_tracks_kmeans,
     compute_cluster_centroid,
 )
-from pytcl.containers.covertree import (
-    CoverTree,
-    CoverTreeNode,
-    CoverTreeResult,
-)
-from pytcl.containers.kd_tree import (
-    BallTree,
-    KDNode,
-    KDTree,
-    NearestNeighborResult,
-)
+from pytcl.containers.covertree import CoverTree, CoverTreeNode, CoverTreeResult
+from pytcl.containers.kd_tree import BallTree, KDNode, KDTree, NearestNeighborResult
 from pytcl.containers.measurement_set import (
     Measurement,
     MeasurementQuery,
@@ -38,18 +48,15 @@ from pytcl.containers.rtree import (
     box_from_points,
     merge_boxes,
 )
-from pytcl.containers.track_list import (
-    TrackList,
-    TrackListStats,
-    TrackQuery,
-)
-from pytcl.containers.vptree import (
-    VPNode,
-    VPTree,
-    VPTreeResult,
-)
+from pytcl.containers.track_list import TrackList, TrackListStats, TrackQuery
+from pytcl.containers.vptree import VPNode, VPTree, VPTreeResult
 
 __all__ = [
+    # Base classes
+    "BaseSpatialIndex",
+    "MetricSpatialIndex",
+    "SpatialQueryResult",
+    "validate_query_input",
     # K-D Tree
     "KDNode",
     "NearestNeighborResult",

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 Any, 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[Any], NDArray[Any]], float]] = None,
+    ):
+        super().__init__(data)
+
+        if metric is None:
+            self.metric = self._euclidean_distance
+        else:
+            self.metric = metric
+
+    @staticmethod
+    def _euclidean_distance(x: NDArray[Any], y: NDArray[Any]) -> 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/cluster_set.py

@@ -8,6 +8,7 @@ that move together (formations, convoys, etc.).
 from __future__ import annotations
 
 from typing import (
+    Any,
     Dict,
     Iterable,
     Iterator,
@@ -312,7 +313,7 @@ class ClusterSet:
         cls,
         tracks: TrackList,
         method: str = "dbscan",
-        **kwargs,
+        **kwargs: Any,
     ) -> ClusterSet:
         """
         Create a ClusterSet by clustering tracks.

pytcl/containers/covertree.py

@@ -11,11 +11,17 @@ References
        neighbor," ICML 2006.
 """
 
-from typing import Callable, List, NamedTuple, Optional, Set, Tuple
+import logging
+from typing import Any, 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,27 +99,24 @@ 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__(
         self,
         data: ArrayLike,
-        metric: Optional[Callable[[NDArray, NDArray], float]] = None,
+        metric: Optional[
+            Callable[[np.ndarray[Any, Any], np.ndarray[Any, Any]], 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 +127,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)."""
@@ -138,7 +143,7 @@ class CoverTree:
             self._distance_cache[key] = self.metric(self.data[i], self.data[j])
         return self._distance_cache[key]
 
-    def _distance_to_point(self, idx: int, query: NDArray) -> float:
+    def _distance_to_point(self, idx: int, query: NDArray[np.floating]) -> float:
         """Distance from data point to query point."""
         return self.metric(self.data[idx], query)
 
@@ -245,11 +250,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 +369,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]] = []