diff --git a/src/pyrecest/filters/discrete_state/__init__.py b/src/pyrecest/filters/discrete_state/__init__.py index b08fd75b9d..1a36ca1763 100644 --- a/src/pyrecest/filters/discrete_state/__init__.py +++ b/src/pyrecest/filters/discrete_state/__init__.py @@ -114,10 +114,11 @@ 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: + scale = float(values.max()) + if scale <= 0.0: raise ValueError(f"{name} must contain positive probability mass") - return values / total + values /= scale + return values / float(values.sum()) def sparse_gaussian_transition_matrix( diff --git a/tests/filters/test_discrete_state_extreme_probabilities.py b/tests/filters/test_discrete_state_extreme_probabilities.py new file mode 100644 index 0000000000..9628c619bd --- /dev/null +++ b/tests/filters/test_discrete_state_extreme_probabilities.py @@ -0,0 +1,39 @@ +import numpy as np + +from pyrecest.filters.discrete_state import ( + discrete_forward_backward, + imm_forward_backward, +) + + +def test_forward_backward_normalizes_extreme_finite_initial_probabilities(): + max_finite = np.finfo(float).max + + result = discrete_forward_backward( + np.zeros((1, 2)), + np.eye(2), + initial_probabilities=np.array([max_finite, max_finite / 2.0]), + ) + + expected = np.array([2.0 / 3.0, 1.0 / 3.0]) + np.testing.assert_allclose(result.filtered_probabilities[0], expected) + np.testing.assert_allclose(result.smoothed_probabilities[0], expected) + assert np.isfinite(result.log_marginal_likelihood) + + +def test_imm_normalizes_extreme_finite_state_and_mode_probabilities(): + max_finite = np.finfo(float).max + + result = imm_forward_backward( + np.zeros((1, 2)), + [np.eye(2), np.eye(2)], + np.eye(2), + initial_state_probabilities=np.array([max_finite, max_finite / 2.0]), + initial_mode_probabilities=np.array([max_finite / 2.0, max_finite]), + ) + + 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_state_probabilities[0], expected_state) + np.testing.assert_allclose(result.filtered_mode_probabilities[0], expected_mode) + np.testing.assert_allclose(result.filtered_joint_probabilities.sum(), 1.0)