codekingpro/portable-devtools

1#cython: language_level=3
2
3from libc.stdint cimport uint32_t
4from cpython.pycapsule cimport PyCapsule_IsValid, PyCapsule_GetPointer
5
6import numpy as np
7cimport numpy as np
8cimport cython
9
10from numpy.random cimport bitgen_t
11from numpy.random import PCG64
12
13np.import_array()
14
15
16@cython.boundscheck(False)
17@cython.wraparound(False)
18def uniform_mean(Py_ssize_t n):
19    cdef Py_ssize_t i
20    cdef bitgen_t *rng
21    cdef const char *capsule_name = "BitGenerator"
22    cdef double[::1] random_values
23    cdef np.ndarray randoms
24
25    x = PCG64()
26    capsule = x.capsule
27    if not PyCapsule_IsValid(capsule, capsule_name):
28        raise ValueError("Invalid pointer to anon_func_state")
29    rng = <bitgen_t *> PyCapsule_GetPointer(capsule, capsule_name)
30    random_values = np.empty(n)
31    # Best practice is to acquire the lock whenever generating random values.
32    # This prevents other threads from modifying the state. Acquiring the lock
33    # is only necessary if the GIL is also released, as in this example.
34    with x.lock, nogil:
35        for i in range(n):
36            random_values[i] = rng.next_double(rng.state)
37    randoms = np.asarray(random_values)
38    return randoms.mean()
39
40
41# This function is declared nogil so it can be used without the GIL below
42cdef uint32_t bounded_uint(uint32_t lb, uint32_t ub, bitgen_t *rng) nogil:
43    cdef uint32_t mask, delta, val
44    mask = delta = ub - lb
45    mask |= mask >> 1
46    mask |= mask >> 2
47    mask |= mask >> 4
48    mask |= mask >> 8
49    mask |= mask >> 16
50
51    val = rng.next_uint32(rng.state) & mask
52    while val > delta:
53        val = rng.next_uint32(rng.state) & mask
54
55    return lb + val
56
57
58@cython.boundscheck(False)
59@cython.wraparound(False)
60def bounded_uints(uint32_t lb, uint32_t ub, Py_ssize_t n):
61    cdef Py_ssize_t i
62    cdef bitgen_t *rng
63    cdef uint32_t[::1] out
64    cdef const char *capsule_name = "BitGenerator"
65
66    x = PCG64()
67    out = np.empty(n, dtype=np.uint32)
68    capsule = x.capsule
69
70    if not PyCapsule_IsValid(capsule, capsule_name):
71        raise ValueError("Invalid pointer to anon_func_state")
72    rng = <bitgen_t *>PyCapsule_GetPointer(capsule, capsule_name)
73
74    with x.lock, nogil:
75        for i in range(n):
76            out[i] = bounded_uint(lb, ub, rng)
77    return np.asarray(out)
78