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