nrl-tracker 1.6.0__py3-none-any.whl → 1.7.1__py3-none-any.whl

pytcl/dynamic_estimation/kalman/constrained.py

@@ -0,0 +1,382 @@
+"""
+Constrained Extended Kalman Filter (CEKF).
+
+Extends the Extended Kalman Filter to enforce constraints on the state
+estimate. Uses Lagrange multiplier method to project onto constraint manifold
+while maintaining positive definite covariance.
+
+References
+----------
+.. [1] Simon, D. (2006). Optimal State Estimation: Kalman, H∞, and Nonlinear
+       Approaches. Wiley-Interscience.
+.. [2] Simon, D. & Simon, D. L. (2010). Constrained Kalman filtering via
+       density function truncation. Journal of Guidance, Control, and Dynamics.
+"""
+
+from typing import Any, Callable, Optional
+
+import numpy as np
+from numpy.typing import ArrayLike, NDArray
+
+from pytcl.dynamic_estimation.kalman.extended import ekf_predict, ekf_update
+from pytcl.dynamic_estimation.kalman.linear import KalmanPrediction, KalmanUpdate
+
+
+class ConstraintFunction:
+    """Base class for state constraints."""
+
+    def __init__(
+        self,
+        g: Callable[[NDArray[Any]], NDArray[Any]],
+        G: Optional[Callable[[NDArray[Any]], NDArray[Any]]] = None,
+        constraint_type: str = "inequality",
+    ):
+        """
+        Define a constraint: g(x) ≤ 0 (inequality) or g(x) = 0 (equality).
+
+        Parameters
+        ----------
+        g : callable
+            Constraint function: g(x) -> scalar or array
+            Inequality: g(x) ≤ 0
+            Equality: g(x) = 0
+        G : callable, optional
+            Jacobian of g with respect to x: ∂g/∂x
+            If None, computed numerically.
+        constraint_type : {'inequality', 'equality'}
+            Constraint type.
+        """
+        self.g = g
+        self.G = G
+        self.constraint_type = constraint_type
+
+    def evaluate(self, x: NDArray[Any]) -> NDArray[Any]:
+        """Evaluate constraint at state x."""
+        return np.atleast_1d(np.asarray(self.g(x), dtype=np.float64))
+
+    def jacobian(self, x: NDArray[Any]) -> NDArray[Any]:
+        """Compute constraint Jacobian at x."""
+        if self.G is not None:
+            return np.atleast_2d(np.asarray(self.G(x), dtype=np.float64))
+        else:
+            # Numerical differentiation
+            eps = 1e-6
+            n = len(x)
+            g_x = self.evaluate(x)
+            m = len(g_x)
+            J = np.zeros((m, n))
+            for i in range(n):
+                x_plus = x.copy()
+                x_plus[i] += eps
+                g_plus = self.evaluate(x_plus)
+                J[:, i] = (g_plus - g_x) / eps
+            return J
+
+    def is_satisfied(self, x: NDArray[Any], tol: float = 1e-6) -> bool:
+        """Check if constraint is satisfied."""
+        g_val = self.evaluate(x)
+        if self.constraint_type == "inequality":
+            return np.all(g_val <= tol)
+        else:  # equality
+            return np.allclose(g_val, 0, atol=tol)
+
+
+class ConstrainedEKF:
+    """
+    Extended Kalman Filter with state constraints.
+
+    Enforces linear and/or nonlinear constraints on state estimate using
+    Lagrange multiplier method with covariance projection.
+
+    Attributes
+    ----------
+    constraints : list of ConstraintFunction
+        List of active constraints.
+    """
+
+    def __init__(self) -> None:
+        """Initialize Constrained EKF."""
+        self.constraints: list[ConstraintFunction] = []
+
+    def add_constraint(self, constraint: ConstraintFunction) -> None:
+        """
+        Add a constraint to the filter.
+
+        Parameters
+        ----------
+        constraint : ConstraintFunction
+            Constraint to enforce.
+        """
+        self.constraints.append(constraint)
+
+    def predict(
+        self,
+        x: ArrayLike,
+        P: ArrayLike,
+        f: Callable[[NDArray[Any]], NDArray[Any]],
+        F: ArrayLike,
+        Q: ArrayLike,
+    ) -> KalmanPrediction:
+        """
+        Constrained EKF prediction step.
+
+        Performs standard EKF prediction (constraints not enforced here,
+        only checked). Constraint enforcement happens in update step.
+
+        Parameters
+        ----------
+        x : array_like
+            Current state estimate, shape (n,).
+        P : array_like
+            Current state covariance, shape (n, n).
+        f : callable
+            Nonlinear state transition function.
+        F : array_like
+            Jacobian of f at current state.
+        Q : array_like
+            Process noise covariance, shape (n, n).
+
+        Returns
+        -------
+        result : KalmanPrediction
+            Predicted state and covariance.
+        """
+        return ekf_predict(x, P, f, F, Q)
+
+    def update(
+        self,
+        x: ArrayLike,
+        P: ArrayLike,
+        z: ArrayLike,
+        h: Callable[[NDArray[Any]], NDArray[Any]],
+        H: ArrayLike,
+        R: ArrayLike,
+    ) -> KalmanUpdate:
+        """
+        Constrained EKF update step.
+
+        Performs standard EKF update, then projects result onto constraint
+        manifold.
+
+        Parameters
+        ----------
+        x : array_like
+            Predicted state estimate, shape (n,).
+        P : array_like
+            Predicted state covariance, shape (n, n).
+        z : array_like
+            Measurement, shape (m,).
+        h : callable
+            Nonlinear measurement function.
+        H : array_like
+            Jacobian of h at current state.
+        R : array_like
+            Measurement noise covariance, shape (m, m).
+
+        Returns
+        -------
+        result : KalmanUpdate
+            Updated state and covariance (constrained).
+        """
+        # Standard EKF update
+        result = ekf_update(x, P, z, h, H, R)
+        x_upd = result.x
+        P_upd = result.P
+
+        # Apply constraint projection
+        if self.constraints:
+            x_upd, P_upd = self._project_onto_constraints(x_upd, P_upd)
+
+        return KalmanUpdate(
+            x=x_upd,
+            P=P_upd,
+            y=result.y,
+            S=result.S,
+            K=result.K,
+            likelihood=result.likelihood,
+        )
+
+    def _project_onto_constraints(
+        self,
+        x: NDArray[Any],
+        P: NDArray[Any],
+        max_iter: int = 10,
+        tol: float = 1e-6,
+    ) -> tuple[NDArray[Any], NDArray[Any]]:
+        """
+        Project state and covariance onto constraint manifold.
+
+        Uses iterative Lagrange multiplier method with covariance
+        projection to enforce constraints while maintaining positive
+        definiteness.
+
+        Parameters
+        ----------
+        x : ndarray
+            State estimate, shape (n,).
+        P : ndarray
+            Covariance matrix, shape (n, n).
+        max_iter : int
+            Maximum iterations for constraint projection.
+        tol : float
+            Convergence tolerance.
+
+        Returns
+        -------
+        x_proj : ndarray
+            Constrained state estimate.
+        P_proj : ndarray
+            Projected covariance.
+        """
+        x_proj = x.copy()
+        P_proj = P.copy()
+
+        # Check which constraints are violated
+        violated: list[ConstraintFunction] = [
+            c for c in self.constraints if not c.is_satisfied(x_proj)
+        ]
+
+        if not violated:
+            return x_proj, P_proj
+
+        # Iterative projection
+        for iteration in range(max_iter):
+            converged = True
+
+            for constraint in violated:
+                g_val = constraint.evaluate(x_proj)
+                G = constraint.jacobian(x_proj)
+
+                # Only process violated constraints
+                if constraint.constraint_type == "inequality":
+                    mask = g_val > tol
+                else:
+                    mask = np.abs(g_val) > tol
+
+                if not np.any(mask):
+                    continue
+
+                converged = False
+
+                # Lagrange multiplier for this constraint
+                # λ = -(G P Gᵀ + μ)⁻¹ G (x - x_ref)
+                # where x_ref is desired state (we use x)
+
+                GP = G @ P_proj
+                GPGt = GP @ G.T
+
+                # Add small regularization for numerical stability
+                mu = np.eye(GPGt.shape[0]) * 1e-6
+
+                try:
+                    GPGt_inv = np.linalg.inv(GPGt + mu)
+                    lam = -GPGt_inv @ (G @ x_proj + g_val)
+
+                    # State correction
+                    x_corr = P_proj @ G.T @ lam
+                    x_proj = x_proj + x_corr
+
+                    # Covariance projection
+                    # P_proj = P - P G^T (G P G^T)^{-1} G P
+                    P_proj = P_proj - GP.T @ GPGt_inv @ GP
+
+                    # Ensure symmetry
+                    P_proj = (P_proj + P_proj.T) / 2
+
+                    # Enforce positive definiteness
+                    eigvals, eigvecs = np.linalg.eigh(P_proj)
+                    if np.any(eigvals < 0):
+                        eigvals[eigvals < 1e-10] = 1e-10
+                        P_proj = eigvecs @ np.diag(eigvals) @ eigvecs.T
+
+                except np.linalg.LinAlgError:
+                    # If inversion fails, use pseudoinverse
+                    GPGt_pinv = np.linalg.pinv(GPGt)
+                    lam = -GPGt_pinv @ (G @ x_proj + g_val)
+                    x_proj = x_proj + P_proj @ G.T @ lam
+
+            if converged:
+                break
+
+        return x_proj, P_proj
+
+
+def constrained_ekf_predict(
+    x: ArrayLike,
+    P: ArrayLike,
+    f: Callable[[NDArray[Any]], NDArray[Any]],
+    F: ArrayLike,
+    Q: ArrayLike,
+) -> KalmanPrediction:
+    """
+    Convenience function for constrained EKF prediction.
+
+    Parameters
+    ----------
+    x : array_like
+        Current state estimate.
+    P : array_like
+        Current covariance.
+    f : callable
+        Nonlinear dynamics function.
+    F : array_like
+        Jacobian of f.
+    Q : array_like
+        Process noise covariance.
+
+    Returns
+    -------
+    result : KalmanPrediction
+        Predicted state and covariance.
+    """
+    cekf = ConstrainedEKF()
+    return cekf.predict(x, P, f, F, Q)
+
+
+def constrained_ekf_update(
+    x: ArrayLike,
+    P: ArrayLike,
+    z: ArrayLike,
+    h: Callable[[NDArray[Any]], NDArray[Any]],
+    H: ArrayLike,
+    R: ArrayLike,
+    constraints: Optional[list[ConstraintFunction]] = None,
+) -> KalmanUpdate:
+    """
+    Convenience function for constrained EKF update.
+
+    Parameters
+    ----------
+    x : array_like
+        Predicted state.
+    P : array_like
+        Predicted covariance.
+    z : array_like
+        Measurement.
+    h : callable
+        Nonlinear measurement function.
+    H : array_like
+        Jacobian of h.
+    R : array_like
+        Measurement noise covariance.
+    constraints : list, optional
+        List of ConstraintFunction objects.
+
+    Returns
+    -------
+    result : KalmanUpdate
+        Updated state and covariance.
+    """
+    cekf = ConstrainedEKF()
+    if constraints:
+        for c in constraints:
+            cekf.add_constraint(c)
+    return cekf.update(x, P, z, h, H, R)
+
+
+__all__ = [
+    "ConstraintFunction",
+    "ConstrainedEKF",
+    "constrained_ekf_predict",
+    "constrained_ekf_update",
+]

