diff options
Diffstat (limited to 'numpy/random/_common.pyx')
-rw-r--r-- | numpy/random/_common.pyx | 977 |
1 files changed, 977 insertions, 0 deletions
diff --git a/numpy/random/_common.pyx b/numpy/random/_common.pyx new file mode 100644 index 000000000..ef1afac7c --- /dev/null +++ b/numpy/random/_common.pyx @@ -0,0 +1,977 @@ +#!python +#cython: wraparound=False, nonecheck=False, boundscheck=False, cdivision=True, language_level=3 +from collections import namedtuple +from cpython cimport PyFloat_AsDouble +import sys +import numpy as np +cimport numpy as np + +from libc.stdint cimport uintptr_t + +__all__ = ['interface'] + +np.import_array() + +interface = namedtuple('interface', ['state_address', 'state', 'next_uint64', + 'next_uint32', 'next_double', + 'bit_generator']) + +cdef double LEGACY_POISSON_LAM_MAX = <double>np.iinfo('l').max - np.sqrt(np.iinfo('l').max)*10 +cdef double POISSON_LAM_MAX = <double>np.iinfo('int64').max - np.sqrt(np.iinfo('int64').max)*10 + +cdef uint64_t MAXSIZE = <uint64_t>sys.maxsize + + +cdef object benchmark(bitgen_t *bitgen, object lock, Py_ssize_t cnt, object method): + """Benchmark command used by BitGenerator""" + cdef Py_ssize_t i + if method==u'uint64': + with lock, nogil: + for i in range(cnt): + bitgen.next_uint64(bitgen.state) + elif method==u'double': + with lock, nogil: + for i in range(cnt): + bitgen.next_double(bitgen.state) + else: + raise ValueError('Unknown method') + + +cdef object random_raw(bitgen_t *bitgen, object lock, object size, object output): + """ + random_raw(self, size=None) + + Return randoms as generated by the underlying PRNG + + Parameters + ---------- + bitgen : BitGenerator + Address of the bit generator struct + lock : Threading.Lock + Lock provided by the bit generator + size : int or tuple of ints, optional + Output shape. If the given shape is, e.g., ``(m, n, k)``, then + ``m * n * k`` samples are drawn. Default is None, in which case a + single value is returned. + output : bool, optional + Output values. Used for performance testing since the generated + values are not returned. + + Returns + ------- + out : uint or ndarray + Drawn samples. + + Notes + ----- + This method directly exposes the the raw underlying pseudo-random + number generator. All values are returned as unsigned 64-bit + values irrespective of the number of bits produced by the PRNG. + + See the class docstring for the number of bits returned. + """ + cdef np.ndarray randoms + cdef uint64_t *randoms_data + cdef Py_ssize_t i, n + + if not output: + if size is None: + with lock: + bitgen.next_raw(bitgen.state) + return None + n = np.asarray(size).sum() + with lock, nogil: + for i in range(n): + bitgen.next_raw(bitgen.state) + return None + + if size is None: + with lock: + return bitgen.next_raw(bitgen.state) + + randoms = <np.ndarray>np.empty(size, np.uint64) + randoms_data = <uint64_t*>np.PyArray_DATA(randoms) + n = np.PyArray_SIZE(randoms) + + with lock, nogil: + for i in range(n): + randoms_data[i] = bitgen.next_raw(bitgen.state) + return randoms + +cdef object prepare_cffi(bitgen_t *bitgen): + """ + Bundles the interfaces to interact with a BitGenerator using cffi + + Parameters + ---------- + bitgen : pointer + A pointer to a BitGenerator instance + + Returns + ------- + interface : namedtuple + The functions required to interface with the BitGenerator using cffi + + * state_address - Memory address of the state struct + * state - pointer to the state struct + * next_uint64 - function pointer to produce 64 bit integers + * next_uint32 - function pointer to produce 32 bit integers + * next_double - function pointer to produce doubles + * bit_generator - pointer to the BitGenerator struct + """ + try: + import cffi + except ImportError: + raise ImportError('cffi cannot be imported.') + + ffi = cffi.FFI() + _cffi = interface(<uintptr_t>bitgen.state, + ffi.cast('void *', <uintptr_t>bitgen.state), + ffi.cast('uint64_t (*)(void *)', <uintptr_t>bitgen.next_uint64), + ffi.cast('uint32_t (*)(void *)', <uintptr_t>bitgen.next_uint32), + ffi.cast('double (*)(void *)', <uintptr_t>bitgen.next_double), + ffi.cast('void *', <uintptr_t>bitgen)) + return _cffi + +cdef object prepare_ctypes(bitgen_t *bitgen): + """ + Bundles the interfaces to interact with a BitGenerator using ctypes + + Parameters + ---------- + bitgen : pointer + A pointer to a BitGenerator instance + + Returns + ------- + interface : namedtuple + The functions required to interface with the BitGenerator using ctypes: + + * state_address - Memory address of the state struct + * state - pointer to the state struct + * next_uint64 - function pointer to produce 64 bit integers + * next_uint32 - function pointer to produce 32 bit integers + * next_double - function pointer to produce doubles + * bit_generator - pointer to the BitGenerator struct + """ + import ctypes + + _ctypes = interface(<uintptr_t>bitgen.state, + ctypes.c_void_p(<uintptr_t>bitgen.state), + ctypes.cast(<uintptr_t>bitgen.next_uint64, + ctypes.CFUNCTYPE(ctypes.c_uint64, + ctypes.c_void_p)), + ctypes.cast(<uintptr_t>bitgen.next_uint32, + ctypes.CFUNCTYPE(ctypes.c_uint32, + ctypes.c_void_p)), + ctypes.cast(<uintptr_t>bitgen.next_double, + ctypes.CFUNCTYPE(ctypes.c_double, + ctypes.c_void_p)), + ctypes.c_void_p(<uintptr_t>bitgen)) + return _ctypes + +cdef double kahan_sum(double *darr, np.npy_intp n): + cdef double c, y, t, sum + cdef np.npy_intp i + sum = darr[0] + c = 0.0 + for i in range(1, n): + y = darr[i] - c + t = sum + y + c = (t-sum) - y + sum = t + return sum + + +cdef object wrap_int(object val, object bits): + """Wraparound to place an integer into the interval [0, 2**bits)""" + mask = ~(~int(0) << bits) + return val & mask + + +cdef np.ndarray int_to_array(object value, object name, object bits, object uint_size): + """Convert a large integer to an array of unsigned integers""" + len = bits // uint_size + value = np.asarray(value) + if uint_size == 32: + dtype = np.uint32 + elif uint_size == 64: + dtype = np.uint64 + else: + raise ValueError('Unknown uint_size') + if value.shape == (): + value = int(value) + upper = int(2)**int(bits) + if value < 0 or value >= upper: + raise ValueError('{name} must be positive and ' + 'less than 2**{bits}.'.format(name=name, bits=bits)) + + out = np.empty(len, dtype=dtype) + for i in range(len): + out[i] = value % 2**int(uint_size) + value >>= int(uint_size) + else: + out = value.astype(dtype) + if out.shape != (len,): + raise ValueError('{name} must have {len} elements when using ' + 'array form'.format(name=name, len=len)) + return out + + +cdef check_output(object out, object dtype, object size): + if out is None: + return + cdef np.ndarray out_array = <np.ndarray>out + if not (np.PyArray_CHKFLAGS(out_array, np.NPY_CARRAY) or + np.PyArray_CHKFLAGS(out_array, np.NPY_FARRAY)): + raise ValueError('Supplied output array is not contiguous, writable or aligned.') + if out_array.dtype != dtype: + raise TypeError('Supplied output array has the wrong type. ' + 'Expected {0}, got {1}'.format(np.dtype(dtype), out_array.dtype)) + if size is not None: + try: + tup_size = tuple(size) + except TypeError: + tup_size = tuple([size]) + if tup_size != out.shape: + raise ValueError('size must match out.shape when used together') + + +cdef object double_fill(void *func, bitgen_t *state, object size, object lock, object out): + cdef random_double_fill random_func = (<random_double_fill>func) + cdef double out_val + cdef double *out_array_data + cdef np.ndarray out_array + cdef np.npy_intp i, n + + if size is None and out is None: + with lock: + random_func(state, 1, &out_val) + return out_val + + if out is not None: + check_output(out, np.float64, size) + out_array = <np.ndarray>out + else: + out_array = <np.ndarray>np.empty(size, np.double) + + n = np.PyArray_SIZE(out_array) + out_array_data = <double *>np.PyArray_DATA(out_array) + with lock, nogil: + random_func(state, n, out_array_data) + return out_array + +cdef object float_fill(void *func, bitgen_t *state, object size, object lock, object out): + cdef random_float_fill random_func = (<random_float_fill>func) + cdef float out_val + cdef float *out_array_data + cdef np.ndarray out_array + cdef np.npy_intp i, n + + if size is None and out is None: + with lock: + random_func(state, 1, &out_val) + return out_val + + if out is not None: + check_output(out, np.float32, size) + out_array = <np.ndarray>out + else: + out_array = <np.ndarray>np.empty(size, np.float32) + + n = np.PyArray_SIZE(out_array) + out_array_data = <float *>np.PyArray_DATA(out_array) + with lock, nogil: + random_func(state, n, out_array_data) + return out_array + +cdef object float_fill_from_double(void *func, bitgen_t *state, object size, object lock, object out): + cdef random_double_0 random_func = (<random_double_0>func) + cdef float *out_array_data + cdef np.ndarray out_array + cdef np.npy_intp i, n + + if size is None and out is None: + with lock: + return <float>random_func(state) + + if out is not None: + check_output(out, np.float32, size) + out_array = <np.ndarray>out + else: + out_array = <np.ndarray>np.empty(size, np.float32) + + n = np.PyArray_SIZE(out_array) + out_array_data = <float *>np.PyArray_DATA(out_array) + with lock, nogil: + for i in range(n): + out_array_data[i] = <float>random_func(state) + return out_array + + +cdef int check_array_constraint(np.ndarray val, object name, constraint_type cons) except -1: + if cons == CONS_NON_NEGATIVE: + if np.any(np.logical_and(np.logical_not(np.isnan(val)), np.signbit(val))): + raise ValueError(name + " < 0") + elif cons == CONS_POSITIVE or cons == CONS_POSITIVE_NOT_NAN: + if cons == CONS_POSITIVE_NOT_NAN and np.any(np.isnan(val)): + raise ValueError(name + " must not be NaN") + elif np.any(np.less_equal(val, 0)): + raise ValueError(name + " <= 0") + elif cons == CONS_BOUNDED_0_1: + if not np.all(np.greater_equal(val, 0)) or \ + not np.all(np.less_equal(val, 1)): + raise ValueError("{0} < 0, {0} > 1 or {0} contains NaNs".format(name)) + elif cons == CONS_BOUNDED_GT_0_1: + if not np.all(np.greater(val, 0)) or not np.all(np.less_equal(val, 1)): + raise ValueError("{0} <= 0, {0} > 1 or {0} contains NaNs".format(name)) + elif cons == CONS_GT_1: + if not np.all(np.greater(val, 1)): + raise ValueError("{0} <= 1 or {0} contains NaNs".format(name)) + elif cons == CONS_GTE_1: + if not np.all(np.greater_equal(val, 1)): + raise ValueError("{0} < 1 or {0} contains NaNs".format(name)) + elif cons == CONS_POISSON: + if not np.all(np.less_equal(val, POISSON_LAM_MAX)): + raise ValueError("{0} value too large".format(name)) + elif not np.all(np.greater_equal(val, 0.0)): + raise ValueError("{0} < 0 or {0} contains NaNs".format(name)) + elif cons == LEGACY_CONS_POISSON: + if not np.all(np.less_equal(val, LEGACY_POISSON_LAM_MAX)): + raise ValueError("{0} value too large".format(name)) + elif not np.all(np.greater_equal(val, 0.0)): + raise ValueError("{0} < 0 or {0} contains NaNs".format(name)) + + return 0 + + +cdef int check_constraint(double val, object name, constraint_type cons) except -1: + cdef bint is_nan + if cons == CONS_NON_NEGATIVE: + if not np.isnan(val) and np.signbit(val): + raise ValueError(name + " < 0") + elif cons == CONS_POSITIVE or cons == CONS_POSITIVE_NOT_NAN: + if cons == CONS_POSITIVE_NOT_NAN and np.isnan(val): + raise ValueError(name + " must not be NaN") + elif val <= 0: + raise ValueError(name + " <= 0") + elif cons == CONS_BOUNDED_0_1: + if not (val >= 0) or not (val <= 1): + raise ValueError("{0} < 0, {0} > 1 or {0} is NaN".format(name)) + elif cons == CONS_BOUNDED_GT_0_1: + if not val >0 or not val <= 1: + raise ValueError("{0} <= 0, {0} > 1 or {0} contains NaNs".format(name)) + elif cons == CONS_GT_1: + if not (val > 1): + raise ValueError("{0} <= 1 or {0} is NaN".format(name)) + elif cons == CONS_GTE_1: + if not (val >= 1): + raise ValueError("{0} < 1 or {0} is NaN".format(name)) + elif cons == CONS_POISSON: + if not (val >= 0): + raise ValueError("{0} < 0 or {0} is NaN".format(name)) + elif not (val <= POISSON_LAM_MAX): + raise ValueError(name + " value too large") + elif cons == LEGACY_CONS_POISSON: + if not (val >= 0): + raise ValueError("{0} < 0 or {0} is NaN".format(name)) + elif not (val <= LEGACY_POISSON_LAM_MAX): + raise ValueError(name + " value too large") + + return 0 + +cdef object cont_broadcast_1(void *func, void *state, object size, object lock, + np.ndarray a_arr, object a_name, constraint_type a_constraint, + object out): + + cdef np.ndarray randoms + cdef double a_val + cdef double *randoms_data + cdef np.broadcast it + cdef random_double_1 f = (<random_double_1>func) + cdef np.npy_intp i, n + + if a_constraint != CONS_NONE: + check_array_constraint(a_arr, a_name, a_constraint) + + if size is not None and out is None: + randoms = <np.ndarray>np.empty(size, np.double) + elif out is None: + randoms = np.PyArray_SimpleNew(np.PyArray_NDIM(a_arr), np.PyArray_DIMS(a_arr), np.NPY_DOUBLE) + else: + randoms = <np.ndarray>out + + randoms_data = <double *>np.PyArray_DATA(randoms) + n = np.PyArray_SIZE(randoms) + it = np.PyArray_MultiIterNew2(randoms, a_arr) + + with lock, nogil: + for i in range(n): + a_val = (<double*>np.PyArray_MultiIter_DATA(it, 1))[0] + randoms_data[i] = f(state, a_val) + + np.PyArray_MultiIter_NEXT(it) + + return randoms + +cdef object cont_broadcast_2(void *func, void *state, object size, object lock, + np.ndarray a_arr, object a_name, constraint_type a_constraint, + np.ndarray b_arr, object b_name, constraint_type b_constraint): + cdef np.ndarray randoms + cdef double a_val, b_val + cdef double *randoms_data + cdef np.broadcast it + cdef random_double_2 f = (<random_double_2>func) + cdef np.npy_intp i, n + + if a_constraint != CONS_NONE: + check_array_constraint(a_arr, a_name, a_constraint) + + if b_constraint != CONS_NONE: + check_array_constraint(b_arr, b_name, b_constraint) + + if size is not None: + randoms = <np.ndarray>np.empty(size, np.double) + else: + it = np.PyArray_MultiIterNew2(a_arr, b_arr) + randoms = <np.ndarray>np.empty(it.shape, np.double) + # randoms = np.PyArray_SimpleNew(it.nd, np.PyArray_DIMS(it), np.NPY_DOUBLE) + + randoms_data = <double *>np.PyArray_DATA(randoms) + n = np.PyArray_SIZE(randoms) + + it = np.PyArray_MultiIterNew3(randoms, a_arr, b_arr) + with lock, nogil: + for i in range(n): + a_val = (<double*>np.PyArray_MultiIter_DATA(it, 1))[0] + b_val = (<double*>np.PyArray_MultiIter_DATA(it, 2))[0] + randoms_data[i] = f(state, a_val, b_val) + + np.PyArray_MultiIter_NEXT(it) + + return randoms + +cdef object cont_broadcast_3(void *func, void *state, object size, object lock, + np.ndarray a_arr, object a_name, constraint_type a_constraint, + np.ndarray b_arr, object b_name, constraint_type b_constraint, + np.ndarray c_arr, object c_name, constraint_type c_constraint): + cdef np.ndarray randoms + cdef double a_val, b_val, c_val + cdef double *randoms_data + cdef np.broadcast it + cdef random_double_3 f = (<random_double_3>func) + cdef np.npy_intp i, n + + if a_constraint != CONS_NONE: + check_array_constraint(a_arr, a_name, a_constraint) + + if b_constraint != CONS_NONE: + check_array_constraint(b_arr, b_name, b_constraint) + + if c_constraint != CONS_NONE: + check_array_constraint(c_arr, c_name, c_constraint) + + if size is not None: + randoms = <np.ndarray>np.empty(size, np.double) + else: + it = np.PyArray_MultiIterNew3(a_arr, b_arr, c_arr) + # randoms = np.PyArray_SimpleNew(it.nd, np.PyArray_DIMS(it), np.NPY_DOUBLE) + randoms = <np.ndarray>np.empty(it.shape, np.double) + + randoms_data = <double *>np.PyArray_DATA(randoms) + n = np.PyArray_SIZE(randoms) + + it = np.PyArray_MultiIterNew4(randoms, a_arr, b_arr, c_arr) + with lock, nogil: + for i in range(n): + a_val = (<double*>np.PyArray_MultiIter_DATA(it, 1))[0] + b_val = (<double*>np.PyArray_MultiIter_DATA(it, 2))[0] + c_val = (<double*>np.PyArray_MultiIter_DATA(it, 3))[0] + randoms_data[i] = f(state, a_val, b_val, c_val) + + np.PyArray_MultiIter_NEXT(it) + + return randoms + +cdef object cont(void *func, void *state, object size, object lock, int narg, + object a, object a_name, constraint_type a_constraint, + object b, object b_name, constraint_type b_constraint, + object c, object c_name, constraint_type c_constraint, + object out): + + cdef np.ndarray a_arr, b_arr, c_arr + cdef double _a = 0.0, _b = 0.0, _c = 0.0 + cdef bint is_scalar = True + check_output(out, np.float64, size) + if narg > 0: + a_arr = <np.ndarray>np.PyArray_FROM_OTF(a, np.NPY_DOUBLE, np.NPY_ALIGNED) + is_scalar = is_scalar and np.PyArray_NDIM(a_arr) == 0 + if narg > 1: + b_arr = <np.ndarray>np.PyArray_FROM_OTF(b, np.NPY_DOUBLE, np.NPY_ALIGNED) + is_scalar = is_scalar and np.PyArray_NDIM(b_arr) == 0 + if narg == 3: + c_arr = <np.ndarray>np.PyArray_FROM_OTF(c, np.NPY_DOUBLE, np.NPY_ALIGNED) + is_scalar = is_scalar and np.PyArray_NDIM(c_arr) == 0 + + if not is_scalar: + if narg == 1: + return cont_broadcast_1(func, state, size, lock, + a_arr, a_name, a_constraint, + out) + elif narg == 2: + return cont_broadcast_2(func, state, size, lock, + a_arr, a_name, a_constraint, + b_arr, b_name, b_constraint) + else: + return cont_broadcast_3(func, state, size, lock, + a_arr, a_name, a_constraint, + b_arr, b_name, b_constraint, + c_arr, c_name, c_constraint) + + if narg > 0: + _a = PyFloat_AsDouble(a) + if a_constraint != CONS_NONE and is_scalar: + check_constraint(_a, a_name, a_constraint) + if narg > 1: + _b = PyFloat_AsDouble(b) + if b_constraint != CONS_NONE: + check_constraint(_b, b_name, b_constraint) + if narg == 3: + _c = PyFloat_AsDouble(c) + if c_constraint != CONS_NONE and is_scalar: + check_constraint(_c, c_name, c_constraint) + + if size is None and out is None: + with lock: + if narg == 0: + return (<random_double_0>func)(state) + elif narg == 1: + return (<random_double_1>func)(state, _a) + elif narg == 2: + return (<random_double_2>func)(state, _a, _b) + elif narg == 3: + return (<random_double_3>func)(state, _a, _b, _c) + + cdef np.npy_intp i, n + cdef np.ndarray randoms + if out is None: + randoms = <np.ndarray>np.empty(size) + else: + randoms = <np.ndarray>out + n = np.PyArray_SIZE(randoms) + + cdef double *randoms_data = <double *>np.PyArray_DATA(randoms) + cdef random_double_0 f0 + cdef random_double_1 f1 + cdef random_double_2 f2 + cdef random_double_3 f3 + + with lock, nogil: + if narg == 0: + f0 = (<random_double_0>func) + for i in range(n): + randoms_data[i] = f0(state) + elif narg == 1: + f1 = (<random_double_1>func) + for i in range(n): + randoms_data[i] = f1(state, _a) + elif narg == 2: + f2 = (<random_double_2>func) + for i in range(n): + randoms_data[i] = f2(state, _a, _b) + elif narg == 3: + f3 = (<random_double_3>func) + for i in range(n): + randoms_data[i] = f3(state, _a, _b, _c) + + if out is None: + return randoms + else: + return out + +cdef object discrete_broadcast_d(void *func, void *state, object size, object lock, + np.ndarray a_arr, object a_name, constraint_type a_constraint): + + cdef np.ndarray randoms + cdef int64_t *randoms_data + cdef np.broadcast it + cdef random_uint_d f = (<random_uint_d>func) + cdef np.npy_intp i, n + + if a_constraint != CONS_NONE: + check_array_constraint(a_arr, a_name, a_constraint) + + if size is not None: + randoms = np.empty(size, np.int64) + else: + # randoms = np.empty(np.shape(a_arr), np.double) + randoms = np.PyArray_SimpleNew(np.PyArray_NDIM(a_arr), np.PyArray_DIMS(a_arr), np.NPY_INT64) + + randoms_data = <int64_t *>np.PyArray_DATA(randoms) + n = np.PyArray_SIZE(randoms) + + it = np.PyArray_MultiIterNew2(randoms, a_arr) + with lock, nogil: + for i in range(n): + a_val = (<double*>np.PyArray_MultiIter_DATA(it, 1))[0] + randoms_data[i] = f(state, a_val) + + np.PyArray_MultiIter_NEXT(it) + + return randoms + +cdef object discrete_broadcast_dd(void *func, void *state, object size, object lock, + np.ndarray a_arr, object a_name, constraint_type a_constraint, + np.ndarray b_arr, object b_name, constraint_type b_constraint): + cdef np.ndarray randoms + cdef int64_t *randoms_data + cdef np.broadcast it + cdef random_uint_dd f = (<random_uint_dd>func) + cdef np.npy_intp i, n + + if a_constraint != CONS_NONE: + check_array_constraint(a_arr, a_name, a_constraint) + if b_constraint != CONS_NONE: + check_array_constraint(b_arr, b_name, b_constraint) + + if size is not None: + randoms = <np.ndarray>np.empty(size, np.int64) + else: + it = np.PyArray_MultiIterNew2(a_arr, b_arr) + randoms = <np.ndarray>np.empty(it.shape, np.int64) + # randoms = np.PyArray_SimpleNew(it.nd, np.PyArray_DIMS(it), np.NPY_INT64) + + randoms_data = <int64_t *>np.PyArray_DATA(randoms) + n = np.PyArray_SIZE(randoms) + + it = np.PyArray_MultiIterNew3(randoms, a_arr, b_arr) + with lock, nogil: + for i in range(n): + a_val = (<double*>np.PyArray_MultiIter_DATA(it, 1))[0] + b_val = (<double*>np.PyArray_MultiIter_DATA(it, 2))[0] + randoms_data[i] = f(state, a_val, b_val) + + np.PyArray_MultiIter_NEXT(it) + + return randoms + +cdef object discrete_broadcast_di(void *func, void *state, object size, object lock, + np.ndarray a_arr, object a_name, constraint_type a_constraint, + np.ndarray b_arr, object b_name, constraint_type b_constraint): + cdef np.ndarray randoms + cdef int64_t *randoms_data + cdef np.broadcast it + cdef random_uint_di f = (<random_uint_di>func) + cdef np.npy_intp i, n + + if a_constraint != CONS_NONE: + check_array_constraint(a_arr, a_name, a_constraint) + + if b_constraint != CONS_NONE: + check_array_constraint(b_arr, b_name, b_constraint) + + if size is not None: + randoms = <np.ndarray>np.empty(size, np.int64) + else: + it = np.PyArray_MultiIterNew2(a_arr, b_arr) + randoms = <np.ndarray>np.empty(it.shape, np.int64) + + randoms_data = <int64_t *>np.PyArray_DATA(randoms) + n = np.PyArray_SIZE(randoms) + + it = np.PyArray_MultiIterNew3(randoms, a_arr, b_arr) + with lock, nogil: + for i in range(n): + a_val = (<double*>np.PyArray_MultiIter_DATA(it, 1))[0] + b_val = (<int64_t*>np.PyArray_MultiIter_DATA(it, 2))[0] + (<int64_t*>np.PyArray_MultiIter_DATA(it, 0))[0] = f(state, a_val, b_val) + + np.PyArray_MultiIter_NEXT(it) + + return randoms + +cdef object discrete_broadcast_iii(void *func, void *state, object size, object lock, + np.ndarray a_arr, object a_name, constraint_type a_constraint, + np.ndarray b_arr, object b_name, constraint_type b_constraint, + np.ndarray c_arr, object c_name, constraint_type c_constraint): + cdef np.ndarray randoms + cdef int64_t *randoms_data + cdef np.broadcast it + cdef random_uint_iii f = (<random_uint_iii>func) + cdef np.npy_intp i, n + + if a_constraint != CONS_NONE: + check_array_constraint(a_arr, a_name, a_constraint) + + if b_constraint != CONS_NONE: + check_array_constraint(b_arr, b_name, b_constraint) + + if c_constraint != CONS_NONE: + check_array_constraint(c_arr, c_name, c_constraint) + + if size is not None: + randoms = <np.ndarray>np.empty(size, np.int64) + else: + it = np.PyArray_MultiIterNew3(a_arr, b_arr, c_arr) + randoms = <np.ndarray>np.empty(it.shape, np.int64) + + randoms_data = <int64_t *>np.PyArray_DATA(randoms) + n = np.PyArray_SIZE(randoms) + + it = np.PyArray_MultiIterNew4(randoms, a_arr, b_arr, c_arr) + with lock, nogil: + for i in range(n): + a_val = (<int64_t*>np.PyArray_MultiIter_DATA(it, 1))[0] + b_val = (<int64_t*>np.PyArray_MultiIter_DATA(it, 2))[0] + c_val = (<int64_t*>np.PyArray_MultiIter_DATA(it, 3))[0] + randoms_data[i] = f(state, a_val, b_val, c_val) + + np.PyArray_MultiIter_NEXT(it) + + return randoms + +cdef object discrete_broadcast_i(void *func, void *state, object size, object lock, + np.ndarray a_arr, object a_name, constraint_type a_constraint): + cdef np.ndarray randoms + cdef int64_t *randoms_data + cdef np.broadcast it + cdef random_uint_i f = (<random_uint_i>func) + cdef np.npy_intp i, n + + if a_constraint != CONS_NONE: + check_array_constraint(a_arr, a_name, a_constraint) + + if size is not None: + randoms = <np.ndarray>np.empty(size, np.int64) + else: + randoms = np.PyArray_SimpleNew(np.PyArray_NDIM(a_arr), np.PyArray_DIMS(a_arr), np.NPY_INT64) + + randoms_data = <int64_t *>np.PyArray_DATA(randoms) + n = np.PyArray_SIZE(randoms) + + it = np.PyArray_MultiIterNew2(randoms, a_arr) + with lock, nogil: + for i in range(n): + a_val = (<int64_t*>np.PyArray_MultiIter_DATA(it, 1))[0] + randoms_data[i] = f(state, a_val) + + np.PyArray_MultiIter_NEXT(it) + + return randoms + +# Needs double <vec>, double-double <vec>, double-int64_t<vec>, int64_t <vec>, int64_t-int64_t-int64_t +cdef object disc(void *func, void *state, object size, object lock, + int narg_double, int narg_int64, + object a, object a_name, constraint_type a_constraint, + object b, object b_name, constraint_type b_constraint, + object c, object c_name, constraint_type c_constraint): + + cdef double _da = 0, _db = 0 + cdef int64_t _ia = 0, _ib = 0, _ic = 0 + cdef bint is_scalar = True + if narg_double > 0: + a_arr = <np.ndarray>np.PyArray_FROM_OTF(a, np.NPY_DOUBLE, np.NPY_ALIGNED) + is_scalar = is_scalar and np.PyArray_NDIM(a_arr) == 0 + if narg_double > 1: + b_arr = <np.ndarray>np.PyArray_FROM_OTF(b, np.NPY_DOUBLE, np.NPY_ALIGNED) + is_scalar = is_scalar and np.PyArray_NDIM(b_arr) == 0 + elif narg_int64 == 1: + b_arr = <np.ndarray>np.PyArray_FROM_OTF(b, np.NPY_INT64, np.NPY_ALIGNED) + is_scalar = is_scalar and np.PyArray_NDIM(b_arr) == 0 + else: + if narg_int64 > 0: + a_arr = <np.ndarray>np.PyArray_FROM_OTF(a, np.NPY_INT64, np.NPY_ALIGNED) + is_scalar = is_scalar and np.PyArray_NDIM(a_arr) == 0 + if narg_int64 > 1: + b_arr = <np.ndarray>np.PyArray_FROM_OTF(b, np.NPY_INT64, np.NPY_ALIGNED) + is_scalar = is_scalar and np.PyArray_NDIM(b_arr) == 0 + if narg_int64 > 2: + c_arr = <np.ndarray>np.PyArray_FROM_OTF(c, np.NPY_INT64, np.NPY_ALIGNED) + is_scalar = is_scalar and np.PyArray_NDIM(c_arr) == 0 + + if not is_scalar: + if narg_int64 == 0: + if narg_double == 1: + return discrete_broadcast_d(func, state, size, lock, + a_arr, a_name, a_constraint) + elif narg_double == 2: + return discrete_broadcast_dd(func, state, size, lock, + a_arr, a_name, a_constraint, + b_arr, b_name, b_constraint) + elif narg_int64 == 1: + if narg_double == 0: + return discrete_broadcast_i(func, state, size, lock, + a_arr, a_name, a_constraint) + elif narg_double == 1: + return discrete_broadcast_di(func, state, size, lock, + a_arr, a_name, a_constraint, + b_arr, b_name, b_constraint) + else: + raise NotImplementedError("No vector path available") + + if narg_double > 0: + _da = PyFloat_AsDouble(a) + if a_constraint != CONS_NONE and is_scalar: + check_constraint(_da, a_name, a_constraint) + + if narg_double > 1: + _db = PyFloat_AsDouble(b) + if b_constraint != CONS_NONE and is_scalar: + check_constraint(_db, b_name, b_constraint) + elif narg_int64 == 1: + _ib = <int64_t>b + if b_constraint != CONS_NONE and is_scalar: + check_constraint(<double>_ib, b_name, b_constraint) + else: + if narg_int64 > 0: + _ia = <int64_t>a + if a_constraint != CONS_NONE and is_scalar: + check_constraint(<double>_ia, a_name, a_constraint) + if narg_int64 > 1: + _ib = <int64_t>b + if b_constraint != CONS_NONE and is_scalar: + check_constraint(<double>_ib, b_name, b_constraint) + if narg_int64 > 2: + _ic = <int64_t>c + if c_constraint != CONS_NONE and is_scalar: + check_constraint(<double>_ic, c_name, c_constraint) + + if size is None: + with lock: + if narg_int64 == 0: + if narg_double == 0: + return (<random_uint_0>func)(state) + elif narg_double == 1: + return (<random_uint_d>func)(state, _da) + elif narg_double == 2: + return (<random_uint_dd>func)(state, _da, _db) + elif narg_int64 == 1: + if narg_double == 0: + return (<random_uint_i>func)(state, _ia) + if narg_double == 1: + return (<random_uint_di>func)(state, _da, _ib) + else: + return (<random_uint_iii>func)(state, _ia, _ib, _ic) + + cdef np.npy_intp i, n + cdef np.ndarray randoms = <np.ndarray>np.empty(size, np.int64) + cdef np.int64_t *randoms_data + cdef random_uint_0 f0 + cdef random_uint_d fd + cdef random_uint_dd fdd + cdef random_uint_di fdi + cdef random_uint_i fi + cdef random_uint_iii fiii + + n = np.PyArray_SIZE(randoms) + randoms_data = <np.int64_t *>np.PyArray_DATA(randoms) + + with lock, nogil: + if narg_int64 == 0: + if narg_double == 0: + f0 = (<random_uint_0>func) + for i in range(n): + randoms_data[i] = f0(state) + elif narg_double == 1: + fd = (<random_uint_d>func) + for i in range(n): + randoms_data[i] = fd(state, _da) + elif narg_double == 2: + fdd = (<random_uint_dd>func) + for i in range(n): + randoms_data[i] = fdd(state, _da, _db) + elif narg_int64 == 1: + if narg_double == 0: + fi = (<random_uint_i>func) + for i in range(n): + randoms_data[i] = fi(state, _ia) + if narg_double == 1: + fdi = (<random_uint_di>func) + for i in range(n): + randoms_data[i] = fdi(state, _da, _ib) + else: + fiii = (<random_uint_iii>func) + for i in range(n): + randoms_data[i] = fiii(state, _ia, _ib, _ic) + + return randoms + + +cdef object cont_broadcast_1_f(void *func, bitgen_t *state, object size, object lock, + np.ndarray a_arr, object a_name, constraint_type a_constraint, + object out): + + cdef np.ndarray randoms + cdef float a_val + cdef float *randoms_data + cdef np.broadcast it + cdef random_float_1 f = (<random_float_1>func) + cdef np.npy_intp i, n + + if a_constraint != CONS_NONE: + check_array_constraint(a_arr, a_name, a_constraint) + + if size is not None and out is None: + randoms = <np.ndarray>np.empty(size, np.float32) + elif out is None: + randoms = np.PyArray_SimpleNew(np.PyArray_NDIM(a_arr), + np.PyArray_DIMS(a_arr), + np.NPY_FLOAT32) + else: + randoms = <np.ndarray>out + + randoms_data = <float *>np.PyArray_DATA(randoms) + n = np.PyArray_SIZE(randoms) + it = np.PyArray_MultiIterNew2(randoms, a_arr) + + with lock, nogil: + for i in range(n): + a_val = (<float*>np.PyArray_MultiIter_DATA(it, 1))[0] + randoms_data[i] = f(state, a_val) + + np.PyArray_MultiIter_NEXT(it) + + return randoms + +cdef object cont_f(void *func, bitgen_t *state, object size, object lock, + object a, object a_name, constraint_type a_constraint, + object out): + + cdef np.ndarray a_arr, b_arr, c_arr + cdef float _a + cdef bint is_scalar = True + cdef int requirements = np.NPY_ALIGNED | np.NPY_FORCECAST + check_output(out, np.float32, size) + a_arr = <np.ndarray>np.PyArray_FROMANY(a, np.NPY_FLOAT32, 0, 0, requirements) + is_scalar = np.PyArray_NDIM(a_arr) == 0 + + if not is_scalar: + return cont_broadcast_1_f(func, state, size, lock, a_arr, a_name, a_constraint, out) + + _a = <float>PyFloat_AsDouble(a) + if a_constraint != CONS_NONE: + check_constraint(_a, a_name, a_constraint) + + if size is None and out is None: + with lock: + return (<random_float_1>func)(state, _a) + + cdef np.npy_intp i, n + cdef np.ndarray randoms + if out is None: + randoms = <np.ndarray>np.empty(size, np.float32) + else: + randoms = <np.ndarray>out + n = np.PyArray_SIZE(randoms) + + cdef float *randoms_data = <float *>np.PyArray_DATA(randoms) + cdef random_float_1 f1 = <random_float_1>func + + with lock, nogil: + for i in range(n): + randoms_data[i] = f1(state, _a) + + if out is None: + return randoms + else: + return out |