Make MEM-RBPF prediction failures atomic (#5446)

* Make MEM-RBPF prediction updates atomic

* Add MEM-RBPF prediction atomicity regressions

* Format MEM-RBPF atomicity regressions

Author: FlorianPfaff

src/pyrecest/filters/mem_rbpf_tracker.py

@@ -325,19 +325,17 @@ def predict_linear(
                 raise ValueError(
                     f"inputs must have shape {expected_shape}, got {inputs.shape}"
                 )
+
+        next_system_matrix = self.system_matrix
         if system_matrix is not None:
-            self.system_matrix = array(system_matrix)
-            self._validate_system_matrix(self.system_matrix)
+            next_system_matrix = array(system_matrix)
+            self._validate_system_matrix(next_system_matrix)
+
+        next_sys_noise = self.sys_noise
         if sys_noise is not None:
-            self.sys_noise = self._as_covariance(
+            next_sys_noise = self._as_covariance(
                 sys_noise, self.state_dim, "sys_noise", require_pd=False
             )
-        self.kinematic_state = self.system_matrix @ self.kinematic_state
-        if inputs is not None:
-            self.kinematic_state = self.kinematic_state + inputs
-        self.covariance = self._symmetrize(
-            self.system_matrix @ self.covariance @ self.system_matrix.T + self.sys_noise
-        )
 
         axis_matrix = eye(2)
         if shape_system_matrix is not None:
@@ -348,26 +346,48 @@ def predict_linear(
                 axis_matrix = shape_system_matrix
             else:
                 raise ValueError("shape_system_matrix must be 3x3 or 2x2")
+
+        next_orientation_process_variance = self.orientation_process_variance
+        next_axis_sys_noise = self.axis_sys_noise
         if shape_sys_noise is not None:
             shape_sys_noise = array(shape_sys_noise)
             if shape_sys_noise.shape == (3, 3):
                 shape_sys_noise = self._as_covariance(
                     shape_sys_noise, 3, "shape_sys_noise", require_pd=False
                 )
-                self.orientation_process_variance = float(shape_sys_noise[0, 0])
-                self.axis_sys_noise = shape_sys_noise[1:, 1:]
+                next_orientation_process_variance = float(shape_sys_noise[0, 0])
+                next_axis_sys_noise = shape_sys_noise[1:, 1:]
             elif shape_sys_noise.shape == (2, 2):
-                self.axis_sys_noise = self._as_covariance(
+                next_axis_sys_noise = self._as_covariance(
                     shape_sys_noise, 2, "shape_sys_noise", require_pd=False
                 )
             else:
                 raise ValueError("shape_sys_noise must be 3x3 or 2x2")
-        self.axis = self.axis @ axis_matrix.T
-        self.axis_covariances = self._symmetrize_stack(
+
+        next_kinematic_state = next_system_matrix @ self.kinematic_state
+        if inputs is not None:
+            next_kinematic_state = next_kinematic_state + inputs
+        next_covariance = self._symmetrize(
+            next_system_matrix @ self.covariance @ next_system_matrix.T
+            + next_sys_noise
+        )
+        next_axis = self.axis @ axis_matrix.T
+        next_axis_covariances = self._symmetrize_stack(
             axis_matrix @ self.axis_covariances @ axis_matrix.T
-            + self.axis_sys_noise.reshape((1, 2, 2))
+            + next_axis_sys_noise.reshape((1, 2, 2))
         )
-        self._apply_axis_floor()
+        if self.axis_floor is not None:
+            next_axis = maximum(next_axis, float(self.axis_floor))
+
+        self.system_matrix = next_system_matrix
+        self.sys_noise = next_sys_noise
+        self.orientation_process_variance = next_orientation_process_variance
+        self.axis_sys_noise = next_axis_sys_noise
+        self.kinematic_state = next_kinematic_state
+        self.covariance = next_covariance
+        self.axis = next_axis
+        self.axis_covariances = next_axis_covariances
+
         if self.log_prior_estimates:
             self.store_prior_estimates()
         if self.log_prior_extents:

tests/filters/test_mem_rbpf_prediction_atomicity.py

@@ -0,0 +1,89 @@
+import numpy as np
+import numpy.testing as npt
+import pytest
+from pyrecest import backend
+from pyrecest.backend import array, diag, eye
+from pyrecest.filters.mem_rbpf_tracker import MEMRBPFTracker
+
+
+pytestmark = pytest.mark.skipif(
+    backend.__backend_name__ == "jax",
+    reason="MEMRBPFTracker is unsupported on JAX.",
+)
+
+
+def _make_tracker():
+    return MEMRBPFTracker(
+        kinematic_state=array([0.0, 0.0, 1.0, -0.5]),
+        covariance=eye(4),
+        shape_state=array([0.2, 2.0, 1.0]),
+        shape_covariance=diag(array([0.05, 0.1, 0.1])),
+        meas_noise_cov=0.05 * eye(2),
+        sys_noise=0.01 * eye(4),
+        shape_sys_noise=diag(array([0.01, 0.01, 0.01])),
+        n_particles=8,
+        resampling_threshold=0,
+        rng=7,
+    )
+
+
+def _to_numpy_copy(value):
+    return np.asarray(backend.to_numpy(value)).copy()
+
+
+def _snapshot(tracker):
+    return {
+        "kinematic_state": _to_numpy_copy(tracker.kinematic_state),
+        "covariance": _to_numpy_copy(tracker.covariance),
+        "system_matrix": _to_numpy_copy(tracker.system_matrix),
+        "sys_noise": _to_numpy_copy(tracker.sys_noise),
+        "axis": _to_numpy_copy(tracker.axis),
+        "axis_covariances": _to_numpy_copy(tracker.axis_covariances),
+        "axis_sys_noise": _to_numpy_copy(tracker.axis_sys_noise),
+        "orientation_process_variance": tracker.orientation_process_variance,
+    }
+
+
+def _assert_snapshot_equal(tracker, snapshot):
+    npt.assert_array_equal(
+        _to_numpy_copy(tracker.kinematic_state), snapshot["kinematic_state"]
+    )
+    npt.assert_array_equal(
+        _to_numpy_copy(tracker.covariance), snapshot["covariance"]
+    )
+    npt.assert_array_equal(
+        _to_numpy_copy(tracker.system_matrix), snapshot["system_matrix"]
+    )
+    npt.assert_array_equal(_to_numpy_copy(tracker.sys_noise), snapshot["sys_noise"])
+    npt.assert_array_equal(_to_numpy_copy(tracker.axis), snapshot["axis"])
+    npt.assert_array_equal(
+        _to_numpy_copy(tracker.axis_covariances), snapshot["axis_covariances"]
+    )
+    npt.assert_array_equal(
+        _to_numpy_copy(tracker.axis_sys_noise), snapshot["axis_sys_noise"]
+    )
+    assert (
+        tracker.orientation_process_variance
+        == snapshot["orientation_process_variance"]
+    )
+
+
+@pytest.mark.parametrize(
+    ("override", "message"),
+    [
+        ({"system_matrix": eye(3)}, "system_matrix"),
+        ({"shape_system_matrix": eye(4)}, "shape_system_matrix"),
+        (
+            {"shape_sys_noise": array([[1.0, 2.0], [2.0, 1.0]])},
+            "shape_sys_noise",
+        ),
+    ],
+)
+def test_predict_validation_failures_leave_tracker_unchanged(override, message):
+    tracker = _make_tracker()
+    snapshot = _snapshot(tracker)
+
+    with pytest.raises(ValueError, match=message):
+        tracker.predict_linear(**override)
+
+    _assert_snapshot_equal(tracker, snapshot)