aboutsummaryrefslogtreecommitdiff
path: root/.venv/lib/python3.12/site-packages/numpy/random/_examples/cython/extending_distributions.pyx
diff options
context:
space:
mode:
Diffstat (limited to '.venv/lib/python3.12/site-packages/numpy/random/_examples/cython/extending_distributions.pyx')
-rw-r--r--.venv/lib/python3.12/site-packages/numpy/random/_examples/cython/extending_distributions.pyx117
1 files changed, 117 insertions, 0 deletions
diff --git a/.venv/lib/python3.12/site-packages/numpy/random/_examples/cython/extending_distributions.pyx b/.venv/lib/python3.12/site-packages/numpy/random/_examples/cython/extending_distributions.pyx
new file mode 100644
index 00000000..d908e92d
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/numpy/random/_examples/cython/extending_distributions.pyx
@@ -0,0 +1,117 @@
+#!/usr/bin/env python3
+#cython: language_level=3
+"""
+This file shows how the to use a BitGenerator to create a distribution.
+"""
+import numpy as np
+cimport numpy as np
+cimport cython
+from cpython.pycapsule cimport PyCapsule_IsValid, PyCapsule_GetPointer
+from libc.stdint cimport uint16_t, uint64_t
+from numpy.random cimport bitgen_t
+from numpy.random import PCG64
+from numpy.random.c_distributions cimport (
+ random_standard_uniform_fill, random_standard_uniform_fill_f)
+
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def uniforms(Py_ssize_t n):
+ """
+ Create an array of `n` uniformly distributed doubles.
+ A 'real' distribution would want to process the values into
+ some non-uniform distribution
+ """
+ cdef Py_ssize_t i
+ cdef bitgen_t *rng
+ cdef const char *capsule_name = "BitGenerator"
+ cdef double[::1] random_values
+
+ x = PCG64()
+ capsule = x.capsule
+ # Optional check that the capsule if from a BitGenerator
+ if not PyCapsule_IsValid(capsule, capsule_name):
+ raise ValueError("Invalid pointer to anon_func_state")
+ # Cast the pointer
+ rng = <bitgen_t *> PyCapsule_GetPointer(capsule, capsule_name)
+ random_values = np.empty(n, dtype='float64')
+ with x.lock, nogil:
+ for i in range(n):
+ # Call the function
+ random_values[i] = rng.next_double(rng.state)
+ randoms = np.asarray(random_values)
+
+ return randoms
+
+# cython example 2
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def uint10_uniforms(Py_ssize_t n):
+ """Uniform 10 bit integers stored as 16-bit unsigned integers"""
+ cdef Py_ssize_t i
+ cdef bitgen_t *rng
+ cdef const char *capsule_name = "BitGenerator"
+ cdef uint16_t[::1] random_values
+ cdef int bits_remaining
+ cdef int width = 10
+ cdef uint64_t buff, mask = 0x3FF
+
+ 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, dtype='uint16')
+ # Best practice is to release GIL and acquire the lock
+ bits_remaining = 0
+ with x.lock, nogil:
+ for i in range(n):
+ if bits_remaining < width:
+ buff = rng.next_uint64(rng.state)
+ random_values[i] = buff & mask
+ buff >>= width
+
+ randoms = np.asarray(random_values)
+ return randoms
+
+# cython example 3
+def uniforms_ex(bit_generator, Py_ssize_t n, dtype=np.float64):
+ """
+ Create an array of `n` uniformly distributed doubles via a "fill" function.
+
+ A 'real' distribution would want to process the values into
+ some non-uniform distribution
+
+ Parameters
+ ----------
+ bit_generator: BitGenerator instance
+ n: int
+ Output vector length
+ dtype: {str, dtype}, optional
+ Desired dtype, either 'd' (or 'float64') or 'f' (or 'float32'). The
+ default dtype value is 'd'
+ """
+ cdef Py_ssize_t i
+ cdef bitgen_t *rng
+ cdef const char *capsule_name = "BitGenerator"
+ cdef np.ndarray randoms
+
+ capsule = bit_generator.capsule
+ # Optional check that the capsule if from a BitGenerator
+ if not PyCapsule_IsValid(capsule, capsule_name):
+ raise ValueError("Invalid pointer to anon_func_state")
+ # Cast the pointer
+ rng = <bitgen_t *> PyCapsule_GetPointer(capsule, capsule_name)
+
+ _dtype = np.dtype(dtype)
+ randoms = np.empty(n, dtype=_dtype)
+ if _dtype == np.float32:
+ with bit_generator.lock:
+ random_standard_uniform_fill_f(rng, n, <float*>np.PyArray_DATA(randoms))
+ elif _dtype == np.float64:
+ with bit_generator.lock:
+ random_standard_uniform_fill(rng, n, <double*>np.PyArray_DATA(randoms))
+ else:
+ raise TypeError('Unsupported dtype %r for random' % _dtype)
+ return randoms
+