Расширение

Генераторы случайных чисел BitGenerators разработаны для расширения с использованием стандартных инструментов для высокопроизводительного Python – numba и Cython. Объект Generator также может использоваться с предоставленными пользователем генераторами BitGenerators, если они экспортируют небольшой набор необходимых функций.

Numba

Numba может использоваться с CTypes или CFFI. Текущая итерация генераторов BitGenerators экспортирует небольшой набор функций через оба интерфейса.

Этот пример демонстрирует, как numba может использоваться для генерации нормальных чисел Бокса-Мюллера с использованием чистого Python-реализации, которая затем компилируется. Случайные числа предоставляются 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')

И CTypes, и CFFI позволяют использовать более сложные распределения непосредственно в Numba после компиляции файла distributions.c в DLL или подобное. Пример использования более сложного распределения находится в папке examples.

Cython

Cython может использоваться для распаковки PyCapsule, предоставленных генератором BitGenerator. Этот пример использует PCG64 и random_gauss_zig, генератор на основе Зиггурата для нормальных распределений, для заполнения массива. Обычные предостережения для написания высокопроизводительного кода с использованием Cython – удаление проверок границ и переполнения, предоставление информации о выравнивании массивов – по-прежнему актуальны.

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

К генератору BitGenerator также можно получить прямой доступ, используя члены базовой структуры RNG.

@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

Эти функции вместе с минимальным файлом настройки включены в папку examples.

Новые базовые генераторы случайных чисел

Generator может использоваться с другими предоставленными пользователем генераторами BitGenerators. Самый простой способ написать новый генератор BitGenerator – изучить файл pyx одного из существующих генераторов BitGenerators. Ключевая структура, которая должна быть предоставлена, – это capsule, которая содержит указатель на структуру типа 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;

которая предоставляет 5 указателей. Первый – это неявный указатель на структуру данных, используемую генераторами BitGenerators. Следующие три – это указатели на функции, которые возвращают следующие 64- и 32-битные беззнаковые целые числа, следующее случайное двойное значение и следующее сырое значение. Эта последняя функция используется для тестирования, и поэтому её можно установить на функцию следующего 64-битного беззнакового целого числа, если она не нужна. Функции внутри Generator используют эту структуру, как в

bitgen_state->next_uint64(bitgen_state->state)