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

pytcl/assignment_algorithms/gating.py

@@ -5,7 +5,7 @@ This module provides gating methods to determine which measurements
 fall within a validation region around predicted track states.
 """
 
-from typing import List, Tuple
+from typing import Any, List, Tuple
 
 import numpy as np
 from numba import njit
@@ -15,8 +15,8 @@ from scipy.stats import chi2
 
 @njit(cache=True, fastmath=True)
 def _mahalanobis_distance_2d(
-    innovation: np.ndarray,
-    S_inv: np.ndarray,
+    innovation: np.ndarray[Any, Any],
+    S_inv: np.ndarray[Any, Any],
 ) -> float:
     """JIT-compiled Mahalanobis distance for 2D innovations."""
     return innovation[0] * (
@@ -26,8 +26,8 @@ def _mahalanobis_distance_2d(
 
 @njit(cache=True, fastmath=True)
 def _mahalanobis_distance_3d(
-    innovation: np.ndarray,
-    S_inv: np.ndarray,
+    innovation: np.ndarray[Any, Any],
+    S_inv: np.ndarray[Any, Any],
 ) -> float:
     """JIT-compiled Mahalanobis distance for 3D innovations."""
     result = 0.0
@@ -39,8 +39,8 @@ def _mahalanobis_distance_3d(
 
 @njit(cache=True, fastmath=True)
 def _mahalanobis_distance_general(
-    innovation: np.ndarray,
-    S_inv: np.ndarray,
+    innovation: np.ndarray[Any, Any],
+    S_inv: np.ndarray[Any, Any],
 ) -> float:
     """JIT-compiled Mahalanobis distance for general dimension."""
     n = len(innovation)
@@ -341,9 +341,9 @@ def compute_gate_volume(
 
 @njit(cache=True, fastmath=True, parallel=False)
 def mahalanobis_batch(
-    innovations: np.ndarray,
-    S_inv: np.ndarray,
-    output: np.ndarray,
+    innovations: np.ndarray[Any, Any],
+    S_inv: np.ndarray[Any, Any],
+    output: np.ndarray[Any, Any],
 ) -> None:
     """
     Compute Mahalanobis distances for a batch of innovations.

pytcl/assignment_algorithms/jpda.py

@@ -9,7 +9,7 @@ This is more sophisticated than GNN which makes hard assignment decisions,
 as JPDA can handle measurement origin uncertainty in cluttered environments.
 """
 
-from typing import List, NamedTuple, Optional, Tuple
+from typing import Any, List, NamedTuple, Optional
 
 import numpy as np
 from numba import njit
@@ -24,7 +24,7 @@ class JPDAResult(NamedTuple):
 
     Attributes
     ----------
-    association_probs : ndarray
+    association_probs : ndarray[Any]
         Association probability matrix of shape (n_tracks, n_measurements + 1).
         association_probs[i, j] is the probability that track i is associated
         with measurement j. The last column (j = n_measurements) represents
@@ -32,9 +32,9 @@ class JPDAResult(NamedTuple):
     marginal_probs : list of ndarray
         List of marginal association probabilities for each track.
         marginal_probs[i][j] = P(measurement j originated from track i).
-    likelihood_matrix : ndarray
+    likelihood_matrix : ndarray[Any]
         Measurement likelihood matrix of shape (n_tracks, n_measurements).
-    gated : ndarray
+    gated : ndarray[Any]
         Boolean matrix indicating which track-measurement pairs passed gating.
     """
 
@@ -53,7 +53,7 @@ class JPDAUpdate(NamedTuple):
         Updated state estimates for each track.
     covariances : list of ndarray
         Updated covariances for each track (includes spread of means).
-    association_probs : ndarray
+    association_probs : ndarray[Any]
         Association probability matrix.
     innovations : list of ndarray
         Combined weighted innovations for each track.