pytcl/dynamic_estimation/kalman/extended.py

@@ -5,7 +5,7 @@ The EKF handles nonlinear dynamics and/or measurements by linearizing
 around the current state estimate using Jacobians.
 """
 
-from typing import Callable
+from typing import Any, Callable
 
 import numpy as np
 from numpy.typing import ArrayLike, NDArray
@@ -16,7 +16,7 @@ from pytcl.dynamic_estimation.kalman.linear import KalmanPrediction, KalmanUpdat
 def ekf_predict(
     x: ArrayLike,
     P: ArrayLike,
-    f: Callable[[NDArray], NDArray],
+    f: Callable[[NDArray[Any]], NDArray[Any]],
     F: ArrayLike,
     Q: ArrayLike,
 ) -> KalmanPrediction:
@@ -86,7 +86,7 @@ def ekf_update(
     x: ArrayLike,
     P: ArrayLike,
     z: ArrayLike,
-    h: Callable[[NDArray], NDArray],
+    h: Callable[[NDArray[Any]], NDArray[Any]],
     H: ArrayLike,
     R: ArrayLike,
 ) -> KalmanUpdate:
@@ -183,7 +183,7 @@ def ekf_update(
 
 
 def numerical_jacobian(
-    f: Callable[[NDArray], NDArray],
+    f: Callable[[NDArray[Any]], NDArray[Any]],
     x: ArrayLike,
     dx: float = 1e-7,
 ) -> NDArray[np.floating]:
@@ -239,7 +239,7 @@ def numerical_jacobian(
 def ekf_predict_auto(
     x: ArrayLike,
     P: ArrayLike,
-    f: Callable[[NDArray], NDArray],
+    f: Callable[[NDArray[Any]], NDArray[Any]],
     Q: ArrayLike,
     dx: float = 1e-7,
 ) -> KalmanPrediction:
@@ -273,7 +273,7 @@ def ekf_update_auto(
     x: ArrayLike,
     P: ArrayLike,
     z: ArrayLike,
-    h: Callable[[NDArray], NDArray],
+    h: Callable[[NDArray[Any]], NDArray[Any]],
     R: ArrayLike,
     dx: float = 1e-7,
 ) -> KalmanUpdate:
@@ -309,8 +309,8 @@ def iterated_ekf_update(
     x: ArrayLike,
     P: ArrayLike,
     z: ArrayLike,
-    h: Callable[[NDArray], NDArray],
-    H_func: Callable[[NDArray], NDArray],
+    h: Callable[[NDArray[Any]], NDArray[Any]],
+    H_func: Callable[[NDArray[Any]], NDArray[Any]],
     R: ArrayLike,
     max_iter: int = 10,
     tol: float = 1e-6,

pytcl/dynamic_estimation/kalman/h_infinity.py

@@ -22,7 +22,7 @@ References
        Proc. IEEE CDC, 1992.
 """
 
