diff --git a/src/pyrecest/filters/mem_rbpf_tracker.py b/src/pyrecest/filters/mem_rbpf_tracker.py index e35db1ec3..2ea64d98a 100644 --- a/src/pyrecest/filters/mem_rbpf_tracker.py +++ b/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: diff --git a/tests/filters/test_mem_rbpf_prediction_atomicity.py b/tests/filters/test_mem_rbpf_prediction_atomicity.py new file mode 100644 index 000000000..c794d560b --- /dev/null +++ b/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)