FEAT: Add jitted function for drawing

quantecon/discrete_rv.py

@@ -11,7 +11,8 @@
 import numpy as np
 from numpy import cumsum
 from numpy.random import uniform
-from .util import check_random_state
+from numba import jit
+
 
 class DiscreteRV:
  """
@@ -21,13 +22,13 @@ class DiscreteRV:
  Parameters
  ----------
  q : array_like(float)
- Nonnegative numbers that sum to 1
+ Nonnegative numbers that sum to 1.
 
  Attributes
  ----------
  q : see Parameters
  Q : array_like(float)
- The cumulative sum of q
+ The cumulative sum of q.
 
  """
 
@@ -58,7 +59,7 @@ def q(self, val):
  self._q = np.asarray(val)
  self.Q = cumsum(val)
 
- def draw(self, k=1, random_state=None):
+ def draw(self, k=None, seed=0):
  """
  Returns k draws from q.
 
@@ -68,20 +69,45 @@ def draw(self, k=1, random_state=None):
  Parameters
  -----------
  k : scalar(int), optional
- Number of draws to be returned
-
- random_state : int or np.random.RandomState, optional
- Random seed (integer) or np.random.RandomState instance to set
- the initial state of the random number generator for
- reproducibility. If None, a randomly initialized RandomState is
- used.
+ Number of draws to be returned.
+ seed : int, optional
+ Random seed (integer) to set the initial state of the random number
+ generator for reproducibility.
 
  Returns
  -------
  array_like(int)
- An array of k independent draws from q
+ An array of k independent draws from q.
 
  """
- random_state = check_random_state(random_state)
+ if k == 1:
+ return discrete_rv(self.Q, size=k, seed=seed)[0]
+ return discrete_rv(self.Q, size=k, seed=seed)
+
+
+@jit(nopython=True)
+def discrete_rv(cum_sum, size=1, seed=0):
+ """
+ Returns `size` daws from `cum_sum` where `cum_sum` is the cumulative sum of
+ the vector of probabilities of a discrete random variable. Optimized using
+ Numba and compilied in nopython mode.
+
+ Parameters
+ -----------
+ cum_sum : array_like(float)
+ Cumulative sum of nonnegative numbers that sum to 1.
+ size : scalar(int), optional
+ Number of draws to be returned.
+ seed : int, optional
+ Random seed (integer) to set the initial state of the random number
+ generator for reproducibility.
+
+ Returns
+ -------
+ array_like(int)
+ An array of k independent draws from q.
+
+ """
+ np.random.seed(seed)
 
-  return self.Q.searchsorted(random_state.uniform(0, 1, size=k))
+ return np.searchsorted(a=cum_sum, v=uniform(0, 1, size=size))