-from typing import NamedTuple, Optional
+from typing import Any, Callable, NamedTuple, Optional
 
 import numpy as np
 from numpy.typing import ArrayLike, NDArray
@@ -379,7 +379,7 @@ def extended_hinf_update(
     x: ArrayLike,
     P: ArrayLike,
     z: ArrayLike,
-    h: callable,
+    h: Callable[[np.ndarray[Any, Any]], np.ndarray[Any, Any]],
     H: ArrayLike,
     R: ArrayLike,
     gamma: float,

pytcl/dynamic_estimation/kalman/square_root.py

@@ -77,7 +77,9 @@ class SRKalmanUpdate(NamedTuple):
     likelihood: float
 
 
-def cholesky_update(S: NDArray, v: NDArray, sign: float = 1.0) -> NDArray:
+def cholesky_update(
+    S: NDArray[np.floating], v: NDArray[np.floating], sign: float = 1.0
+) -> NDArray[np.floating]:
     """
     Rank-1 Cholesky update/downdate.
 
@@ -148,7 +150,11 @@ def cholesky_update(S: NDArray, v: NDArray, sign: float = 1.0) -> NDArray:
     return S
 
 
-def qr_update(S_x: NDArray, S_noise: NDArray, F: Optional[NDArray] = None) -> NDArray:
+def qr_update(
+    S_x: NDArray[np.floating],
+    S_noise: NDArray[np.floating],
+    F: Optional[NDArray[np.floating]] = None,
+) -> NDArray[np.floating]:
     """
     QR-based covariance square root update.
 

pytcl/dynamic_estimation/kalman/sr_ukf.py

@@ -14,7 +14,7 @@ References
        Estimation," Proceedings of the IEEE, 2004.
 """
 
-from typing import Callable
+from typing import Any, Callable
 
 import numpy as np
 import scipy.linalg
@@ -30,7 +30,7 @@ from pytcl.dynamic_estimation.kalman.square_root import (
 def sr_ukf_predict(
     x: ArrayLike,
     S: ArrayLike,
-    f: Callable,
+    f: Callable[[np.ndarray[Any, Any]], np.ndarray[Any, Any]],
     S_Q: ArrayLike,
     alpha: float = 1e-3,
     beta: float = 2.0,
@@ -143,7 +143,7 @@ def sr_ukf_update(
     x: ArrayLike,
     S: ArrayLike,
     z: ArrayLike,
-    h: Callable,
+    h: Callable[[np.ndarray[Any, Any]], np.ndarray[Any, Any]],
     S_R: ArrayLike,
     alpha: float = 1e-3,
     beta: float = 2.0,

pytcl/dynamic_estimation/kalman/ud_filter.py

@@ -41,7 +41,7 @@ class UDState(NamedTuple):
     D: NDArray[np.floating]
 
 
-def ud_factorize(P: ArrayLike) -> tuple[NDArray, NDArray]:
+def ud_factorize(P: ArrayLike) -> tuple[NDArray[np.floating], NDArray[np.floating]]:
     """
     Compute U-D factorization of a symmetric positive definite matrix.
 
@@ -91,7 +91,7 @@ def ud_factorize(P: ArrayLike) -> tuple[NDArray, NDArray]:
     return U, D
 
 
-def ud_reconstruct(U: ArrayLike, D: ArrayLike) -> NDArray:
+def ud_reconstruct(U: ArrayLike, D: ArrayLike) -> NDArray[np.floating]:
     """
     Reconstruct covariance matrix from U-D factors.
 
@@ -128,7 +128,7 @@ def ud_predict(
     D: ArrayLike,
     F: ArrayLike,
     Q: ArrayLike,
-) -> tuple[NDArray, NDArray, NDArray]:
+) -> tuple[NDArray[np.floating], NDArray[np.floating], NDArray[np.floating]]:
     """
     U-D filter prediction step.
 
@@ -193,7 +193,7 @@ def ud_update_scalar(
     z: float,
     h: ArrayLike,
     r: float,
-) -> tuple[NDArray, NDArray, NDArray]:
+) -> tuple[NDArray[np.floating], NDArray[np.floating], NDArray[np.floating]]:
     """
     U-D filter scalar measurement update (Bierman's algorithm).
 
@@ -290,7 +290,13 @@ def ud_update(
     z: ArrayLike,
     H: ArrayLike,
     R: ArrayLike,
-) -> tuple[NDArray, NDArray, NDArray, NDArray, float]:
+) -> tuple[
+    NDArray[np.floating],
+    NDArray[np.floating],
+    NDArray[np.floating],
+    NDArray[np.floating],
+    float,
+]:
     """
     U-D filter vector measurement update.
 

