FEAT: Add random.draw

quantecon/__init__.py

@@ -15,7 +15,7 @@
 
 #-Objects-#
 from .compute_fp import compute_fixed_point
-from .discrete_rv import DiscreteRV, random_choice_scalar, random_choice
+from .discrete_rv import DiscreteRV
 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,7 @@
 import numpy as np
 from numpy import cumsum
 from numpy.random import uniform
-from numba import jit
+from .util import check_random_state
 
 
 class DiscreteRV:
@@ -26,7 +26,7 @@ class DiscreteRV:
 
     Attributes
     ----------
-    q : see Parameters
+    q : see Parameters.
     Q : array_like(float)
         The cumulative sum of q.
 
@@ -59,7 +59,7 @@ def q(self, val):
         self._q = np.asarray(val)
         self.Q = cumsum(val)
 
-    def draw(self, k=None, seed=None):
+    def draw(self, k=1, random_state=None):
         """
         Returns k draws from q.
 
@@ -69,87 +69,21 @@ def draw(self, k=None, seed=None):
         Parameters
         -----------
         k : scalar(int), optional
-            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.
+            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.
 
         Returns
         -------
         array_like(int)
-            An array of k independent draws from q.
+            An array of k independent draws from q
 
         """
-        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)
+        random_state = check_random_state(random_state)
 
-    return np.searchsorted(a=cum_sum, v=uniform(0, 1, size=size))
+        return self.Q.searchsorted(random_state.uniform(0, 1, size=k),
+                                   side='right')

quantecon/random/__init__.py

@@ -13,4 +13,4 @@
    1. AR1 Function
 """
 
-from .utilities import probvec, sample_without_replacement
+from .utilities import probvec, sample_without_replacement, draw

quantecon/random/tests/test_utilities.py

@@ -7,10 +7,11 @@
 sample_without_replacement
 
 """
+import numbers
 import numpy as np
-from numpy.testing import assert_array_equal, assert_raises
-from nose.tools import eq_
-from quantecon.random import probvec, sample_without_replacement
+from numpy.testing import assert_array_equal, assert_allclose, assert_raises
+from nose.tools import eq_, ok_
+from quantecon.random import probvec, sample_without_replacement, draw
 
 
 # probvec #
@@ -64,6 +65,39 @@ def test_sample_without_replacement_value_error():
     assert_raises(ValueError, sample_without_replacement, 2, 3)
 
 
+# draw #
+
+class TestDraw:
+    def setUp(self):
+        self.pmf = np.array([0.4, 0.1, 0.5])
+        self.cdf = np.cumsum(self.pmf)
+        self.n = len(self.pmf)
+
+    def test_return_types(self):
+        out = draw(self.cdf)
+        ok_(isinstance(out, numbers.Integral))
+
+        size = 10
+        out = draw(self.cdf, size)
+        eq_(out.shape, (size,))
+
+    def test_return_values(self):
+        out = draw(self.cdf)
+        ok_(out in range(self.n))
+
+        size = 10
+        out = draw(self.cdf, size)
+        ok_(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)
+
+
 if __name__ == '__main__':
     import sys
     import nose

quantecon/random/utilities.py

@@ -3,9 +3,9 @@
 """
 
 import numpy as np
-from numba import jit, guvectorize
+from numba import jit, guvectorize, generated_jit, types
 
-from ..util import check_random_state
+from ..util import check_random_state, searchsorted
 
 
 # Generating Arrays and Vectors #
@@ -161,3 +161,46 @@ def sample_without_replacement(n, k, num_trials=None, random_state=None):
         return result[0]
     else:
         return result
+
+
+@generated_jit(nopython=True)
+def draw(cdf, size=None):
+    """
+    Generate a random sample according to the cumulative distribution
+    given by `cdf`. Jit-complied by Numba in nopython mode.
+
+    Parameters
+    ----------
+    cdf : array_like(float, ndim=1)
+        Array containing the cumulative distribution.
+
+    size : scalar(int), optional(default=None)
+        Size of the sample. If an integer is supplied, an ndarray of
+        `size` independent draws is returned; otherwise, a single draw
+        is returned as a scalar.
+
+    Returns
+    -------
+    scalar(int) or ndarray(int, ndim=1)
+
+    Examples
+    --------
+    >>> cdf = np.cumsum([0.4, 0.6])
+    >>> qe.random.draw(cdf)
+    1
+    >>> qe.random.draw(cdf, 10)
+    array([1, 0, 1, 0, 1, 0, 0, 0, 1, 0])
+
+    """
+    if isinstance(size, types.Integer):
+        def draw_impl(cdf, size):
+            rs = np.random.random_sample(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):
+            r = np.random.random_sample()
+            return searchsorted(cdf, r)
+    return draw_impl

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, seed=5)
+        draws = DiscreteRV(x).draw(k=10, random_state=5)
 
         expected_output = np.array([1, 2, 1, 2, 1, 1, 2, 1, 1, 1])