5 files changed, 318 insertions, 0 deletions
diff --git a/numpy/random/_examples/cython/extending.pyx b/numpy/random/_examples/cython/extending.pyx
new file mode 100644
index 000000000..2a866648d
--- /dev/null
+++ b/numpy/random/_examples/cython/extending.pyx
@@ -0,0 +1,78 @@
+#!/usr/bin/env python
+#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._bit_generator 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 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 declated 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)
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
diff --git a/numpy/random/_examples/cython/setup.py b/numpy/random/_examples/cython/setup.py
new file mode 100644
index 000000000..315527a2d
--- /dev/null
+++ b/numpy/random/_examples/cython/setup.py
@@ -0,0 +1,32 @@
+#!/usr/bin/env python3
+"""
+Build the demos
+
+Usage: python setup.py build_ext -i
+"""
+
+import numpy as np
+from distutils.core import setup
+from Cython.Build import cythonize
+from setuptools.extension import Extension
+from os.path import join, abspath, dirname
+
+curpath = abspath(dirname(__file__))
+
+extending = Extension("extending",
+                      sources=[join(curpath, 'extending.pyx')],
+                      include_dirs=[
+                            np.get_include(),
+                            join(curpath, '..', '..')
+                        ],
+                      )
+distributions = Extension("extending_distributions",
+                          sources=[join(curpath, 'extending_distributions.pyx'),
+                                  ],
+                          include_dirs=[np.get_include()])
+
+extensions = [extending, distributions]
+
+setup(
+    ext_modules=cythonize(extensions)
+)
diff --git a/numpy/random/_examples/numba/extending.py b/numpy/random/_examples/numba/extending.py
new file mode 100644
index 000000000..0d240596b
--- /dev/null
+++ b/numpy/random/_examples/numba/extending.py
@@ -0,0 +1,84 @@
+import numpy as np
+import numba as nb
+
+from numpy.random import PCG64
+from timeit import timeit
+
+bit_gen = PCG64()
+next_d = bit_gen.cffi.next_double
+state_addr = bit_gen.cffi.state_address
+
+def normals(n, state):
+    out = np.empty(n)
+    for i in range((n + 1) // 2):
+        x1 = 2.0 * next_d(state) - 1.0
+        x2 = 2.0 * next_d(state) - 1.0
+        r2 = x1 * x1 + x2 * x2
+        while r2 >= 1.0 or r2 == 0.0:
+            x1 = 2.0 * next_d(state) - 1.0
+            x2 = 2.0 * next_d(state) - 1.0
+            r2 = x1 * x1 + x2 * x2
+        f = np.sqrt(-2.0 * np.log(r2) / r2)
+        out[2 * i] = f * x1
+        if 2 * i + 1 < n:
+            out[2 * i + 1] = f * x2
+    return out
+
+# Compile using Numba
+normalsj = nb.jit(normals, nopython=True)
+# Must use state address not state with numba
+n = 10000
+
+def numbacall():
+    return normalsj(n, state_addr)
+
+rg = np.random.Generator(PCG64())
+
+def numpycall():
+    return rg.normal(size=n)
+
+# Check that the functions work
+r1 = numbacall()
+r2 = numpycall()
+assert r1.shape == (n,)
+assert r1.shape == r2.shape
+
+t1 = timeit(numbacall, number=1000)
+print('{:.2f} secs for {} PCG64 (Numba/PCG64) gaussian randoms'.format(t1, n))
+t2 = timeit(numpycall, number=1000)
+print('{:.2f} secs for {} PCG64 (NumPy/PCG64) gaussian randoms'.format(t2, n))
+
+# example 2
+
+next_u32 = bit_gen.ctypes.next_uint32
+ctypes_state = bit_gen.ctypes.state
+
+@nb.jit(nopython=True)
+def bounded_uint(lb, ub, state):
+    mask = delta = ub - lb
+    mask |= mask >> 1
+    mask |= mask >> 2
+    mask |= mask >> 4
+    mask |= mask >> 8
+    mask |= mask >> 16
+
+    val = next_u32(state) & mask
+    while val > delta:
+        val = next_u32(state) & mask
+
+    return lb + val
+
+
+print(bounded_uint(323, 2394691, ctypes_state.value))
+
+
+@nb.jit(nopython=True)
+def bounded_uints(lb, ub, n, state):
+    out = np.empty(n, dtype=np.uint32)
+    for i in range(n):
+        out[i] = bounded_uint(lb, ub, state)
+
+
+bounded_uints(323, 2394691, 10000000, ctypes_state.value)
+
+
diff --git a/numpy/random/_examples/numba/extending_distributions.py b/numpy/random/_examples/numba/extending_distributions.py
new file mode 100644
index 000000000..9233ccced
--- /dev/null
+++ b/numpy/random/_examples/numba/extending_distributions.py
@@ -0,0 +1,61 @@
+r"""
+On *nix, execute in randomgen/src/distributions
+
+export PYTHON_INCLUDE=#path to Python's include folder, usually \
+    ${PYTHON_HOME}/include/python${PYTHON_VERSION}m
+export NUMPY_INCLUDE=#path to numpy's include folder, usually \
+    ${PYTHON_HOME}/lib/python${PYTHON_VERSION}/site-packages/numpy/core/include
+gcc -shared -o libdistributions.so -fPIC distributions.c \
+    -I${NUMPY_INCLUDE} -I${PYTHON_INCLUDE}
+mv libdistributions.so ../../examples/numba/
+
+On Windows
+
+rem PYTHON_HOME is setup dependent, this is an example
+set PYTHON_HOME=c:\Anaconda
+cl.exe /LD .\distributions.c -DDLL_EXPORT \
+    -I%PYTHON_HOME%\lib\site-packages\numpy\core\include \
+    -I%PYTHON_HOME%\include %PYTHON_HOME%\libs\python36.lib
+move distributions.dll ../../examples/numba/
+"""
+import os
+
+import numba as nb
+import numpy as np
+from cffi import FFI
+
+from numpy.random import PCG64
+
+ffi = FFI()
+if os.path.exists('./distributions.dll'):
+    lib = ffi.dlopen('./distributions.dll')
+elif os.path.exists('./libdistributions.so'):
+    lib = ffi.dlopen('./libdistributions.so')
+else:
+    raise RuntimeError('Required DLL/so file was not found.')
+
+ffi.cdef("""
+double random_gauss_zig(void *bitgen_state);
+""")
+x = PCG64()
+xffi = x.cffi
+bit_generator = xffi.bit_generator
+
+random_gauss_zig = lib.random_gauss_zig
+
+
+def normals(n, bit_generator):
+    out = np.empty(n)
+    for i in range(n):
+        out[i] = random_gauss_zig(bit_generator)
+    return out
+
+
+normalsj = nb.jit(normals, nopython=True)
+
+# Numba requires a memory address for void *
+# Can also get address from x.ctypes.bit_generator.value
+bit_generator_address = int(ffi.cast('uintptr_t', bit_generator))
+
+norm = normalsj(1000, bit_generator_address)
+print(norm[:12])