8000 BUG: Add parameter check to negative_binomial by Androp0v · Pull Request #21005 · numpy/numpy · GitHub
[go: up one dir, main page]

Skip to content

BUG: Add parameter check to negative_binomial #21005

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Mar 13, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 51 additions & 2 deletions numpy/random/_generator.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ from numpy.core.multiarray import normalize_axis_index

from .c_distributions cimport *
from libc cimport string
from libc.math cimport sqrt
from libc.stdint cimport (uint8_t, uint16_t, uint32_t, uint64_t,
int32_t, int64_t, INT64_MAX, SIZE_MAX)
from ._bounded_integers cimport (_rand_bool, _rand_int32, _rand_int64,
Expand Down Expand Up @@ -2997,6 +2998,22 @@ cdef class Generator:
then the probability distribution of the number of non-"1"s that
appear before the third "1" is a negative binomial distribution.

Because this method internally calls ``Generator.poisson`` with an
intermediate random value, a ValueError is raised when the choice of
:math:`n` and :math:`p` would result in the mean + 10 sigma of the sampled
intermediate distribution exceeding the max acceptable value of the
``Generator.poisson`` method. This happens when :math:`p` is too low
(a lot of failures happen for every success) and :math:`n` is too big (
a lot of sucesses are allowed).
Therefore, the :math:`n` and :math:`p` values must satisfy the constraint:

.. math:: n\\frac{1-p}{p}+10n\\sqrt{n}\\frac{1-p}{p}<2^{63}-1-10\\sqrt{2^{63}-1},

Where the left side of the equation is the derived mean + 10 sigma of
a sample from the gamma distribution internally used as the :math:`lam`
parameter of a poisson sample, and the right side of the equation is
the constraint for maximum value of :math:`lam` in ``Generator.poisson``.

References
----------
.. [1] Weisstein, Eric W. "Negative Binomial Distribution." From
Expand All @@ -3021,9 +3038,41 @@ cdef class Generator:
... print(i, "wells drilled, probability of one success =", probability)

"""

cdef bint is_scalar = True
cdef double *_dn
cdef double *_dp
cdef double _dmax_lam

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_DOUBLE, 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, see Generator.negative_binomial Notes")

else:
_dn = <double*>np.PyArray_DATA(n_arr)
_dp = <double*>np.PyArray_DATA(p_arr)

check_constraint(_dn[0], 'n', CONS_POSITIVE_NOT_NAN)
check_constraint(_dp[0], '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[0]) / _dp[0] * (_dn[0] + 10 * sqrt(_dn[0]))
if _dmax_lam > POISSON_LAM_MAX:
raise ValueError("n too large or p too small, see Generator.negative_binomial Notes")

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_arr, 'n', CONS_NONE,
p_arr, 'p', CONS_NONE,
0.0, '', CONS_NONE)

def poisson(self, lam=1.0, size=None):
Expand Down
7 changes: 7 additions & 0 deletions numpy/random/tests/test_generator_mt19937.py
Original file line number Diff line number Diff line change
Expand Up @@ -1485,6 +1485,13 @@ def test_negative_binomial_p0_exception(self):
with assert_raises(ValueError):
x = random.negative_binomial(1, 0)

def test_negative_binomial_invalid_p_n_combination(self):
# Verify that values of p and n that would result in an overflow
# or infinite loop raise an exception.
with np.errstate(invalid='ignore'):
assert_raises(ValueError, random.negative_binomial, 2**62, 0.1)
assert_raises(ValueError, random.negative_binomial, [2**62], [0.1])

def test_noncentral_chisquare(self):
random = Generator(MT19937(self.seed))
actual = random.noncentral_chisquare(df=5, nonc=5, size=(3, 2))
Expand Down
0