pytcl/dynamic_estimation/kalman/unscented.py

@@ -5,7 +5,7 @@ The UKF uses the unscented transform to propagate the mean and covariance
 through nonlinear functions without requiring Jacobian computation.
 """
 
-from typing import Callable, NamedTuple, Optional, Tuple
+from typing import Any, Callable, NamedTuple, Optional, Tuple
 
 import numpy as np
 from numpy.typing import ArrayLike, NDArray
@@ -226,7 +226,7 @@ def unscented_transform(
 def ukf_predict(
     x: ArrayLike,
     P: ArrayLike,
-    f: Callable[[NDArray], NDArray],
+    f: Callable[[NDArray[Any]], NDArray[Any]],
     Q: ArrayLike,
     alpha: float = 1e-3,
     beta: float = 2.0,
@@ -292,7 +292,7 @@ def ukf_update(
     x: ArrayLike,
     P: ArrayLike,
     z: ArrayLike,
-    h: Callable[[NDArray], NDArray],
+    h: Callable[[NDArray[Any]], NDArray[Any]],
     R: ArrayLike,
     alpha: float = 1e-3,
     beta: float = 2.0,
@@ -382,7 +382,9 @@ def ukf_update(
     )
 
 
-def ckf_spherical_cubature_points(n: int) -> Tuple[NDArray, NDArray]:
+def ckf_spherical_cubature_points(
+    n: int,
+) -> tuple[NDArray[np.floating], NDArray[np.floating]]:
     """
     Generate cubature points for Cubature Kalman Filter.
 
