nrl-tracker 1.7.0__py3-none-any.whl → 1.7.3__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/base.py

@@ -8,7 +8,7 @@ Cover trees.
 
 import logging
 from abc import ABC, abstractmethod
-from typing import Callable, List, NamedTuple, Optional
+from typing import Any, Callable, List, NamedTuple, Optional
 
 import numpy as np
 from numpy.typing import ArrayLike, NDArray
@@ -155,7 +155,7 @@ class MetricSpatialIndex(BaseSpatialIndex):
     def __init__(
         self,
         data: ArrayLike,
-        metric: Optional[Callable[[NDArray, NDArray], float]] = None,
+        metric: Optional[Callable[[NDArray[Any], NDArray[Any]], float]] = None,
     ):
         super().__init__(data)
 
@@ -165,7 +165,7 @@ class MetricSpatialIndex(BaseSpatialIndex):
             self.metric = metric
 
     @staticmethod
-    def _euclidean_distance(x: NDArray, y: NDArray) -> float:
+    def _euclidean_distance(x: NDArray[Any], y: NDArray[Any]) -> float:
         """Default Euclidean distance metric."""
         return float(np.sqrt(np.sum((x - y) ** 2)))
 

pytcl/containers/cluster_set.py

@@ -7,7 +7,17 @@ that move together (formations, convoys, etc.).
 
 from __future__ import annotations
 
-from typing import Dict, Iterable, Iterator, List, NamedTuple, Optional, Tuple, Union
+from typing import (
+    Any,
+    Dict,
+    Iterable,
+    Iterator,
+    List,
+    NamedTuple,
+    Optional,
+    Tuple,
+    Union,
+)
 
 import numpy as np
 from numpy.typing import ArrayLike, NDArray
@@ -303,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

@@ -12,7 +12,7 @@ References
 """
 
 import logging
-from typing import Callable, List, NamedTuple, Optional, Set, Tuple
+from typing import Any, Callable, List, NamedTuple, Optional, Set, Tuple
 
 import numpy as np
 from numpy.typing import ArrayLike, NDArray
@@ -109,7 +109,9 @@ class CoverTree(MetricSpatialIndex):
     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,
     ):
         super().__init__(data, metric)
@@ -141,7 +143,7 @@ class CoverTree(MetricSpatialIndex):
             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)
 

pytcl/containers/rtree.py

@@ -654,7 +654,7 @@ class RTree:
         query = np.asarray(query_point, dtype=np.float64)
         neighbors: List[Tuple[float, int]] = []
 
-        def min_dist_to_box(point: NDArray, bbox: BoundingBox) -> float:
+        def min_dist_to_box(point: NDArray[np.floating], bbox: BoundingBox) -> float:
             """Minimum distance from point to bounding box."""
             clamped = np.clip(point, bbox.min_coords, bbox.max_coords)
             return float(np.sqrt(np.sum((point - clamped) ** 2)))

pytcl/containers/vptree.py

@@ -12,7 +12,7 @@ References
 """
 
 import logging
-from typing import Callable, List, NamedTuple, Optional, Tuple
+from typing import Any, Callable, List, NamedTuple, Optional, Tuple
 
 import numpy as np
 from numpy.typing import ArrayLike, NDArray
@@ -105,7 +105,9 @@ class VPTree(MetricSpatialIndex):
     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,
     ):
         super().__init__(data, metric)
 

pytcl/coordinate_systems/conversions/geodetic.py

