Temporary rebase staging for #4295

src/pyrecest/distributions/circle/piecewise_constant_distribution.py

@@ -4,6 +4,7 @@
 import numpy as np
 import pyrecest.backend
 from pyrecest.backend import (
+    amax,
     arange,
     array,
     exp,
@@ -144,11 +145,12 @@ def __init__(self, w):
         if any(bool(weight < 0.0) for weight in w):
             raise ValueError("Weights must be nonnegative")
 
-        mean_weight = mean(w)
-        if not bool(mean_weight > 0.0):
+        max_weight = amax(w)
+        if not bool(max_weight > 0.0):
             raise ValueError("Weights must have positive total mass")
 
-        self.w = w / (mean_weight * 2.0 * pi)
+        scaled_weights = w / max_weight
+        self.w = scaled_weights / (mean(scaled_weights) * 2.0 * pi)
 
     def pdf(self, xs):
         """Evaluate the pdf at each point in xs.

tests/distributions/test_piecewise_constant_weight_overflow.py

@@ -0,0 +1,21 @@
+import numpy as np
+import numpy.testing as npt
+
+from pyrecest.backend import array, to_numpy
+from pyrecest.distributions import PiecewiseConstantDistribution
+
+
+def test_max_finite_weights_normalize_without_overflow():
+    probe = np.asarray(to_numpy(array([0.0], dtype=float)))
+    maximum = np.finfo(probe.dtype).max
+
+    with np.errstate(over="raise", invalid="raise", divide="raise"):
+        distribution = PiecewiseConstantDistribution(
+            array([maximum, maximum / 2.0], dtype=float)
+        )
+
+    actual = np.asarray(to_numpy(distribution.w), dtype=float)
+    expected = np.array([2.0, 1.0]) / (3.0 * np.pi)
+
+    npt.assert_allclose(actual, expected, rtol=1e-6, atol=0.0)
+    npt.assert_allclose(np.mean(actual) * 2.0 * np.pi, 1.0, rtol=1e-6)