Skip to content

Commit 929fb49

Browse files
kc611kc611
committed
Added Generator().choice() and Generator.bytes()
1 parent af575e6 commit 929fb49

3 files changed

Lines changed: 225 additions & 4 deletions

File tree

‎numba/np/random/generator_methods.py‎

Lines changed: 47 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,8 @@
1919
from numba.np.random.random_methods import \
2020
(random_bounded_uint64_fill, random_bounded_uint32_fill,
2121
random_bounded_uint16_fill, random_bounded_uint8_fill,
22-
random_bounded_bool_fill, _randint_arg_check)
23-
22+
random_bounded_bool_fill, _randint_arg_check, _choice)
23+
from numba.core.extending import register_jitable
2424

2525
registry = Registry('generator_methods')
2626

@@ -123,6 +123,51 @@ def impl(inst, low, high=None, size=None,
123123
return impl
124124

125125

126+
# Overload the Generator().choice()
127+
@overload_method(types.NumPyRandomGeneratorType, 'choice')
128+
def NumPyRandomGeneratorType_choice(inst, a, size=None, replace=True,
129+
p=None, axis=0, shuffle=True):
130+
131+
if isinstance(a, types.Integer):
132+
@register_jitable
133+
def array_maker(a):
134+
return np.arange(a)
135+
else:
136+
@register_jitable
137+
def array_maker(a):
138+
return a
139+
if isinstance(size, types.Omitted):
140+
size = size.value
141+
142+
if isinstance(size, (types.NoneType,)) or size is None:
143+
def impl(inst, a, size=None, replace=True,
144+
p=None, axis=0, shuffle=True):
145+
return _choice(inst, array_maker(a), 1,
146+
replace, p, axis, shuffle)[0]
147+
return impl
148+
else:
149+
def impl(inst, a, size=None, replace=True,
150+
p=None, axis=0, shuffle=True):
151+
return _choice(inst, array_maker(a), size,
152+
replace, p, axis, shuffle)
153+
return impl
154+
155+
156+
# Fails right now since .tobytes() method of Numpy arrays
157+
# is not yet supported in Numba
158+
# Overload the Generator().bytes()
159+
160+
# @overload_method(types.NumPyRandomGeneratorType, 'bytes')
161+
# def NumPyRandomGeneratorType_bytes(inst, length):
162+
# def impl(inst, length):
163+
# _size = ((length - 1) // 4 + 1)
164+
# # Interpret the uint32s as little-endian to convert them to bytes
165+
# # consistently.
166+
# return inst.integers(0, 4294967296, size=_size, dtype=np.uint32)\
167+
# .astype('<u4').tobytes()[:length]
168+
# return impl
169+
170+
126171
# Overload the Generator().random()
127172
@overload_method(types.NumPyRandomGeneratorType, 'random')
128173
def NumPyRandomGeneratorType_random(inst, size=None, dtype=np.float64):

‎numba/np/random/random_methods.py‎

Lines changed: 158 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -347,3 +347,161 @@ def _randint_arg_check(low, high, endpoint, lower_bound, upper_bound):
347347
raise ValueError("low is greater than high in given interval")
348348
rng = (high - low)
349349
return low, rng
350+
351+
352+
@register_jitable
353+
def _shuffle_int(bitgen, n, first, data):
354+
for i in range(n, first, -1):
355+
j = random_bounded_uint64_fill(bitgen, 0, i, None, 1, np.uint64)[0]
356+
data[i], data[j] = data[j], data[i]
357+
return data
358+
359+
360+
@register_jitable
361+
def _choice(bitgen_inst, a, size, replace, p, axis, shuffle):
362+
pop_size = a.shape[axis]
363+
364+
if p is not None:
365+
atol = np.sqrt(np.finfo(np.float64).eps)
366+
if isinstance(p, np.ndarray):
367+
if np.issubdtype(p.dtype, np.floating):
368+
atol = max(atol, np.sqrt(np.finfo(p.dtype).eps))
369+
370+
flat_size = np.prod(np.array(size))
371+
372+
# Actual sampling
373+
if replace:
374+
if p is not None:
375+
cdf = p.cumsum()
376+
cdf /= cdf[-1]
377+
uniform_samples = bitgen_inst.random(size)
378+
idx = cdf.searchsorted(uniform_samples, side='right')
379+
# searchsorted returns a scalar
380+
idx = np.array(idx, copy=False, dtype=np.int64)
381+
else:
382+
idx = bitgen_inst.integers(0, pop_size, size=size, dtype=np.int64)
383+
else:
384+
if flat_size > pop_size:
385+
raise ValueError("Cannot take a larger sample than "
386+
"population when replace is False")
387+
elif flat_size < 0:
388+
raise ValueError("negative dimensions are not allowed")
389+
390+
if p is not None:
391+
if np.count_nonzero(p > 0) < flat_size:
392+
raise ValueError("Fewer non-zero entries in p than size")
393+
n_uniq = 0
394+
p = p.copy()
395+
found = np.zeros(size, dtype=np.int64)
396+
flat_found = found.ravel()
397+
while n_uniq < flat_size:
398+
x = bitgen_inst.random((flat_size - n_uniq,))
399+
if n_uniq > 0:
400+
p[flat_found[0:n_uniq]] = 0
401+
cdf = np.cumsum(p)
402+
cdf /= cdf[-1]
403+
new = cdf.searchsorted(x, side='right')
404+
_, unique_indices = np.unique(new, return_index=True)
405+
unique_indices.sort()
406+
new = new.take(unique_indices)
407+
flat_found[n_uniq:n_uniq + new.size] = new
408+
n_uniq += new.size
409+
idx = found
410+
else:
411+
size_i = flat_size
412+
pop_size_i = pop_size
413+
# This is a heuristic tuning. should be improvable
414+
if shuffle:
415+
cutoff = 50
416+
else:
417+
cutoff = 20
418+
if pop_size_i > 10000 and (size_i > (pop_size_i // cutoff)):
419+
# Tail shuffle size elements
420+
idx = np.arange(0, pop_size_i)
421+
idx = _shuffle_int(bitgen_inst.bit_generator, pop_size_i,
422+
max(pop_size_i - size_i, 1), idx)
423+
# Copy to allow potentially large array backing idx to be gc
424+
idx = idx[(pop_size - flat_size):].copy()
425+
else:
426+
# Floyd's algorithm
427+
idx = np.empty(int(flat_size), dtype=np.int64)
428+
# smallest power of 2 larger than 1.2 * flat_size
429+
set_size = (1.2 * size_i)
430+
mask = gen_mask(int(set_size))
431+
set_size = 1 + mask
432+
hash_set = np.full(set_size, UINT64_MAX - 1, np.uint64)
433+
for j in range(pop_size_i - size_i, pop_size_i):
434+
val = random_bounded_uint64_fill(bitgen_inst.bit_generator,
435+
0, j, None, 1, np.int64)[0]
436+
loc = val & mask
437+
while hash_set[loc] != UINT64_MAX - 1 \
438+
and hash_set[loc] != val:
439+
loc = (loc + 1) & mask
440+
if hash_set[loc] == UINT64_MAX - 1:
441+
hash_set[loc] = val
442+
idx[int(j - pop_size_i + size_i)] = val
443+
else: # we need to insert j instead
444+
loc = j & mask
445+
while hash_set[loc] != UINT64_MAX - 1:
446+
loc = (loc + 1) & mask
447+
hash_set[loc] = j
448+
idx[int(j - pop_size_i + size_i)] = j
449+
if shuffle:
450+
idx = _shuffle_int(bitgen_inst.bit_generator,
451+
size_i, 1, idx)
452+
453+
return idx
454+
455+
456+
def random_interval(bitgen, _max):
457+
if (_max == 0):
458+
return 0
459+
460+
mask = _max
461+
462+
# /* Smallest bit mask >= max */
463+
mask |= mask >> 1
464+
mask |= mask >> 2
465+
mask |= mask >> 4
466+
mask |= mask >> 8
467+
mask |= mask >> 16
468+
mask |= mask >> 32
469+
470+
# /* Search a random value in [0..mask] <= max */
471+
if (_max <= 0xffffffff):
472+
value = next_uint32(bitgen)
473+
while (value & mask) > _max:
474+
value = next_uint32(bitgen)
475+
else:
476+
value = next_uint64(bitgen)
477+
while (value & mask) > _max:
478+
value = next_uint64(bitgen)
479+
480+
return value
481+
482+
483+
def _shuffle(bitgen, x, axis=0):
484+
485+
if isinstance(x, np.ndarray):
486+
if x.size == 0:
487+
return x
488+
489+
x = np.swapaxes(x, 0, axis)
490+
buf = np.empty_like(x[0, ...])
491+
for i in range(len(x) - 1, 0, -1):
492+
j = random_interval(bitgen, i)
493+
if i == j:
494+
continue
495+
buf[...] = x[j]
496+
x[j] = x[i]
497+
x[i] = buf
498+
else:
499+
if axis != 0:
500+
raise NotImplementedError("Axis argument is only supported "
501+
"on ndarray objects")
502+
503+
for i in range(x.size - 1, 0, -1):
504+
j = random_interval(bitgen, i)
505+
x[i], x[j] = x[j], x[i]
506+
507+
return x

‎numba/tests/test_np_randomgen.py‎

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,8 @@
77

88
class TestRandomGenerators(TestCase):
99
def check_numpy_parity(self, distribution_func,
10-
bitgen_instance=None, seed=10):
10+
bitgen_instance=None, seed=10,
11+
test_sizes=None):
1112
distribution_func = numba.njit(distribution_func)
1213
if bitgen_instance is None:
1314
numba_rng_instance = np.random.default_rng(seed=seed)
@@ -16,8 +17,12 @@ def check_numpy_parity(self, distribution_func,
1617
numba_rng_instance = Generator(bitgen_instance(seed))
1718
numpy_rng_instance = Generator(bitgen_instance(seed))
1819

20+
if test_sizes is None:
21+
# Default sizes for testing
22+
test_sizes = [None, (), (100,), (10,20,30)]
23+
1924
# Check parity for different size cases
20-
for size in [None, (), (100,), (10,20,30)]:
25+
for size in test_sizes:
2126
numba_res = distribution_func(numba_rng_instance, size)
2227
numpy_res = distribution_func.py_func(numpy_rng_instance,
2328
size)
@@ -44,6 +49,19 @@ def test_integers(self):
4449
self.check_numpy_parity(dist_func)
4550
self.check_numpy_parity(dist_func, bitgen_instance=MT19937)
4651

52+
def test_choice(self):
53+
dist_func = lambda x, size:x.choice(100, size=size)
54+
self.check_numpy_parity(dist_func)
55+
self.check_numpy_parity(dist_func, bitgen_instance=MT19937)
56+
57+
# Fails right now since .tobytes() method of Numpy arrays
58+
# is not yet supported in Numba
59+
# def test_bytes(self):
60+
# dist_func = lambda x, size:x.bytes(length=size)
61+
# self.check_numpy_parity(dist_func, test_sizes=[1, 1000])
62+
# self.check_numpy_parity(dist_func, bitgen_instance=MT19937,
63+
# test_sizes=[1, 1000])
64+
4765
def test_random(self):
4866
dist_func = lambda x, size:x.random(size=size)
4967
self.check_numpy_parity(dist_func)

0 commit comments

Comments
 (0)