Stabilize discrete-state prior normalization

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__"]

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, :],
+    )