@@ -66,8 +66,8 @@ class JPDAUpdate(NamedTuple):
 
 
 def compute_measurement_likelihood(
-    innovation: NDArray,
-    innovation_cov: NDArray,
+    innovation: NDArray[Any],
+    innovation_cov: NDArray[Any],
     detection_prob: float = 1.0,
 ) -> float:
     """
@@ -75,9 +75,9 @@ def compute_measurement_likelihood(
 
     Parameters
     ----------
-    innovation : ndarray
+    innovation : ndarray[Any]
         Measurement innovation (residual), shape (m,).
-    innovation_cov : ndarray
+    innovation_cov : ndarray[Any]
         Innovation covariance, shape (m, m).
     detection_prob : float
         Probability of detection (Pd).
@@ -102,14 +102,14 @@ def compute_measurement_likelihood(
 
 
 def compute_likelihood_matrix(
-    track_states: List[NDArray],
-    track_covariances: List[NDArray],
-    measurements: NDArray,
-    H: NDArray,
-    R: NDArray,
+    track_states: list[NDArray[Any]],
+    track_covariances: list[NDArray[Any]],
+    measurements: NDArray[Any],
+    H: NDArray[Any],
+    R: NDArray[Any],
     detection_prob: float = 1.0,
     gate_threshold: Optional[float] = None,
-) -> Tuple[NDArray, NDArray]:
+) -> tuple[NDArray[Any], NDArray[Any]]:
     """
     Compute likelihood matrix for all track-measurement pairs.
 
