Added UnscentedKalmanFilter

Author: FlorianPfaff

pyrecest/filters/unscented_kalman_filter.py

@@ -0,0 +1,122 @@
+from bayesian_filters.kalman import UnscentedKalmanFilter as BayesianFiltersUKF
+from bayesian_filters.kalman import MerweScaledSigmaPoints
+
+from .manifold_mixins import EuclideanFilterMixin
+from pyrecest.distributions import GaussianDistribution
+
+# pylint: disable=redefined-builtin,no-name-in-module,no-member
+from pyrecest.backend import atleast_1d
+import numbers
+from typing import Callable
+
+class UnscentedKalmanFilter(EuclideanFilterMixin):
+    def __init__(self, initial_state: GaussianDistribution | tuple,
+                 dt: numbers.Real = 1,
+                 fx: Callable = lambda x, dt: x,
+                 hx: Callable = lambda x: x,
+                 points=None,
+    ):
+        if isinstance(initial_state, GaussianDistribution):
+            dim_x = initial_state.dim
+        elif isinstance(initial_state, tuple) and len(initial_state) == 2:
+            dim_x = len(initial_state[0])
+        else:
+            raise ValueError(
+                "initial_state must be a GaussianDistribution or a tuple of (mean, covariance)"
+            )
+            
+        if points is None:
+            # Standard settings for Gaussian approximations
+            points = MerweScaledSigmaPoints(dim_x, alpha=0.001, beta=2, kappa=0)
+            
+        # Initialize filterpy UKF
+        # Note: We initialize dim_z to dim_x as a default, but this can be 
+        # overridden dynamically in update() by providing R and hx
+        self._filter_state = BayesianFiltersUKF(dim_x=dim_x, dim_z=dim_x, dt=dt, hx=hx, fx=fx, points=points)
+        
+        # Set initial state
+        if isinstance(initial_state, GaussianDistribution):
+            self._filter_state.x = initial_state.mu
+            self._filter_state.P = initial_state.C
+        else:
+            self._filter_state.x = initial_state[0]
+            self._filter_state.P = initial_state[1]
+
+        self._filter_state.x_prior = self._filter_state.x.copy()
+        self._filter_state.P_prior = self._filter_state.P.copy()
+
+    @property
+    def filter_state(
+        self,
+    ) -> (
+        GaussianDistribution | tuple
+    ):  # It can only return GaussianDistribution, this just serves to prevent mypy linter warnings
+        return GaussianDistribution(self._filter_state.x, self._filter_state.P)
+
+    @filter_state.setter
+    def filter_state(
+        self, new_state: GaussianDistribution | tuple
+    ):
+        """
+        Set the filter state.
+
+        :param new_state: Provide GaussianDistribution or mean and covariance as state.
+        """
+        if isinstance(new_state, GaussianDistribution):
+            self._filter_state.x = new_state.mu
+            self._filter_state.P = new_state.C
+        elif isinstance(new_state, tuple) and len(new_state) == 2:
+            self._filter_state.x = new_state[0]
+            self._filter_state.P = new_state[1]
+        else:
+            raise ValueError(
+                "new_state must be a GaussianDistribution or a tuple of (mean, covariance)"
+            )
+
+    def predict_nonlinear(self, fx, sys_noise_cov, dt=None, **fx_args):
+        """
+        :param fx: Function with signature fx(x, dt, **fx_args)
+        :param sys_noise_cov: Process noise matrix Q
+        """
+        # FIX: FilterPy UKF uses member variable Q, not an argument to predict()
+        self._filter_state.Q = sys_noise_cov
+        self._filter_state.predict(dt=dt, fx=fx, **fx_args)
+
+    def update_nonlinear(self, measurement, hx, cov_mat_meas, **hx_args):
+        # Update allows R to be passed as argument
+        self._filter_state.update(z=measurement, hx=hx, R=cov_mat_meas, **hx_args)
+
+    def predict_identity(self, sys_noise_cov, dt=None):
+        self._filter_state.Q = sys_noise_cov
+        def fx(x, dt):
+            return x
+
+        self._filter_state.predict(dt=dt, fx=fx)
+
+    def predict_linear(self, system_matrix, sys_noise_cov, sys_input=None, dt=None):
+        self._filter_state.Q = sys_noise_cov
+        
+        if sys_input is None:
+            # F * x
+            fx = lambda x, _dt: system_matrix @ x
+        else:
+            # F * x + B * u (Assuming B is Identity based on your previous code)
+            # ideally sys_input should already be B*u or you should pass B explicitly
+            fx = lambda x, _dt: system_matrix @ x + sys_input
+
+        self._filter_state.predict(dt=dt, fx=fx)
+
+    def update_identity(self, meas, meas_cov):
+        # h(x) = x
+        hx = lambda x: x
+        self._filter_state.update(z=atleast_1d(meas), R=meas_cov, hx=hx)
+
+    def update_linear(self, measurement, measurement_matrix, cov_mat_meas):
+        # h(x) = H * x
+        # Note: hx must return a 1D array for filterpy
+        hx = lambda x: measurement_matrix @ x
+        
+        self._filter_state.update(z=measurement, R=cov_mat_meas, hx=hx)
+
+    def get_estimate(self):
+        return GaussianDistribution(self._filter_state.x, self._filter_state.P)

