Reject undefined axial Kalman mean updates

src/pyrecest/filters/axial_kalman_filter.py

@@ -8,6 +8,9 @@
 from .abstract_axial_filter import AbstractAxialFilter
 
 
+_MIN_NORMALIZABLE_MEAN_NORM = 1e-12
+
+
 def _is_complex_array(value):
     """Return whether a NumPy/JAX array or PyTorch tensor has complex dtype."""
     dtype = getattr(value, "dtype", None)
@@ -128,7 +131,16 @@ def update_identity(self, gauss_v, z):
         mu_new = self._filter_state.mu + K @ (z - self._filter_state.mu)
         C_new = (eye(d) - K) @ self._filter_state.C
 
-        mu_new = mu_new / linalg.norm(mu_new)  # enforce unit vector
+        mu_new_norm = linalg.norm(mu_new)
+        if not bool(isfinite(mu_new_norm)):
+            raise ValueError(
+                "Axial Kalman update produced a non-finite posterior mean."
+            )
+        if not bool(mu_new_norm > _MIN_NORMALIZABLE_MEAN_NORM):
+            raise ValueError(
+                "Axial Kalman update produced an undefined zero-length posterior mean."
+            )
+        mu_new = mu_new / mu_new_norm  # enforce unit vector
         self._filter_state = GaussianDistribution(mu_new, C_new, check_validity=False)
 
     def get_point_estimate(self):

tests/filters/test_axial_kalman_filter_zero_mean.py

@@ -0,0 +1,43 @@
+import unittest
+
+import numpy.testing as npt
+
+import pyrecest.backend
+from pyrecest.backend import array
+from pyrecest.distributions import GaussianDistribution
+from pyrecest.filters.axial_kalman_filter import AxialKalmanFilter
+
+
+class TestAxialKalmanFilterZeroMean(unittest.TestCase):
+    @unittest.skipIf(
+        pyrecest.backend.__backend_name__ == "pytorch",
+        reason="Not supported on this backend",  # pylint: disable=no-member
+    )
+    def test_update_rejects_zero_length_posterior_mean_atomically(self):
+        inv_sqrt_two = 2.0**-0.5
+        state_cov = array([[5.0, 2.0], [2.0, 1.0]])
+        noise_cov = array(
+            [
+                [2.0 - inv_sqrt_two, 1.0],
+                [1.0, (1.0 + inv_sqrt_two) / 2.0],
+            ]
+        )
+
+        axial_filter = AxialKalmanFilter()
+        axial_filter.filter_state = GaussianDistribution(
+            array([1.0, 0.0]), state_cov
+        )
+        prior_mu = axial_filter.filter_state.mu.copy()
+        prior_cov = axial_filter.filter_state.C.copy()
+        noise = GaussianDistribution(array([1.0, 0.0]), noise_cov)
+        measurement = array([inv_sqrt_two, inv_sqrt_two])
+
+        with self.assertRaisesRegex(ValueError, "zero-length posterior mean"):
+            axial_filter.update_identity(noise, measurement)
+
+        npt.assert_array_equal(axial_filter.filter_state.mu, prior_mu)
+        npt.assert_array_equal(axial_filter.filter_state.C, prior_cov)
+
+
+if __name__ == "__main__":
+    unittest.main()