From 45af3a50d146c400c016259afd213c76459312e8 Mon Sep 17 00:00:00 2001
From: QBatista <quentin.batista@mail.mcgill.ca>
Date: Wed, 7 Mar 2018 11:10:16 +1100
Subject: [PATCH] FEAT: Add jitted function for drawing

---
 quantecon/discrete_rv.py | 54 +++++++++++++++++++++++++++++-----------
 1 file changed, 40 insertions(+), 14 deletions(-)

diff --git a/quantecon/discrete_rv.py b/quantecon/discrete_rv.py
index 7cc012666..5ce0605fa 100644
--- a/quantecon/discrete_rv.py
+++ b/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))