ENH/BUG: Allow multinomial to check pvals with other float types

bashtage · bashtage · commit c1e3e114ffd3 · 2021-02-18T16:39:14.000Z
Add additional check when original input is an array that does not have dtype double closes numpy#8317
diff --git a/numpy/random/_generator.pyx b/numpy/random/_generator.pyx
@@ -3726,6 +3726,7 @@ cdef class Generator:
         cdef np.npy_intp d, i, sz, offset
         cdef np.ndarray parr, mnarr, on, temp_arr
         cdef double *pix
+        cdef double max_sum
         cdef int64_t *mnix
         cdef int64_t ni
         cdef np.broadcast it
@@ -3736,8 +3737,16 @@ cdef class Generator:
             pvals, np.NPY_DOUBLE, 1, 1, np.NPY_ARRAY_ALIGNED | np.NPY_ARRAY_C_CONTIGUOUS)
         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")
+        max_sum = 1.0 + 1e-12
+        if kahan_sum(pix, d-1) > max_sum:
+            # Further checks to handle case where the pvals is an array with
+            # a dtype that differs from double. Comparison is slow, but should
+            # almost never be hit when pvals is valid
+            if (not isinstance(pvals, np.ndarray)
+                    or pvals.dtype == float
+                    or not np.issubdtype(pvals.dtype, np.floating)
+                    or pvals[:-1].sum() > pvals.dtype.type(max_sum)):
+                raise ValueError(f"sum(pvals[:-1]) > 1.0")
 
         if np.PyArray_NDIM(on) != 0: # vector
             check_array_constraint(on, 'n', CONS_NON_NEGATIVE)
diff --git a/numpy/random/tests/test_generator_mt19937.py b/numpy/random/tests/test_generator_mt19937.py
@@ -142,6 +142,20 @@ def test_multidimensional_pvals(self):
         assert_raises(ValueError, random.multinomial, 10, [[[0], [1]], [[1], [0]]])
         assert_raises(ValueError, random.multinomial, 10, np.array([[0, 1], [1, 0]]))
 
+    def test_multinomial_pvals_float32(self):
+        x = np.array([9.9e-01, 9.9e-01, 1.0e-09, 1.0e-09, 1.0e-09, 1.0e-09,
+                      1.0e-09, 1.0e-09, 1.0e-09, 1.0e-09], dtype=np.float32)
+        pvals = x / x.sum()
+        random = Generator(MT19937(1432985819))
+        result = random.multinomial(1, pvals)
+
+        random = Generator(MT19937(1432985819))
+        pvals = pvals.astype(float)
+        assert_raises(ValueError, random.multinomial, 1, pvals)
+        pvals = pvals / pvals.sum()
+        expected = random.multinomial(1, pvals)
+        assert_array_equal(result, expected)
+
 
 class TestMultivariateHypergeometric:
 