@@ -418,7 +420,7 @@ def ckf_spherical_cubature_points(n: int) -> Tuple[NDArray, NDArray]:
 def ckf_predict(
     x: ArrayLike,
     P: ArrayLike,
-    f: Callable[[NDArray], NDArray],
+    f: Callable[[NDArray[Any]], NDArray[Any]],
     Q: ArrayLike,
 ) -> KalmanPrediction:
     """
@@ -487,7 +489,7 @@ def ckf_update(
     x: ArrayLike,
     P: ArrayLike,
     z: ArrayLike,
-    h: Callable[[NDArray], NDArray],
+    h: Callable[[NDArray[Any]], NDArray[Any]],
     R: ArrayLike,
 ) -> KalmanUpdate:
     """

pytcl/dynamic_estimation/particle_filters/bootstrap.py

@@ -5,7 +5,7 @@ This module provides particle filtering algorithms for nonlinear/non-Gaussian
 state estimation.
 """
 
-from typing import Callable, NamedTuple, Optional, Tuple
+from typing import Any, Callable, NamedTuple, Optional, Tuple
 
 import numpy as np
 from numba import njit
@@ -102,9 +102,9 @@ def resample_systematic(
 
 @njit(cache=True)
 def _resample_residual_deterministic(
-    particles: np.ndarray,
-    floor_Nw: np.ndarray,
-) -> Tuple[np.ndarray, int]:
+    particles: np.ndarray[Any, Any],
+    floor_Nw: np.ndarray[Any, Any],
+) -> Tuple[np.ndarray[Any, Any], int]:
     """JIT-compiled deterministic copy portion of residual resampling."""
     N = particles.shape[0]
     n = particles.shape[1]
@@ -196,8 +196,8 @@ def effective_sample_size(weights: NDArray[np.floating]) -> float:
 
 def bootstrap_pf_predict(
     particles: NDArray[np.floating],
-    f: Callable[[NDArray], NDArray],
-    Q_sample: Callable[[int, Optional[np.random.Generator]], NDArray],
+    f: Callable[[NDArray[Any]], NDArray[Any]],
+    Q_sample: Callable[[int, Optional[np.random.Generator]], NDArray[Any]],
     rng: Optional[np.random.Generator] = None,
 ) -> NDArray[np.floating]:
     """
