Extending

The BitGenerators 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 BitGenerators 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 BitGenerators all export a small set of functions through both interfaces.

This example shows how numba can be used to produce Box-Muller normals using a pure Python implementation which is then compiled. The random numbers are provided by ctypes.next_double.

from numpy.random import PCG64
import numpy as np
import numba as nb

x = PCG64()
f = x.ctypes.next_double
s = x.ctypes.state
state_addr = x.ctypes.state_address

def normals(n, state):
    out = np.empty(n)
    for i in range((n+1)//2):
        x1 = 2.0*f(state) - 1.0
        x2 = 2.0*f(state) - 1.0
        r2 = x1*x1 + x2*x2
        while r2 >= 1.0 or r2 == 0.0:
            x1 = 2.0*f(state) - 1.0
            x2 = 2.0*f(state) - 1.0
            r2 = x1*x1 + x2*x2
        g = np.sqrt(-2.0*np.log(r2)/r2)
        out[2*i] = g*x1
        if 2*i+1 < n:
            out[2*i+1] = g*x2
    return out

# Compile using Numba
print(normals(10, s).var())
# Warm up
normalsj = nb.jit(normals, nopython=True)
# Must use state address not state with numba
normalsj(1, state_addr)
%timeit normalsj(1000000, state_addr)
print('1,000,000 Box-Muller (numba/PCG64) randoms')
%timeit np.random.standard_normal(1000000)
print('1,000,000 Box-Muller (NumPy) 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 folder.

Cython

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

import numpy as np
cimport numpy as np
cimport cython
from cpython.pycapsule cimport PyCapsule_IsValid, PyCapsule_GetPointer
from numpy.random.common cimport *
from numpy.random.distributions cimport random_gauss_zig
from numpy.random import PCG64


@cython.boundscheck(False)
@cython.wraparound(False)
def normals_zig(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)
    # Best practice is to release GIL and acquire the lock
    with x.lock, nogil:
        for i in range(n):
            random_values[i] = random_gauss_zig(rng)
    randoms = np.asarray(random_values)
    return randoms

The BitGenerator can also be directly accessed using the members of the basic RNG structure.

@cython.boundscheck(False)
@cython.wraparound(False)
def 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
    # 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)
    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

These functions along with a minimal setup file are included in the examples folder.

New Basic RNGs

Generator can be used with other user-provided BitGenerators. The simplest way to write a new BitGenerator is to examine the pyx file of one of the existing BitGenerators. 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 BitGenerators. 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)