Validate UKF process-noise covariance shapes

src/pyrecest/filters/_ukf.py

@@ -92,6 +92,20 @@ def predict(self, fx=None, dt=None, **fx_args):
         if dt is None:
             dt = self._model.dt
 
+        dim_x = self._model.dim_x
+        process_noise_covariance = asarray(self.Q, dtype=float64)
+        if len(process_noise_covariance.shape) == 0 or (
+            len(process_noise_covariance.shape) == 1
+            and process_noise_covariance.shape[0] == 1
+            and dim_x == 1
+        ):
+            process_noise_covariance = reshape(process_noise_covariance, (1, 1))
+        if process_noise_covariance.shape != (dim_x, dim_x):
+            raise ValueError(
+                "process noise covariance Q has shape "
+                f"{process_noise_covariance.shape}, expected {(dim_x, dim_x)}"
+            )
+
         points = self._model.points
         sigmas = points.sigma_points(self.x, self.P)
         n_sigmas = sigmas.shape[0]
@@ -104,10 +118,10 @@ def predict(self, fx=None, dt=None, **fx_args):
             for i in range(n_sigmas)
         ]
         for sigma_f in sigma_rows:
-            if sigma_f.shape[0] != self._model.dim_x:
+            if sigma_f.shape[0] != dim_x:
                 raise ValueError(
                     "transition function must return vectors with state dimension "
-                    f"{self._model.dim_x}; got {sigma_f.shape[0]}"
+                    f"{dim_x}; got {sigma_f.shape[0]}"
                 )
         sigmas_f = stack(sigma_rows)
 
@@ -116,11 +130,11 @@ def predict(self, fx=None, dt=None, **fx_args):
 
         x_pred = einsum("i,ij->j", Wm, sigmas_f)
 
-        P_pred = zeros((self._model.dim_x, self._model.dim_x))
+        P_pred = zeros((dim_x, dim_x))
         for i in range(n_sigmas):
             d = expand_dims(sigmas_f[i] - x_pred, -1)
             P_pred = P_pred + Wc[i] * (d @ transpose(d))
-        P_pred = P_pred + asarray(self.Q, dtype=float64)
+        P_pred = P_pred + process_noise_covariance
         P_pred = 0.5 * (P_pred + transpose(P_pred))
 
         self.x = x_pred

tests/filters/test_unscented_kalman_filter_process_noise_shape.py

@@ -0,0 +1,44 @@
+"""Regression tests for UKF process-noise covariance shape handling."""
+
+import unittest
+
+import numpy.testing as npt
+
+# pylint: disable=no-name-in-module,no-member
+import pyrecest.backend
+from pyrecest.backend import array
+from pyrecest.distributions import GaussianDistribution
+from pyrecest.filters.unscented_kalman_filter import UnscentedKalmanFilter
+
+
+@unittest.skipIf(
+    pyrecest.backend.__backend_name__ in ("pytorch", "jax"),
+    reason="UnscentedKalmanFilter is not supported on this backend",
+)
+class UnscentedKalmanFilterProcessNoiseShapeTest(unittest.TestCase):
+    def test_predict_rejects_vector_process_noise_without_mutating_state(self):
+        initial_mean = array([0.5, -0.25])
+        initial_covariance = array([[1.2, 0.3], [0.3, 0.8]])
+        ukf = UnscentedKalmanFilter(
+            GaussianDistribution(initial_mean, initial_covariance)
+        )
+
+        with self.assertRaisesRegex(ValueError, "process noise covariance Q"):
+            ukf.predict_identity(array([0.4, 0.2]))
+
+        npt.assert_allclose(ukf.get_point_estimate(), initial_mean)
+        npt.assert_allclose(ukf.filter_state.covariance(), initial_covariance)
+
+    def test_predict_accepts_length_one_process_noise_for_one_dimensional_state(self):
+        ukf = UnscentedKalmanFilter(
+            GaussianDistribution(array([0.5]), array([[1.25]]))
+        )
+
+        ukf.predict_identity(array([0.5]))
+
+        npt.assert_allclose(ukf.get_point_estimate(), array([0.5]))
+        npt.assert_allclose(ukf.filter_state.covariance(), array([[1.75]]))
+
+
+if __name__ == "__main__":
+    unittest.main()