codekingpro/portable-devtools

1from timeit import timeit
2
3import numba as nb
4
5import numpy as np
6from numpy.random import PCG64
7
8bit_gen = PCG64()
9next_d = bit_gen.cffi.next_double
10state_addr = bit_gen.cffi.state_address
11
12def normals(n, state):
13    out = np.empty(n)
14    for i in range((n + 1) // 2):
15        x1 = 2.0 * next_d(state) - 1.0
16        x2 = 2.0 * next_d(state) - 1.0
17        r2 = x1 * x1 + x2 * x2
18        while r2 >= 1.0 or r2 == 0.0:
19            x1 = 2.0 * next_d(state) - 1.0
20            x2 = 2.0 * next_d(state) - 1.0
21            r2 = x1 * x1 + x2 * x2
22        f = np.sqrt(-2.0 * np.log(r2) / r2)
23        out[2 * i] = f * x1
24        if 2 * i + 1 < n:
25            out[2 * i + 1] = f * x2
26    return out
27
28
29# Compile using Numba
30normalsj = nb.jit(normals, nopython=True)
31# Must use state address not state with numba
32n = 10000
33
34def numbacall():
35    return normalsj(n, state_addr)
36
37
38rg = np.random.Generator(PCG64())
39
40def numpycall():
41    return rg.normal(size=n)
42
43
44# Check that the functions work
45r1 = numbacall()
46r2 = numpycall()
47assert r1.shape == (n,)
48assert r1.shape == r2.shape
49
50t1 = timeit(numbacall, number=1000)
51print(f'{t1:.2f} secs for {n} PCG64 (Numba/PCG64) gaussian randoms')
52t2 = timeit(numpycall, number=1000)
53print(f'{t2:.2f} secs for {n} PCG64 (NumPy/PCG64) gaussian randoms')
54
55# example 2
56
57next_u32 = bit_gen.ctypes.next_uint32
58ctypes_state = bit_gen.ctypes.state
59
60@nb.jit(nopython=True)
61def bounded_uint(lb, ub, state):
62    mask = delta = ub - lb
63    mask |= mask >> 1
64    mask |= mask >> 2
65    mask |= mask >> 4
66    mask |= mask >> 8
67    mask |= mask >> 16
68
69    val = next_u32(state) & mask
70    while val > delta:
71        val = next_u32(state) & mask
72
73    return lb + val
74
75
76print(bounded_uint(323, 2394691, ctypes_state.value))
77
78
79@nb.jit(nopython=True)
80def bounded_uints(lb, ub, n, state):
81    out = np.empty(n, dtype=np.uint32)
82    for i in range(n):
83        out[i] = bounded_uint(lb, ub, state)
84
85
86bounded_uints(323, 2394691, 10000000, ctypes_state.value)
87