diff --git a/pyproject.toml b/pyproject.toml index 8797dad4..d24fd0d5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,7 +23,7 @@ keywords = [ dynamic = ["description", "version"] requires-python = ">=3.7" dependencies = [ - 'numba>=0.49.0', + 'numba>=0.56.0', 'numpy>=1.17.0', 'requests', 'scipy>=1.5.0', diff --git a/quantecon/random/tests/test_utilities.py b/quantecon/random/tests/test_utilities.py index 430c7a1c..30020a90 100644 --- a/quantecon/random/tests/test_utilities.py +++ b/quantecon/random/tests/test_utilities.py @@ -1,17 +1,19 @@ """ -Tests for util/random.py +Tests for random/utilities.py Functions --------- probvec sample_without_replacement +draw """ import numbers import numpy as np +import pytest from numpy.testing import (assert_array_equal, assert_allclose, assert_raises, assert_) -from numba import njit +from numba import config, njit, TypingError from quantecon.random import probvec, sample_without_replacement, draw @@ -73,6 +75,39 @@ def draw_jitted(cdf, size=None): return draw(cdf, size) +@njit +def draw_jitted_rng(cdf, size, rng): + return draw(cdf, size, rng) + + +@njit +def draw_jitted_rng_kw(cdf, rng): + return draw(cdf, rng=rng) + + +@njit +def draw_jitted_explicit_none(cdf, size): + return draw(cdf, size, None) + + +@njit +def draw_jitted_fwd(cdf, size=None, rng=None): + # Forwards its own defaults, which reach the overload as types.none + return draw(cdf, size, rng) + + +@njit +def draw_jitted_seeded(cdf, size, seed): + np.random.seed(seed) + return draw(cdf, size) + + +@njit +def draw_jitted_optional_rng(cdf, size, rng, flag): + rng_or_none = rng if flag else None + return draw(cdf, size, rng_or_none) + + class TestDraw: def setup_method(self): self.pmf = np.array([0.4, 0.1, 0.5]) @@ -109,6 +144,140 @@ def test_lln(self): atol = 1e-2 assert_allclose(pmf_computed, self.pmf, atol=atol) + # rng: a Generator behaves the same on both paths # + + def test_python_jitted_agree(self): + # The contract that keeps the two implementations in step: for + # the same Generator seed they must return the same values. + for seed in [0, 1234, 20260801]: + for size in [1, 10, 1000]: + out_py = draw(self.cdf, size, np.random.default_rng(seed)) + out_jit = draw_jitted_rng(self.cdf, size, + np.random.default_rng(seed)) + assert_array_equal(out_py, out_jit) + + # Scalar draws are compared as arrays: the Python path + # returns np.int64 and the jitted path a Python int. + out_py = draw(self.cdf, None, np.random.default_rng(seed)) + out_jit = draw_jitted_rng(self.cdf, None, + np.random.default_rng(seed)) + assert_array_equal(np.asarray(out_py), np.asarray(out_jit)) + + def test_generator_reproducible(self): + for size in [None, 10]: + for func in [draw, draw_jitted_rng]: + out0 = func(self.cdf, size, np.random.default_rng(1234)) + out1 = func(self.cdf, size, np.random.default_rng(1234)) + assert_array_equal(np.asarray(out0), np.asarray(out1)) + + def test_generator_state_advances(self): + # The Generator must be mutated in place by the jitted call, not + # copied at the boundary, so that successive calls continue the + # stream rather than restarting it. + rng = np.random.default_rng(1234) + parts = [draw_jitted_rng(self.cdf, 5, rng), + draw_jitted_rng(self.cdf, 5, rng), + draw(self.cdf, 5, rng)] + expected = draw(self.cdf, 15, np.random.default_rng(1234)) + assert_array_equal(np.concatenate(parts), expected) + + def test_consumes_exactly_one_draw_per_variate(self): + # An extra variate taken at the *end* of a call is invisible to + # the value comparisons above, but it desynchronises a shared + # Generator for every later consumer. Complements + # test_generator_state_advances, which covers the in-place + # mutation half of the same contract. + for func in [draw, draw_jitted_rng]: + for size in [None, 1, 10]: + n = 1 if size is None else size + rng = np.random.default_rng(1234) + func(self.cdf, size, rng) + ref = np.random.default_rng(1234) + ref.random(n) + assert_array_equal(rng.random(4), ref.random(4)) + + def test_generator_keyword_form_in_jit(self): + out_py = draw(self.cdf, rng=np.random.default_rng(1234)) + out_jit = draw_jitted_rng_kw(self.cdf, np.random.default_rng(1234)) + assert_array_equal(np.asarray(out_py), np.asarray(out_jit)) + + # rng=None: unchanged behaviour on both paths # + + def test_none_matches_legacy_global(self): + for size in [None, 10]: + np.random.seed(99) + out = draw(self.cdf, size) + np.random.seed(99) + r = np.random.random(size) if size is not None \ + else np.random.random() + expected = np.searchsorted(self.cdf, r, side='right') + assert_array_equal(np.asarray(out), np.asarray(expected)) + + def test_none_spellings_all_compile(self): + # Omitted, explicitly None, and forwarded from a jitted caller's + # own default are distinct numba types, and each must reach the + # np.random branch of the overload. + size = 10 + for out in [draw_jitted(self.cdf, size), + draw_jitted_explicit_none(self.cdf, size), + draw_jitted_fwd(self.cdf, size), + draw_jitted_fwd(self.cdf, size, None)]: + assert_(out.shape == (size,)) + assert_(np.isin(out, range(self.n)).all()) + + for out in [draw_jitted(self.cdf), draw_jitted_fwd(self.cdf)]: + assert_(out in range(self.n)) + + def test_jitted_np_random_seed_still_reproducible(self): + out0 = draw_jitted_seeded(self.cdf, 10, 1234) + out1 = draw_jitted_seeded(self.cdf, 10, 1234) + assert_array_equal(out0, out1) + + # rng: compile-time rejections, jitted path only # + # + # The None | Generator contract is enforced only by the overload: + # the pure-Python body does no input checking and its behaviour for + # other inputs is unspecified, so nothing here calls `draw` + # directly, and all three tests are skipped when NUMBA_DISABLE_JIT + # routes the jitted wrappers to the Python body. + + @pytest.mark.skipif(config.DISABLE_JIT, + reason='requires nopython compilation') + def test_int_seed_raises_in_jit(self): + # Stricter than SPEC 7 by design: nopython mode cannot construct + # a generator from a seed. + with pytest.raises(TypingError) as excinfo: + draw_jitted_rng(self.cdf, 10, 1234) + # Numba nests the message in its report of the candidate + # implementations it rejected. Assert on tokens that can + # only come from draw's own message, not from the caller's + # source line that numba echoes alongside it. + msg = str(excinfo.value) + assert_('quantecon.random.draw' in msg) + assert_('np.random.default_rng' in msg) + + @pytest.mark.skipif(config.DISABLE_JIT, + reason='requires nopython compilation') + def test_randomstate_raises_in_jit(self): + # In nopython mode a RandomState cannot be typed at all, so it + # is rejected during argument typing and the overload never + # runs. Assert only the exception type, never the message. + assert_raises(TypingError, draw_jitted_rng, self.cdf, 10, + np.random.RandomState(1234)) + + @pytest.mark.skipif(config.DISABLE_JIT, + reason='requires nopython compilation') + def test_optional_generator_raises_in_jit(self): + assert_raises(TypingError, draw_jitted_optional_rng, self.cdf, 10, + np.random.default_rng(1234), True) + try: + draw_jitted_optional_rng(self.cdf, 10, + np.random.default_rng(1234), True) + except TypingError as e: + # A token unique to the Optional hint, so that this branch + # cannot be satisfied by the int-seed message. + assert_('could not prove' in str(e)) + @njit def draw_jitted_w_o_size(n): diff --git a/quantecon/random/utilities.py b/quantecon/random/utilities.py index bd309b06..9ca73ece 100644 --- a/quantecon/random/utilities.py +++ b/quantecon/random/utilities.py @@ -5,6 +5,7 @@ import numpy as np from numba import guvectorize, types +from numba import TypingError from numba.extending import overload from ..util import check_random_state @@ -170,10 +171,10 @@ def _sample_without_replacement(n, r, out): # Pure python implementation that will run if the JIT compiler is disabled -def draw(cdf, size=None): +def draw(cdf, size=None, rng=None): """ Generate a random sample according to the cumulative distribution - given by `cdf`. Jit-complied by Numba in nopython mode. + given by `cdf`. JIT-compiled by Numba in nopython mode. Parameters ---------- @@ -185,40 +186,145 @@ def draw(cdf, size=None): `size` independent draws is returned; otherwise, a single draw is returned as a scalar. + rng : np.random.Generator, optional(default=None) + Random number generator to draw from. Must be a + `np.random.Generator` or None; in particular, integer seeds and + `np.random.RandomState` are not accepted. If None, the global + random state is used (see Notes). + Returns ------- scalar(int) or ndarray(int, ndim=1) + Notes + ----- + `draw` is intended primarily for use inside jit-compiled functions. + Pass a `np.random.Generator` as `rng` whenever the draws should be + reproducible; the generator is consumed in place, so its state + advances across successive calls. + + `rng` accepts only None or a `Generator` because Numba's nopython + mode can neither construct a generator from a seed nor represent + `np.random.RandomState`. A jit-compiled caller passing anything + else gets a compile-time `numba.TypingError`. Build a generator + outside any jit-compiled function with `np.random.default_rng(seed)` + and pass it in. + + If `rng` is None, a jit-compiled caller draws from Numba's own + internal random state, which is seeded by calling `np.random.seed` + inside a jit-compiled function; seeding NumPy's global random state + from Python has no effect on it. + + A single `Generator` must not be shared across the iterations of a + `numba.prange` loop. Its state is updated in place, so concurrent + iterations race and can consume the same underlying random number; + the draws are then neither reproducible nor independent. Give each + iteration its own generator, spawned with `Generator.spawn` outside + the jit-compiled function -- `spawn` itself cannot be called in + nopython mode. + Examples -------- + >>> import numpy as np + >>> import quantecon as qe >>> 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]) + >>> rng = np.random.default_rng(1234) + >>> qe.random.draw(cdf, 10, rng=rng) + array([1, 0, 1, 0, 0, 0, 0, 0, 1, 0]) + >>> qe.random.draw(cdf, rng=np.random.default_rng(1234)) + np.int64(1) + + Inside a jit-compiled function: + + >>> from numba import njit + >>> @njit + ... def draw_jitted(cdf, size, rng): + ... return qe.random.draw(cdf, size, rng) + >>> draw_jitted(cdf, 5, np.random.default_rng(1234)) + array([1, 0, 1, 0, 0]) """ + if rng is None: + rng = np.random if isinstance(size, int): - rs = np.random.random(size) + rs = rng.random(size) out = np.searchsorted(cdf, rs, side='right') return out else: - r = np.random.random() + r = rng.random() return np.searchsorted(cdf, r, side='right') +def _is_no_rng(numba_type): + """ + Return True if `numba_type`, as seen from inside the `draw` + overload, means that `rng` was not supplied. + + Numba spells "no value" in more than one way depending on the call + site, and each must be recognised: an omitted argument, as in + ``draw(cdf, 10)``, arrives as the *Python* object `None`, while an + explicit `None`, as in ``draw(cdf, 10, None)``, and a `None` default + forwarded by a jit-compiled caller both arrive as + `numba.types.none`. `numba.types.Omitted` is not currently observed + in an `@overload` body but is handled for the benefit of other Numba + typing templates. + + """ + return (numba_type is None or + isinstance(numba_type, types.NoneType) or + (isinstance(numba_type, types.Omitted) and + numba_type.value is None)) + + # Overload for the `draw` function +# +# The implementations below must return the same values as the pure +# Python body above: the same random numbers drawn in the same order, +# and the same `searchsorted` result. The sized branches keep the +# hand-written loop rather than the array `np.searchsorted` used in the +# Python body -- Numba supports both and they agree exactly, but the +# loop is measurably faster. `TestDraw.test_python_jitted_agree` is what +# holds the two paths together. @overload(draw) -def ol_draw(cdf, size=None): - if isinstance(size, types.Integer): - def draw_impl(cdf, size=None): - rs = np.random.random(size) - out = np.empty(size, dtype=np.int_) - for i in range(size): - out[i] = np.searchsorted(cdf, rs[i], side='right') - return out +def ol_draw(cdf, size=None, rng=None): + if isinstance(rng, types.NumPyRandomGeneratorType): + if isinstance(size, types.Integer): + def draw_impl(cdf, size=None, rng=None): + rs = rng.random(size) + out = np.empty(size, dtype=np.int_) + for i in range(size): + out[i] = np.searchsorted(cdf, rs[i], side='right') + return out + else: + def draw_impl(cdf, size=None, rng=None): + r = rng.random() + return np.searchsorted(cdf, r, side='right') + elif _is_no_rng(rng): + if isinstance(size, types.Integer): + def draw_impl(cdf, size=None, rng=None): + rs = np.random.random(size) + out = np.empty(size, dtype=np.int_) + for i in range(size): + out[i] = np.searchsorted(cdf, rs[i], side='right') + return out + else: + def draw_impl(cdf, size=None, rng=None): + r = np.random.random() + return np.searchsorted(cdf, r, side='right') else: - def draw_impl(cdf, size=None): - r = np.random.random() - return np.searchsorted(cdf, r, side='right') + if (isinstance(rng, types.Optional) and + isinstance(rng.type, types.NumPyRandomGeneratorType)): + hint = ('Numba could not prove that `rng` is a Generator ' + 'rather than None at this call site; hoist the None ' + 'case out of the branch that reaches `draw`.') + else: + hint = ('Integer seeds and np.random.RandomState are not ' + 'accepted, because nopython mode cannot construct a ' + 'generator. Build one outside the jit-compiled ' + 'function with np.random.default_rng(seed) and pass ' + 'that in.') + raise TypingError( + 'quantecon.random.draw: `rng` must be None or an ' + f'np.random.Generator; got {rng}. {hint}' + ) return draw_impl