Source code for pytcl.assignment_algorithms.gating

"""
Gating functions for data association in target tracking.

This module provides gating methods to determine which measurements
fall within a validation region around predicted track states.
"""

from typing import Any, List, Tuple

import numpy as np
from numba import njit
from numpy.typing import ArrayLike, NDArray
from scipy.stats import chi2

from pytcl.diagnostics import diagnostics_enabled, logger


@njit(cache=True, fastmath=True)
def _mahalanobis_distance_2d(
    innovation: np.ndarray[Any, Any],
    S_inv: np.ndarray[Any, Any],
) -> float:
    """JIT-compiled Mahalanobis distance for 2D innovations."""
    return innovation[0] * (
        S_inv[0, 0] * innovation[0] + S_inv[0, 1] * innovation[1]
    ) + innovation[1] * (S_inv[1, 0] * innovation[0] + S_inv[1, 1] * innovation[1])


@njit(cache=True, fastmath=True)
def _mahalanobis_distance_3d(
    innovation: np.ndarray[Any, Any],
    S_inv: np.ndarray[Any, Any],
) -> float:
    """JIT-compiled Mahalanobis distance for 3D innovations."""
    result = 0.0
    for i in range(3):
        for j in range(3):
            result += innovation[i] * S_inv[i, j] * innovation[j]
    return result


@njit(cache=True, fastmath=True)
def _mahalanobis_distance_general(
    innovation: np.ndarray[Any, Any],
    S_inv: np.ndarray[Any, Any],
) -> float:
    """JIT-compiled Mahalanobis distance for general dimension."""
    n = len(innovation)
    result = 0.0
    for i in range(n):
        for j in range(n):
            result += innovation[i] * S_inv[i, j] * innovation[j]
    return result


@njit(cache=True, fastmath=True)
def _invert_2x2(S: np.ndarray[Any, Any]) -> "tuple[np.ndarray[Any, Any], float]":
    """Closed-form inverse of a 2x2 matrix; no LAPACK dispatch involved.

    Returns (inverse, determinant). Caller must check determinant != 0
    before using the inverse -- for an exactly singular S the returned
    `inv` is uninitialized garbage (division by zero is skipped rather
    than executed, since numba's fastmath float division by exact 0.0
    raises ZeroDivisionError instead of producing inf); the caller must
    fall back to `np.linalg.solve` in that case, not use this result.
    """
    det = S[0, 0] * S[1, 1] - S[0, 1] * S[1, 0]
    inv = np.empty((2, 2))
    if det != 0.0:
        inv[0, 0] = S[1, 1] / det
        inv[0, 1] = -S[0, 1] / det
        inv[1, 0] = -S[1, 0] / det
        inv[1, 1] = S[0, 0] / det
    return inv, det


@njit(cache=True, fastmath=True)
def _invert_3x3(S: np.ndarray[Any, Any]) -> "tuple[np.ndarray[Any, Any], float]":
    """Closed-form (adjugate/cofactor) inverse of a 3x3 matrix.

    No LAPACK dispatch involved. Returns (inverse, determinant). Caller must
    check determinant != 0 before using the inverse -- for an exactly
    singular S the returned `inv` is uninitialized garbage (see
    `_invert_2x2` for why: division by zero is skipped, not executed).
    """
    a, b, c = S[0, 0], S[0, 1], S[0, 2]
    d, e, f = S[1, 0], S[1, 1], S[1, 2]
    g, h, i = S[2, 0], S[2, 1], S[2, 2]

    A = e * i - f * h
    B = f * g - d * i
    C = d * h - e * g
    D = c * h - b * i
    E = a * i - c * g
    F = b * g - a * h
    G = b * f - c * e
    H = c * d - a * f
    I = a * e - b * d  # noqa: E741

    det = a * A + b * B + c * C

    inv = np.empty((3, 3))
    if det != 0.0:
        inv[0, 0], inv[0, 1], inv[0, 2] = A / det, D / det, G / det
        inv[1, 0], inv[1, 1], inv[1, 2] = B / det, E / det, H / det
        inv[2, 0], inv[2, 1], inv[2, 2] = C / det, F / det, I / det
    return inv, det


