BUG/ENH: Fix zipf changes missed in NumPy

bashtage · bashtage · commit 09bdc6eeb6ee · 2019-04-12T14:18:39.000+01:00
Fix zipf changes missed in NumPy
Enable 0 as valid input for hypergeometric
diff --git a/numpy/random/generator.pyx b/numpy/random/generator.pyx
@@ -3642,7 +3642,7 @@ cdef class RandomGenerator:
         x.shape = tuple(final_shape)
         return x
 
-    def multinomial(self, np.npy_intp n, object pvals, size=None):
+    def multinomial(self, object n, object pvals, size=None):
         """
         multinomial(n, pvals, size=None)
 
@@ -3741,11 +3741,12 @@ cdef class RandomGenerator:
         on = <np.ndarray>np.PyArray_FROM_OTF(n, np.NPY_INT64, np.NPY_ALIGNED)
         parr = <np.ndarray>np.PyArray_FROM_OTF(pvals, np.NPY_DOUBLE, np.NPY_ALIGNED)
         pix = <double*>np.PyArray_DATA(parr)
-
+        check_array_constraint(parr, 'pvals', CONS_BOUNDED_0_1)
         if kahan_sum(pix, d-1) > (1.0 + 1e-12):
             raise ValueError("sum(pvals[:-1]) > 1.0")
 
         if np.PyArray_NDIM(on) != 0: # vector
+            check_array_constraint(on, 'n', CONS_NON_NEGATIVE)
             if size is None:
                 it = np.PyArray_MultiIterNew1(on)
             else:
@@ -3779,6 +3780,7 @@ cdef class RandomGenerator:
         mnix = <int64_t*>np.PyArray_DATA(mnarr)
         sz = np.PyArray_SIZE(mnarr)
         ni = n
+        check_constraint(ni, 'n', CONS_NON_NEGATIVE)
         offset = 0
         with self.lock, nogil:
             for i in range(sz // d):
diff --git a/numpy/random/mtrand.pyx b/numpy/random/mtrand.pyx
@@ -3799,7 +3799,7 @@ cdef class RandomState:
         d = len(pvals)
         parr = <np.ndarray>np.PyArray_FROM_OTF(pvals, np.NPY_DOUBLE, np.NPY_ALIGNED)
         pix = <double*>np.PyArray_DATA(parr)
-
+        check_array_constraint(parr, 'pvals', CONS_BOUNDED_0_1)
         if kahan_sum(pix, d-1) > (1.0 + 1e-12):
             raise ValueError("sum(pvals[:-1]) > 1.0")
 
@@ -3816,6 +3816,7 @@ cdef class RandomState:
         mnix = <int64_t*>np.PyArray_DATA(mnarr)
         sz = np.PyArray_SIZE(mnarr)
         ni = n
+        check_constraint(ni, 'n', CONS_NON_NEGATIVE)
         offset = 0
         with self.lock, nogil:
             for i in range(sz // d):
diff --git a/numpy/random/src/distributions/distributions.c b/numpy/random/src/distributions/distributions.c
@@ -1048,25 +1048,31 @@ int64_t random_geometric(brng_t *brng_state, double p) {
 }
 
 int64_t random_zipf(brng_t *brng_state, double a) {
-  double T, U, V;
-  int64_t X;
   double am1, b;
 
   am1 = a - 1.0;
   b = pow(2.0, am1);
-  do {
-    U = 1.0 - next_double(brng_state);
-    V = next_double(brng_state);
-    X = (int64_t)floor(pow(U, -1.0 / am1));
-    /* The real result may be above what can be represented in a int64.
-     * It will get casted to -sys.maxint-1. Since this is
-     * a straightforward rejection algorithm, we can just reject this value
-     * in the rejection condition below. This function then models a Zipf
+  while (1) {
+    double T, U, V, X;
+
+    U = 1.0 - random_double(brng_state);
+    V = random_double(brng_state);
+    X = floor(pow(U, -1.0 / am1));
+    /*
+     * The real result may be above what can be represented in a signed
+     * long. Since this is a straightforward rejection algorithm, we can
+     * just reject this value. This function then models a Zipf
      * distribution truncated to sys.maxint.
      */
+    if (X > LONG_MAX || X < 1.0) {
+      continue;
+    }
+
     T = pow(1.0 + 1.0 / X, am1);
-  } while (((V * X * (T - 1.0) / (b - 1.0)) > (T / b)) || X < 1);
-  return X;
+    if (V * X * (T - 1.0) / (b - 1.0) <= T / b) {
+      return (long)X;
+    }
+  }
 }
 
 double random_triangular(brng_t *brng_state, double left, double mode,
diff --git a/numpy/random/tests/test_generator_mt19937.py b/numpy/random/tests/test_generator_mt19937.py
@@ -90,6 +90,11 @@ def test_size(self):
 
     def test_invalid_prob(self):
         assert_raises(ValueError, random.multinomial, 100, [1.1, 0.2])
+        assert_raises(ValueError, random.multinomial, 100, [-.1, 0.9])
+
+    def test_invalid_n(self):
+        assert_raises(ValueError, random.multinomial, -1, [0.8, 0.2])
+        assert_raises(ValueError, random.multinomial, [-1] * 10, [0.8, 0.2])
 
 
 class TestSetState(object):
@@ -804,8 +809,7 @@ def test_geometric_exceptions(self):
         assert_raises(ValueError, random.geometric, [1.1] * 10)
         assert_raises(ValueError, random.geometric, -0.1)
         assert_raises(ValueError, random.geometric, [-0.1] * 10)
-        with suppress_warnings() as sup:
-            sup.record(RuntimeWarning)
+        with np.errstate(invalid='ignore'):
             assert_raises(ValueError, random.geometric, np.nan)
             assert_raises(ValueError, random.geometric, [np.nan] * 10)
 
@@ -888,8 +892,7 @@ def test_logseries(self):
         assert_array_equal(actual, desired)
 
     def test_logseries_exceptions(self):
-        with suppress_warnings() as sup:
-            sup.record(RuntimeWarning)
+        with np.errstate(invalid='ignore'):
             assert_raises(ValueError, random.logseries, np.nan)
             assert_raises(ValueError, random.logseries, [np.nan] * 10)
 
@@ -964,8 +967,7 @@ def test_negative_binomial(self):
         assert_array_equal(actual, desired)
 
     def test_negative_binomial_exceptions(self):
-        with suppress_warnings() as sup:
-            sup.record(RuntimeWarning)
+        with np.errstate(invalid='ignore'):
             assert_raises(ValueError, random.negative_binomial, 100, np.nan)
             assert_raises(ValueError, random.negative_binomial, 100,
                           [np.nan] * 10)
@@ -1046,8 +1048,7 @@ def test_poisson_exceptions(self):
         assert_raises(ValueError, random.poisson, [lamneg] * 10)
         assert_raises(ValueError, random.poisson, lambig)
         assert_raises(ValueError, random.poisson, [lambig] * 10)
-        with suppress_warnings() as sup:
-            sup.record(RuntimeWarning)
+        with np.errstate(invalid='ignore'):
             assert_raises(ValueError, random.poisson, np.nan)
             assert_raises(ValueError, random.poisson, [np.nan] * 10)
 
@@ -1850,7 +1851,7 @@ def test_logseries(self):
         assert_raises(ValueError, logseries, bad_p_two * 3)
 
     def test_multinomial(self):
-        random.seed(self.seed)
+        random.brng.seed(self.seed)
         actual = random.multinomial([5, 20], [1 / 6.] * 6, size=(3, 2))
         desired = np.array([[[1, 1, 1, 1, 0, 1],
                              [4, 5, 1, 4, 3, 3]],
@@ -1860,7 +1861,7 @@ def test_multinomial(self):
                              [3, 2, 3, 4, 2, 6]]], dtype=np.int64)
         assert_array_equal(actual, desired)
 
-        random.seed(self.seed)
+        random.brng.seed(self.seed)
         actual = random.multinomial([5, 20], [1 / 6.] * 6)
         desired = np.array([[1, 1, 1, 1, 0, 1],
                             [4, 5, 1, 4, 3, 3]], dtype=np.int64)
diff --git a/numpy/random/tests/test_randomstate.py b/numpy/random/tests/test_randomstate.py
@@ -112,6 +112,10 @@ def test_size(self):
 
     def test_invalid_prob(self):
         assert_raises(ValueError, random.multinomial, 100, [1.1, 0.2])
+        assert_raises(ValueError, random.multinomial, 100, [-.1, 0.9])
+
+    def test_invalid_n(self):
+        assert_raises(ValueError, random.multinomial, -1, [0.8, 0.2])
 
 
 class TestSetState(object):
