Merge pull request #27046 from WarrenWeckesser/zipf-large-param

charris · web-flow · commit e214a1271fd9 · 2024-07-25T20:36:21.000-06:00
BUG: random: prevent zipf from hanging when parameter is large.
diff --git a/numpy/random/src/distributions/distributions.c b/numpy/random/src/distributions/distributions.c
@@ -1000,6 +1000,15 @@ int64_t random_geometric(bitgen_t *bitgen_state, double p) {
 RAND_INT_TYPE random_zipf(bitgen_t *bitgen_state, double a) {
   double am1, b;
 
+  if (a >= 1025) {
+    /*
+     * If a exceeds 1025, the calculation of b will overflow and the loop
+     * will not terminate.  It is safe to simply return 1 here, because the
+     * probability of generating a value greater than 1 in this case is
+     * less than 3e-309.
+     */
+    return (RAND_INT_TYPE) 1;
+  }
   am1 = a - 1.0;
   b = pow(2.0, am1);
   while (1) {
diff --git a/numpy/random/tests/test_generator_mt19937_regressions.py b/numpy/random/tests/test_generator_mt19937_regressions.py
@@ -163,3 +163,10 @@ def test_geometric_tiny_prob(self):
         # is 0.9999999999907766, so we expect the result to be all 2**63-1.
         assert_array_equal(self.mt19937.geometric(p=1e-30, size=3),
                            np.iinfo(np.int64).max)
+
+    def test_zipf_large_parameter(self):
+        # Regression test for part of gh-9829: a call such as rng.zipf(10000)
+        # would hang.
+        n = 8
+        sample = self.mt19937.zipf(10000, size=n)
+        assert_array_equal(sample, np.ones(n, dtype=np.int64))
