RFC: Replace @generated_jit with @overload in random.draw

To avoid `NumbaDeprecationWarning`

Author: oyamad

quantecon/random/tests/test_utilities.py

@@ -11,6 +11,7 @@
 import numpy as np
 from numpy.testing import (assert_array_equal, assert_allclose, assert_raises,
                            assert_)
+from numba import njit
 from quantecon.random import probvec, sample_without_replacement, draw
 
 
@@ -67,32 +68,43 @@ def test_sample_without_replacement_value_error():
 
 # draw #
 
+@njit
+def draw_jitted(cdf, size=None):
+    return draw(cdf, size)
+
+
 class TestDraw:
     def setup_method(self):
         self.pmf = np.array([0.4, 0.1, 0.5])
         self.cdf = np.cumsum(self.pmf)
         self.n = len(self.pmf)
+        self.draw_funcs = [draw, draw_jitted]
 
     def test_return_types(self):
-        out = draw(self.cdf)
-        assert_(isinstance(out, numbers.Integral))
+        for func in self.draw_funcs:
+            out = func(self.cdf)
+            assert_(isinstance(out, numbers.Integral))
 
         size = 10
-        out = draw(self.cdf, size)
-        assert_(out.shape == (size,))
+        for func in self.draw_funcs:
+            out = func(self.cdf, size)
+            assert_(out.shape == (size,))
 
     def test_return_values(self):
-        out = draw(self.cdf)
-        assert_(out in range(self.n))
+        for func in self.draw_funcs:
+            out = func(self.cdf)
+            assert_(out in range(self.n))
 
         size = 10
-        out = draw(self.cdf, size)
-        assert_(np.isin(out, range(self.n)).all())
+        for func in self.draw_funcs:
+            out = func(self.cdf, size)
+            assert_(np.isin(out, range(self.n)).all())
 
     def test_lln(self):
         size = 1000000
-        out = draw(self.cdf, size)
-        hist, bin_edges = np.histogram(out, bins=self.n, density=True)
-        pmf_computed = hist * np.diff(bin_edges)
-        atol = 1e-2
-        assert_allclose(pmf_computed, self.pmf, atol=atol)
+        for func in self.draw_funcs:
+            out = func(self.cdf, size)
+            hist, bin_edges = np.histogram(out, bins=self.n, density=True)
+            pmf_computed = hist * np.diff(bin_edges)
+            atol = 1e-2
+            assert_allclose(pmf_computed, self.pmf, atol=atol)

quantecon/random/utilities.py

@@ -4,8 +4,8 @@
 """
 
 import numpy as np
-from numba import guvectorize, generated_jit, types
-
+from numba import guvectorize, types
+from numba.extending import overload
 from ..util import check_random_state, searchsorted
 
 
@@ -169,7 +169,7 @@ def _sample_without_replacement(n, r, out):
         pool[idx] = pool[n-j-1]
 
 
-@generated_jit(nopython=True)
+# Pure python implementation that will run if the JIT compiler is disabled
 def draw(cdf, size=None):
     """
     Generate a random sample according to the cumulative distribution
@@ -198,6 +198,20 @@ def draw(cdf, size=None):
     array([1, 0, 1, 0, 1, 0, 0, 0, 1, 0])
 
     """
+    if isinstance(size, int):
+        rs = np.random.random(size)
+        out = np.empty(size, dtype=np.int_)
+        for i in range(size):
+            out[i] = searchsorted(cdf, rs[i])
+        return out
+    else:
+        r = np.random.random()
+        return searchsorted(cdf, r)
+
+
+# Overload for the `draw` function
+@overload(draw)
+def ol_draw(cdf, size):
     if isinstance(size, types.Integer):
         def draw_impl(cdf, size):
             rs = np.random.random(size)