ENH: Enable 0 samples in hypergeometric

Enable Hypergeometric to work with 0 samples xref numpy/numpy#9237
bashtage · Apr 10, 2019 · 699256c · 699256c
1 parent 9a32519
commit 699256c
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 5 deletions.
diff --git a/randomgen/generator.pyx b/randomgen/generator.pyx
@@ -3448,14 +3448,14 @@ cdef class RandomGenerator:
             return disc(&random_hypergeometric, self._brng, size, self.lock, 0, 3,
                         lngood, 'ngood', CONS_NON_NEGATIVE,
                         lnbad, 'nbad', CONS_NON_NEGATIVE,
-                        lnsample, 'nsample', CONS_GTE_1)
+                        lnsample, 'nsample', CONS_NON_NEGATIVE)
 
         if np.any(np.less(np.add(ongood, onbad), onsample)):
             raise ValueError("ngood + nbad < nsample")
         return discrete_broadcast_iii(&random_hypergeometric, self._brng, size, self.lock,
                                       ongood, 'ngood', CONS_NON_NEGATIVE,
                                       onbad, 'nbad', CONS_NON_NEGATIVE,
-                                      onsample, 'nsample', CONS_GTE_1)
+                                      onsample, 'nsample', CONS_NON_NEGATIVE)
 
     def logseries(self, p, size=None):
         """

diff --git a/randomgen/src/distributions/distributions.c b/randomgen/src/distributions/distributions.c
@@ -1179,8 +1179,10 @@ int64_t random_hypergeometric(brng_t *brng_state, int64_t good, int64_t bad,
                               int64_t sample) {
   if (sample > 10) {
     return random_hypergeometric_hrua(brng_state, good, bad, sample);
-  } else {
+  } else if (sample > 0) {
     return random_hypergeometric_hyp(brng_state, good, bad, sample);
+  } else {
+  return 0;
   }
 }
 

diff --git a/randomgen/tests/test_generator_mt19937.py b/randomgen/tests/test_generator_mt19937.py
@@ -1832,7 +1832,7 @@ def test_hypergeometric(self):
         nsample = [2]
         bad_ngood = [-1]
         bad_nbad = [-2]
-        bad_nsample_one = [0]
+        bad_nsample_one = [-1]
         bad_nsample_two = [4]
         hypergeom = random.hypergeometric
         desired = np.array([1, 1, 1])
@@ -1863,7 +1863,7 @@ def test_hypergeometric(self):
 
         assert_raises(ValueError, hypergeom, -1, 10, 20)
         assert_raises(ValueError, hypergeom, 10, -1, 20)
-        assert_raises(ValueError, hypergeom, 10, 10, 0)
+        assert_raises(ValueError, hypergeom, 10, 10, -1)
         assert_raises(ValueError, hypergeom, 10, 10, 25)
 
     def test_logseries(self):