diff --git a/src/pyrecest/filters/discrete_state/__init__.py b/src/pyrecest/filters/discrete_state/__init__.py index edd913205b..40d0c47143 100644 --- a/src/pyrecest/filters/discrete_state/__init__.py +++ b/src/pyrecest/filters/discrete_state/__init__.py @@ -283,10 +283,19 @@ def _validated_probability_vector( mask = _module_globals["_coerce_valid_state_mask"](valid_state_mask, n_entries) if mask is not None: values[~mask] = 0.0 - total = float(values.sum()) - if total <= 0.0: + + # Summing finite, non-negative weights directly can overflow even though the + # normalized probability vector is perfectly well-defined. Scale by the + # largest entry first so both very large and subnormal priors preserve their + # relative mass without overflowing or flushing the total to zero. + value_scale = float(np.max(values)) + if value_scale <= 0.0: + raise ValueError(f"{name} must contain positive probability mass") + scaled_values = values / value_scale + scaled_total = float(scaled_values.sum()) + if not np.isfinite(scaled_total) or scaled_total <= 0.0: raise ValueError(f"{name} must contain positive probability mass") - return values / total + return scaled_values / scaled_total @wraps(_original_sticky_mode_transition_matrix) @@ -394,6 +403,11 @@ def sparse_gaussian_transition_matrix( _module_globals["mode_transition_matrix"] = sticky_mode_transition_matrix _module_globals["sparse_gaussian_transition_matrix"] = sparse_gaussian_transition_matrix _module_globals["_normalize_probability_vector"] = _validated_probability_vector +# ``runpy.run_path`` returns a copy of the executed module globals. Functions +# retain the original execution dictionary, so patch that dictionary as well. +_original_discrete_forward_backward.__globals__["_normalize_probability_vector"] = ( + _validated_probability_vector +) for name in _module_globals["__all__"]: globals()[name] = _module_globals[name] -__all__ = _module_globals["__all__"] +__all__ = _module_globals["__all__"] \ No newline at end of file diff --git a/tests/filters/test_discrete_state_probability_normalization_overflow.py b/tests/filters/test_discrete_state_probability_normalization_overflow.py new file mode 100644 index 0000000000..7c37ba51e1 --- /dev/null +++ b/tests/filters/test_discrete_state_probability_normalization_overflow.py @@ -0,0 +1,43 @@ +import numpy as np + +from pyrecest.filters.discrete_state import ( + discrete_forward_backward, + imm_forward_backward, +) + + +def test_forward_backward_normalizes_finite_priors_without_overflow(): + huge = np.finfo(float).max + initial = np.array([huge, huge / 2.0]) + + result = discrete_forward_backward( + np.zeros((1, 2)), + np.eye(2), + initial_probabilities=initial, + ) + + np.testing.assert_allclose( + result.filtered_probabilities[0], + np.array([2.0 / 3.0, 1.0 / 3.0]), + ) + + +def test_imm_normalizes_finite_state_and_mode_priors_without_overflow(): + huge = np.finfo(float).max + state_prior = np.array([huge, huge / 2.0]) + mode_prior = np.array([huge / 2.0, huge]) + + result = imm_forward_backward( + np.zeros((1, 2)), + [np.eye(2), np.eye(2)], + np.eye(2), + initial_state_probabilities=state_prior, + initial_mode_probabilities=mode_prior, + ) + + expected_state = np.array([2.0 / 3.0, 1.0 / 3.0]) + expected_mode = np.array([1.0 / 3.0, 2.0 / 3.0]) + np.testing.assert_allclose( + result.filtered_joint_probabilities[0], + expected_mode[:, None] * expected_state[None, :], + )