aboutsummaryrefslogtreecommitdiff
path: root/.venv/lib/python3.12/site-packages/numpy/random/_examples/cython/extending.pyx
diff options
context:
space:
mode:
Diffstat (limited to '.venv/lib/python3.12/site-packages/numpy/random/_examples/cython/extending.pyx')
-rw-r--r--.venv/lib/python3.12/site-packages/numpy/random/_examples/cython/extending.pyx78
1 files changed, 78 insertions, 0 deletions
diff --git a/.venv/lib/python3.12/site-packages/numpy/random/_examples/cython/extending.pyx b/.venv/lib/python3.12/site-packages/numpy/random/_examples/cython/extending.pyx
new file mode 100644
index 00000000..30efd744
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/numpy/random/_examples/cython/extending.pyx
@@ -0,0 +1,78 @@
+#!/usr/bin/env python3
+#cython: language_level=3
+
+from libc.stdint cimport uint32_t
+from cpython.pycapsule cimport PyCapsule_IsValid, PyCapsule_GetPointer
+
+import numpy as np
+cimport numpy as np
+cimport cython
+
+from numpy.random cimport bitgen_t
+from numpy.random import PCG64
+
+np.import_array()
+
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def uniform_mean(Py_ssize_t n):
+ cdef Py_ssize_t i
+ cdef bitgen_t *rng
+ cdef const char *capsule_name = "BitGenerator"
+ cdef double[::1] random_values
+ cdef np.ndarray randoms
+
+ x = PCG64()
+ capsule = x.capsule
+ if not PyCapsule_IsValid(capsule, capsule_name):
+ raise ValueError("Invalid pointer to anon_func_state")
+ rng = <bitgen_t *> PyCapsule_GetPointer(capsule, capsule_name)
+ random_values = np.empty(n)
+ # Best practice is to acquire the lock whenever generating random values.
+ # This prevents other threads from modifying the state. Acquiring the lock
+ # is only necessary if the GIL is also released, as in this example.
+ with x.lock, nogil:
+ for i in range(n):
+ random_values[i] = rng.next_double(rng.state)
+ randoms = np.asarray(random_values)
+ return randoms.mean()
+
+
+# This function is declared nogil so it can be used without the GIL below
+cdef uint32_t bounded_uint(uint32_t lb, uint32_t ub, bitgen_t *rng) nogil:
+ cdef uint32_t mask, delta, val
+ mask = delta = ub - lb
+ mask |= mask >> 1
+ mask |= mask >> 2
+ mask |= mask >> 4
+ mask |= mask >> 8
+ mask |= mask >> 16
+
+ val = rng.next_uint32(rng.state) & mask
+ while val > delta:
+ val = rng.next_uint32(rng.state) & mask
+
+ return lb + val
+
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def bounded_uints(uint32_t lb, uint32_t ub, Py_ssize_t n):
+ cdef Py_ssize_t i
+ cdef bitgen_t *rng
+ cdef uint32_t[::1] out
+ cdef const char *capsule_name = "BitGenerator"
+
+ x = PCG64()
+ out = np.empty(n, dtype=np.uint32)
+ capsule = x.capsule
+
+ if not PyCapsule_IsValid(capsule, capsule_name):
+ raise ValueError("Invalid pointer to anon_func_state")
+ rng = <bitgen_t *>PyCapsule_GetPointer(capsule, capsule_name)
+
+ with x.lock, nogil:
+ for i in range(n):
+ out[i] = bounded_uint(lb, ub, rng)
+ return np.asarray(out)