FEAT: Add jitted function for drawing

quantecon/__init__.py

@@ -15,7 +15,7 @@
 
 #-Objects-#
 from .compute_fp import compute_fixed_point
-from .discrete_rv import DiscreteRV
+from .discrete_rv import DiscreteRV, random_choice_scalar, random_choice
 from .ecdf import ECDF
 from .estspec import smooth, periodogram, ar_periodogram
 # from .game_theory import <objects-here> 							#Place Holder if we wish to promote any general objects to the qe namespace.

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=None):
         """
         Returns k draws from q.
 
@@ -68,20 +69,87 @@ 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 : scalar(int), optional
+            Random seed (integer) to set the initial state of the random number
+            generator for reproducibility. If None, a seed is randomly
+            generated.
 
         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 seed is None:
+            seed = np.random.randint(10**18)
+
+        if k is None:
+            return random_choice_scalar(self._q, seed=seed, cum_sum=self.Q)
+        else:
+            return random_choice(self._q, seed=seed, cum_sum=self.Q, size=k)
+
+
+@jit(nopython=True)
+def random_choice_scalar(p_vals, seed, cum_sum=None):
+    """
+    Returns one draw from `p_vals`. Optimized using Numba and compilied in
+    nopython mode.
+
+    Parameters
+    -----------
+    p_vals : array_like(float)
+        Nonnegative numbers that sum to 1.
+    seed : scalar(int)
+        Random seed (integer) to set the initial state of the random number
+        generator for reproducibility.
+    cum_sum : array_like(float), optional
+        The cumulative sum of p_vals.
+
+    Returns
+    -------
+    scalar(int)
+       One draw from p_vals.
+
+    """
+    np.random.seed(seed)
+
+    if cum_sum is None:
+        cum_sum = cumsum(p_vals)
+
+    return np.searchsorted(a=cum_sum, v=uniform(0, 1))
+
+
+@jit(nopython=True)
+def random_choice(p_vals, seed, cum_sum=None, size=1):
+    """
+    Returns `size` draws from `p_vals`. Optimized using Numba and compilied in
+    nopython mode.
+
+    For each such draw, the value i is returned with probability
+    p_vals[i].
+
+    Parameters
+    -----------
+    p_vals : array_like(float)
+        Nonnegative numbers that sum to 1.
+    seed : scalar(int)
+        Random seed (integer) to set the initial state of the random number
+        generator for reproducibility.
+    cum_sum : array_like(float), optional
+        The cumulative sum of p_vals.
+    size : scalar(int), optional
+        Number of draws to be returned.
+
+    Returns
+    -------
+    array_like(int)
+        An array of k independent draws from p_vals.
+
+    """
+    np.random.seed(seed)
+
+    if cum_sum is None:
+        cum_sum = cumsum(p_vals)
 
-        return self.Q.searchsorted(random_state.uniform(0, 1, size=k))
+    return np.searchsorted(a=cum_sum, v=uniform(0, 1, size=size))

quantecon/tests/test_discrete_rv.py

@@ -60,7 +60,7 @@ def test_draw_with_seed(self):
         x = np.array([0.03326189, 0.60713005, 0.84514831, 0.28493183,
                       0.12393182, 0.35308009, 0.70371579, 0.81728178,
                       0.21294538, 0.05358209])
-        draws = DiscreteRV(x).draw(k=10, random_state=5)
+        draws = DiscreteRV(x).draw(k=10, seed=5)
 
         expected_output = np.array([1, 2, 1, 2, 1, 1, 2, 1, 1, 1])