@@ -166,7 +166,7 @@ def ecef2geodetic(
 
         # Altitude
         # Use cos_lat when available, otherwise use sin_lat with guard against division by zero
-        with np.errstate(divide='ignore', invalid='ignore'):
+        with np.errstate(divide="ignore", invalid="ignore"):
             alt = np.where(
                 np.abs(cos_lat) > 1e-10,
                 p / cos_lat - N,
@@ -651,14 +651,30 @@ def ecef2sez(
     # Z = cos(lat)*cos(lon)*dX + cos(lat)*sin(lon)*dY + sin(lat)*dZ
 
     if delta_ecef.ndim == 1:
-        s = -sin_lat * cos_lon * delta_ecef[0] - sin_lat * sin_lon * delta_ecef[1] + cos_lat * delta_ecef[2]
+        s = (
+            -sin_lat * cos_lon * delta_ecef[0]
+            - sin_lat * sin_lon * delta_ecef[1]
+            + cos_lat * delta_ecef[2]
+        )
         e = -sin_lon * delta_ecef[0] + cos_lon * delta_ecef[1]
-        z = cos_lat * cos_lon * delta_ecef[0] + cos_lat * sin_lon * delta_ecef[1] + sin_lat * delta_ecef[2]
+        z = (
+            cos_lat * cos_lon * delta_ecef[0]
+            + cos_lat * sin_lon * delta_ecef[1]
+            + sin_lat * delta_ecef[2]
+        )
         return np.array([s, e, z], dtype=np.float64)
     else:
-        s = -sin_lat * cos_lon * delta_ecef[0, :] - sin_lat * sin_lon * delta_ecef[1, :] + cos_lat * delta_ecef[2, :]
+        s = (
+            -sin_lat * cos_lon * delta_ecef[0, :]
+            - sin_lat * sin_lon * delta_ecef[1, :]
+            + cos_lat * delta_ecef[2, :]
+        )
         e = -sin_lon * delta_ecef[0, :] + cos_lon * delta_ecef[1, :]
-        z = cos_lat * cos_lon * delta_ecef[0, :] + cos_lat * sin_lon * delta_ecef[1, :] + sin_lat * delta_ecef[2, :]
+        z = (
+            cos_lat * cos_lon * delta_ecef[0, :]
+            + cos_lat * sin_lon * delta_ecef[1, :]
+            + sin_lat * delta_ecef[2, :]
+        )
         return np.array([s, e, z], dtype=np.float64)
 
 
@@ -714,8 +730,16 @@ def sez2ecef(
     else:
         if sez.shape[0] != 3:
             sez = sez.T
-        dX = -sin_lat * cos_lon * sez[0, :] - sin_lon * sez[1, :] + cos_lat * cos_lon * sez[2, :]
-        dY = -sin_lat * sin_lon * sez[0, :] + cos_lon * sez[1, :] + cos_lat * sin_lon * sez[2, :]
+        dX = (
+            -sin_lat * cos_lon * sez[0, :]
+            - sin_lon * sez[1, :]
+            + cos_lat * cos_lon * sez[2, :]
+        )
+        dY = (
+            -sin_lat * sin_lon * sez[0, :]
+            + cos_lon * sez[1, :]
+            + cos_lat * sin_lon * sez[2, :]
+        )
         dZ = cos_lat * sez[0, :] + sin_lat * sez[2, :]
         return ecef_ref[:, np.newaxis] + np.array([dX, dY, dZ], dtype=np.float64)
 

pytcl/coordinate_systems/jacobians/jacobians.py

@@ -6,7 +6,7 @@ coordinate transformations, essential for error propagation in tracking
 filters (e.g., converting measurement covariances between coordinate systems).
 """
 
-from typing import Literal
+from typing import Callable, Literal
 
 import numpy as np
 from numpy.typing import ArrayLike, NDArray
@@ -431,7 +431,7 @@ def cross_covariance_transform(
 
 
 def numerical_jacobian(
-    func,
+    func: Callable[[ArrayLike], ArrayLike],
     x: ArrayLike,
     dx: float = 1e-7,
 ) -> NDArray[np.floating]:

pytcl/coordinate_systems/projections/__init__.py

@@ -27,7 +27,7 @@ Examples
 """
 
 from pytcl.coordinate_systems.projections.projections import (
-    WGS84_A,  # Constants; Result types; Azimuthal Equidistant; UTM; Lambert Conformal Conic; Mercator; Stereographic; Transverse Mercator
+    WGS84_A,  # Constants; Result types
 )
 from pytcl.coordinate_systems.projections.projections import (
     WGS84_B,

pytcl/coordinate_systems/projections/projections.py

@@ -25,7 +25,7 @@ References
        nanometers." Journal of Geodesy 85.8 (2011): 475-485.
 """
 
-from typing import NamedTuple, Optional, Tuple
+from typing import Any, NamedTuple, Optional, Tuple
 
 import numpy as np
 from numpy.typing import NDArray
@@ -1253,7 +1253,7 @@ def geodetic2utm_batch(
     lats: NDArray[np.floating],
     lons: NDArray[np.floating],
     zone: Optional[int] = None,
-) -> Tuple[NDArray[np.floating], NDArray[np.floating], NDArray[np.int_], NDArray]:
+) -> tuple[NDArray[np.floating], NDArray[np.floating], NDArray[np.intp], NDArray[Any]]:
     """
     Batch convert geodetic coordinates to UTM.
 

pytcl/coordinate_systems/rotations/rotations.py

@@ -6,7 +6,7 @@ representations including rotation matrices, quaternions, Euler angles,
 axis-angle, and Rodrigues parameters.
 """
 
-from typing import Tuple
+from typing import Any, Tuple
 
 import numpy as np
 from numba import njit
@@ -14,7 +14,7 @@ from numpy.typing import ArrayLike, NDArray
 
 
 @njit(cache=True, fastmath=True)
-def _rotx_inplace(angle: float, R: np.ndarray) -> None:
+def _rotx_inplace(angle: float, R: np.ndarray[Any, Any]) -> None:
     """JIT-compiled rotation about x-axis (fills existing matrix)."""
     c = np.cos(angle)
     s = np.sin(angle)
@@ -30,7 +30,7 @@ def _rotx_inplace(angle: float, R: np.ndarray) -> None:
 
 
 @njit(cache=True, fastmath=True)
-def _roty_inplace(angle: float, R: np.ndarray) -> None:
+def _roty_inplace(angle: float, R: np.ndarray[Any, Any]) -> None:
     """JIT-compiled rotation about y-axis (fills existing matrix)."""
     c = np.cos(angle)
     s = np.sin(angle)
@@ -46,7 +46,7 @@ def _roty_inplace(angle: float, R: np.ndarray) -> None:
 
 
 @njit(cache=True, fastmath=True)
-def _rotz_inplace(angle: float, R: np.ndarray) -> None:
+def _rotz_inplace(angle: float, R: np.ndarray[Any, Any]) -> None:
     """JIT-compiled rotation about z-axis (fills existing matrix)."""
     c = np.cos(angle)
     s = np.sin(angle)
@@ -62,7 +62,9 @@ def _rotz_inplace(angle: float, R: np.ndarray) -> None:
 
 
 @njit(cache=True, fastmath=True)
-def _euler_zyx_to_rotmat(yaw: float, pitch: float, roll: float, R: np.ndarray) -> None:
+def _euler_zyx_to_rotmat(
+    yaw: float, pitch: float, roll: float, R: np.ndarray[Any, Any]
+) -> None:
     """JIT-compiled ZYX Euler angles to rotation matrix."""
     cy = np.cos(yaw)
     sy = np.sin(yaw)
@@ -84,7 +86,9 @@ def _euler_zyx_to_rotmat(yaw: float, pitch: float, roll: float, R: np.ndarray) -
 
 
 @njit(cache=True, fastmath=True)
-def _matmul_3x3(A: np.ndarray, B: np.ndarray, C: np.ndarray) -> None:
+def _matmul_3x3(
+    A: np.ndarray[Any, Any], B: np.ndarray[Any, Any], C: np.ndarray[Any, Any]
+) -> None:
     """JIT-compiled 3x3 matrix multiplication C = A @ B."""
     for i in range(3):
         for j in range(3):

pytcl/core/validation.py

@@ -28,7 +28,7 @@ def validate_array(
     arr: ArrayLike,
     name: str = "array",
     *,
-    dtype: type | np.dtype | None = None,
+    dtype: type | np.dtype[Any] | None = None,
     ndim: int | tuple[int, ...] | None = None,
     shape: tuple[int | None, ...] | None = None,
     min_ndim: int | None = None,
@@ -415,7 +415,7 @@ def validate_same_shape(*arrays: ArrayLike, names: Sequence[str] | None = None)
 def validated_array_input(
     param_name: str,
     *,
-    dtype: type | np.dtype | None = None,
+    dtype: type | np.dtype[Any] | None = None,
     ndim: int | tuple[int, ...] | None = None,
     shape: tuple[int | None, ...] | None = None,
     finite: bool = False,
@@ -516,7 +516,7 @@ class ArraySpec:
     def __init__(
         self,
         *,
-        dtype: type | np.dtype | None = None,
+        dtype: type | np.dtype[Any] | None = None,
         ndim: int | tuple[int, ...] | None = None,
         shape: tuple[int | None, ...] | None = None,
         min_ndim: int | None = None,

pytcl/dynamic_estimation/__init__.py

@@ -14,6 +14,14 @@ This module provides filtering and smoothing algorithms for state estimation:
 # Import submodules for easy access
 from pytcl.dynamic_estimation import kalman, particle_filters
 
+# Gaussian Sum Filter
+from pytcl.dynamic_estimation.gaussian_sum_filter import (
+    GaussianComponent,
+    GaussianSumFilter,
+    gaussian_sum_filter_predict,
+    gaussian_sum_filter_update,
+)
+
 # IMM estimator
 from pytcl.dynamic_estimation.imm import (
     IMMEstimator,
@@ -25,14 +33,6 @@ from pytcl.dynamic_estimation.imm import (
     imm_update,
 )
 
-# Gaussian Sum Filter
-from pytcl.dynamic_estimation.gaussian_sum_filter import (
-    GaussianComponent,
-    GaussianSumFilter,
-    gaussian_sum_filter_predict,
-    gaussian_sum_filter_update,
-)
-
 # Information filter
 from pytcl.dynamic_estimation.information_filter import (
     InformationFilterResult,
@@ -93,14 +93,6 @@ from pytcl.dynamic_estimation.kalman import (
     unscented_transform,
 )
 
-# Rao-Blackwellized Particle Filter
-from pytcl.dynamic_estimation.rbpf import (
-    RBPFFilter,
-    RBPFParticle,
-    rbpf_predict,
-    rbpf_update,
-)
-
 # Particle filters
 from pytcl.dynamic_estimation.particle_filters import (
     ParticleState,
@@ -117,6 +109,14 @@ from pytcl.dynamic_estimation.particle_filters import (
     resample_systematic,
 )
 
+# Rao-Blackwellized Particle Filter
+from pytcl.dynamic_estimation.rbpf import (
+    RBPFFilter,
+    RBPFParticle,
+    rbpf_predict,
+    rbpf_update,
+)
+
 # Smoothers
 from pytcl.dynamic_estimation.smoothers import (
     FixedLagResult,

pytcl/dynamic_estimation/gaussian_sum_filter.py

@@ -31,8 +31,8 @@ from pytcl.dynamic_estimation.kalman.extended import ekf_predict, ekf_update
 class GaussianComponent(NamedTuple):
     """Single Gaussian component in mixture."""
 
-    x: NDArray  # State estimate
-    P: NDArray  # Covariance
+    x: NDArray[np.floating]  # State estimate
+    P: NDArray[np.floating]  # Covariance
     w: float  # Weight (probability)
 
 
@@ -108,15 +108,13 @@ class GaussianSumFilter:
                 x = x0.copy()
             else:
                 # Slight perturbation for diversity
-                x = x0 + np.random.randn(x0.shape[0]) * np.sqrt(
-                    np.diag(P0)
-                ) * 0.1
+                x = x0 + np.random.randn(x0.shape[0]) * np.sqrt(np.diag(P0)) * 0.1
 
             self.components.append(GaussianComponent(x=x, P=P0.copy(), w=weight))
 
     def predict(
         self,
-        f: Callable,
+        f: Callable[[NDArray[np.floating]], NDArray[np.floating]],
         F: ArrayLike,
         Q: ArrayLike,
     ) -> None:
@@ -138,16 +136,14 @@ class GaussianSumFilter:
         for comp in self.components:
             # EKF predict for each component
             pred = ekf_predict(comp.x, comp.P, f, F, Q)
-            new_components.append(
-                GaussianComponent(x=pred.x, P=pred.P, w=comp.w)
-            )
+            new_components.append(GaussianComponent(x=pred.x, P=pred.P, w=comp.w))
 
         self.components = new_components
 
     def update(
         self,
         z: ArrayLike,
-        h: Callable,
+        h: Callable[[NDArray[np.floating]], NDArray[np.floating]],
         H: ArrayLike,
         R: ArrayLike,
     ) -> None:
@@ -179,9 +175,7 @@ class GaussianSumFilter:
             # Likelihood from this measurement
             likelihood = upd.likelihood
 
-            updated_components.append(
-                GaussianComponent(x=upd.x, P=upd.P, w=comp.w)
-            )
+            updated_components.append(GaussianComponent(x=upd.x, P=upd.P, w=comp.w))
             likelihoods.append(likelihood)
 
         # Adapt weights based on measurement likelihood
@@ -210,15 +204,11 @@ class GaussianSumFilter:
 
     def _prune_components(self) -> None:
         """Remove components with weight below threshold."""
-        self.components = [
-            c for c in self.components if c.w >= self.prune_threshold
-        ]
+        self.components = [c for c in self.components if c.w >= self.prune_threshold]
 
         if len(self.components) == 0:
             # Failsafe: keep best component
-            self.components = [
-                max(self.components, key=lambda c: c.w)
-            ]
+            self.components = [max(self.components, key=lambda c: c.w)]
 
         # Renormalize weights
         total_weight = sum(c.w for c in self.components)
@@ -232,13 +222,11 @@ class GaussianSumFilter:
         """Merge similar components to keep count manageable."""
         while len(self.components) > self.max_components:
             # Find pair with smallest KL divergence
-            best_i, best_j, best_kl = 0, 1, float('inf')
+            best_i, best_j, best_kl = 0, 1, float("inf")
 
             for i in range(len(self.components)):
                 for j in range(i + 1, len(self.components)):
-                    kl = self._kl_divergence(
-                        self.components[i], self.components[j]
-                    )
+                    kl = self._kl_divergence(self.components[i], self.components[j])
                     if kl < best_kl:
                         best_kl = kl
                         best_i = best_j = i
@@ -261,8 +249,7 @@ class GaussianSumFilter:
                 dx_i = ci.x - x_new
                 dx_j = cj.x - x_new
                 P_new += (
-                    ci.w * np.outer(dx_i, dx_i)
-                    + cj.w * np.outer(dx_j, dx_j)
+                    ci.w * np.outer(dx_i, dx_i) + cj.w * np.outer(dx_j, dx_j)
                 ) / w_new
 
                 # Create merged component
@@ -270,7 +257,8 @@ class GaussianSumFilter:
 
                 # Replace with merged, remove old
                 self.components = [
-                    c for i, c in enumerate(self.components)
+                    c
+                    for i, c in enumerate(self.components)
                     if i != best_i and i != best_j
                 ]
                 self.components.append(merged)
@@ -300,9 +288,7 @@ class GaussianSumFilter:
 
         try:
             P2_inv = np.linalg.inv(c2.P)
-            logdet_ratio = np.linalg.slogdet(c2.P)[1] - np.linalg.slogdet(
-                c1.P
-            )[1]
+            logdet_ratio = np.linalg.slogdet(c2.P)[1] - np.linalg.slogdet(c1.P)[1]
 
             trace_term = np.trace(P2_inv @ c1.P)
             quad_term = dx @ P2_inv @ dx
@@ -313,7 +299,7 @@ class GaussianSumFilter:
             # Singular matrix, return large KL
             return 1e6
 
-    def estimate(self) -> tuple[NDArray, NDArray]:
+    def estimate(self) -> tuple[NDArray[np.floating], NDArray[np.floating]]:
         """Get overall state estimate (weighted mean and covariance).
 
         Returns
@@ -356,7 +342,7 @@ class GaussianSumFilter:
 
 def gaussian_sum_filter_predict(
     components: List[GaussianComponent],
-    f: Callable,
+    f: Callable[[NDArray[np.floating]], NDArray[np.floating]],
     F: ArrayLike,
     Q: ArrayLike,
 ) -> List[GaussianComponent]:
@@ -384,9 +370,7 @@ def gaussian_sum_filter_predict(
     new_components = []
     for comp in components:
         pred = ekf_predict(comp.x, comp.P, f, F, Q)
-        new_components.append(
-            GaussianComponent(x=pred.x, P=pred.P, w=comp.w)
-        )
+        new_components.append(GaussianComponent(x=pred.x, P=pred.P, w=comp.w))
 
     return new_components
 
@@ -394,7 +378,7 @@ def gaussian_sum_filter_predict(
 def gaussian_sum_filter_update(
     components: List[GaussianComponent],
     z: ArrayLike,
-    h: Callable,
+    h: Callable[[NDArray[np.floating]], NDArray[np.floating]],
     H: ArrayLike,
     R: ArrayLike,
 ) -> List[GaussianComponent]:
@@ -429,9 +413,7 @@ def gaussian_sum_filter_update(
         upd = ekf_update(comp.x, comp.P, z, h, H, R)
         likelihood = upd.likelihood
 
-        updated_components.append(
-            GaussianComponent(x=upd.x, P=upd.P, w=comp.w)
-        )
+        updated_components.append(GaussianComponent(x=upd.x, P=upd.P, w=comp.w))
         likelihoods.append(likelihood)
 
     # Adapt weights