ENH: Extend random.draw to accept a random number generator

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',

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):

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