FIX: random.draw: Add =None

Fix #704

quantecon/random/tests/test_utilities.py

@@ -108,3 +108,14 @@ def test_lln(self):
  pmf_computed = hist * np.diff(bin_edges)
  atol = 1e-2
  assert_allclose(pmf_computed, self.pmf, atol=atol)
+
+
+@njit
+def draw_jitted_w_o_size(n):
+ cdf = np.linspace(1/n, 1, n)
+ return draw(cdf)
+
+
+def test_draw_jitted_w_o_size():
+ n = 3
+ assert_(draw_jitted_w_o_size(n) in range(n))

quantecon/random/utilities.py

@@ -211,16 +211,16 @@ def draw(cdf, size=None):
 
 # Overload for the `draw` function
 @overload(draw)
-def ol_draw(cdf, size):
+def ol_draw(cdf, size=None):
  if isinstance(size, types.Integer):
- def draw_impl(cdf, size):
+ def draw_impl(cdf, size=None):
  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:
- def draw_impl(cdf, size):
+ def draw_impl(cdf, size=None):
  r = np.random.random()
  return searchsorted(cdf, r)
  return draw_impl