diff --git a/quantecon/game_theory/random.py b/quantecon/game_theory/random.py index c7e5285fd..9cf1d564c 100644 --- a/quantecon/game_theory/random.py +++ b/quantecon/game_theory/random.py @@ -202,7 +202,7 @@ def _random_mixed_actions(out, random_state): x[0] = 1 else: r = random_state.random(size=n-1) - _probvec_cpu(r, x) + _probvec_cpu()(r, x) return out diff --git a/quantecon/game_theory/vertex_enumeration.py b/quantecon/game_theory/vertex_enumeration.py index a811950fa..494e7ec80 100644 --- a/quantecon/game_theory/vertex_enumeration.py +++ b/quantecon/game_theory/vertex_enumeration.py @@ -111,8 +111,12 @@ def vertex_enumeration_gen(g, qhull_options=None): g.players[1-i], idx=i, qhull_options=qhull_options ) for i in range(N)] - labelings_bits_tup = \ - tuple(_ints_arr_to_bits(brps[i].labelings) for i in range(N)) + labelings_bits_tup = tuple( + _ints_arr_to_bits( + brps[i].labelings, + np.empty(brps[i].labelings.shape[0], dtype=np.uint64) + ) for i in range(N) + ) equations_tup = tuple(brps[i].equations for i in range(N)) trans_recips = tuple(brps[i].trans_recip for i in range(N)) @@ -302,33 +306,37 @@ def __init__(self, opponent_player, idx=0, qhull_options=None): self.trans_recip = trans_recip -@guvectorize(['(i4[:], u8[:])'], '(m)->()', nopython=True, cache=True) +@guvectorize('(m)->()', nopython=True, cache=True) def _ints_arr_to_bits(ints_arr, out): """ Convert an array of integers representing the set bits into the corresponding integer. - Compiled as a ufunc by Numba's `@guvectorize`: if the input is a - 2-dim array with shape[0]=K, the function returns a 1-dim array of - K converted integers. + Compiled as a dynamic (signature-less) gufunc by Numba's + `@guvectorize` on first call: if the input is a 2-dim array with + shape[0]=K, the function fills the 1-dim output array `out` with + K converted integers. The output array must be supplied by the + caller, with dtype np.uint64. Parameters ---------- - ints_arr : ndarray(int32, ndim=1) - Array of distinct integers from 0, ..., 63. + ints_arr : ndarray(int, ndim=2) + Array of distinct integers from 0, ..., 63 in each row. + + out : ndarray(np.uint64, ndim=1) + Output array, of length ints_arr.shape[0]. Returns ------- - np.uint64 - Integer with set bits represented by the input integers. + out : ndarray(np.uint64, ndim=1) + Array of integers with set bits represented by the rows of + the input. Examples -------- - >>> ints_arr = np.array([0, 1, 2], dtype=np.int32) - >>> _ints_arr_to_bits(ints_arr) - np.uint64(7) - >>> ints_arr2d = np.array([[0, 1, 2], [3, 0, 1]], dtype=np.int32) - >>> _ints_arr_to_bits(ints_arr2d) + >>> ints_arr = np.array([[0, 1, 2], [3, 0, 1]], dtype=np.int32) + >>> out = np.empty(ints_arr.shape[0], dtype=np.uint64) + >>> _ints_arr_to_bits(ints_arr, out) array([ 7, 11], dtype=uint64) """ diff --git a/quantecon/random/utilities.py b/quantecon/random/utilities.py index c9dae5785..863775cb4 100644 --- a/quantecon/random/utilities.py +++ b/quantecon/random/utilities.py @@ -3,6 +3,8 @@ """ +import functools + import numpy as np from numba import guvectorize, types from numba import TypingError @@ -60,9 +62,9 @@ def probvec(m, k, random_state=None, parallel=True): # Parse Parallel Option # if parallel: - _probvec_parallel(r, x) + _probvec_parallel()(r, x) else: - _probvec_cpu(r, x) + _probvec_cpu()(r, x) return x @@ -92,14 +94,41 @@ def _probvec(r, out): # pragma: no cover out[i] = r[i] - r[i-1] out[n] = 1 - r[n-1] -_probvec_parallel = guvectorize( - ['(f8[:], f8[:])'], '(n), (k)', nopython=True, target='parallel', - cache=True - )(_probvec) -_probvec_cpu = guvectorize( - ['(f8[:], f8[:])'], '(n), (k)', nopython=True, target='cpu', - cache=True - )(_probvec) +# The two gufuncs of `_probvec` are constructed on first call rather than +# at import, since an explicit signature makes `guvectorize` compile +# eagerly. The explicit signature is required by the parallel target, and +# is retained for the cpu target as well: an explicit-signature gufunc and +# a dynamic (signature-less) gufunc of the same kernel must not be mixed +# with `cache=True`, as they share a cache entry that is not compatible. + +_PROBVEC_SIGNATURE = ['(f8[:], f8[:])'] +_PROBVEC_LAYOUT = '(n), (k)' + + +@functools.cache +def _probvec_parallel(): + """ + Return the parallel-target gufunc of `_probvec`, compiling it on + first call. + + """ + return guvectorize( + _PROBVEC_SIGNATURE, _PROBVEC_LAYOUT, nopython=True, + target='parallel', cache=True + )(_probvec) + + +@functools.cache +def _probvec_cpu(): + """ + Return the cpu-target gufunc of `_probvec`, compiling it on first + call. + + """ + return guvectorize( + _PROBVEC_SIGNATURE, _PROBVEC_LAYOUT, nopython=True, + target='cpu', cache=True + )(_probvec) def sample_without_replacement(n, k, num_trials=None, random_state=None): @@ -155,16 +184,18 @@ def sample_without_replacement(n, k, num_trials=None, random_state=None): random_state = check_random_state(random_state) r = random_state.random(size=size) - result = _sample_without_replacement(n, r) + result = np.empty(size, dtype=np.int64) + _sample_without_replacement(n, r, result) return result -@guvectorize(['(i8, f8[:], i8[:])'], '(),(k)->(k)', nopython=True, cache=True) +@guvectorize('(),(k)->(k)', nopython=True, cache=True) def _sample_without_replacement(n, r, out): """ - Main body of `sample_without_replacement`. To be compiled as a ufunc - by guvectorize of Numba. + Main body of `sample_without_replacement`. Compiled as a dynamic + (signature-less) gufunc by Numba's `@guvectorize` on first call; + the output array `out` must be supplied by the caller. """ k = r.shape[0]