118 lines
3.8 KiB
Cython
118 lines
3.8 KiB
Cython
|
#!/usr/bin/env python3
|
||
|
|
#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)
|
||
|
|
|
||
|
|
|
||
|
|
@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
|
||
|
|
|
||
|
|
# cython example 2
|
||
|
|
@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 example 3
|
||
|
|
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
|
||
|
|
|