| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384 |
- import numpy as np
- import numba as nb
- from numpy.random import PCG64
- from timeit import timeit
- bit_gen = PCG64()
- next_d = bit_gen.cffi.next_double
- state_addr = bit_gen.cffi.state_address
- def normals(n, state):
- out = np.empty(n)
- for i in range((n + 1) // 2):
- x1 = 2.0 * next_d(state) - 1.0
- x2 = 2.0 * next_d(state) - 1.0
- r2 = x1 * x1 + x2 * x2
- while r2 >= 1.0 or r2 == 0.0:
- x1 = 2.0 * next_d(state) - 1.0
- x2 = 2.0 * next_d(state) - 1.0
- r2 = x1 * x1 + x2 * x2
- f = np.sqrt(-2.0 * np.log(r2) / r2)
- out[2 * i] = f * x1
- if 2 * i + 1 < n:
- out[2 * i + 1] = f * x2
- return out
- # Compile using Numba
- normalsj = nb.jit(normals, nopython=True)
- # Must use state address not state with numba
- n = 10000
- def numbacall():
- return normalsj(n, state_addr)
- rg = np.random.Generator(PCG64())
- def numpycall():
- return rg.normal(size=n)
- # Check that the functions work
- r1 = numbacall()
- r2 = numpycall()
- assert r1.shape == (n,)
- assert r1.shape == r2.shape
- t1 = timeit(numbacall, number=1000)
- print(f'{t1:.2f} secs for {n} PCG64 (Numba/PCG64) gaussian randoms')
- t2 = timeit(numpycall, number=1000)
- print(f'{t2:.2f} secs for {n} PCG64 (NumPy/PCG64) gaussian randoms')
- # example 2
- next_u32 = bit_gen.ctypes.next_uint32
- ctypes_state = bit_gen.ctypes.state
- @nb.jit(nopython=True)
- def bounded_uint(lb, ub, state):
- mask = delta = ub - lb
- mask |= mask >> 1
- mask |= mask >> 2
- mask |= mask >> 4
- mask |= mask >> 8
- mask |= mask >> 16
- val = next_u32(state) & mask
- while val > delta:
- val = next_u32(state) & mask
- return lb + val
- print(bounded_uint(323, 2394691, ctypes_state.value))
- @nb.jit(nopython=True)
- def bounded_uints(lb, ub, n, state):
- out = np.empty(n, dtype=np.uint32)
- for i in range(n):
- out[i] = bounded_uint(lb, ub, state)
- bounded_uints(323, 2394691, 10000000, ctypes_state.value)
|