@@ -242,7 +242,7 @@ def bootstrap_pf_update(
     particles: NDArray[np.floating],
     weights: NDArray[np.floating],
     z: ArrayLike,
-    likelihood_func: Callable[[NDArray, NDArray], float],
+    likelihood_func: Callable[[NDArray[Any], NDArray[Any]], float],
 ) -> Tuple[NDArray[np.floating], float]:
     """
     Bootstrap particle filter update step.
@@ -328,9 +328,9 @@ def bootstrap_pf_step(
     particles: NDArray[np.floating],
     weights: NDArray[np.floating],
     z: ArrayLike,
-    f: Callable[[NDArray], NDArray],
-    h: Callable[[NDArray], NDArray],
-    Q_sample: Callable[[int, Optional[np.random.Generator]], NDArray],
+    f: Callable[[NDArray[Any]], NDArray[Any]],
+    h: Callable[[NDArray[Any]], NDArray[Any]],
+    Q_sample: Callable[[int, Optional[np.random.Generator]], NDArray[Any]],
     R: ArrayLike,
     resample_threshold: float = 0.5,
     resample_method: str = "systematic",
@@ -378,7 +378,7 @@ def bootstrap_pf_step(
     particles_pred = bootstrap_pf_predict(particles, f, Q_sample, rng)
 
     # Update
-    def likelihood_func(z, x):
+    def likelihood_func(z: NDArray[Any], x: NDArray[Any]) -> Any:
         z_pred = h(x)
         return gaussian_likelihood(z, z_pred, R)
 
@@ -426,10 +426,10 @@ def particle_mean(
 
 @njit(cache=True)
 def _particle_covariance_core(
-    particles: np.ndarray,
-    weights: np.ndarray,
-    mean: np.ndarray,
-) -> np.ndarray:
+    particles: np.ndarray[Any, Any],
+    weights: np.ndarray[Any, Any],
+    mean: np.ndarray[Any, Any],
+) -> np.ndarray[Any, Any]:
     """JIT-compiled core for particle covariance computation."""
     N = particles.shape[0]
     n = particles.shape[1]