BUG: random: Fix long delays/hangs with zipf(a) when a near 1.

WarrenWeckesser · WarrenWeckesser · commit f5a5e04555ce · 2024-07-26T00:05:50.000-04:00
Closes gh-9829.
diff --git a/numpy/random/src/distributions/distributions.c b/numpy/random/src/distributions/distributions.c
@@ -998,7 +998,7 @@ int64_t random_geometric(bitgen_t *bitgen_state, double p) {
 }
 
 RAND_INT_TYPE random_zipf(bitgen_t *bitgen_state, double a) {
-  double am1, b;
+  double am1, b, Umin;
 
   if (a >= 1025) {
     /*
@@ -1011,10 +1011,21 @@ RAND_INT_TYPE random_zipf(bitgen_t *bitgen_state, double a) {
   }
   am1 = a - 1.0;
   b = pow(2.0, am1);
+  /*
+   * In the while loop, X is generated from the uniform distribution (Umin, 1].
+   * Values below Umin would result in X being rejected because it is too
+   * large, so there is no point in including them in the distribution of U.
+   */
+  Umin = pow(RAND_INT_MAX, -am1);
   while (1) {
-    double T, U, V, X;
+    double U01, T, U, V, X;
 
-    U = 1.0 - next_double(bitgen_state);
+    /*
+     * U is sampled from (Umin, 1]. Note that Umin might be 0, and we don't
+     * want U to be 0.
+     */
+    U01 = next_double(bitgen_state);
+    U = U01*Umin + (1 - U01);
     V = next_double(bitgen_state);
     X = floor(pow(U, -1.0 / am1));
     /*
diff --git a/numpy/random/tests/test_generator_mt19937_regressions.py b/numpy/random/tests/test_generator_mt19937_regressions.py
@@ -170,3 +170,14 @@ def test_zipf_large_parameter(self):
         n = 8
         sample = self.mt19937.zipf(10000, size=n)
         assert_array_equal(sample, np.ones(n, dtype=np.int64))
+
+    def test_zipf_a_near_1(self):
+        # Regression test for gh-9829: a call such as rng.zipf(1.0000000000001)
+        # would hang.
+        n = 100000
+        sample = self.mt19937.zipf(1.0000000000001, size=n)
+        # Not much of a test, but let's do something more than verify that
+        # it doesn't hang.  Certainly for a monotonically decreasing
+        # discrete distribution truncated to signed 64 bit integers, more
+        # than half should be less than 2**62.
+        assert np.count_nonzero(sample < 2**62) > n/2
