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