# extending.pyx

    #!/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)

previous

setup.py

next

extending_distributions.pyx