[docs] def mahalanobis_distance( innovation: ArrayLike, innovation_covariance: ArrayLike, ) -> float: """ Compute the squared Mahalanobis distance. The Mahalanobis distance measures how many standard deviations a point is from the center of a distribution. Parameters ---------- innovation : array_like Innovation (measurement residual) vector of shape (m,). innovation_covariance : array_like Innovation covariance matrix of shape (m, m). Returns ------- float Squared Mahalanobis distance. Examples -------- >>> innovation = np.array([1.0, 0.5]) >>> S = np.array([[2.0, 0.0], [0.0, 1.0]]) >>> d2 = mahalanobis_distance(innovation, S) >>> d2 0.75 Notes ----- The squared Mahalanobis distance is defined as: d^2 = (z - z_pred)^T @ S^{-1} @ (z - z_pred) where S is the innovation covariance matrix. Dispatch (measured; see task-C2-report.md in the v2.5.0-region-lcd-perf campaign for the full before/after numbers and the behavior-equality bounds): - n == 2 or n == 3: closed-form njit kernels (`_mahalanobis_distance_2d`/ `_3d`, fed by `_invert_2x2`/`_invert_3x3`) that avoid `np.linalg.solve`'s generic LAPACK dispatch entirely (measured ~4x faster than the generic path at these sizes). - 1 <= n <= 10 otherwise: `np.linalg.inv` + the njit `_mahalanobis_distance_general` quadratic-form kernel, measured ~10-25% faster than `np.linalg.solve` in this range (a single `inv` call plus a tight njit loop beats `solve`'s per-call LAPACK dispatch overhead here; the constant-overhead advantage disappears and reverses for n > 12, where `inv`'s O(n^3) cost with a larger constant than solving one right-hand side starts to dominate -- hence the n <= 10 cutoff, comfortably inside the measured-winning range). - n > 10, or an exactly-singular 2D/3D covariance: the original generic `np.linalg.solve` path. Every branch raises the same `numpy.linalg.LinAlgError` on an exactly singular covariance. """ nu = np.asarray(innovation, dtype=np.float64) S = np.asarray(innovation_covariance, dtype=np.float64) n = nu.shape[0] if n == 2: S_inv, det = _invert_2x2(S) if det != 0.0: return float(_mahalanobis_distance_2d(nu, S_inv)) elif n == 3: S_inv, det = _invert_3x3(S) if det != 0.0: return float(_mahalanobis_distance_3d(nu, S_inv)) elif n <= 10: S_inv = np.linalg.inv(S) return float(_mahalanobis_distance_general(nu, S_inv)) # Use solve instead of inverse for numerical stability S_inv_nu = np.linalg.solve(S, nu) return float(nu @ S_inv_nu)
[docs] def ellipsoidal_gate( innovation: ArrayLike, innovation_covariance: ArrayLike, gate_threshold: float, ) -> bool: """ Test if a measurement passes an ellipsoidal gate. The ellipsoidal gate defines a validation region based on the chi-squared distribution of the squared Mahalanobis distance. Parameters ---------- innovation : array_like Innovation vector of shape (m,). innovation_covariance : array_like Innovation covariance matrix of shape (m, m). gate_threshold : float Gate threshold (chi-squared value). Common values: - 9.21 for 99% probability with 2 measurements - 11.34 for 99% probability with 3 measurements - 16.27 for 99% probability with 4 measurements Returns ------- bool True if measurement passes the gate (is inside the ellipsoid). Examples -------- >>> innovation = np.array([1.0, 0.5]) >>> S = np.array([[2.0, 0.0], [0.0, 1.0]]) >>> ellipsoidal_gate(innovation, S, gate_threshold=9.21) True See Also -------- chi2_gate_threshold : Compute threshold from probability. """ d2 = mahalanobis_distance(innovation, innovation_covariance) return d2 <= gate_threshold
[docs] def chi2_gate_threshold( probability: float, num_dimensions: int, ) -> float: """ Compute chi-squared gate threshold for a given probability. Parameters ---------- probability : float Gate probability (e.g., 0.99 for 99% of true measurements to pass). num_dimensions : int Measurement dimension (degrees of freedom). Returns ------- float Chi-squared threshold value. Examples -------- >>> round(chi2_gate_threshold(0.99, 2), 6) # 2D measurement, 99% probability 9.21034 >>> round(chi2_gate_threshold(0.99, 3), 6) # 3D measurement, 99% probability 11.344867 """ return float(chi2.ppf(probability, df=num_dimensions))
[docs] def rectangular_gate( innovation: ArrayLike, innovation_covariance: ArrayLike, num_sigmas: float = 3.0, ) -> bool: """ Test if a measurement passes a rectangular gate. The rectangular gate defines a validation region as a hypercube based on the marginal standard deviations. Parameters ---------- innovation : array_like Innovation vector of shape (m,). innovation_covariance : array_like Innovation covariance matrix of shape (m, m). num_sigmas : float, optional Number of standard deviations for gate bounds (default: 3.0). Returns ------- bool True if measurement passes the gate. Examples -------- >>> innovation = np.array([1.0, 0.5]) >>> S = np.array([[4.0, 0.0], [0.0, 1.0]]) >>> rectangular_gate(innovation, S, num_sigmas=3.0) True Notes ----- Rectangular gating is computationally cheaper but less tight than ellipsoidal gating. It may pass more false measurements. """ nu = np.asarray(innovation, dtype=np.float64) S = np.asarray(innovation_covariance, dtype=np.float64) # Extract marginal standard deviations sigmas = np.sqrt(np.diag(S)) # Check if all components are within bounds return bool(np.all(np.abs(nu) <= num_sigmas * sigmas))
[docs] def gate_measurements( predicted_measurement: ArrayLike, innovation_covariance: ArrayLike, measurements: ArrayLike, gate_threshold: float, gate_type: str = "ellipsoidal", ) -> Tuple[NDArray[np.intp], NDArray[np.float64]]: """ Gate multiple measurements against a predicted track state. Parameters ---------- predicted_measurement : array_like Predicted measurement of shape (m,). innovation_covariance : array_like Innovation covariance matrix of shape (m, m). measurements : array_like Array of measurements of shape (n_meas, m). gate_threshold : float Gate threshold. For ellipsoidal gates, this is the chi-squared value. For rectangular gates, this is the number of sigmas. gate_type : str, optional Type of gate: "ellipsoidal" or "rectangular" (default: "ellipsoidal"). Returns ------- valid_indices : ndarray Indices of measurements that pass the gate. distances : ndarray Squared Mahalanobis distances for valid measurements. Examples -------- >>> z_pred = np.array([0.0, 0.0]) >>> S = np.eye(2) >>> measurements = np.array([[0.5, 0.5], [5.0, 5.0], [1.0, -1.0]]) >>> valid_idx, dists = gate_measurements(z_pred, S, measurements, 9.21) >>> valid_idx array([0, 2]) Notes ----- This function efficiently gates multiple measurements against a single track prediction, which is common in multi-target tracking. """ z_pred = np.asarray(predicted_measurement, dtype=np.float64) S = np.asarray(innovation_covariance, dtype=np.float64) Z = np.asarray(measurements, dtype=np.float64) if Z.ndim == 1: Z = Z.reshape(1, -1) n_meas = Z.shape[0] valid_indices: List[int] = [] distances: List[float] = [] for i in range(n_meas): innovation = Z[i] - z_pred if gate_type == "ellipsoidal": d2 = mahalanobis_distance(innovation, S) if d2 <= gate_threshold: valid_indices.append(i) distances.append(d2) elif gate_type == "rectangular": if rectangular_gate(innovation, S, num_sigmas=gate_threshold): # For rectangular gate, still compute Mahalanobis distance for ranking d2 = mahalanobis_distance(innovation, S) valid_indices.append(i) distances.append(d2) else: raise ValueError(f"Unknown gate type: {gate_type}") if diagnostics_enabled(): n_rejected = n_meas - len(valid_indices) if n_rejected: logger.bind(site="gating").debug( "gate_measurements: rejected {} of {} measurement(s), threshold={:.2f}", n_rejected, n_meas, gate_threshold, ) return ( np.array(valid_indices, dtype=np.intp), np.array(distances, dtype=np.float64), )
[docs] def compute_gate_volume( innovation_covariance: ArrayLike, gate_threshold: float, ) -> float: """ Compute the volume of an ellipsoidal gate. Parameters ---------- innovation_covariance : array_like Innovation covariance matrix of shape (m, m). gate_threshold : float Chi-squared gate threshold. Returns ------- float Volume of the ellipsoidal gate region. Notes ----- The gate volume is used in probabilistic data association methods to compute the clutter density. For an m-dimensional ellipsoid with threshold gamma: V = c_m * sqrt(det(S)) * gamma^(m/2) where c_m is the volume of the unit hypersphere in m dimensions. Examples -------- Compute gate volume for a 2D measurement with 99% gate probability: >>> import numpy as np >>> from scipy.stats import chi2 >>> S = np.array([[4.0, 0.0], [0.0, 1.0]]) # innovation covariance >>> gate_prob = 0.99 >>> threshold = chi2.ppf(gate_prob, df=2) >>> volume = compute_gate_volume(S, threshold) >>> volume > 0 True See Also -------- ellipsoidal_gate : Test if measurement passes gate. mahalanobis_distance : Compute distance used in gating. """ S = np.asarray(innovation_covariance, dtype=np.float64) m = S.shape[0] # Volume of unit hypersphere in m dimensions # c_m = pi^(m/2) / Gamma(m/2 + 1) from scipy.special import gamma as gamma_func c_m = np.pi ** (m / 2) / gamma_func(m / 2 + 1) # Gate volume det_S = np.linalg.det(S) volume = c_m * np.sqrt(det_S) * gate_threshold ** (m / 2) return float(volume)
[docs] @njit(cache=True, fastmath=True, parallel=False) def mahalanobis_batch( 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. JIT-compiled for performance. Computes squared Mahalanobis distances for multiple innovations against a single covariance matrix. Parameters ---------- innovations : ndarray Innovations of shape (n_measurements, dim). S_inv : ndarray Inverse of innovation covariance matrix of shape (dim, dim). output : ndarray Output array of shape (n_measurements,) to store distances. """ n_meas = innovations.shape[0] dim = innovations.shape[1] for i in range(n_meas): result = 0.0 for j in range(dim): for k in range(dim): result += innovations[i, j] * S_inv[j, k] * innovations[i, k] output[i] = result
__all__ = [ "mahalanobis_distance", "mahalanobis_batch", "ellipsoidal_gate", "rectangular_gate", "gate_measurements", "chi2_gate_threshold", "compute_gate_volume", ]