Skip to main content
`Ctrl`+`K`

# Extending

The `BitGenerator`s have been designed to be extendable using standard tools for high-performance Python – numba and Cython. The `Generator` object can also be used with user-provided `BitGenerator`s as long as these export a small set of required functions.

## Numba

Numba can be used with either CTypes or CFFI. The current iteration of the `BitGenerator`s all export a small set of functions through both interfaces.

This example shows how Numba can be used to produce Gaussian samples using a pure Python implementation which is then compiled. The random numbers are provided by `ctypes.next_double`.

    from timeit import timeit

    import numba as nb

    import numpy as np
    from numpy.random import PCG64

    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(f'{t1:.2f} secs for {n} PCG64 (Numba/PCG64) gaussian randoms')
    t2 = timeit(numpycall, number=1000)
    print(f'{t2:.2f} secs for {n} PCG64 (NumPy/PCG64) gaussian randoms')

Both CTypes and CFFI allow the more complicated distributions to be used directly in Numba after compiling the file distributions.c into a `DLL` or `so`. An example showing the use of a more complicated distribution is in the Examples section below.

## Cython

Cython can be used to unpack the `PyCapsule` provided by a `BitGenerator`. This example uses `PCG64` and the example from above. The usual caveats for writing high-performance code using Cython – removing bounds checks and wrap around, providing array alignment information – still apply.

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

    np.import_array()

    @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

The `BitGenerator` can also be directly accessed using the members of the `bitgen_t` struct.

    @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 can be used to directly access the functions in `numpy/random/c_distributions.pxd`. This requires linking with the `npyrandom` library located in `numpy/random/lib`.

    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

See Extending numpy.random via Cython for the complete listings of these examples and a minimal `setup.py` to build the c-extension modules.

## CFFI

CFFI can be used to directly access the functions in `include/numpy/random/distributions.h`. Some “massaging” of the header file is required:

    """
    Use cffi to access any of the underlying C functions from distributions.h
    """
    import os

    import cffi

    import numpy as np

    from .parse import parse_distributions_h

    ffi = cffi.FFI()

    inc_dir = os.path.join(np.get_include(), 'numpy')

    # Basic numpy types
    ffi.cdef('''
        typedef intptr_t npy_intp;
        typedef unsigned char npy_bool;

    ''')

    parse_distributions_h(ffi, inc_dir)

Once the header is parsed by `ffi.cdef`, the functions can be accessed directly from the `_generator` shared object, using the `BitGenerator.cffi` interface.

    lib = ffi.dlopen(np.random._generator.__file__)

    # Compare the distributions.h random_standard_normal_fill to
    # Generator.standard_random
    bit_gen = np.random.PCG64()
    rng = np.random.Generator(bit_gen)
    state = bit_gen.state

    interface = rng.bit_generator.cffi
    n = 100
    vals_cffi = ffi.new('double[%d]' % n)
    lib.random_standard_normal_fill(interface.bit_generator, n, vals_cffi)

    # reset the state
    bit_gen.state = state

    vals = rng.standard_normal(n)

    for i in range(n):
        assert vals[i] == vals_cffi[i]

## New BitGenerators

`Generator` can be used with user-provided `BitGenerator`s. The simplest way to write a new `BitGenerator` is to examine the pyx file of one of the existing `BitGenerator`s. The key structure that must be provided is the `capsule` which contains a `PyCapsule` to a struct pointer of type `bitgen_t`,

    typedef struct bitgen {
      void *state;
      uint64_t (*next_uint64)(void *st);
      uint32_t (*next_uint32)(void *st);
      double (*next_double)(void *st);
      uint64_t (*next_raw)(void *st);
    } bitgen_t;

which provides 5 pointers. The first is an opaque pointer to the data structure used by the `BitGenerator`s. The next three are function pointers which return the next 64- and 32-bit unsigned integers, the next random double and the next raw value. This final function is used for testing and so can be set to the next 64-bit unsigned integer function if not needed. Functions inside `Generator` use this structure as in

    bitgen_state->next_uint64(bitgen_state->state)

## Examples

- Numba
- CFFI + Numba
- Cython
  - meson.build
  - extending.pyx
  - extending_distributions.pyx
- CFFI