Add parameter check in negative_binomial generator to avoid infinite loop for large values

1 files changed, 35 insertions, 2 deletions

diff --git a/numpy/random/_generator.pyx b/numpy/random/_generator.pyx
index d7c1879e7..c30c8d984 100644
--- a/numpy/random/_generator.pyx
+++ b/numpy/random/_generator.pyx
@@ -3021,9 +3021,42 @@ cdef class Generator:
         ...    print(i, "wells drilled, probability of one success =", probability)
 
         """
+
+        cdef bint is_scalar = True
+
+        p_arr = <np.ndarray>np.PyArray_FROM_OTF(p, np.NPY_DOUBLE, np.NPY_ALIGNED)
+        is_scalar = is_scalar and np.PyArray_NDIM(p_arr) == 0
+        n_arr = <np.ndarray>np.PyArray_FROM_OTF(n, np.NPY_INT64, np.NPY_ALIGNED)
+        is_scalar = is_scalar and np.PyArray_NDIM(n_arr) == 0
+
+        if not is_scalar:
+            check_array_constraint(n_arr, 'n', CONS_POSITIVE_NOT_NAN)
+            check_array_constraint(p_arr, 'p', CONS_BOUNDED_GT_0_1)
+            # Check that the choice of negative_binomial parameters won't result in a
+            # call to the poisson distribution function with a value of lam too large.
+            max_lam_arr = (1 - p_arr) / p_arr * (n_arr + 10 * np.sqrt(n_arr))
+            if np.any(np.greater(max_lam_arr, POISSON_LAM_MAX)):
+                raise ValueError("n too large or p too small")
+
+            return disc(&random_negative_binomial, &self._bitgen, size, self.lock, 2, 0,
+                        n, '', CONS_NONE,
+                        p, '', CONS_NONE,
+                        0.0, '', CONS_NONE)
+
+        _dp = PyFloat_AsDouble(p)
+        _in = PyFloat_AsDouble(n)
+
+        check_constraint(<double>_in, 'n', CONS_POSITIVE_NOT_NAN)
+        check_constraint(_dp, 'p', CONS_BOUNDED_GT_0_1)
+        # Check that the choice of negative_binomial parameters won't result in a
+        # call to the poisson distribution function with a value of lam too large.
+        _dmax_lam = (1 - _dp) / _dp * (_in + 10 * np.sqrt(_in))
+        if _dmax_lam > POISSON_LAM_MAX:
+            raise ValueError("n too large or p too small")
+
         return disc(&random_negative_binomial, &self._bitgen, size, self.lock, 2, 0,
-                    n, 'n', CONS_POSITIVE_NOT_NAN,
-                    p, 'p', CONS_BOUNDED_GT_0_1,
+                    n, 'n', CONS_NONE,
+                    p, 'p', CONS_NONE,
                     0.0, '', CONS_NONE)
 
     def poisson(self, lam=1.0, size=None):