Source code for pytcl.dynamic_estimation.kalman.unscented

"""
Unscented Kalman Filter (UKF) implementation.

The UKF uses the unscented transform to propagate the mean and covariance
through nonlinear functions without requiring Jacobian computation.
"""

from typing import Any, Callable, NamedTuple, Optional, Tuple

import numpy as np
from numpy.typing import ArrayLike, NDArray

from pytcl.dynamic_estimation.kalman.linear import KalmanPrediction, KalmanUpdate


[docs] class SigmaPoints(NamedTuple): """Sigma points and weights for unscented transform. Attributes ---------- points : ndarray Sigma points, shape (2n+1, n). Wm : ndarray Weights for mean computation. Wc : ndarray Weights for covariance computation. """ points: NDArray[np.floating] Wm: NDArray[np.floating] Wc: NDArray[np.floating]
[docs] def sigma_points_merwe( x: ArrayLike, P: ArrayLike, alpha: float = 1e-3, beta: float = 2.0, kappa: float = 0.0, ) -> SigmaPoints: """ Generate sigma points using Van der Merwe's scaled method. Parameters ---------- x : array_like State mean, shape (n,). P : array_like State covariance, shape (n, n). alpha : float, optional Spread of sigma points (default: 1e-3). Small values keep points close to mean. beta : float, optional Prior knowledge of distribution (default: 2.0 for Gaussian). kappa : float, optional Secondary scaling parameter (default: 0, or 3-n). Returns ------- result : SigmaPoints Sigma points and weights. Examples -------- >>> import numpy as np >>> x = np.array([0.0, 0.0]) >>> P = np.eye(2) >>> sp = sigma_points_merwe(x, P) >>> sp.points.shape (5, 2) Notes ----- Generates 2n+1 sigma points: X_0 = x X_i = x + sqrt((n+lambda)*P)_i, i = 1..n X_{i+n} = x - sqrt((n+lambda)*P)_i, i = 1..n where lambda = alpha^2 * (n + kappa) - n. References ---------- - Van der Merwe, R., "Sigma-Point Kalman Filters for Probabilistic Inference in Dynamic State-Space Models", PhD Thesis, 2004. """ x = np.asarray(x, dtype=np.float64).flatten() P = np.asarray(P, dtype=np.float64) n = len(x) lambda_ = alpha**2 * (n + kappa) - n # Compute matrix square root try: sqrt_P = np.linalg.cholesky((n + lambda_) * P) except np.linalg.LinAlgError: # Fall back to eigendecomposition for near-singular P eigvals, eigvecs = np.linalg.eigh(P) eigvals = np.maximum(eigvals, 1e-10) sqrt_P = eigvecs @ np.diag(np.sqrt((n + lambda_) * eigvals)) # Generate sigma points points = np.zeros((2 * n + 1, n), dtype=np.float64) points[0] = x for i in range(n): points[i + 1] = x + sqrt_P[:, i] points[n + i + 1] = x - sqrt_P[:, i] # Compute weights Wm = np.full(2 * n + 1, 1 / (2 * (n + lambda_)), dtype=np.float64) Wm[0] = lambda_ / (n + lambda_) Wc = Wm.copy() Wc[0] = lambda_ / (n + lambda_) + (1 - alpha**2 + beta) return SigmaPoints(points=points, Wm=Wm, Wc=Wc)
[docs] def sigma_points_julier( x: ArrayLike, P: ArrayLike, kappa: float = 0.0, ) -> SigmaPoints: """ Generate sigma points using Julier's original method. Parameters ---------- x : array_like State mean. P : array_like State covariance. kappa : float, optional Scaling parameter (default: 0, typical: 3-n). Returns ------- result : SigmaPoints Sigma points and weights. Examples -------- >>> import numpy as np >>> x = np.array([1.0, 2.0]) >>> P = np.eye(2) * 0.1 >>> sp = sigma_points_julier(x, P, kappa=1.0) >>> sp.points.shape # 2*n+1 = 5 points for n=2 (5, 2) >>> np.allclose(sp.Wm.sum(), 1.0) True Notes ----- Julier's method is a special case of Merwe's with alpha=1, beta=0. """ x = np.asarray(x, dtype=np.float64).flatten() P = np.asarray(P, dtype=np.float64) n = len(x) # Compute matrix square root try: sqrt_P = np.linalg.cholesky((n + kappa) * P) except np.linalg.LinAlgError: eigvals, eigvecs = np.linalg.eigh(P) eigvals = np.maximum(eigvals, 1e-10) sqrt_P = eigvecs @ np.diag(np.sqrt((n + kappa) * eigvals)) # Generate sigma points points = np.zeros((2 * n + 1, n), dtype=np.float64) points[0] = x for i in range(n): points[i + 1] = x + sqrt_P[:, i] points[n + i + 1] = x - sqrt_P[:, i] # Weights (same for mean and covariance in Julier's method) W = np.full(2 * n + 1, 1 / (2 * (n + kappa)), dtype=np.float64) W[0] = kappa / (n + kappa) return SigmaPoints(points=points, Wm=W, Wc=W)
[docs] def unscented_transform( sigmas: NDArray[np.floating], Wm: NDArray[np.floating], Wc: NDArray[np.floating], noise_cov: Optional[ArrayLike] = None, ) -> Tuple[NDArray[np.floating], NDArray[np.floating]]: """ Compute mean and covariance from transformed sigma points. Parameters ---------- sigmas : ndarray Transformed sigma points, shape (2n+1, m). Wm : ndarray Weights for mean. Wc : ndarray Weights for covariance. noise_cov : array_like, optional Additive noise covariance. Returns ------- mean : ndarray Weighted mean, shape (m,). cov : ndarray Weighted covariance, shape (m, m). Examples -------- >>> import numpy as np >>> # Generate sigma points >>> x = np.array([0.0, 0.0]) >>> P = np.eye(2) >>> sp = sigma_points_merwe(x, P) >>> # Pass through nonlinear function >>> def f(x): ... return np.array([x[0]**2, x[1]]) >>> sigmas_f = np.array([f(sp.points[i]) for i in range(len(sp.points))]) >>> mean, cov = unscented_transform(sigmas_f, sp.Wm, sp.Wc) >>> mean.shape (2,) """ # Weighted mean mean = np.sum(Wm[:, np.newaxis] * sigmas, axis=0) # Weighted covariance (vectorized: avoids loop over sigma points) residuals = sigmas - mean # Compute weighted outer products in one operation: (W * residuals)^T @ residuals weighted_residuals = np.sqrt(np.abs(Wc))[:, np.newaxis] * residuals # Handle negative weights (e.g., from Merwe scaling) by adjusting sign cov = weighted_residuals.T @ weighted_residuals # Correct for any negative weights (subtract their contribution twice to flip sign) neg_mask = Wc < 0 if np.any(neg_mask): neg_residuals = residuals[neg_mask] neg_weights = -Wc[neg_mask] for i, (w, r) in enumerate(zip(neg_weights, neg_residuals)): cov -= 2 * w * np.outer(r, r) if noise_cov is not None: cov += np.asarray(noise_cov, dtype=np.float64) # Ensure symmetry cov = (cov + cov.T) / 2 return mean, cov
[docs] def ukf_predict( x: ArrayLike, P: ArrayLike, f: Callable[[NDArray[Any]], NDArray[Any]], Q: ArrayLike, alpha: float = 1e-3, beta: float = 2.0, kappa: float = 0.0, ) -> KalmanPrediction: """ Unscented Kalman filter prediction 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(x) -> x_next. Q : array_like Process noise covariance, shape (n, n). alpha : float, optional UKF spread parameter. beta : float, optional UKF distribution parameter. kappa : float, optional UKF scaling parameter. Returns ------- result : KalmanPrediction Predicted state and covariance. Examples -------- >>> import numpy as np >>> def f(x): ... return np.array([x[0] + x[1], x[1]]) >>> x = np.array([0.0, 1.0]) >>> P = np.eye(2) * 0.1 >>> Q = np.eye(2) * 0.01 >>> pred = ukf_predict(x, P, f, Q) See Also -------- ukf_update : UKF measurement update. sigma_points_merwe : Sigma point generation. """ x = np.asarray(x, dtype=np.float64).flatten() P = np.asarray(P, dtype=np.float64) Q = np.asarray(Q, dtype=np.float64) # Generate sigma points sp = sigma_points_merwe(x, P, alpha, beta, kappa) # Transform sigma points through dynamics sigmas_f = np.array([f(s) for s in sp.points], dtype=np.float64) # Compute predicted mean and covariance x_pred, P_pred = unscented_transform(sigmas_f, sp.Wm, sp.Wc, Q) return KalmanPrediction(x=x_pred, P=P_pred)
[docs] def ukf_update( x: ArrayLike, P: ArrayLike, z: ArrayLike, h: Callable[[NDArray[Any]], NDArray[Any]], R: ArrayLike, alpha: float = 1e-3, beta: float = 2.0, kappa: float = 0.0, ) -> KalmanUpdate: """ Unscented Kalman filter update step. 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(x) -> z. R : array_like Measurement noise covariance, shape (m, m). alpha : float, optional UKF spread parameter. beta : float, optional UKF distribution parameter. kappa : float, optional UKF scaling parameter. Returns ------- result : KalmanUpdate Updated state, covariance, and innovation statistics. Examples -------- Update with a range-bearing measurement: >>> import numpy as np >>> # State: [x, y], range-bearing measurement >>> def h_rb(x): ... r = np.sqrt(x[0]**2 + x[1]**2) ... theta = np.arctan2(x[1], x[0]) ... return np.array([r, theta]) >>> x = np.array([100.0, 50.0]) >>> P = np.eye(2) * 10.0 >>> z = np.array([112.0, 0.46]) # measured range and bearing >>> R = np.diag([1.0, 0.01]) # measurement noise >>> upd = ukf_update(x, P, z, h_rb, R) >>> upd.x.shape (2,) See Also -------- ukf_predict : UKF prediction step. """ x = np.asarray(x, dtype=np.float64).flatten() P = np.asarray(P, dtype=np.float64) z = np.asarray(z, dtype=np.float64).flatten() R = np.asarray(R, dtype=np.float64) m = len(z) # Generate sigma points sp = sigma_points_merwe(x, P, alpha, beta, kappa) # Transform sigma points through measurement function sigmas_h = np.array([h(s) for s in sp.points], dtype=np.float64) # Predicted measurement mean and covariance z_pred, S = unscented_transform(sigmas_h, sp.Wm, sp.Wc, R) # Cross-covariance between state and measurement (vectorized) x_residuals = sp.points - x z_residuals = sigmas_h - z_pred # Weighted cross-covariance: sum of Wc[i] * outer(x_res[i], z_res[i]) Pxz = (sp.Wc[:, np.newaxis] * x_residuals).T @ z_residuals # Kalman gain K = np.linalg.solve(S.T, Pxz.T).T # Innovation y = z - z_pred # Updated state and covariance x_upd = x + K @ y P_upd = P - K @ S @ K.T # Ensure symmetry P_upd = (P_upd + P_upd.T) / 2 # Likelihood det_S = np.linalg.det(S) if det_S > 0: mahal_sq = y @ np.linalg.solve(S, y) likelihood = np.exp(-0.5 * mahal_sq) / np.sqrt((2 * np.pi) ** m * det_S) else: likelihood = 0.0 return KalmanUpdate( x=x_upd, P=P_upd, y=y, S=S, K=K, likelihood=likelihood, )
[docs] def ckf_spherical_cubature_points( n: int, ) -> tuple[NDArray[np.floating], NDArray[np.floating]]: """ Generate cubature points for Cubature Kalman Filter. Parameters ---------- n : int State dimension. Returns ------- points : ndarray Unit cubature points, shape (2n, n). weights : ndarray Cubature weights. Examples -------- >>> import numpy as np >>> points, weights = ckf_spherical_cubature_points(3) >>> points.shape # 2n = 6 points for n=3 (6, 3) >>> np.allclose(weights.sum(), 1.0) True Notes ----- The CKF uses a third-degree spherical-radial cubature rule with 2n points at ±sqrt(n) along each axis. See Also -------- pytcl.mathematical_functions.numerical_integration.cubature_points.second_order_cubature_points : A structurally different (n+2)-point rule (the scaled unscented transformation's minimal spherical simplex points), exact through degree 2 only -- fewer points than this 2n-point rule, but without this rule's degree-3 exactness. Not an alternate parameterization of this rule. """ points = np.zeros((2 * n, n), dtype=np.float64) sqrt_n = np.sqrt(n) for i in range(n): points[i, i] = sqrt_n points[n + i, i] = -sqrt_n weights = np.full(2 * n, 1 / (2 * n), dtype=np.float64) return points, weights
[docs] def ckf_predict( x: ArrayLike, P: ArrayLike, f: Callable[[NDArray[Any]], NDArray[Any]], Q: ArrayLike, *, points: Optional[ArrayLike] = None, weights: Optional[ArrayLike] = None, ) -> KalmanPrediction: """ Cubature Kalman filter prediction step. The CKF uses spherical-radial cubature for numerical integration. Its 2n-point rule avoids the negative center weight the standard UKF parameterization produces for n > 3, which is the basis of the CKF's high-dimensional robustness argument (Arasaratnam & Haykin, 2009) -- not a general accuracy guarantee over the UKF. Parameters ---------- x : array_like Current state estimate. P : array_like Current state covariance. f : callable Nonlinear state transition function. Q : array_like Process noise covariance. points : array_like, optional Custom unit cubature points, shape ``(m, n)``, in the N(0, I) convention (i.e. the points that would be used to integrate against a standard normal, before scaling by the Cholesky factor of ``P``). Must be provided together with ``weights``. If omitted, the default third-degree spherical-radial rule (``ckf_spherical_cubature_points``) is used. weights : array_like, optional Weights matching ``points``, shape ``(m,)``, summing to 1. May contain negative values (e.g. higher-degree rules). Returns ------- result : KalmanPrediction Predicted state and covariance. Examples -------- >>> import numpy as np >>> # Linear dynamics: x_k+1 = F @ x_k >>> def f(x): ... F = np.array([[1, 1], [0, 1]]) ... return F @ x >>> x = np.array([0.0, 1.0]) >>> P = np.eye(2) * 0.1 >>> Q = np.eye(2) * 0.01 >>> pred = ckf_predict(x, P, f, Q) >>> pred.x # Should be approximately [1, 1] array([1., 1.]) See Also -------- pytcl.mathematical_functions.numerical_integration.cubature_points : Higher-degree Gaussian cubature rules (e.g. ``fifth_order_cubature_points``) that can be passed via ``points``/``weights``. References ---------- - Arasaratnam, I. and Haykin, S., "Cubature Kalman Filters", IEEE Trans. Automatic Control, 2009. """ x = np.asarray(x, dtype=np.float64).flatten() P = np.asarray(P, dtype=np.float64) Q = np.asarray(Q, dtype=np.float64) n = len(x) # Matrix square root try: sqrt_P = np.linalg.cholesky(P) except np.linalg.LinAlgError: eigvals, eigvecs = np.linalg.eigh(P) eigvals = np.maximum(eigvals, 1e-10) sqrt_P = eigvecs @ np.diag(np.sqrt(eigvals)) # Generate cubature points if (points is None) != (weights is None): raise ValueError("points and weights must be provided together") if points is None: unit_pts, weights_arr = ckf_spherical_cubature_points(n) else: unit_pts = np.asarray(points, dtype=np.float64) weights_arr = np.asarray(weights, dtype=np.float64) if unit_pts.ndim != 2 or unit_pts.shape[1] != n: raise ValueError( f"points shape {unit_pts.shape} incompatible with state dimension {n}" ) if weights_arr.shape != (unit_pts.shape[0],): raise ValueError( f"weights shape {weights_arr.shape} does not match " f"{unit_pts.shape[0]} points" ) cubature_pts = x + (sqrt_P @ unit_pts.T).T # Transform through dynamics transformed = np.array([f(pt) for pt in cubature_pts], dtype=np.float64) # Predicted mean x_pred = np.sum(weights_arr[:, np.newaxis] * transformed, axis=0) # Predicted covariance (vectorized; sign-safe for negative weights) residuals = transformed - x_pred P_pred = residuals.T @ (weights_arr[:, np.newaxis] * residuals) + Q P_pred = (P_pred + P_pred.T) / 2 return KalmanPrediction(x=x_pred, P=P_pred)
[docs] def ckf_update( x: ArrayLike, P: ArrayLike, z: ArrayLike, h: Callable[[NDArray[Any]], NDArray[Any]], R: ArrayLike, *, points: Optional[ArrayLike] = None, weights: Optional[ArrayLike] = None, ) -> KalmanUpdate: """ Cubature Kalman filter update step. Parameters ---------- x : array_like Predicted state estimate. P : array_like Predicted state covariance. z : array_like Measurement. h : callable Nonlinear measurement function. R : array_like Measurement noise covariance. points : array_like, optional Custom unit cubature points, shape ``(m, n)``, in the N(0, I) convention (i.e. the points that would be used to integrate against a standard normal, before scaling by the Cholesky factor of ``P``). Must be provided together with ``weights``. If omitted, the default third-degree spherical-radial rule (``ckf_spherical_cubature_points``) is used. weights : array_like, optional Weights matching ``points``, shape ``(m,)``, summing to 1. May contain negative values (e.g. higher-degree rules). Returns ------- result : KalmanUpdate Updated state and covariance. Examples -------- >>> import numpy as np >>> # Position measurement h(x) = x[0] >>> def h(x): ... return np.array([x[0]]) >>> x = np.array([3.0, 1.0]) # predicted [position, velocity] >>> P = np.eye(2) * 0.5 >>> z = np.array([3.1]) # measurement >>> R = np.array([[0.1]]) >>> upd = ckf_update(x, P, z, h, R) >>> upd.x.shape (2,) See Also -------- pytcl.mathematical_functions.numerical_integration.cubature_points : Higher-degree Gaussian cubature rules (e.g. ``fifth_order_cubature_points``) that can be passed via ``points``/``weights``. """ x = np.asarray(x, dtype=np.float64).flatten() P = np.asarray(P, dtype=np.float64) z = np.asarray(z, dtype=np.float64).flatten() R = np.asarray(R, dtype=np.float64) n = len(x) m = len(z) # Matrix square root try: sqrt_P = np.linalg.cholesky(P) except np.linalg.LinAlgError: eigvals, eigvecs = np.linalg.eigh(P) eigvals = np.maximum(eigvals, 1e-10) sqrt_P = eigvecs @ np.diag(np.sqrt(eigvals)) # Generate cubature points if (points is None) != (weights is None): raise ValueError("points and weights must be provided together") if points is None: unit_pts, weights_arr = ckf_spherical_cubature_points(n) else: unit_pts = np.asarray(points, dtype=np.float64) weights_arr = np.asarray(weights, dtype=np.float64) if unit_pts.ndim != 2 or unit_pts.shape[1] != n: raise ValueError( f"points shape {unit_pts.shape} incompatible with state dimension {n}" ) if weights_arr.shape != (unit_pts.shape[0],): raise ValueError( f"weights shape {weights_arr.shape} does not match " f"{unit_pts.shape[0]} points" ) cubature_pts = x + (sqrt_P @ unit_pts.T).T # Transform through measurement function transformed = np.array([h(pt) for pt in cubature_pts], dtype=np.float64) # Predicted measurement z_pred = np.sum(weights_arr[:, np.newaxis] * transformed, axis=0) # Innovation covariance (vectorized; sign-safe for negative weights) z_residuals = transformed - z_pred S = z_residuals.T @ (weights_arr[:, np.newaxis] * z_residuals) + R # Cross-covariance (vectorized) x_residuals = cubature_pts - x Pxz = (weights_arr[:, np.newaxis] * x_residuals).T @ z_residuals # Kalman gain K = np.linalg.solve(S.T, Pxz.T).T # Innovation y = z - z_pred # Update x_upd = x + K @ y P_upd = P - K @ S @ K.T P_upd = (P_upd + P_upd.T) / 2 # Likelihood det_S = np.linalg.det(S) if det_S > 0: mahal_sq = y @ np.linalg.solve(S, y) likelihood = np.exp(-0.5 * mahal_sq) / np.sqrt((2 * np.pi) ** m * det_S) else: likelihood = 0.0 return KalmanUpdate( x=x_upd, P=P_upd, y=y, S=S, K=K, likelihood=likelihood, )
__all__ = [ "SigmaPoints", "sigma_points_merwe", "sigma_points_julier", "unscented_transform", "ukf_predict", "ukf_update", "ckf_spherical_cubature_points", "ckf_predict", "ckf_update", ]