Reject non-finite UKF vectors and keep failed predictions atomic

src/pyrecest/filters/_ukf.py

@@ -7,12 +7,14 @@
 from copy import deepcopy
 
 # pylint: disable=no-name-in-module,no-member
+from pyrecest.backend import all as backend_all
 from pyrecest.backend import (
     asarray,
     einsum,
     expand_dims,
     eye,
     float64,
+    isfinite,
     linalg,
     reshape,
     stack,
@@ -36,14 +38,16 @@
 
 
 def _as_vector(value, description):
-    """Return scalar or vector input without silently flattening matrices."""
+    """Return a finite scalar or vector without silently flattening matrices."""
     value = asarray(value, dtype=float64)
     if len(value.shape) == 0:
-        return reshape(value, (1,))
-    if len(value.shape) != 1:
+        value = reshape(value, (1,))
+    elif len(value.shape) != 1:
         raise ValueError(
             f"{description} must be scalar or one-dimensional; got shape {value.shape}"
         )
+    if not bool(backend_all(isfinite(value))):
+        raise ValueError(f"{description} must contain only finite values")
     return value
 
 
@@ -137,13 +141,16 @@ def predict(self, fx=None, dt=None, **fx_args):
         P_pred = P_pred + process_noise_covariance
         P_pred = 0.5 * (P_pred + transpose(P_pred))
 
-        self.x = x_pred
-        self.P = P_pred
         # ``sigmas_f`` represents only the deterministic transition.  Process
         # noise is added analytically above, so regenerate the cached sigma
         # points from the complete predicted covariance before the measurement
         # update; otherwise additive process noise is ignored in Pz and Pxz.
-        self._sigmas_f = points.sigma_points(self.x, self.P)
+        # Build the cache before committing x/P so failed covariance validation
+        # cannot leave the filter in a partially-updated state.
+        predicted_sigmas = points.sigma_points(x_pred, P_pred)
+        self.x = x_pred
+        self.P = P_pred
+        self._sigmas_f = predicted_sigmas
 
     def _innovation_matrices(  # pylint: disable=too-many-positional-arguments
         self, sigmas_f, sigmas_h, z_pred, R, Wc

tests/filters/test_ukf_nonfinite_atomicity.py

@@ -0,0 +1,89 @@
+import unittest
+
+import numpy as np
+import numpy.testing as npt
+
+from pyrecest.backend import array, zeros
+from pyrecest.filters._ukf import _UKFModel, UnscentedKalmanFilter
+from pyrecest.sampling.sigma_points import MerweScaledSigmaPoints
+
+
+class UKFNonfiniteAtomicityTest(unittest.TestCase):
+    @staticmethod
+    def _make_filter(*, fx=None, hx=None):
+        points = MerweScaledSigmaPoints(n=1, alpha=1.0, beta=2.0, kappa=0.0)
+        model = _UKFModel(
+            dim_x=1,
+            dim_z=1,
+            dt=1.0,
+            hx=(lambda x: x) if hx is None else hx,
+            fx=(lambda x, _dt: x) if fx is None else fx,
+            points=points,
+        )
+        ukf = UnscentedKalmanFilter(model)
+        ukf.x = array([0.25])
+        ukf.P = array([[0.5]])
+        ukf.Q = zeros((1, 1))
+        ukf.R = array([[0.1]])
+        return ukf
+
+    @staticmethod
+    def _assert_state_unchanged(ukf, x_before, p_before):
+        npt.assert_allclose(np.asarray(ukf.x), x_before)
+        npt.assert_allclose(np.asarray(ukf.P), p_before)
+
+    def test_predict_rejects_nonfinite_transition_without_state_change(self):
+        ukf = self._make_filter(
+            fx=lambda _x, _dt: array([float("nan")]),
+        )
+        x_before = np.asarray(ukf.x).copy()
+        p_before = np.asarray(ukf.P).copy()
+
+        with self.assertRaisesRegex(
+            ValueError,
+            "transition function output must contain only finite values",
+        ):
+            ukf.predict()
+
+        self._assert_state_unchanged(ukf, x_before, p_before)
+
+    def test_update_rejects_nonfinite_measurement_without_state_change(self):
+        ukf = self._make_filter()
+        x_before = np.asarray(ukf.x).copy()
+        p_before = np.asarray(ukf.P).copy()
+
+        with self.assertRaisesRegex(
+            ValueError,
+            "measurement z must contain only finite values",
+        ):
+            ukf.update(array([float("nan")]))
+
+        self._assert_state_unchanged(ukf, x_before, p_before)
+
+    def test_update_rejects_nonfinite_model_output_without_state_change(self):
+        ukf = self._make_filter(hx=lambda _x: array([float("inf")]))
+        x_before = np.asarray(ukf.x).copy()
+        p_before = np.asarray(ukf.P).copy()
+
+        with self.assertRaisesRegex(
+            ValueError,
+            "measurement function output must contain only finite values",
+        ):
+            ukf.update(array([0.0]))
+
+        self._assert_state_unchanged(ukf, x_before, p_before)
+
+    def test_failed_predicted_covariance_validation_is_atomic(self):
+        ukf = self._make_filter()
+        ukf.Q = array([[float("nan")]])
+        x_before = np.asarray(ukf.x).copy()
+        p_before = np.asarray(ukf.P).copy()
+
+        with self.assertRaisesRegex(ValueError, "P must contain only finite values"):
+            ukf.predict()
+
+        self._assert_state_unchanged(ukf, x_before, p_before)
+
+
+if __name__ == "__main__":
+    unittest.main()