diff --git a/quantecon/__init__.py b/quantecon/__init__.py index a0bcf7484..d18e84870 100644 --- a/quantecon/__init__.py +++ b/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 #Place Holder if we wish to promote any general objects to the qe namespace. diff --git a/quantecon/discrete_rv.py b/quantecon/discrete_rv.py index 8972e67d8..5030cf93e 100644 --- a/quantecon/discrete_rv.py +++ b/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') diff --git a/quantecon/random/__init__.py b/quantecon/random/__init__.py index 33c35e3a8..3d5f67415 100644 --- a/quantecon/random/__init__.py +++ b/quantecon/random/__init__.py @@ -13,4 +13,4 @@ 1. AR1 Function """ -from .utilities import probvec, sample_without_replacement \ No newline at end of file +from .utilities import probvec, sample_without_replacement, draw diff --git a/quantecon/random/tests/test_utilities.py b/quantecon/random/tests/test_utilities.py index dccd88e7a..c51fbd62f 100644 --- a/quantecon/random/tests/test_utilities.py +++ b/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 diff --git a/quantecon/random/utilities.py b/quantecon/random/utilities.py index e34c9fd5e..299077a47 100644 --- a/quantecon/random/utilities.py +++ b/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 diff --git a/quantecon/tests/test_discrete_rv.py b/quantecon/tests/test_discrete_rv.py index bc1a169eb..281bbd848 100644 --- a/quantecon/tests/test_discrete_rv.py +++ b/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])