1 files changed, 63 insertions, 0 deletions
diff --git a/numpy/random/_examples/cython/extending_distributions.pyx b/numpy/random/_examples/cython/extending_distributions.pyx
new file mode 100644
index 000000000..d17da45c1
--- /dev/null
+++ b/numpy/random/_examples/cython/extending_distributions.pyx
@@ -0,0 +1,63 @@
+#!/usr/bin/env python
+#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 numpy.random._bit_generator cimport bitgen_t
+from numpy.random import PCG64
+
+
+@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 uint16_uniforms(Py_ssize_t n):
+    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
+    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='uint32')
+    # Best practice is to release GIL and acquire the lock
+    with x.lock, nogil:
+        for i in range(n):
+            random_values[i] = rng.next_uint32(rng.state)
+    randoms = np.asarray(random_values)
+    return randoms