@@ -119,11 +119,11 @@ def compute_likelihood_matrix(
         State estimates for each track.
     track_covariances : list of ndarray
         Covariances for each track.
-    measurements : ndarray
+    measurements : ndarray[Any]
         Measurements, shape (n_meas, m).
-    H : ndarray
+    H : ndarray[Any]
         Measurement matrix, shape (m, n).
-    R : ndarray
+    R : ndarray[Any]
         Measurement noise covariance, shape (m, m).
     detection_prob : float
         Probability of detection.
@@ -132,9 +132,9 @@ def compute_likelihood_matrix(
 
     Returns
     -------
-    likelihood_matrix : ndarray
+    likelihood_matrix : ndarray[Any]
         Likelihood values, shape (n_tracks, n_meas).
-    gated : ndarray
+    gated : ndarray[Any]
         Boolean gating matrix, shape (n_tracks, n_meas).
     """
     n_tracks = len(track_states)
@@ -163,11 +163,11 @@ def compute_likelihood_matrix(
 
 
 def jpda_probabilities(
-    likelihood_matrix: NDArray,
-    gated: NDArray,
+    likelihood_matrix: NDArray[Any],
+    gated: NDArray[Any],
     detection_prob: float = 1.0,
     clutter_density: float = 1e-6,
-) -> NDArray:
+) -> NDArray[Any]:
     """
     Compute JPDA association probabilities.
 
@@ -176,9 +176,9 @@ def jpda_probabilities(
 
     Parameters
     ----------
-    likelihood_matrix : ndarray
+    likelihood_matrix : ndarray[Any]
         Likelihood values, shape (n_tracks, n_meas).
-    gated : ndarray
+    gated : ndarray[Any]
         Boolean gating matrix, shape (n_tracks, n_meas).
     detection_prob : float
         Probability of detection (Pd).
@@ -187,7 +187,7 @@ def jpda_probabilities(
 
     Returns
     -------
-    beta : ndarray
+    beta : ndarray[Any]
         Association probability matrix, shape (n_tracks, n_meas + 1).
         beta[i, j] = P(measurement j is from track i) for j < n_meas.
         beta[i, n_meas] = P(track i has no measurement).
@@ -218,11 +218,11 @@ def jpda_probabilities(
 
 
 def _jpda_exact(
-    likelihood_matrix: NDArray,
-    gated: NDArray,
+    likelihood_matrix: NDArray[Any],
+    gated: NDArray[Any],
     detection_prob: float,
     clutter_density: float,
-) -> NDArray:
+) -> NDArray[Any]:
     """
     Exact JPDA computation via hypothesis enumeration.
 
@@ -241,8 +241,8 @@ def _jpda_exact(
     def generate_hypotheses(
         meas_idx: int,
         current_assignment: List[int],
-        used_tracks: set,
-    ):
+        used_tracks: set[Any],
+    ) -> Any:
         """Recursively generate valid hypotheses."""
         if meas_idx == n_meas:
             yield current_assignment.copy()
@@ -268,11 +268,11 @@ def _jpda_exact(
     hypothesis_probs = []
     hypothesis_assignments = []
 
-    for assignment in generate_hypotheses(0, [], set()):
+    for assignment in generate_hypotheses(0, [], set[Any]()):
         # Compute probability of this hypothesis
         prob = 1.0
 
-        detected_tracks = set()
+        detected_tracks = set[Any]()
         for j, track_idx in enumerate(assignment):
             if track_idx == -1:
                 # Measurement j is clutter
@@ -301,7 +301,7 @@ def _jpda_exact(
     for h_idx, (assignment, prob) in enumerate(
         zip(hypothesis_assignments, hypothesis_probs)
     ):
-        detected_tracks = set()
+        detected_tracks = set[Any]()
         for j, track_idx in enumerate(assignment):
             if track_idx >= 0:
                 beta[track_idx, j] += prob
@@ -317,11 +317,11 @@ def _jpda_exact(
 
 @njit(cache=True)
 def _jpda_approximate_core(
-    likelihood_matrix: np.ndarray,
-    gated: np.ndarray,
+    likelihood_matrix: np.ndarray[Any, Any],
+    gated: np.ndarray[Any, Any],
     detection_prob: float,
     clutter_density: float,
-) -> np.ndarray:
+) -> np.ndarray[Any, Any]:
     """JIT-compiled core of approximate JPDA computation."""
     n_tracks = likelihood_matrix.shape[0]
     n_meas = likelihood_matrix.shape[1]
@@ -369,11 +369,11 @@ def _jpda_approximate_core(
 
 
 def _jpda_approximate(
-    likelihood_matrix: NDArray,
-    gated: NDArray,
+    likelihood_matrix: NDArray[Any],
+    gated: NDArray[Any],
     detection_prob: float,
     clutter_density: float,
-) -> NDArray:
+) -> NDArray[Any]:
     """
     Approximate JPDA using parametric approach.
 

pytcl/assignment_algorithms/nd_assignment.py

@@ -0,0 +1,379 @@
+"""
+N-dimensional assignment algorithms (4D and higher).
+
+This module extends the 3D assignment solver to arbitrary dimensions,
+enabling more complex assignment scenarios such as:
+- 4D: Measurements × Tracks × Hypotheses × Sensors
+- 5D+: Additional dimensions for time frames, maneuver classes, etc.
+
+The module provides a unified interface for solving high-dimensional
+assignment problems using generalized relaxation methods.
+
+References
+----------
+.. [1] Poore, A. B., "Multidimensional Assignment Problem and Data
+       Association," IEEE Transactions on Aerospace and Electronic Systems,
+       2013.
+.. [2] Cramer, R. D., et al., "The Emerging Role of Chemical Similarity in
+       Drug Discovery," Perspectives in Drug Discovery and Design, 2003.
+"""
+
+from typing import NamedTuple, Optional, Tuple
+
+import numpy as np
+from numpy.typing import NDArray
+
+
+class AssignmentNDResult(NamedTuple):
+    """Result of an N-dimensional assignment problem.
+
+    Attributes
+    ----------
+    assignments : ndarray
+        Array of shape (n_assignments, n_dimensions) containing assigned
+        index tuples. Each row is an n-tuple of indices.
+    cost : float
+        Total assignment cost.
+    converged : bool
+        Whether the algorithm converged (for iterative methods).
+    n_iterations : int
+        Number of iterations used (for iterative methods).
+    gap : float
+        Optimality gap (upper_bound - lower_bound) for relaxation methods.
+    """
+
+    assignments: NDArray[np.intp]
+    cost: float
+    converged: bool
+    n_iterations: int
+    gap: float
+
+
+def validate_cost_tensor(cost_tensor: NDArray[np.float64]) -> Tuple[int, ...]:
+    """
+    Validate cost tensor and return dimensions.
+
+    Parameters
+    ----------
+    cost_tensor : ndarray
+        Cost tensor of arbitrary dimension.
+
+    Returns
+    -------
+    dims : tuple
+        Dimensions of the cost tensor.
+
+    Raises
+    ------
+    ValueError
+        If tensor has fewer than 2 dimensions.
+    """
+    if cost_tensor.ndim < 2:
+        raise ValueError(
+            f"Cost tensor must have at least 2 dimensions, got {cost_tensor.ndim}"
+        )
+
+    return cost_tensor.shape
+
+
+def greedy_assignment_nd(
+    cost_tensor: NDArray[np.float64],
+    max_assignments: Optional[int] = None,
+) -> AssignmentNDResult:
+    """
+    Greedy solver for N-dimensional assignment.
+
+    Selects minimum-cost tuples in order until no more valid assignments
+    exist (no dimension index is repeated).
+
+    Parameters
+    ----------
+    cost_tensor : ndarray
+        Cost tensor of shape (n1, n2, ..., nk).
+    max_assignments : int, optional
+        Maximum number of assignments to find (default: min(dimensions)).
+
+    Returns
+    -------
+    AssignmentNDResult
+        Assignments, total cost, and algorithm info.
+
+    Notes
+    -----
+    Greedy assignment is fast O(n log n) but not optimal. Used as
+    heuristic or starting solution for optimization methods.
+    """
+    dims = cost_tensor.shape
+    n_dims = len(dims)
+
+    if max_assignments is None:
+        max_assignments = min(dims)
+
+    # Flatten tensor with index mapping
+    flat_costs = cost_tensor.ravel()
+    sorted_indices = np.argsort(flat_costs)
+
+    assignments: list[tuple[int, ...]] = []
+    used_indices: list[set[int]] = [set() for _ in range(n_dims)]
+
+    for flat_idx in sorted_indices:
+        if len(assignments) >= max_assignments:
+            break
+
+        # Convert flat index to multi-dimensional index
+        multi_idx = np.unravel_index(flat_idx, dims)
+
+        # Check if any dimension index is already used
+        conflict = False
+        for d, idx in enumerate(multi_idx):
+            if idx in used_indices[d]:
+                conflict = True
+                break
+
+        if not conflict:
+            assignments.append(multi_idx)
+            for d, idx in enumerate(multi_idx):
+                used_indices[d].add(idx)
+
+    assignments_array = np.array(assignments, dtype=np.intp)
+    if assignments_array.size > 0:
+        total_cost = float(np.sum(cost_tensor[tuple(assignments_array.T)]))
+    else:
+        total_cost = 0.0
+
+    return AssignmentNDResult(
+        assignments=assignments_array,
+        cost=total_cost,
+        converged=True,
+        n_iterations=1,
+        gap=0.0,  # Greedy doesn't compute lower bound
+    )
+
+
+def relaxation_assignment_nd(
+    cost_tensor: NDArray[np.float64],
+    max_iterations: int = 100,
+    tolerance: float = 1e-6,
+    verbose: bool = False,
+) -> AssignmentNDResult:
+    """
+    Lagrangian relaxation solver for N-dimensional assignment.
+
+    Uses iterative subgradient optimization on Lagrange multipliers
+    to tighten the lower bound and find good solutions.
+
+    Parameters
+    ----------
+    cost_tensor : ndarray
+        Cost tensor of shape (n1, n2, ..., nk).
+    max_iterations : int, optional
+        Maximum iterations (default 100).
+    tolerance : float, optional
+        Convergence tolerance for gap (default 1e-6).
+    verbose : bool, optional
+        Print iteration info (default False).
+
+    Returns
+    -------
+    AssignmentNDResult
+        Assignments, total cost, convergence info, and optimality gap.
+
+    Notes
+    -----
+    The relaxation approach:
+    1. Maintain Lagrange multipliers for each dimension
+    2. Solve relaxed problem (select best entries per tuple)
+    3. Update multipliers based on constraint violations
+    4. Iterate until convergence or gap tolerance met
+
+    This guarantees a lower bound on optimal cost and often finds
+    near-optimal or optimal solutions.
+    """
+    dims = cost_tensor.shape
+    n_dims = len(dims)
+
+    # Initialize Lagrange multipliers (one per dimension per index)
+    lambdas = [np.zeros(dim) for dim in dims]
+
+    best_cost = np.inf
+    best_assignments = None
+    lower_bound = -np.inf
+
+    for iteration in range(max_iterations):
+        # Compute relaxed costs: original - Lagrange penalty
+        relaxed_cost = cost_tensor.copy()
+        for d in range(n_dims):
+            # Reshape lambda[d] to broadcast correctly
+            shape = [1] * n_dims
+            shape[d] = dims[d]
+            relaxed_cost = relaxed_cost - lambdas[d].reshape(shape)
+
+        # Solve relaxed problem: greedy on relaxed costs
+        result_relaxed = greedy_assignment_nd(relaxed_cost)
+
+        # Compute lower bound from relaxed solution
+        lower_bound = result_relaxed.cost + sum(
+            np.sum(lambdas[d]) for d in range(n_dims)
+        )
+
+        # Extract solution from relaxed problem
+        if len(result_relaxed.assignments) > 0:
+            actual_cost = float(
+                np.sum(cost_tensor[tuple(result_relaxed.assignments.T)])
+            )
+
+            if actual_cost < best_cost:
+                best_cost = actual_cost
+                best_assignments = result_relaxed.assignments
+
+        # Compute constraint violations and update multipliers
+        violations = [np.zeros(dim) for dim in dims]
+
+        for assignment in result_relaxed.assignments:
+            for d, idx in enumerate(assignment):
+                violations[d][idx] += 1
+
+        # Subgradient descent on multipliers
+        step_size = 1.0 / (iteration + 1)
+        for d in range(n_dims):
+            lambdas[d] -= step_size * (violations[d] - 1.0)
+
+        # Compute gap
+        gap = best_cost - lower_bound if best_cost != np.inf else np.inf
+
+        if verbose:
+            print(
+                f"Iter {iteration+1}: LB={lower_bound:.4f}, UB={best_cost:.4f}, "
+                f"Gap={gap:.6f}"
+            )
+
+        if gap < tolerance:
+            if verbose:
+                print(f"Converged at iteration {iteration+1}")
+            break
+
+    if best_assignments is None:
+        best_assignments = np.empty((0, n_dims), dtype=np.intp)
+        best_cost = 0.0
+
+    gap = best_cost - lower_bound if best_cost != np.inf else np.inf
+
+    return AssignmentNDResult(
+        assignments=best_assignments,
+        cost=best_cost,
+        converged=gap < tolerance,
+        n_iterations=iteration + 1,
+        gap=gap,
+    )
+
+
+def auction_assignment_nd(
+    cost_tensor: NDArray[np.float64],
+    max_iterations: int = 100,
+    epsilon: float = 0.01,
+    verbose: bool = False,
+) -> AssignmentNDResult:
+    """
+    Auction algorithm for N-dimensional assignment.
+
+    Inspired by the classical auction algorithm for 2D assignment,
+    adapted to higher dimensions. Objects bid for assignments based
+    on relative costs.
+
+    Parameters
+    ----------
+    cost_tensor : ndarray
+        Cost tensor of shape (n1, n2, ..., nk).
+    max_iterations : int, optional
+        Maximum iterations (default 100).
+    epsilon : float, optional
+        Bid increment (default 0.01). Larger epsilon → fewer iterations,
+        worse solution; smaller epsilon → more iterations, better solution.
+    verbose : bool, optional
+        Print iteration info (default False).
+
+    Returns
+    -------
+    AssignmentNDResult
+        Assignments, total cost, convergence info, gap estimate.
+
+    Notes
+    -----
+    The algorithm maintains a "price" for each index and allows bidding
+    (price adjustment) to maximize value. Converges to epsilon-optimal
+    solution in finite iterations.
+    """
+    dims = cost_tensor.shape
+    n_dims = len(dims)
+
+    # Initialize prices (one per dimension per index)
+    prices = [np.zeros(dim) for dim in dims]
+
+    for iteration in range(max_iterations):
+        # Compute profit: cost - price penalty
+        profit = cost_tensor.copy()
+        for d in range(n_dims):
+            shape = [1] * n_dims
+            shape[d] = dims[d]
+            profit = profit - prices[d].reshape(shape)
+
+        # Find best assignment at current prices (greedy)
+        result = greedy_assignment_nd(profit)
+
+        if len(result.assignments) == 0:
+            break
+
+        # Update prices: increase price for "in-demand" indices
+        demands = [np.zeros(dim) for dim in dims]
+        for assignment in result.assignments:
+            for d, idx in enumerate(assignment):
+                demands[d][idx] += 1
+
+        for d in range(n_dims):
+            prices[d] += epsilon * (demands[d] - 1.0)
+
+        if verbose and (iteration + 1) % 10 == 0:
+            actual_cost = float(np.sum(cost_tensor[tuple(result.assignments.T)]))
+            print(f"Iter {iteration+1}: Cost={actual_cost:.4f}")
+
+    # Final solution
+    result = greedy_assignment_nd(cost_tensor)
+
+    return AssignmentNDResult(
+        assignments=result.assignments,
+        cost=result.cost,
+        converged=True,
+        n_iterations=iteration + 1,
+        gap=0.0,  # Auction algorithm doesn't track gap formally
+    )
+
+
+def detect_dimension_conflicts(
+    assignments: NDArray[np.intp],
+    dims: Tuple[int, ...],
+) -> bool:
+    """
+    Check if assignments violate dimension uniqueness.
+
+    For valid assignment, each index should appear at most once per dimension.
+
+    Parameters
+    ----------
+    assignments : ndarray
+        Array of shape (n_assignments, n_dimensions) with assignments.
+    dims : tuple
+        Dimensions of the cost tensor.
+
+    Returns
+    -------
+    has_conflicts : bool
+        True if any index appears more than once in any dimension.
+    """
+    n_dims = len(dims)
+
+    for d in range(n_dims):
+        indices_in_dim = assignments[:, d]
+        if len(indices_in_dim) != len(np.unique(indices_in_dim)):
+            return True
+
+    return False