Aluode/PerceptionLabPortable
0
1"""
2Implement the random and np.random module functions.
3"""
4
5
6import math
7import random
8
9import numpy as np
10
11from llvmlite import ir
12
13from numba.core.cgutils import is_nonelike, is_empty_tuple
14from numba.core.extending import intrinsic, overload, register_jitable
15from numba.core.imputils import Registry
16from numba.core.typing import signature
17from numba.core import types, cgutils
18from numba.core.errors import NumbaTypeError
19from numba.np.random._constants import LONG_MAX
20
21registry = Registry('randomimpl')
22lower = registry.lower
23
24int32_t = ir.IntType(32)
25int64_t = ir.IntType(64)
26def const_int(x):
27 return ir.Constant(int32_t, x)
28double = ir.DoubleType()
29
30N = 624
31N_const = ir.Constant(int32_t, N)
32
33
34# This is the same struct as rnd_state_t in _random.c.
35rnd_state_t = ir.LiteralStructType([
36 # index
37 int32_t,
38 # mt[N]
39 ir.ArrayType(int32_t, N),
40 # has_gauss
41 int32_t,
42 # gauss
43 double,
44 # is_initialized
45 int32_t,
46 ])
47rnd_state_ptr_t = ir.PointerType(rnd_state_t)
48
49
50def get_state_ptr(context, builder, name):
51 """
52 Get a pointer to the given thread-local random state
53 (depending on *name*: "py" or "np").
54 If the state isn't initialized, it is lazily initialized with
55 system entropy.
56 """
57 assert name in ('py', 'np', 'internal')
58 func_name = "numba_get_%s_random_state" % name
59 fnty = ir.FunctionType(rnd_state_ptr_t, ())
60 fn = cgutils.get_or_insert_function(builder.module, fnty, func_name)
61 # These two attributes allow LLVM to hoist the function call
62 # outside of loops.
63 fn.attributes.add('readnone')
64 fn.attributes.add('nounwind')
65 return builder.call(fn, ())
66
67def get_py_state_ptr(context, builder):
68 """
69 Get a pointer to the thread-local Python random state.
70 """
71 return get_state_ptr(context, builder, 'py')
72
73def get_np_state_ptr(context, builder):
74 """
75 Get a pointer to the thread-local Numpy random state.
76 """
77 return get_state_ptr(context, builder, 'np')
78
79def get_internal_state_ptr(context, builder):
80 """
81 Get a pointer to the thread-local internal random state.
82 """
83 return get_state_ptr(context, builder, 'internal')
84
85# Accessors
86def get_index_ptr(builder, state_ptr):
87 return cgutils.gep_inbounds(builder, state_ptr, 0, 0)
88
89def get_array_ptr(builder, state_ptr):
90 return cgutils.gep_inbounds(builder, state_ptr, 0, 1)
91
92def get_has_gauss_ptr(builder, state_ptr):
93 return cgutils.gep_inbounds(builder, state_ptr, 0, 2)
94
95def get_gauss_ptr(builder, state_ptr):
96 return cgutils.gep_inbounds(builder, state_ptr, 0, 3)
97
98def get_rnd_shuffle(builder):
99 """
100 Get the internal function to shuffle the MT taste.
101 """
102 fnty = ir.FunctionType(ir.VoidType(), (rnd_state_ptr_t,))
103 fn = cgutils.get_or_insert_function(builder.function.module, fnty,
104 "numba_rnd_shuffle")
105 fn.args[0].add_attribute("nocapture")
106 return fn
107
108
109def get_next_int32(context, builder, state_ptr):
110 """
111 Get the next int32 generated by the PRNG at *state_ptr*.
112 """
113 idxptr = get_index_ptr(builder, state_ptr)
114 idx = builder.load(idxptr)
115 need_reshuffle = builder.icmp_unsigned('>=', idx, N_const)
116 with cgutils.if_unlikely(builder, need_reshuffle):
117 fn = get_rnd_shuffle(builder)
118 builder.call(fn, (state_ptr,))
119 builder.store(const_int(0), idxptr)
120 idx = builder.load(idxptr)
121 array_ptr = get_array_ptr(builder, state_ptr)
122 y = builder.load(cgutils.gep_inbounds(builder, array_ptr, 0, idx))
123 idx = builder.add(idx, const_int(1))
124 builder.store(idx, idxptr)
125 # Tempering
126 y = builder.xor(y, builder.lshr(y, const_int(11)))
127 y = builder.xor(y, builder.and_(builder.shl(y, const_int(7)),
128 const_int(0x9d2c5680)))
129 y = builder.xor(y, builder.and_(builder.shl(y, const_int(15)),
130 const_int(0xefc60000)))
131 y = builder.xor(y, builder.lshr(y, const_int(18)))
132 return y
133
134def get_next_double(context, builder, state_ptr):
135 """
136 Get the next double generated by the PRNG at *state_ptr*.
137 """
138 # a = rk_random(state) >> 5, b = rk_random(state) >> 6;
139 a = builder.lshr(get_next_int32(context, builder, state_ptr), const_int(5))
140 b = builder.lshr(get_next_int32(context, builder, state_ptr), const_int(6))
141
142 # return (a * 67108864.0 + b) / 9007199254740992.0;
143 a = builder.uitofp(a, double)
144 b = builder.uitofp(b, double)
145 return builder.fdiv(
146 builder.fadd(b, builder.fmul(a, ir.Constant(double, 67108864.0))),
147 ir.Constant(double, 9007199254740992.0))
148
149def get_next_int(context, builder, state_ptr, nbits, is_numpy):
150 """
151 Get the next integer with width *nbits*.
152 """
153 c32 = ir.Constant(nbits.type, 32)
154 def get_shifted_int(nbits):
155 shift = builder.sub(c32, nbits)
156 y = get_next_int32(context, builder, state_ptr)
157
158 # This truncation/extension is safe because 0 < nbits <= 64
159 if nbits.type.width < y.type.width:
160 shift = builder.zext(shift, y.type)
161 elif nbits.type.width > y.type.width:
162 shift = builder.trunc(shift, y.type)
163
164 if is_numpy:
165 # Use the last N bits, to match np.random
166 mask = builder.not_(ir.Constant(y.type, 0))
167 mask = builder.lshr(mask, shift)
168 return builder.and_(y, mask)
169 else:
170 # Use the first N bits, to match CPython random
171 return builder.lshr(y, shift)
172
173 ret = cgutils.alloca_once_value(builder, ir.Constant(int64_t, 0))
174
175 is_32b = builder.icmp_unsigned('<=', nbits, c32)
176 with builder.if_else(is_32b) as (ifsmall, iflarge):
177 with ifsmall:
178 low = get_shifted_int(nbits)
179 builder.store(builder.zext(low, int64_t), ret)
180 with iflarge:
181 # XXX This assumes nbits <= 64
182 if is_numpy:
183 # Get the high bits first to match np.random
184 high = get_shifted_int(builder.sub(nbits, c32))
185 low = get_next_int32(context, builder, state_ptr)
186 if not is_numpy:
187 # Get the high bits second to match CPython random
188 high = get_shifted_int(builder.sub(nbits, c32))
189 total = builder.add(
190 builder.zext(low, int64_t),
191 builder.shl(builder.zext(high, int64_t),
192 ir.Constant(int64_t, 32)))
193 builder.store(total, ret)
194
195 return builder.load(ret)
196
197
198@overload(random.seed)
199def seed_impl(a):
200 if isinstance(a, types.Integer):
201 fn = register_jitable(_seed_impl('py'))
202 def impl(a):
203 return fn(a)
204 return impl
205
206
207@overload(np.random.seed)
208def seed_impl(seed):
209 if isinstance(seed, types.Integer):
210 return _seed_impl('np')
211
212
213def _seed_impl(state_type):
214 @intrinsic
215 def _impl(typingcontext, seed):
216 def codegen(context, builder, sig, args):
217 seed_value, = args
218 fnty = ir.FunctionType(ir.VoidType(), (rnd_state_ptr_t, int32_t))
219 fn = cgutils.get_or_insert_function(builder.function.module, fnty,
220 'numba_rnd_init')
221 builder.call(fn, (get_state_ptr(context, builder, state_type),
222 seed_value))
223 return context.get_constant(types.none, None)
224 return signature(types.void, types.uint32), codegen
225 return lambda seed: _impl(seed)
226
227
228@overload(random.random)
229def random_impl():
230 @intrinsic
231 def _impl(typingcontext):
232 def codegen(context, builder, sig, args):
233 state_ptr = get_state_ptr(context, builder, "py")
234 return get_next_double(context, builder, state_ptr)
235 return signature(types.double), codegen
236 return lambda: _impl()
237
238
239@overload(np.random.random)
240@overload(np.random.random_sample)
241@overload(np.random.sample)
242@overload(np.random.ranf)
243def random_impl0():
244 @intrinsic
245 def _impl(typingcontext):
246 def codegen(context, builder, sig, args):
247 state_ptr = get_state_ptr(context, builder, "np")
248 return get_next_double(context, builder, state_ptr)
249 return signature(types.float64), codegen
250 return lambda: _impl()
251
252
253@overload(np.random.random)
254@overload(np.random.random_sample)
255@overload(np.random.sample)
256@overload(np.random.ranf)
257def random_impl1(size=None):
258 if is_nonelike(size):
259 return lambda size=None: np.random.random()
260 if is_empty_tuple(size):
261 # Handle size = ()
262 return lambda size=None: np.array(np.random.random())
263 if isinstance(size, types.Integer) or (isinstance(size, types.UniTuple)
264 and isinstance(size.dtype,
265 types.Integer)):
266 def _impl(size=None):
267 out = np.empty(size)
268 out_flat = out.flat
269 for idx in range(out.size):
270 out_flat[idx] = np.random.random()
271 return out
272 return _impl
273
274
275@overload(random.gauss)
276@overload(random.normalvariate)
277def gauss_impl(mu, sigma):
278 if isinstance(mu, (types.Float, types.Integer)) and isinstance(
279 sigma, (types.Float, types.Integer)):
280 @intrinsic
281 def _impl(typingcontext, mu, sigma):
282 loc_preprocessor = _double_preprocessor(mu)
283 scale_preprocessor = _double_preprocessor(sigma)
284 return signature(types.float64, mu, sigma),\
285 _gauss_impl("py", loc_preprocessor, scale_preprocessor)
286 return lambda mu, sigma: _impl(mu, sigma)
287
288
289@overload(np.random.standard_normal)
290@overload(np.random.normal)
291def np_gauss_impl0():
292 return lambda: np.random.normal(0.0, 1.0)
293
294
295@overload(np.random.normal)
296def np_gauss_impl1(loc):
297 if isinstance(loc, (types.Float, types.Integer)):
298 return lambda loc: np.random.normal(loc, 1.0)
299
300
301@overload(np.random.normal)
302def np_gauss_impl2(loc, scale):
303 if isinstance(loc, (types.Float, types.Integer)) and isinstance(
304 scale, (types.Float, types.Integer)):
305 @intrinsic
306 def _impl(typingcontext, loc, scale):
307 loc_preprocessor = _double_preprocessor(loc)
308 scale_preprocessor = _double_preprocessor(scale)
309 return signature(types.float64, loc, scale),\
310 _gauss_impl("np", loc_preprocessor, scale_preprocessor)
311 return lambda loc, scale: _impl(loc, scale)
312
313
314@overload(np.random.standard_normal)
315def standard_normal_impl1(size):
316 if is_nonelike(size):
317 return lambda size: np.random.standard_normal()
318 if is_empty_tuple(size):
319 # Handle size = ()
320 return lambda size: np.array(np.random.standard_normal())
321 if isinstance(size, types.Integer) or (isinstance(size, types.UniTuple) and
322 isinstance(size.dtype,
323 types.Integer)):
324 def _impl(size):
325 out = np.empty(size)
326 out_flat = out.flat
327 for idx in range(out.size):
328 out_flat[idx] = np.random.standard_normal()
329 return out
330 return _impl
331
332
333@overload(np.random.normal)
334def np_gauss_impl3(loc, scale, size):
335 if (isinstance(loc, (types.Float, types.Integer)) and isinstance(
336 scale, (types.Float, types.Integer)) and
337 is_nonelike(size)):
338 return lambda loc, scale, size: np.random.normal(loc, scale)
339 if (isinstance(loc, (types.Float, types.Integer)) and isinstance(
340 scale, (types.Float, types.Integer)) and
341 is_empty_tuple(size)):
342 # Handle size = ()
343 return lambda loc, scale, size: np.array(np.random.normal(loc, scale))
344 if (isinstance(loc, (types.Float, types.Integer)) and isinstance(
345 scale, (types.Float, types.Integer)) and
346 (isinstance(size, types.Integer) or (isinstance(size, types.UniTuple)
347 and isinstance(size.dtype,
348 types.Integer)))):
349 def _impl(loc, scale, size):
350 out = np.empty(size)
351 out_flat = out.flat
352 for idx in range(out.size):
353 out_flat[idx] = np.random.normal(loc, scale)
354 return out
355 return _impl
356
357
358def _gauss_pair_impl(_random):
359 def compute_gauss_pair():
360 """
361 Compute a pair of numbers on the normal distribution.
362 """
363 while True:
364 x1 = 2.0 * _random() - 1.0
365 x2 = 2.0 * _random() - 1.0
366 r2 = x1*x1 + x2*x2
367 if r2 < 1.0 and r2 != 0.0:
368 break
369
370 # Box-Muller transform
371 f = math.sqrt(-2.0 * math.log(r2) / r2)
372 return f * x1, f * x2
373 return compute_gauss_pair
374
375
376def _gauss_impl(state, loc_preprocessor, scale_preprocessor):
377 def _impl(context, builder, sig, args):
378 # The type for all computations (either float or double)
379 ty = sig.return_type
380 llty = context.get_data_type(ty)
381 _random = {"py": random.random,
382 "np": np.random.random}[state]
383
384 state_ptr = get_state_ptr(context, builder, state)
385
386 ret = cgutils.alloca_once(builder, llty, name="result")
387
388 gauss_ptr = get_gauss_ptr(builder, state_ptr)
389 has_gauss_ptr = get_has_gauss_ptr(builder, state_ptr)
390 has_gauss = cgutils.is_true(builder, builder.load(has_gauss_ptr))
391 with builder.if_else(has_gauss) as (then, otherwise):
392 with then:
393 # if has_gauss: return it
394 builder.store(builder.load(gauss_ptr), ret)
395 builder.store(const_int(0), has_gauss_ptr)
396 with otherwise:
397 # if not has_gauss: compute a pair of numbers using the Box-Muller
398 # transform; keep one and return the other
399 pair = context.compile_internal(builder,
400 _gauss_pair_impl(_random),
401 signature(types.UniTuple(ty, 2)),
402 ())
403
404 first, second = cgutils.unpack_tuple(builder, pair, 2)
405 builder.store(first, gauss_ptr)
406 builder.store(second, ret)
407 builder.store(const_int(1), has_gauss_ptr)
408
409 mu, sigma = args
410 return builder.fadd(loc_preprocessor(builder, mu),
411 builder.fmul(scale_preprocessor(builder, sigma),
412 builder.load(ret)))
413 return _impl
414
415
416def _double_preprocessor(value):
417 ty = ir.types.DoubleType()
418
419 if isinstance(value, types.Integer):
420 if value.signed:
421 return lambda builder, v: builder.sitofp(v, ty)
422 else:
423 return lambda builder, v: builder.uitofp(v, ty)
424 elif isinstance(value, types.Float):
425 if value.bitwidth != 64:
426 return lambda builder, v: builder.fpext(v, ty)
427 else:
428 return lambda _builder, v: v
429 else:
430 raise NumbaTypeError("Cannot convert {} to floating point type" % value)
431
432
433@overload(random.getrandbits)
434def getrandbits_impl(k):
435 if isinstance(k, types.Integer):
436 @intrinsic
437 def _impl(typingcontext, k):
438 def codegen(context, builder, sig, args):
439 nbits, = args
440
441 too_large = builder.icmp_unsigned(">=", nbits, const_int(65))
442 too_small = builder.icmp_unsigned("==", nbits, const_int(0))
443 with cgutils.if_unlikely(builder, builder.or_(too_large,
444 too_small)):
445 msg = "getrandbits() limited to 64 bits"
446 context.call_conv.return_user_exc(builder, OverflowError,
447 (msg,))
448 state_ptr = get_state_ptr(context, builder, "py")
449 return get_next_int(context, builder, state_ptr, nbits, False)
450 return signature(types.uint64, k), codegen
451 return lambda k: _impl(k)
452
453
454def _randrange_impl(context, builder, start, stop, step, ty, signed, state):
455 state_ptr = get_state_ptr(context, builder, state)
456 zero = ir.Constant(ty, 0)
457 one = ir.Constant(ty, 1)
458 nptr = cgutils.alloca_once(builder, ty, name="n")
459
460 # n = stop - start
461 builder.store(builder.sub(stop, start), nptr)
462
463 with builder.if_then(builder.icmp_signed('<', step, zero)):
464 # n = (n + step + 1) // step
465 w = builder.add(builder.add(builder.load(nptr), step), one)
466 n = builder.sdiv(w, step)
467 builder.store(n, nptr)
468 with builder.if_then(builder.icmp_signed('>', step, one)):
469 # n = (n + step - 1) // step
470 w = builder.sub(builder.add(builder.load(nptr), step), one)
471 n = builder.sdiv(w, step)
472 builder.store(n, nptr)
473
474 n = builder.load(nptr)
475 with cgutils.if_unlikely(builder, builder.icmp_signed('<=', n, zero)):
476 # n <= 0
477 msg = "empty range for randrange()"
478 context.call_conv.return_user_exc(builder, ValueError, (msg,))
479
480 fnty = ir.FunctionType(ty, [ty, cgutils.true_bit.type])
481 fn = cgutils.get_or_insert_function(builder.function.module, fnty,
482 "llvm.ctlz.%s" % ty)
483 # Since the upper bound is exclusive, we need to subtract one before
484 # calculating the number of bits. This leads to a special case when
485 # n == 1; there's only one possible result, so we don't need bits from
486 # the PRNG. This case is handled separately towards the end of this
487 # function. CPython's implementation is simpler and just runs another
488 # iteration of the while loop when the resulting number is too large
489 # instead of subtracting one, to avoid needing to handle a special
490 # case. Thus, we only perform this subtraction for the NumPy case.
491 nm1 = builder.sub(n, one) if state == "np" else n
492 nbits = builder.trunc(builder.call(fn, [nm1, cgutils.true_bit]), int32_t)
493 nbits = builder.sub(ir.Constant(int32_t, ty.width), nbits)
494
495 rptr = cgutils.alloca_once(builder, ty, name="r")
496
497 def get_num():
498 bbwhile = builder.append_basic_block("while")
499 bbend = builder.append_basic_block("while.end")
500 builder.branch(bbwhile)
501
502 builder.position_at_end(bbwhile)
503 r = get_next_int(context, builder, state_ptr, nbits, state == "np")
504 r = builder.trunc(r, ty)
505 too_large = builder.icmp_signed('>=', r, n)
506 builder.cbranch(too_large, bbwhile, bbend)
507
508 builder.position_at_end(bbend)
509 builder.store(r, rptr)
510
511 if state == "np":
512 # Handle n == 1 case, per previous comment.
513 with builder.if_else(builder.icmp_signed('==', n, one)) as (is_one, is_not_one):
514 with is_one:
515 builder.store(zero, rptr)
516 with is_not_one:
517 get_num()
518 else:
519 get_num()
520
521 return builder.add(start, builder.mul(builder.load(rptr), step))
522
523
524@overload(random.randrange)
525def randrange_impl_1(start):
526 if isinstance(start, types.Integer):
527 return lambda start: random.randrange(0, start, 1)
528
529
530@overload(random.randrange)
531def randrange_impl_2(start, stop):
532 if isinstance(start, types.Integer) and isinstance(stop, types.Integer):
533 return lambda start, stop: random.randrange(start, stop, 1)
534
535
536def _randrange_preprocessor(bitwidth, ty):
537 if ty.bitwidth != bitwidth:
538 return (ir.IRBuilder.sext if ty.signed
539 else ir.IRBuilder.zext)
540 else:
541 return lambda _builder, v, _ty: v
542
543
544@overload(random.randrange)
545def randrange_impl_3(start, stop, step):
546 if (isinstance(start, types.Integer) and isinstance(stop, types.Integer) and
547 isinstance(step, types.Integer)):
548 signed = max(start.signed, stop.signed, step.signed)
549 bitwidth = max(start.bitwidth, stop.bitwidth, step.bitwidth)
550 int_ty = types.Integer.from_bitwidth(bitwidth, signed)
551 llvm_type = ir.IntType(bitwidth)
552
553 start_preprocessor = _randrange_preprocessor(bitwidth, start)
554 stop_preprocessor = _randrange_preprocessor(bitwidth, stop)
555 step_preprocessor = _randrange_preprocessor(bitwidth, step)
556
557 @intrinsic
558 def _impl(typingcontext, start, stop, step):
559 def codegen(context, builder, sig, args):
560 start, stop, step = args
561
562 start = start_preprocessor(builder, start, llvm_type)
563 stop = stop_preprocessor(builder, stop, llvm_type)
564 step = step_preprocessor(builder, step, llvm_type)
565 return _randrange_impl(context, builder, start, stop, step,
566 llvm_type, signed, 'py')
567 return signature(int_ty, start, stop, step), codegen
568 return lambda start, stop, step: _impl(start, stop, step)
569
570
571@overload(random.randint)
572def randint_impl_1(a, b):
573 if isinstance(a, types.Integer) and isinstance(b, types.Integer):
574 return lambda a, b: random.randrange(a, b + 1, 1)
575
576
577@overload(np.random.randint)
578def np_randint_impl_1(low):
579 if isinstance(low, types.Integer):
580 return lambda low: np.random.randint(0, low)
581
582
583@overload(np.random.randint)
584def np_randint_impl_2(low, high):
585 if isinstance(low, types.Integer) and isinstance(high, types.Integer):
586 signed = max(low.signed, high.signed)
587 bitwidth = max(low.bitwidth, high.bitwidth)
588 int_ty = types.Integer.from_bitwidth(bitwidth, signed)
589 llvm_type = ir.IntType(bitwidth)
590
591 start_preprocessor = _randrange_preprocessor(bitwidth, low)
592 stop_preprocessor = _randrange_preprocessor(bitwidth, high)
593
594 @intrinsic
595 def _impl(typingcontext, low, high):
596 def codegen(context, builder, sig, args):
597 start, stop = args
598
599 start = start_preprocessor(builder, start, llvm_type)
600 stop = stop_preprocessor(builder, stop, llvm_type)
601 step = ir.Constant(llvm_type, 1)
602 return _randrange_impl(context, builder, start, stop, step,
603 llvm_type, signed, 'np')
604 return signature(int_ty, low, high), codegen
605 return lambda low, high: _impl(low, high)
606
607
608@overload(np.random.randint)
609def np_randint_impl_3(low, high, size):
610 if (isinstance(low, types.Integer) and isinstance(high, types.Integer) and
611 is_nonelike(size)):
612 return lambda low, high, size: np.random.randint(low, high)
613 if (isinstance(low, types.Integer) and isinstance(high, types.Integer) and
614 is_empty_tuple(size)):
615 # Handle size = ()
616 return lambda low, high, size: np.array(np.random.randint(low, high))
617 if (isinstance(low, types.Integer) and isinstance(high, types.Integer) and
618 (isinstance(size, types.Integer) or (isinstance(size, types.UniTuple)
619 and isinstance(size.dtype,
620 types.Integer)))):
621 bitwidth = max(low.bitwidth, high.bitwidth)
622 result_type = getattr(np, f'int{bitwidth}')
623
624 def _impl(low, high, size):
625 out = np.empty(size, dtype=result_type)
626 out_flat = out.flat
627 for idx in range(out.size):
628 out_flat[idx] = np.random.randint(low, high)
629 return out
630 return _impl
631
632
633@overload(np.random.uniform)
634def np_uniform_impl0():
635 return lambda: np.random.uniform(0.0, 1.0)
636
637
638@overload(random.uniform)
639def uniform_impl2(a, b):
640 if isinstance(a, (types.Float, types.Integer)) and isinstance(
641 b, (types.Float, types.Integer)):
642 @intrinsic
643 def _impl(typingcontext, a, b):
644 low_preprocessor = _double_preprocessor(a)
645 high_preprocessor = _double_preprocessor(b)
646 return signature(types.float64, a, b), uniform_impl(
647 'py', low_preprocessor, high_preprocessor)
648 return lambda a, b: _impl(a, b)
649
650
651@overload(np.random.uniform)
652def np_uniform_impl2(low, high):
653 if isinstance(low, (types.Float, types.Integer)) and isinstance(
654 high, (types.Float, types.Integer)):
655 @intrinsic
656 def _impl(typingcontext, low, high):
657 low_preprocessor = _double_preprocessor(low)
658 high_preprocessor = _double_preprocessor(high)
659 return signature(types.float64, low, high), uniform_impl(
660 'np', low_preprocessor, high_preprocessor)
661 return lambda low, high: _impl(low, high)
662
663
664def uniform_impl(state, a_preprocessor, b_preprocessor):
665 def impl(context, builder, sig, args):
666 state_ptr = get_state_ptr(context, builder, state)
667 a, b = args
668 a = a_preprocessor(builder, a)
669 b = b_preprocessor(builder, b)
670 width = builder.fsub(b, a)
671 r = get_next_double(context, builder, state_ptr)
672 return builder.fadd(a, builder.fmul(width, r))
673 return impl
674
675
676@overload(np.random.uniform)
677def np_uniform_impl3(low, high, size):
678 if (isinstance(low, (types.Float, types.Integer)) and isinstance(
679 high, (types.Float, types.Integer)) and
680 is_nonelike(size)):
681 return lambda low, high, size: np.random.uniform(low, high)
682 if (isinstance(low, (types.Float, types.Integer)) and isinstance(
683 high, (types.Float, types.Integer)) and
684 is_empty_tuple(size)):
685 # When calling np.random.uniform with size = (), the returned value isn't a
686 # float like when size = None. Instead, it's an array of shape ()
687 return lambda low, high, size: np.array(np.random.uniform(low, high))
688 if (isinstance(low, (types.Float, types.Integer)) and isinstance(
689 high, (types.Float, types.Integer)) and
690 (isinstance(size, types.Integer) or (isinstance(size, types.UniTuple)
691 and isinstance(size.dtype,
692 types.Integer)))):
693 def _impl(low, high, size):
694 out = np.empty(size)
695 out_flat = out.flat
696 for idx in range(out.size):
697 out_flat[idx] = np.random.uniform(low, high)
698 return out
699 return _impl
700
701
702@overload(random.triangular)
703def triangular_impl_2(low, high):
704 def _impl(low, high):
705 u = random.random()
706 c = 0.5
707 if u > c:
708 u = 1.0 - u
709 low, high = high, low
710 return low + (high - low) * math.sqrt(u * c)
711
712 if isinstance(low, (types.Float, types.Integer)) and isinstance(
713 high, (types.Float, types.Integer)):
714 return _impl
715
716
717@overload(random.triangular)
718def triangular_impl_3(low, high, mode):
719 if (isinstance(low, (types.Float, types.Integer)) and isinstance(
720 high, (types.Float, types.Integer)) and
721 isinstance(mode, (types.Float, types.Integer))):
722 def _impl(low, high, mode):
723 if high == low:
724 return low
725 u = random.random()
726 c = (mode - low) / (high - low)
727 if u > c:
728 u = 1.0 - u
729 c = 1.0 - c
730 low, high = high, low
731 return low + (high - low) * math.sqrt(u * c)
732
733 return _impl
734
735
736@overload(np.random.triangular)
737def triangular_impl_3(left, mode, right):
738 if (isinstance(left, (types.Float, types.Integer)) and isinstance(
739 mode, (types.Float, types.Integer)) and
740 isinstance(right, (types.Float, types.Integer))):
741 def _impl(left, mode, right):
742 if right == left:
743 return left
744 u = np.random.random()
745 c = (mode - left) / (right - left)
746 if u > c:
747 u = 1.0 - u
748 c = 1.0 - c
749 left, right = right, left
750 return left + (right - left) * math.sqrt(u * c)
751
752 return _impl
753
754
755@overload(np.random.triangular)
756def triangular_impl(left, mode, right, size=None):
757 if is_nonelike(size):
758 return lambda left, mode, right, size=None: np.random.triangular(left,
759 mode,
760 right)
761 if is_empty_tuple(size):
762 # Handle size = ()
763 return lambda left, mode, right, size=None: np.array(
764 np.random.triangular(left, mode, right)
765 )
766 if (isinstance(size, types.Integer) or (isinstance(size, types.UniTuple) and
767 isinstance(size.dtype,
768 types.Integer))):
769 def _impl(left, mode, right, size=None):
770 out = np.empty(size)
771 out_flat = out.flat
772 for idx in range(out.size):
773 out_flat[idx] = np.random.triangular(left, mode, right)
774 return out
775 return _impl
776
777
778@overload(random.gammavariate)
779def gammavariate_impl(alpha, beta):
780 if isinstance(alpha, (types.Float, types.Integer)) and isinstance(
781 beta, (types.Float, types.Integer)):
782 return _gammavariate_impl(random.random)
783
784
785@overload(np.random.standard_gamma)
786@overload(np.random.gamma)
787def ol_np_random_gamma1(shape):
788 if isinstance(shape, (types.Float, types.Integer)):
789 return lambda shape: np.random.gamma(shape, 1.0)
790
791
792@overload(np.random.gamma)
793def ol_np_random_gamma2(shape, scale):
794 if isinstance(shape, (types.Float, types.Integer)) and isinstance(
795 scale, (types.Float, types.Integer)):
796 fn = register_jitable(_gammavariate_impl(np.random.random))
797 def impl(shape, scale):
798 return fn(shape, scale)
799 return impl
800
801
802def _gammavariate_impl(_random):
803 def _impl(alpha, beta):
804 """Gamma distribution. Taken from CPython.
805 """
806 SG_MAGICCONST = 1.0 + math.log(4.5)
807 # alpha > 0, beta > 0, mean is alpha*beta, variance is alpha*beta**2
808
809 # Warning: a few older sources define the gamma distribution in terms
810 # of alpha > -1.0
811 if alpha <= 0.0 or beta <= 0.0:
812 raise ValueError('gammavariate: alpha and beta must be > 0.0')
813
814 if alpha > 1.0:
815 # Uses R.C.H. Cheng, "The generation of Gamma
816 # variables with non-integral shape parameters",
817 # Applied Statistics, (1977), 26, No. 1, p71-74
818 ainv = math.sqrt(2.0 * alpha - 1.0)
819 bbb = alpha - math.log(4.0)
820 ccc = alpha + ainv
821
822 while 1:
823 u1 = _random()
824 if not 1e-7 < u1 < .9999999:
825 continue
826 u2 = 1.0 - _random()
827 v = math.log(u1/(1.0-u1))/ainv
828 x = alpha*math.exp(v)
829 z = u1*u1*u2
830 r = bbb+ccc*v-x
831 if r + SG_MAGICCONST - 4.5*z >= 0.0 or r >= math.log(z):
832 return x * beta
833
834 elif alpha == 1.0:
835 # expovariate(1)
836
837 # Adjust due to cpython
838 # commit 63d152232e1742660f481c04a811f824b91f6790
839 return -math.log(1.0 - _random()) * beta
840
841 else: # alpha is between 0 and 1 (exclusive)
842 # Uses ALGORITHM GS of Statistical Computing - Kennedy & Gentle
843 while 1:
844 u = _random()
845 b = (math.e + alpha)/math.e
846 p = b*u
847 if p <= 1.0:
848 x = p ** (1.0/alpha)
849 else:
850 x = -math.log((b-p)/alpha)
851 u1 = _random()
852 if p > 1.0:
853 if u1 <= x ** (alpha - 1.0):
854 break
855 elif u1 <= math.exp(-x):
856 break
857 return x * beta
858 return _impl
859
860
861@overload(np.random.gamma)
862def gamma_impl(shape, scale, size):
863 if is_nonelike(size):
864 return lambda shape, scale, size: np.random.gamma(shape, scale)
865 if is_empty_tuple(size):
866 # Handle size = ()
867 return lambda shape, scale, size: np.array(np.random.gamma(shape, scale))
868 if isinstance(size, types.Integer) or (isinstance(size, types.UniTuple) and
869 isinstance(size.dtype,
870 types.Integer)):
871 def _impl(shape, scale, size):
872 out = np.empty(size)
873 out_flat = out.flat
874 for idx in range(out.size):
875 out_flat[idx] = np.random.gamma(shape, scale)
876 return out
877 return _impl
878
879
880@overload(np.random.standard_gamma)
881def standard_gamma_impl(shape, size):
882 if is_nonelike(size):
883 return lambda shape, size: np.random.standard_gamma(shape)
884 if is_empty_tuple(size):
885 # Handle size = ()
886 return lambda shape, size: np.array(np.random.standard_gamma(shape))
887 if (isinstance(size, types.Integer) or (isinstance(size, types.UniTuple)
888 and isinstance(size.dtype,
889 types.Integer))):
890 def _impl(shape, size):
891 out = np.empty(size)
892 out_flat = out.flat
893 for idx in range(out.size):
894 out_flat[idx] = np.random.standard_gamma(shape)
895 return out
896 return _impl
897
898
899@overload(random.betavariate)
900def betavariate_impl(alpha, beta):
901 if isinstance(alpha, (types.Float, types.Integer)) and isinstance(
902 beta, (types.Float, types.Integer)):
903 return _betavariate_impl(random.gammavariate)
904
905
906@overload(np.random.beta)
907def ol_np_random_beta(a, b):
908 if isinstance(a, (types.Float, types.Integer)) and isinstance(
909 b, (types.Float, types.Integer)):
910 fn = register_jitable(_betavariate_impl(np.random.gamma))
911 def impl(a, b):
912 return fn(a, b)
913 return impl
914
915
916def _betavariate_impl(gamma):
917 def _impl(alpha, beta):
918 """Beta distribution. Taken from CPython.
919 """
920 # This version due to Janne Sinkkonen, and matches all the std
921 # texts (e.g., Knuth Vol 2 Ed 3 pg 134 "the beta distribution").
922 y = gamma(alpha, 1.)
923 if y == 0.0:
924 return 0.0
925 else:
926 return y / (y + gamma(beta, 1.))
927 return _impl
928
929
930@overload(np.random.beta)
931def beta_impl(a, b, size):
932 if is_nonelike(size):
933 return lambda a, b, size: np.random.beta(a, b)
934 if is_empty_tuple(size):
935 # When calling np.random.beta with size = (), the returned value isn't a
936 # float like when size = None. Instead, it's an array of shape ()
937 return lambda a, b, size: np.array(np.random.beta(a, b))
938 if (isinstance(size, types.Integer) or (isinstance(size, types.UniTuple)
939 and isinstance(size.dtype,
940 types.Integer))):
941 def _impl(a, b, size):
942 out = np.empty(size)
943 out_flat = out.flat
944 for idx in range(out.size):
945 out_flat[idx] = np.random.beta(a, b)
946 return out
947 return _impl
948
949
950@overload(random.expovariate)
951def expovariate_impl(lambd):
952 if isinstance(lambd, types.Float):
953 def _impl(lambd):
954 """Exponential distribution. Taken from CPython.
955 """
956 # lambd: rate lambd = 1/mean
957 # ('lambda' is a Python reserved word)
958
959 # we use 1-random() instead of random() to preclude the
960 # possibility of taking the log of zero.
961 return -math.log(1.0 - random.random()) / lambd
962
963 return _impl
964
965
966@overload(np.random.exponential)
967def exponential_impl(scale):
968 if isinstance(scale, (types.Float, types.Integer)):
969 def _impl(scale):
970 return -math.log(1.0 - np.random.random()) * scale
971 return _impl
972
973
974@overload(np.random.exponential)
975def exponential_impl(scale, size):
976 if is_nonelike(size):
977 return lambda scale, size: np.random.exponential(scale)
978 if is_empty_tuple(size):
979 # Handle size = ()
980 return lambda scale, size: np.array(np.random.exponential(scale))
981 if (isinstance(size, types.Integer) or (isinstance(size, types.UniTuple) and
982 isinstance(size.dtype,
983 types.Integer))):
984 def _impl(scale, size):
985 out = np.empty(size)
986 out_flat = out.flat
987 for idx in range(out.size):
988 out_flat[idx] = np.random.exponential(scale)
989 return out
990 return _impl
991
992
993@overload(np.random.standard_exponential)
994@overload(np.random.exponential)
995def exponential_impl():
996 def _impl():
997 return -math.log(1.0 - np.random.random())
998 return _impl
999
1000
1001@overload(np.random.standard_exponential)
1002def standard_exponential_impl(size):
1003 if is_nonelike(size):
1004 return lambda size: np.random.standard_exponential()
1005 if is_empty_tuple(size):
1006 # Handle size = ()
1007 return lambda size: np.array(np.random.standard_exponential())
1008 if (isinstance(size, types.Integer) or
1009 (isinstance(size, types.UniTuple) and isinstance(size.dtype,
1010 types.Integer))
1011 ):
1012 def _impl(size):
1013 out = np.empty(size)
1014 out_flat = out.flat
1015 for idx in range(out.size):
1016 out_flat[idx] = np.random.standard_exponential()
1017 return out
1018 return _impl
1019
1020
1021@overload(np.random.lognormal)
1022def np_lognormal_impl0():
1023 return lambda: np.random.lognormal(0.0, 1.0)
1024
1025
1026@overload(np.random.lognormal)
1027def np_log_normal_impl1(mean):
1028 if isinstance(mean, (types.Float, types.Integer)):
1029 return lambda mean: np.random.lognormal(mean, 1.0)
1030
1031
1032@overload(np.random.lognormal)
1033def np_log_normal_impl2(mean, sigma):
1034 if isinstance(mean, (types.Float, types.Integer)) and isinstance(
1035 sigma, (types.Float, types.Integer)):
1036 fn = register_jitable(_lognormvariate_impl(np.random.normal))
1037 return lambda mean, sigma: fn(mean, sigma)
1038
1039
1040@overload(np.random.lognormal)
1041def lognormal_impl(mean, sigma, size):
1042 if is_nonelike(size):
1043 return lambda mean, sigma, size: np.random.lognormal(mean, sigma)
1044 if is_empty_tuple(size):
1045 # Handle size = ()
1046 return lambda mean, sigma, size: np.array(np.random.lognormal(mean, sigma))
1047 if (isinstance(size, types.Integer) or (isinstance(size, types.UniTuple) and
1048 isinstance(size.dtype,
1049 types.Integer))):
1050 def _impl(mean, sigma, size):
1051 out = np.empty(size)
1052 out_flat = out.flat
1053 for idx in range(out.size):
1054 out_flat[idx] = np.random.lognormal(mean, sigma)
1055 return out
1056 return _impl
1057
1058
1059@overload(random.lognormvariate)
1060def lognormvariate_impl(mu, sigma):
1061 if isinstance(mu, types.Float) and isinstance(sigma, types.Float):
1062 fn = register_jitable(_lognormvariate_impl(random.gauss))
1063 return lambda mu, sigma: fn(mu, sigma)
1064
1065
1066def _lognormvariate_impl(_gauss):
1067 return lambda mu, sigma: math.exp(_gauss(mu, sigma))
1068
1069
1070@overload(random.paretovariate)
1071def paretovariate_impl(alpha):
1072 if isinstance(alpha, types.Float):
1073 def _impl(alpha):
1074 """Pareto distribution. Taken from CPython."""
1075 # Jain, pg. 495
1076 u = 1.0 - random.random()
1077 return 1.0 / u ** (1.0/alpha)
1078
1079 return _impl
1080
1081
1082@overload(np.random.pareto)
1083def pareto_impl(a):
1084 if isinstance(a, types.Float):
1085 def _impl(a):
1086 # Same as paretovariate() - 1.
1087 u = 1.0 - np.random.random()
1088 return 1.0 / u ** (1.0/a) - 1
1089
1090 return _impl
1091
1092
1093@overload(np.random.pareto)
1094def pareto_impl(a, size):
1095 if is_nonelike(size):
1096 return lambda a, size: np.random.pareto(a)
1097 if is_empty_tuple(size):
1098 # Handle size = ()
1099 return lambda a, size: np.array(np.random.pareto(a))
1100 if (isinstance(size, types.Integer) or (isinstance(size, types.UniTuple) and
1101 isinstance(size.dtype,
1102 types.Integer))):
1103 def _impl(a, size):
1104 out = np.empty(size)
1105 out_flat = out.flat
1106 for idx in range(out.size):
1107 out_flat[idx] = np.random.pareto(a)
1108 return out
1109 return _impl
1110
1111
1112@overload(random.weibullvariate)
1113def weibullvariate_impl(alpha, beta):
1114 if isinstance(alpha, (types.Float, types.Integer)) and isinstance(
1115 beta, (types.Float, types.Integer)):
1116 def _impl(alpha, beta):
1117 """Weibull distribution. Taken from CPython."""
1118 # Jain, pg. 499; bug fix courtesy Bill Arms
1119 u = 1.0 - random.random()
1120 return alpha * (-math.log(u)) ** (1.0/beta)
1121
1122 return _impl
1123
1124
1125@overload(np.random.weibull)
1126def weibull_impl(a):
1127 if isinstance(a, (types.Float, types.Integer)):
1128 def _impl(a):
1129 # Same as weibullvariate(1.0, a)
1130 u = 1.0 - np.random.random()
1131 return (-math.log(u)) ** (1.0/a)
1132
1133 return _impl
1134
1135
1136@overload(np.random.weibull)
1137def weibull_impl2(a, size):
1138 if is_nonelike(size):
1139 return lambda a, size: np.random.weibull(a)
1140 if is_empty_tuple(size):
1141 # Handle size = ()
1142 return lambda a, size: np.array(np.random.weibull(a))
1143 if (isinstance(size, types.Integer) or (isinstance(size, types.UniTuple) and
1144 isinstance(size.dtype,
1145 types.Integer))):
1146 def _impl(a, size):
1147 out = np.empty(size)
1148 out_flat = out.flat
1149 for idx in range(out.size):
1150 out_flat[idx] = np.random.weibull(a)
1151 return out
1152 return _impl
1153
1154
1155@overload(random.vonmisesvariate)
1156def vonmisesvariate_impl(mu, kappa):
1157 if isinstance(mu, types.Float) and isinstance(kappa, types.Float):
1158 return _vonmisesvariate_impl(random.random)
1159
1160
1161@overload(np.random.vonmises)
1162def vonmisesvariate_impl(mu, kappa):
1163 if isinstance(mu, types.Float) and isinstance(kappa, types.Float):
1164 return _vonmisesvariate_impl(np.random.random)
1165
1166
1167def _vonmisesvariate_impl(_random):
1168 def _impl(mu, kappa):
1169 """Circular data distribution. Taken from CPython.
1170 Note the algorithm in Python 2.6 and Numpy is different:
1171 http://bugs.python.org/issue17141
1172 """
1173 # mu: mean angle (in radians between 0 and 2*pi)
1174 # kappa: concentration parameter kappa (>= 0)
1175 # if kappa = 0 generate uniform random angle
1176
1177 # Based upon an algorithm published in: Fisher, N.I.,
1178 # "Statistical Analysis of Circular Data", Cambridge
1179 # University Press, 1993.
1180
1181 # Thanks to Magnus Kessler for a correction to the
1182 # implementation of step 4.
1183 if kappa <= 1e-6:
1184 return 2.0 * math.pi * _random()
1185
1186 s = 0.5 / kappa
1187 r = s + math.sqrt(1.0 + s * s)
1188
1189 while 1:
1190 u1 = _random()
1191 z = math.cos(math.pi * u1)
1192
1193 d = z / (r + z)
1194 u2 = _random()
1195 if u2 < 1.0 - d * d or u2 <= (1.0 - d) * math.exp(d):
1196 break
1197
1198 q = 1.0 / r
1199 f = (q + z) / (1.0 + q * z)
1200 u3 = _random()