1#!/usr/bin/env python3
2#cython: language_level=3
3"""
4This file shows how the to use a BitGenerator to create a distribution.
5"""
6import numpy as np
7cimport numpy as np
8cimport cython
9from cpython.pycapsule cimport PyCapsule_IsValid, PyCapsule_GetPointer
10from libc.stdint cimport uint16_t, uint64_t
11from numpy.random cimport bitgen_t
12from numpy.random import PCG64
13from numpy.random.c_distributions cimport (
14      random_standard_uniform_fill, random_standard_uniform_fill_f)
15
16
17@cython.boundscheck(False)
18@cython.wraparound(False)
19def uniforms(Py_ssize_t n):
20    """
21    Create an array of `n` uniformly distributed doubles.
22    A 'real' distribution would want to process the values into
23    some non-uniform distribution
24    """
25    cdef Py_ssize_t i
26    cdef bitgen_t *rng
27    cdef const char *capsule_name = "BitGenerator"
28    cdef double[::1] random_values
29
30    x = PCG64()
31    capsule = x.capsule
32    # Optional check that the capsule if from a BitGenerator
33    if not PyCapsule_IsValid(capsule, capsule_name):
34        raise ValueError("Invalid pointer to anon_func_state")
35    # Cast the pointer
36    rng = <bitgen_t *> PyCapsule_GetPointer(capsule, capsule_name)
37    random_values = np.empty(n, dtype='float64')
38    with x.lock, nogil:
39        for i in range(n):
40            # Call the function
41            random_values[i] = rng.next_double(rng.state)
42    randoms = np.asarray(random_values)
43
44    return randoms
45
46# cython example 2
47@cython.boundscheck(False)
48@cython.wraparound(False)
49def uint10_uniforms(Py_ssize_t n):
50    """Uniform 10 bit integers stored as 16-bit unsigned integers"""
51    cdef Py_ssize_t i
52    cdef bitgen_t *rng
53    cdef const char *capsule_name = "BitGenerator"
54    cdef uint16_t[::1] random_values
55    cdef int bits_remaining
56    cdef int width = 10
57    cdef uint64_t buff, mask = 0x3FF
58
59    x = PCG64()
60    capsule = x.capsule
61    if not PyCapsule_IsValid(capsule, capsule_name):
62        raise ValueError("Invalid pointer to anon_func_state")
63    rng = <bitgen_t *> PyCapsule_GetPointer(capsule, capsule_name)
64    random_values = np.empty(n, dtype='uint16')
65    # Best practice is to release GIL and acquire the lock
66    bits_remaining = 0
67    with x.lock, nogil:
68        for i in range(n):
69            if bits_remaining < width:
70                buff = rng.next_uint64(rng.state)
71            random_values[i] = buff & mask
72            buff >>= width
73
74    randoms = np.asarray(random_values)
75    return randoms
76
77# cython example 3
78def uniforms_ex(bit_generator, Py_ssize_t n, dtype=np.float64):
79    """
80    Create an array of `n` uniformly distributed doubles via a "fill" function.
81
82    A 'real' distribution would want to process the values into
83    some non-uniform distribution
84
85    Parameters
86    ----------
87    bit_generator: BitGenerator instance
88    n: int
89        Output vector length
90    dtype: {str, dtype}, optional
91        Desired dtype, either 'd' (or 'float64') or 'f' (or 'float32'). The
92        default dtype value is 'd'
93    """
94    cdef Py_ssize_t i
95    cdef bitgen_t *rng
96    cdef const char *capsule_name = "BitGenerator"
97    cdef np.ndarray randoms
98
99    capsule = bit_generator.capsule
100    # Optional check that the capsule if from a BitGenerator
101    if not PyCapsule_IsValid(capsule, capsule_name):
102        raise ValueError("Invalid pointer to anon_func_state")
103    # Cast the pointer
104    rng = <bitgen_t *> PyCapsule_GetPointer(capsule, capsule_name)
105
106    _dtype = np.dtype(dtype)
107    randoms = np.empty(n, dtype=_dtype)
108    if _dtype == np.float32:
109        with bit_generator.lock:
110            random_standard_uniform_fill_f(rng, n, <float*>np.PyArray_DATA(randoms))
111    elif _dtype == np.float64:
112        with bit_generator.lock:
113            random_standard_uniform_fill(rng, n, <double*>np.PyArray_DATA(randoms))
114    else:
115        raise TypeError('Unsupported dtype %r for random' % _dtype)
116    return randoms
117
118