pyrecest/tests/filters/test_unscented_kalman_filter.py

@@ -0,0 +1,41 @@
+import unittest
+from pyrecest.filters.unscented_kalman_filter import UnscentedKalmanFilter
+from pyrecest.distributions import GaussianDistribution
+import copy
+from pyrecest.backend import array, diag, eye, allclose
+import numpy.testing as npt
+
+class UnscentedKalmanFilterTest(unittest.TestCase):
+
+    def test_initialization(self):
+        filter_custom = UnscentedKalmanFilter(GaussianDistribution(array([1.0]), array([[10000.0]])))
+        npt.assert_allclose(filter_custom.get_point_estimate(), 1.0)
+    
+    def test_initialization_gauss(self):
+        filter_custom = UnscentedKalmanFilter(GaussianDistribution(array([4.0]), array([[10000.0]])))
+        npt.assert_allclose(filter_custom.get_point_estimate(), 4)
+
+    def test_update_linear_1d(self):
+        kf = UnscentedKalmanFilter(GaussianDistribution(array([0.0]), array([[1.0]])))
+        kf.update_identity(array([3.0]), array([[1.0]]))
+        npt.assert_allclose(kf.get_point_estimate(), 1.5)
+
+    def test_update_linear_2d(self):
+        filter_add = UnscentedKalmanFilter(GaussianDistribution(array([0.0, 1.0]), diag(array([1.0, 2.0]))))
+        filter_id = copy.deepcopy(filter_add)
+        gauss = GaussianDistribution(array([1.0, 0.0]), diag(array([2.0, 1.0])))
+        filter_add.update_linear(gauss.mu, eye(2), gauss.C)
+        filter_id.update_identity(gauss.mu, gauss.C)
+        self.assertTrue(allclose(filter_add.get_point_estimate(), filter_id.get_point_estimate()))
+        self.assertTrue(allclose(filter_add.filter_state.covariance(), filter_id.filter_state.covariance()))
+
+    def test_predict_linear_2d(self):
+        kf = UnscentedKalmanFilter(GaussianDistribution(array([0.0, 1.0]), diag(array(array([1.0, 2.0])))))
+        kf.predict_linear(diag(array([1.0, 2.0])), diag(array([2.0, 1.0])))
+        self.assertTrue(allclose(kf.get_point_estimate(), array([0.0, 2.0])))
+        self.assertTrue(allclose(kf.filter_state.covariance(), diag(array([3.0, 9.0]))))
+        kf.predict_linear(diag(array([1.0, 2.0])), diag(array([2.0, 1.0])), array([2.0, -2.0]))
+        self.assertTrue(allclose(kf.get_point_estimate(), array([2.0, 2.0])))
+
+if __name__ == "__main__":
+    unittest.main()