diff --git a/quantecon/game_theory/vertex_enumeration.py b/quantecon/game_theory/vertex_enumeration.py index 236feea5..dbe3beae 100644 --- a/quantecon/game_theory/vertex_enumeration.py +++ b/quantecon/game_theory/vertex_enumeration.py @@ -296,25 +296,33 @@ def __init__(self, opponent_player, idx=0, qhull_options=None): self.trans_recip = trans_recip -@guvectorize(['(i4[:], u8[:])'], '(m)->()', nopython=True, cache=True) -def _ints_arr_to_bits(ints_arr, out): +def _ints_arr_to_bits_kernel(ints_arr, out): # pragma: no cover + m = ints_arr.shape[0] + out[0] = 0 + for i in range(m): + out[0] |= np.uint64(1) << np.uint64(ints_arr[i]) + + +_ints_arr_to_bits_gufunc = None + + +def _ints_arr_to_bits(ints_arr): """ 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. + If the input is a 2-dim array with shape[0]=K, returns a 1-dim array + of K converted integers. Parameters ---------- - ints_arr : ndarray(int32, ndim=1) + ints_arr : ndarray(int32, ndim=1 or 2) Array of distinct integers from 0, ..., 63. Returns ------- - np.uint64 - Integer with set bits represented by the input integers. + np.uint64 or ndarray(uint64) + Integer(s) with set bits represented by the input integers. Examples -------- @@ -326,10 +334,12 @@ def _ints_arr_to_bits(ints_arr, out): array([ 7, 11], dtype=uint64) """ - m = ints_arr.shape[0] - out[0] = 0 - for i in range(m): - out[0] |= np.uint64(1) << np.uint64(ints_arr[i]) + global _ints_arr_to_bits_gufunc + if _ints_arr_to_bits_gufunc is None: + _ints_arr_to_bits_gufunc = guvectorize( + ['(i4[:], u8[:])'], '(m)->()', nopython=True, cache=True + )(_ints_arr_to_bits_kernel) + return _ints_arr_to_bits_gufunc(ints_arr) @jit(nopython=True, cache=True) diff --git a/quantecon/random/utilities.py b/quantecon/random/utilities.py index 9ca73ece..97ec19ba 100644 --- a/quantecon/random/utilities.py +++ b/quantecon/random/utilities.py @@ -89,14 +89,26 @@ 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) +_probvec_parallel_gufunc = None +_probvec_cpu_gufunc = None + + +def _probvec_parallel(r, out): + global _probvec_parallel_gufunc + if _probvec_parallel_gufunc is None: + _probvec_parallel_gufunc = guvectorize( + ['(f8[:], f8[:])'], '(n),(k)', nopython=True, + target='parallel', cache=True)(_probvec) + _probvec_parallel_gufunc(r, out) + + +def _probvec_cpu(r, out): + global _probvec_cpu_gufunc + if _probvec_cpu_gufunc is None: + _probvec_cpu_gufunc = guvectorize( + ['(f8[:], f8[:])'], '(n),(k)', nopython=True, + target='cpu', cache=True)(_probvec) + _probvec_cpu_gufunc(r, out) def sample_without_replacement(n, k, num_trials=None, random_state=None): @@ -153,10 +165,9 @@ def sample_without_replacement(n, k, num_trials=None, random_state=None): return result -@guvectorize(['(i8, f8[:], i8[:])'], '(),(k)->(k)', nopython=True, cache=True) -def _sample_without_replacement(n, r, out): +def _sample_without_replacement_kernel(n, r, out): # pragma: no cover """ - Main body of `sample_without_replacement`. To be complied as a ufunc + Main body of `sample_without_replacement`. To be compiled as a ufunc by guvectorize of Numba. """ @@ -170,6 +181,18 @@ def _sample_without_replacement(n, r, out): pool[idx] = pool[n-j-1] +_swr_gufunc = None + + +def _sample_without_replacement(n, r): + global _swr_gufunc + if _swr_gufunc is None: + _swr_gufunc = guvectorize( + ['(i8, f8[:], i8[:])'], '(),(k)->(k)', nopython=True, cache=True + )(_sample_without_replacement_kernel) + return _swr_gufunc(n, r) + + # Pure python implementation that will run if the JIT compiler is disabled def draw(cdf, size=None, rng=None): """