1#!python
2#cython: wraparound=False, nonecheck=False, boundscheck=False, cdivision=True, language_level=3
3from collections import namedtuple
4from cpython cimport PyFloat_AsDouble
5import sys
6import numpy as np
7cimport numpy as np
8
9from libc.stdint cimport uintptr_t
10
11__all__ = ['interface']
12
13np.import_array()
14
15interface = namedtuple('interface', ['state_address', 'state', 'next_uint64',
16                                     'next_uint32', 'next_double',
17                                     'bit_generator'])
18
19cdef double LEGACY_POISSON_LAM_MAX = <double>np.iinfo('l').max - np.sqrt(np.iinfo('l').max)*10
20cdef double POISSON_LAM_MAX = <double>np.iinfo('int64').max - np.sqrt(np.iinfo('int64').max)*10
21
22cdef uint64_t MAXSIZE = <uint64_t>sys.maxsize
23
24
25cdef object benchmark(bitgen_t *bitgen, object lock, Py_ssize_t cnt, object method):
26    """Benchmark command used by BitGenerator"""
27    cdef Py_ssize_t i
28    if method==u'uint64':
29        with lock, nogil:
30            for i in range(cnt):
31                bitgen.next_uint64(bitgen.state)
32    elif method==u'double':
33        with lock, nogil:
34            for i in range(cnt):
35                bitgen.next_double(bitgen.state)
36    else:
37        raise ValueError('Unknown method')
38
39
40cdef object random_raw(bitgen_t *bitgen, object lock, object size, object output):
41    """
42    random_raw(self, size=None)
43
44    Return randoms as generated by the underlying PRNG
45
46    Parameters
47    ----------
48    bitgen : BitGenerator
49        Address of the bit generator struct
50    lock : Threading.Lock
51        Lock provided by the bit generator
52    size : int or tuple of ints, optional
53        Output shape.  If the given shape is, e.g., ``(m, n, k)``, then
54        ``m * n * k`` samples are drawn.  Default is None, in which case a
55        single value is returned.
56    output : bool, optional
57        Output values.  Used for performance testing since the generated
58        values are not returned.
59
60    Returns
61    -------
62    out : uint or ndarray
63        Drawn samples.
64
65    Notes
66    -----
67    This method directly exposes the the raw underlying pseudo-random
68    number generator. All values are returned as unsigned 64-bit
69    values irrespective of the number of bits produced by the PRNG.
70
71    See the class docstring for the number of bits returned.
72    """
73    cdef np.ndarray randoms
74    cdef uint64_t *randoms_data
75    cdef Py_ssize_t i, n
76
77    if not output:
78        if size is None:
79            with lock:
80                bitgen.next_raw(bitgen.state)
81            return None
82        n = np.asarray(size).sum()
83        with lock, nogil:
84            for i in range(n):
85                bitgen.next_raw(bitgen.state)
86        return None
87
88    if size is None:
89        with lock:
90            return bitgen.next_raw(bitgen.state)
91
92    randoms = <np.ndarray>np.empty(size, np.uint64)
93    randoms_data = <uint64_t*>np.PyArray_DATA(randoms)
94    n = np.PyArray_SIZE(randoms)
95
96    with lock, nogil:
97        for i in range(n):
98            randoms_data[i] = bitgen.next_raw(bitgen.state)
99    return randoms
100
101cdef object prepare_cffi(bitgen_t *bitgen):
102    """
103    Bundles the interfaces to interact with a BitGenerator using cffi
104
105    Parameters
106    ----------
107    bitgen : pointer
108        A pointer to a BitGenerator instance
109
110    Returns
111    -------
112    interface : namedtuple
113        The functions required to interface with the BitGenerator using cffi
114
115        * state_address - Memory address of the state struct
116        * state - pointer to the state struct
117        * next_uint64 - function pointer to produce 64 bit integers
118        * next_uint32 - function pointer to produce 32 bit integers
119        * next_double - function pointer to produce doubles
120        * bit_generator - pointer to the BitGenerator struct
121    """
122    try:
123        import cffi
124    except ImportError:
125        raise ImportError('cffi cannot be imported.')
126
127    ffi = cffi.FFI()
128    _cffi = interface(<uintptr_t>bitgen.state,
129                      ffi.cast('void *', <uintptr_t>bitgen.state),
130                      ffi.cast('uint64_t (*)(void *)', <uintptr_t>bitgen.next_uint64),
131                      ffi.cast('uint32_t (*)(void *)', <uintptr_t>bitgen.next_uint32),
132                      ffi.cast('double (*)(void *)', <uintptr_t>bitgen.next_double),
133                      ffi.cast('void *', <uintptr_t>bitgen))
134    return _cffi
135
136cdef object prepare_ctypes(bitgen_t *bitgen):
137    """
138    Bundles the interfaces to interact with a BitGenerator using ctypes
139
140    Parameters
141    ----------
142    bitgen : pointer
143        A pointer to a BitGenerator instance
144
145    Returns
146    -------
147    interface : namedtuple
148        The functions required to interface with the BitGenerator using ctypes:
149
150        * state_address - Memory address of the state struct
151        * state - pointer to the state struct
152        * next_uint64 - function pointer to produce 64 bit integers
153        * next_uint32 - function pointer to produce 32 bit integers
154        * next_double - function pointer to produce doubles
155        * bit_generator - pointer to the BitGenerator struct
156    """
157    import ctypes
158
159    _ctypes = interface(<uintptr_t>bitgen.state,
160                        ctypes.c_void_p(<uintptr_t>bitgen.state),
161                        ctypes.cast(<uintptr_t>bitgen.next_uint64,
162                                    ctypes.CFUNCTYPE(ctypes.c_uint64,
163                                                     ctypes.c_void_p)),
164                        ctypes.cast(<uintptr_t>bitgen.next_uint32,
165                                    ctypes.CFUNCTYPE(ctypes.c_uint32,
166                                                     ctypes.c_void_p)),
167                        ctypes.cast(<uintptr_t>bitgen.next_double,
168                                    ctypes.CFUNCTYPE(ctypes.c_double,
169                                                     ctypes.c_void_p)),
170                        ctypes.c_void_p(<uintptr_t>bitgen))
171    return _ctypes
172
173cdef double kahan_sum(double *darr, np.npy_intp n):
174    cdef double c, y, t, sum
175    cdef np.npy_intp i
176    sum = darr[0]
177    c = 0.0
178    for i in range(1, n):
179        y = darr[i] - c
180        t = sum + y
181        c = (t-sum) - y
182        sum = t
183    return sum
184
185
186cdef object wrap_int(object val, object bits):
187    """Wraparound to place an integer into the interval [0, 2**bits)"""
188    mask = ~(~int(0) << bits)
189    return val & mask
190
191
192cdef np.ndarray int_to_array(object value, object name, object bits, object uint_size):
193    """Convert a large integer to an array of unsigned integers"""
194    len = bits // uint_size
195    value = np.asarray(value)
196    if uint_size == 32:
197        dtype = np.uint32
198    elif uint_size == 64:
199        dtype = np.uint64
200    else:
201        raise ValueError('Unknown uint_size')
202    if value.shape == ():
203        value = int(value)
204        upper = int(2)**int(bits)
205        if value < 0 or value >= upper:
206            raise ValueError('{name} must be positive and '
207                             'less than 2**{bits}.'.format(name=name, bits=bits))
208
209        out = np.empty(len, dtype=dtype)
210        for i in range(len):
211            out[i] = value % 2**int(uint_size)
212            value >>= int(uint_size)
213    else:
214        out = value.astype(dtype)
215        if out.shape != (len,):
216            raise ValueError('{name} must have {len} elements when using '
217                             'array form'.format(name=name, len=len))
218    return out
219
220
221cdef validate_output_shape(iter_shape, np.ndarray output):
222    cdef np.npy_intp *shape
223    cdef ndim, i
224    cdef bint error
225    dims = np.PyArray_DIMS(output)
226    ndim = np.PyArray_NDIM(output)
227    output_shape = tuple((dims[i] for i in range(ndim)))
228    if iter_shape != output_shape:
229        raise ValueError(
230            f"Output size {output_shape} is not compatible with broadcast "
231            f"dimensions of inputs {iter_shape}."
232        )
233
234
235cdef check_output(object out, object dtype, object size):
236    if out is None:
237        return
238    cdef np.ndarray out_array = <np.ndarray>out
239    if not (np.PyArray_CHKFLAGS(out_array, np.NPY_CARRAY) or
240            np.PyArray_CHKFLAGS(out_array, np.NPY_FARRAY)):
241        raise ValueError('Supplied output array is not contiguous, writable or aligned.')
242    if out_array.dtype != dtype:
243        raise TypeError('Supplied output array has the wrong type. '
244                        'Expected {0}, got {1}'.format(np.dtype(dtype), out_array.dtype))
245    if size is not None:
246        try:
247            tup_size = tuple(size)
248        except TypeError:
249            tup_size = tuple([size])
250        if tup_size != out.shape:
251            raise ValueError('size must match out.shape when used together')
252
253
254cdef object double_fill(void *func, bitgen_t *state, object size, object lock, object out):
255    cdef random_double_fill random_func = (<random_double_fill>func)
256    cdef double out_val
257    cdef double *out_array_data
258    cdef np.ndarray out_array
259    cdef np.npy_intp i, n
260
261    if size is None and out is None:
262        with lock:
263            random_func(state, 1, &out_val)
264            return out_val
265
266    if out is not None:
267        check_output(out, np.float64, size)
268        out_array = <np.ndarray>out
269    else:
270        out_array = <np.ndarray>np.empty(size, np.double)
271
272    n = np.PyArray_SIZE(out_array)
273    out_array_data = <double *>np.PyArray_DATA(out_array)
274    with lock, nogil:
275        random_func(state, n, out_array_data)
276    return out_array
277
278cdef object float_fill(void *func, bitgen_t *state, object size, object lock, object out):
279    cdef random_float_fill random_func = (<random_float_fill>func)
280    cdef float out_val
281    cdef float *out_array_data
282    cdef np.ndarray out_array
283    cdef np.npy_intp i, n
284
285    if size is None and out is None:
286        with lock:
287            random_func(state, 1, &out_val)
288            return out_val
289
290    if out is not None:
291        check_output(out, np.float32, size)
292        out_array = <np.ndarray>out
293    else:
294        out_array = <np.ndarray>np.empty(size, np.float32)
295
296    n = np.PyArray_SIZE(out_array)
297    out_array_data = <float *>np.PyArray_DATA(out_array)
298    with lock, nogil:
299        random_func(state, n, out_array_data)
300    return out_array
301
302cdef object float_fill_from_double(void *func, bitgen_t *state, object size, object lock, object out):
303    cdef random_double_0 random_func = (<random_double_0>func)
304    cdef float *out_array_data
305    cdef np.ndarray out_array
306    cdef np.npy_intp i, n
307
308    if size is None and out is None:
309        with lock:
310            return <float>random_func(state)
311
312    if out is not None:
313        check_output(out, np.float32, size)
314        out_array = <np.ndarray>out
315    else:
316        out_array = <np.ndarray>np.empty(size, np.float32)
317
318    n = np.PyArray_SIZE(out_array)
319    out_array_data = <float *>np.PyArray_DATA(out_array)
320    with lock, nogil:
321        for i in range(n):
322            out_array_data[i] = <float>random_func(state)
323    return out_array
324
325
326cdef int check_array_constraint(np.ndarray val, object name, constraint_type cons) except -1:
327    if cons == CONS_NON_NEGATIVE:
328        if np.any(np.logical_and(np.logical_not(np.isnan(val)), np.signbit(val))):
329            raise ValueError(name + " < 0")
330    elif cons == CONS_POSITIVE or cons == CONS_POSITIVE_NOT_NAN:
331        if cons == CONS_POSITIVE_NOT_NAN and np.any(np.isnan(val)):
332            raise ValueError(name + " must not be NaN")
333        elif np.any(np.less_equal(val, 0)):
334            raise ValueError(name + " <= 0")
335    elif cons == CONS_BOUNDED_0_1:
336        if not np.all(np.greater_equal(val, 0)) or \
337                not np.all(np.less_equal(val, 1)):
338            raise ValueError("{0} < 0, {0} > 1 or {0} contains NaNs".format(name))
339    elif cons == CONS_BOUNDED_GT_0_1:
340        if not np.all(np.greater(val, 0)) or not np.all(np.less_equal(val, 1)):
341            raise ValueError("{0} <= 0, {0} > 1 or {0} contains NaNs".format(name))
342    elif cons == CONS_GT_1:
343        if not np.all(np.greater(val, 1)):
344            raise ValueError("{0} <= 1 or {0} contains NaNs".format(name))
345    elif cons == CONS_GTE_1:
346        if not np.all(np.greater_equal(val, 1)):
347            raise ValueError("{0} < 1 or {0} contains NaNs".format(name))
348    elif cons == CONS_POISSON:
349        if not np.all(np.less_equal(val, POISSON_LAM_MAX)):
350            raise ValueError("{0} value too large".format(name))
351        elif not np.all(np.greater_equal(val, 0.0)):
352            raise ValueError("{0} < 0 or {0} contains NaNs".format(name))
353    elif cons == LEGACY_CONS_POISSON:
354        if not np.all(np.less_equal(val, LEGACY_POISSON_LAM_MAX)):
355            raise ValueError("{0} value too large".format(name))
356        elif not np.all(np.greater_equal(val, 0.0)):
357            raise ValueError("{0} < 0 or {0} contains NaNs".format(name))
358
359    return 0
360
361
362cdef int check_constraint(double val, object name, constraint_type cons) except -1:
363    cdef bint is_nan
364    if cons == CONS_NON_NEGATIVE:
365        if not np.isnan(val) and np.signbit(val):
366            raise ValueError(name + " < 0")
367    elif cons == CONS_POSITIVE or cons == CONS_POSITIVE_NOT_NAN:
368        if cons == CONS_POSITIVE_NOT_NAN and np.isnan(val):
369            raise ValueError(name + " must not be NaN")
370        elif val <= 0:
371            raise ValueError(name + " <= 0")
372    elif cons == CONS_BOUNDED_0_1:
373        if not (val >= 0) or not (val <= 1):
374            raise ValueError("{0} < 0, {0} > 1 or {0} is NaN".format(name))
375    elif cons == CONS_BOUNDED_GT_0_1:
376        if not val >0 or not val <= 1:
377            raise ValueError("{0} <= 0, {0} > 1 or {0} contains NaNs".format(name))
378    elif cons == CONS_GT_1:
379        if not (val > 1):
380            raise ValueError("{0} <= 1 or {0} is NaN".format(name))
381    elif cons == CONS_GTE_1:
382        if not (val >= 1):
383            raise ValueError("{0} < 1 or {0} is NaN".format(name))
384    elif cons == CONS_POISSON:
385        if not (val >= 0):
386            raise ValueError("{0} < 0 or {0} is NaN".format(name))
387        elif not (val <= POISSON_LAM_MAX):
388            raise ValueError(name + " value too large")
389    elif cons == LEGACY_CONS_POISSON:
390        if not (val >= 0):
391            raise ValueError("{0} < 0 or {0} is NaN".format(name))
392        elif not (val <= LEGACY_POISSON_LAM_MAX):
393            raise ValueError(name + " value too large")
394
395    return 0
396
397cdef object cont_broadcast_1(void *func, void *state, object size, object lock,
398                             np.ndarray a_arr, object a_name, constraint_type a_constraint,
399                             object out):
400
401    cdef np.ndarray randoms
402    cdef double a_val
403    cdef double *randoms_data
404    cdef np.broadcast it
405    cdef random_double_1 f = (<random_double_1>func)
406    cdef np.npy_intp i, n
407
408    if a_constraint != CONS_NONE:
409        check_array_constraint(a_arr, a_name, a_constraint)
410
411    if size is not None and out is None:
412        randoms = <np.ndarray>np.empty(size, np.double)
413    elif out is None:
414        randoms = np.PyArray_SimpleNew(np.PyArray_NDIM(a_arr), np.PyArray_DIMS(a_arr), np.NPY_DOUBLE)
415    else:
416        randoms = <np.ndarray>out
417
418    randoms_data = <double *>np.PyArray_DATA(randoms)
419    n = np.PyArray_SIZE(randoms)
420    it = np.PyArray_MultiIterNew2(randoms, a_arr)
421    validate_output_shape(it.shape, randoms)
422
423    with lock, nogil:
424        for i in range(n):
425            a_val = (<double*>np.PyArray_MultiIter_DATA(it, 1))[0]
426            randoms_data[i] = f(state, a_val)
427
428            np.PyArray_MultiIter_NEXT(it)
429
430    return randoms
431
432cdef object cont_broadcast_2(void *func, void *state, object size, object lock,
433                 np.ndarray a_arr, object a_name, constraint_type a_constraint,
434                 np.ndarray b_arr, object b_name, constraint_type b_constraint):
435    cdef np.ndarray randoms
436    cdef double a_val, b_val
437    cdef double *randoms_data
438    cdef np.broadcast it
439    cdef random_double_2 f = (<random_double_2>func)
440    cdef np.npy_intp i, n
441
442    if a_constraint != CONS_NONE:
443        check_array_constraint(a_arr, a_name, a_constraint)
444
445    if b_constraint != CONS_NONE:
446        check_array_constraint(b_arr, b_name, b_constraint)
447
448    if size is not None:
449        randoms = <np.ndarray>np.empty(size, np.double)
450    else:
451        it = np.PyArray_MultiIterNew2(a_arr, b_arr)
452        randoms = <np.ndarray>np.empty(it.shape, np.double)
453        # randoms = np.PyArray_SimpleNew(it.nd, np.PyArray_DIMS(it), np.NPY_DOUBLE)
454
455    randoms_data = <double *>np.PyArray_DATA(randoms)
456    n = np.PyArray_SIZE(randoms)
457
458    it = np.PyArray_MultiIterNew3(randoms, a_arr, b_arr)
459    validate_output_shape(it.shape, randoms)
460
461    with lock, nogil:
462        for i in range(n):
463            a_val = (<double*>np.PyArray_MultiIter_DATA(it, 1))[0]
464            b_val = (<double*>np.PyArray_MultiIter_DATA(it, 2))[0]
465            randoms_data[i] = f(state, a_val, b_val)
466
467            np.PyArray_MultiIter_NEXT(it)
468
469    return randoms
470
471cdef object cont_broadcast_3(void *func, void *state, object size, object lock,
472                             np.ndarray a_arr, object a_name, constraint_type a_constraint,
473                             np.ndarray b_arr, object b_name, constraint_type b_constraint,
474                             np.ndarray c_arr, object c_name, constraint_type c_constraint):
475    cdef np.ndarray randoms
476    cdef double a_val, b_val, c_val
477    cdef double *randoms_data
478    cdef np.broadcast it
479    cdef random_double_3 f = (<random_double_3>func)
480    cdef np.npy_intp i, n
481
482    if a_constraint != CONS_NONE:
483        check_array_constraint(a_arr, a_name, a_constraint)
484
485    if b_constraint != CONS_NONE:
486        check_array_constraint(b_arr, b_name, b_constraint)
487
488    if c_constraint != CONS_NONE:
489        check_array_constraint(c_arr, c_name, c_constraint)
490
491    if size is not None:
492        randoms = <np.ndarray>np.empty(size, np.double)
493    else:
494        it = np.PyArray_MultiIterNew3(a_arr, b_arr, c_arr)
495        # randoms = np.PyArray_SimpleNew(it.nd, np.PyArray_DIMS(it), np.NPY_DOUBLE)
496        randoms = <np.ndarray>np.empty(it.shape, np.double)
497
498    randoms_data = <double *>np.PyArray_DATA(randoms)
499    n = np.PyArray_SIZE(randoms)
500
501    it = np.PyArray_MultiIterNew4(randoms, a_arr, b_arr, c_arr)
502    validate_output_shape(it.shape, randoms)
503
504    with lock, nogil:
505        for i in range(n):
506            a_val = (<double*>np.PyArray_MultiIter_DATA(it, 1))[0]
507            b_val = (<double*>np.PyArray_MultiIter_DATA(it, 2))[0]
508            c_val = (<double*>np.PyArray_MultiIter_DATA(it, 3))[0]
509            randoms_data[i] = f(state, a_val, b_val, c_val)
510
511            np.PyArray_MultiIter_NEXT(it)
512
513    return randoms
514
515cdef object cont(void *func, void *state, object size, object lock, int narg,
516                 object a, object a_name, constraint_type a_constraint,
517                 object b, object b_name, constraint_type b_constraint,
518                 object c, object c_name, constraint_type c_constraint,
519                 object out):
520
521    cdef np.ndarray a_arr, b_arr, c_arr
522    cdef double _a = 0.0, _b = 0.0, _c = 0.0
523    cdef bint is_scalar = True
524    check_output(out, np.float64, size)
525    if narg > 0:
526        a_arr = <np.ndarray>np.PyArray_FROM_OTF(a, np.NPY_DOUBLE, np.NPY_ALIGNED)
527        is_scalar = is_scalar and np.PyArray_NDIM(a_arr) == 0
528    if narg > 1:
529        b_arr = <np.ndarray>np.PyArray_FROM_OTF(b, np.NPY_DOUBLE, np.NPY_ALIGNED)
530        is_scalar = is_scalar and np.PyArray_NDIM(b_arr) == 0
531    if narg == 3:
532        c_arr = <np.ndarray>np.PyArray_FROM_OTF(c, np.NPY_DOUBLE, np.NPY_ALIGNED)
533        is_scalar = is_scalar and np.PyArray_NDIM(c_arr) == 0
534
535    if not is_scalar:
536        if narg == 1:
537            return cont_broadcast_1(func, state, size, lock,
538                                    a_arr, a_name, a_constraint,
539                                    out)
540        elif narg == 2:
541            return cont_broadcast_2(func, state, size, lock,
542                                    a_arr, a_name, a_constraint,
543                                    b_arr, b_name, b_constraint)
544        else:
545            return cont_broadcast_3(func, state, size, lock,
546                                    a_arr, a_name, a_constraint,
547                                    b_arr, b_name, b_constraint,
548                                    c_arr, c_name, c_constraint)
549
550    if narg > 0:
551        _a = PyFloat_AsDouble(a)
552        if a_constraint != CONS_NONE and is_scalar:
553            check_constraint(_a, a_name, a_constraint)
554    if narg > 1:
555        _b = PyFloat_AsDouble(b)
556        if b_constraint != CONS_NONE:
557            check_constraint(_b, b_name, b_constraint)
558    if narg == 3:
559        _c = PyFloat_AsDouble(c)
560        if c_constraint != CONS_NONE and is_scalar:
561            check_constraint(_c, c_name, c_constraint)
562
563    if size is None and out is None:
564        with lock:
565            if narg == 0:
566                return (<random_double_0>func)(state)
567            elif narg == 1:
568                return (<random_double_1>func)(state, _a)
569            elif narg == 2:
570                return (<random_double_2>func)(state, _a, _b)
571            elif narg == 3:
572                return (<random_double_3>func)(state, _a, _b, _c)
573
574    cdef np.npy_intp i, n
575    cdef np.ndarray randoms
576    if out is None:
577        randoms = <np.ndarray>np.empty(size)
578    else:
579        randoms = <np.ndarray>out
580    n = np.PyArray_SIZE(randoms)
581
582    cdef double *randoms_data = <double *>np.PyArray_DATA(randoms)
583    cdef random_double_0 f0
584    cdef random_double_1 f1
585    cdef random_double_2 f2
586    cdef random_double_3 f3
587
588    with lock, nogil:
589        if narg == 0:
590            f0 = (<random_double_0>func)
591            for i in range(n):
592                randoms_data[i] = f0(state)
593        elif narg == 1:
594            f1 = (<random_double_1>func)
595            for i in range(n):
596                randoms_data[i] = f1(state, _a)
597        elif narg == 2:
598            f2 = (<random_double_2>func)
599            for i in range(n):
600                randoms_data[i] = f2(state, _a, _b)
601        elif narg == 3:
602            f3 = (<random_double_3>func)
603            for i in range(n):
604                randoms_data[i] = f3(state, _a, _b, _c)
605
606    if out is None:
607        return randoms
608    else:
609        return out
610
611cdef object discrete_broadcast_d(void *func, void *state, object size, object lock,
612                                 np.ndarray a_arr, object a_name, constraint_type a_constraint):
613
614    cdef np.ndarray randoms
615    cdef int64_t *randoms_data
616    cdef np.broadcast it
617    cdef random_uint_d f = (<random_uint_d>func)
618    cdef np.npy_intp i, n
619
620    if a_constraint != CONS_NONE:
621        check_array_constraint(a_arr, a_name, a_constraint)
622
623    if size is not None:
624        randoms = np.empty(size, np.int64)
625    else:
626        # randoms = np.empty(np.shape(a_arr), np.double)
627        randoms = np.PyArray_SimpleNew(np.PyArray_NDIM(a_arr), np.PyArray_DIMS(a_arr), np.NPY_INT64)
628
629    randoms_data = <int64_t *>np.PyArray_DATA(randoms)
630    n = np.PyArray_SIZE(randoms)
631
632    it = np.PyArray_MultiIterNew2(randoms, a_arr)
633    validate_output_shape(it.shape, randoms)
634
635    with lock, nogil:
636        for i in range(n):
637            a_val = (<double*>np.PyArray_MultiIter_DATA(it, 1))[0]
638            randoms_data[i] = f(state, a_val)
639
640            np.PyArray_MultiIter_NEXT(it)
641
642    return randoms
643
644cdef object discrete_broadcast_dd(void *func, void *state, object size, object lock,
645                                  np.ndarray a_arr, object a_name, constraint_type a_constraint,
646                                  np.ndarray b_arr, object b_name, constraint_type b_constraint):
647    cdef np.ndarray randoms
648    cdef int64_t *randoms_data
649    cdef np.broadcast it
650    cdef random_uint_dd f = (<random_uint_dd>func)
651    cdef np.npy_intp i, n
652
653    if a_constraint != CONS_NONE:
654        check_array_constraint(a_arr, a_name, a_constraint)
655    if b_constraint != CONS_NONE:
656        check_array_constraint(b_arr, b_name, b_constraint)
657
658    if size is not None:
659        randoms = <np.ndarray>np.empty(size, np.int64)
660    else:
661        it = np.PyArray_MultiIterNew2(a_arr, b_arr)
662        randoms = <np.ndarray>np.empty(it.shape, np.int64)
663        # randoms = np.PyArray_SimpleNew(it.nd, np.PyArray_DIMS(it), np.NPY_INT64)
664
665    randoms_data = <int64_t *>np.PyArray_DATA(randoms)
666    n = np.PyArray_SIZE(randoms)
667
668    it = np.PyArray_MultiIterNew3(randoms, a_arr, b_arr)
669    validate_output_shape(it.shape, randoms)
670
671    with lock, nogil:
672        for i in range(n):
673            a_val = (<double*>np.PyArray_MultiIter_DATA(it, 1))[0]
674            b_val = (<double*>np.PyArray_MultiIter_DATA(it, 2))[0]
675            randoms_data[i] = f(state, a_val, b_val)
676
677            np.PyArray_MultiIter_NEXT(it)
678
679    return randoms
680
681cdef object discrete_broadcast_di(void *func, void *state, object size, object lock,
682                                  np.ndarray a_arr, object a_name, constraint_type a_constraint,
683                                  np.ndarray b_arr, object b_name, constraint_type b_constraint):
684    cdef np.ndarray randoms
685    cdef int64_t *randoms_data
686    cdef np.broadcast it
687    cdef random_uint_di f = (<random_uint_di>func)
688    cdef np.npy_intp i, n
689
690    if a_constraint != CONS_NONE:
691        check_array_constraint(a_arr, a_name, a_constraint)
692
693    if b_constraint != CONS_NONE:
694        check_array_constraint(b_arr, b_name, b_constraint)
695
696    if size is not None:
697        randoms = <np.ndarray>np.empty(size, np.int64)
698    else:
699        it = np.PyArray_MultiIterNew2(a_arr, b_arr)
700        randoms = <np.ndarray>np.empty(it.shape, np.int64)
701
702    randoms_data = <int64_t *>np.PyArray_DATA(randoms)
703    n = np.PyArray_SIZE(randoms)
704
705    it = np.PyArray_MultiIterNew3(randoms, a_arr, b_arr)
706    validate_output_shape(it.shape, randoms)
707
708    with lock, nogil:
709        for i in range(n):
710            a_val = (<double*>np.PyArray_MultiIter_DATA(it, 1))[0]
711            b_val = (<int64_t*>np.PyArray_MultiIter_DATA(it, 2))[0]
712            (<int64_t*>np.PyArray_MultiIter_DATA(it, 0))[0] = f(state, a_val, b_val)
713
714            np.PyArray_MultiIter_NEXT(it)
715
716    return randoms
717
718cdef object discrete_broadcast_iii(void *func, void *state, object size, object lock,
719                                   np.ndarray a_arr, object a_name, constraint_type a_constraint,
720                                   np.ndarray b_arr, object b_name, constraint_type b_constraint,
721                                   np.ndarray c_arr, object c_name, constraint_type c_constraint):
722    cdef np.ndarray randoms
723    cdef int64_t *randoms_data
724    cdef np.broadcast it
725    cdef random_uint_iii f = (<random_uint_iii>func)
726    cdef np.npy_intp i, n
727
728    if a_constraint != CONS_NONE:
729        check_array_constraint(a_arr, a_name, a_constraint)
730
731    if b_constraint != CONS_NONE:
732        check_array_constraint(b_arr, b_name, b_constraint)
733
734    if c_constraint != CONS_NONE:
735        check_array_constraint(c_arr, c_name, c_constraint)
736
737    if size is not None:
738        randoms = <np.ndarray>np.empty(size, np.int64)
739    else:
740        it = np.PyArray_MultiIterNew3(a_arr, b_arr, c_arr)
741        randoms = <np.ndarray>np.empty(it.shape, np.int64)
742
743    randoms_data = <int64_t *>np.PyArray_DATA(randoms)
744    n = np.PyArray_SIZE(randoms)
745
746    it = np.PyArray_MultiIterNew4(randoms, a_arr, b_arr, c_arr)
747    validate_output_shape(it.shape, randoms)
748
749    with lock, nogil:
750        for i in range(n):
751            a_val = (<int64_t*>np.PyArray_MultiIter_DATA(it, 1))[0]
752            b_val = (<int64_t*>np.PyArray_MultiIter_DATA(it, 2))[0]
753            c_val = (<int64_t*>np.PyArray_MultiIter_DATA(it, 3))[0]
754            randoms_data[i] = f(state, a_val, b_val, c_val)
755
756            np.PyArray_MultiIter_NEXT(it)
757
758    return randoms
759
760cdef object discrete_broadcast_i(void *func, void *state, object size, object lock,
761                                 np.ndarray a_arr, object a_name, constraint_type a_constraint):
762    cdef np.ndarray randoms
763    cdef int64_t *randoms_data
764    cdef np.broadcast it
765    cdef random_uint_i f = (<random_uint_i>func)
766    cdef np.npy_intp i, n
767
768    if a_constraint != CONS_NONE:
769        check_array_constraint(a_arr, a_name, a_constraint)
770
771    if size is not None:
772        randoms = <np.ndarray>np.empty(size, np.int64)
773    else:
774        randoms = np.PyArray_SimpleNew(np.PyArray_NDIM(a_arr), np.PyArray_DIMS(a_arr), np.NPY_INT64)
775
776    randoms_data = <int64_t *>np.PyArray_DATA(randoms)
777    n = np.PyArray_SIZE(randoms)
778
779    it = np.PyArray_MultiIterNew2(randoms, a_arr)
780    validate_output_shape(it.shape, randoms)
781
782    with lock, nogil:
783        for i in range(n):
784            a_val = (<int64_t*>np.PyArray_MultiIter_DATA(it, 1))[0]
785            randoms_data[i] = f(state, a_val)
786
787            np.PyArray_MultiIter_NEXT(it)
788
789    return randoms
790
791# Needs double <vec>, double-double <vec>, double-int64_t<vec>, int64_t <vec>, int64_t-int64_t-int64_t
792cdef object disc(void *func, void *state, object size, object lock,
793                 int narg_double, int narg_int64,
794                 object a, object a_name, constraint_type a_constraint,
795                 object b, object b_name, constraint_type b_constraint,
796                 object c, object c_name, constraint_type c_constraint):
797
798    cdef double _da = 0, _db = 0
799    cdef int64_t _ia = 0, _ib = 0, _ic = 0
800    cdef bint is_scalar = True
801    if narg_double > 0:
802        a_arr = <np.ndarray>np.PyArray_FROM_OTF(a, np.NPY_DOUBLE, np.NPY_ALIGNED)
803        is_scalar = is_scalar and np.PyArray_NDIM(a_arr) == 0
804        if narg_double > 1:
805            b_arr = <np.ndarray>np.PyArray_FROM_OTF(b, np.NPY_DOUBLE, np.NPY_ALIGNED)
806            is_scalar = is_scalar and np.PyArray_NDIM(b_arr) == 0
807        elif narg_int64 == 1:
808            b_arr = <np.ndarray>np.PyArray_FROM_OTF(b, np.NPY_INT64, np.NPY_ALIGNED)
809            is_scalar = is_scalar and np.PyArray_NDIM(b_arr) == 0
810    else:
811        if narg_int64 > 0:
812            a_arr = <np.ndarray>np.PyArray_FROM_OTF(a, np.NPY_INT64, np.NPY_ALIGNED)
813            is_scalar = is_scalar and np.PyArray_NDIM(a_arr) == 0
814        if narg_int64 > 1:
815            b_arr = <np.ndarray>np.PyArray_FROM_OTF(b, np.NPY_INT64, np.NPY_ALIGNED)
816            is_scalar = is_scalar and np.PyArray_NDIM(b_arr) == 0
817        if narg_int64 > 2:
818            c_arr = <np.ndarray>np.PyArray_FROM_OTF(c, np.NPY_INT64, np.NPY_ALIGNED)
819            is_scalar = is_scalar and np.PyArray_NDIM(c_arr) == 0
820
821    if not is_scalar:
822        if narg_int64 == 0:
823            if narg_double == 1:
824                return discrete_broadcast_d(func, state, size, lock,
825                                            a_arr, a_name, a_constraint)
826            elif narg_double == 2:
827                return discrete_broadcast_dd(func, state, size, lock,
828                                             a_arr, a_name, a_constraint,
829                                             b_arr, b_name, b_constraint)
830        elif narg_int64 == 1:
831            if narg_double == 0:
832                return discrete_broadcast_i(func, state, size, lock,
833                                            a_arr, a_name, a_constraint)
834            elif narg_double == 1:
835                return discrete_broadcast_di(func, state, size, lock,
836                                             a_arr, a_name, a_constraint,
837                                             b_arr, b_name, b_constraint)
838        else:
839            raise NotImplementedError("No vector path available")
840
841    if narg_double > 0:
842        _da = PyFloat_AsDouble(a)
843        if a_constraint != CONS_NONE and is_scalar:
844            check_constraint(_da, a_name, a_constraint)
845
846        if narg_double > 1:
847            _db = PyFloat_AsDouble(b)
848            if b_constraint != CONS_NONE and is_scalar:
849                check_constraint(_db, b_name, b_constraint)
850        elif narg_int64 == 1:
851            _ib = <int64_t>b
852            if b_constraint != CONS_NONE and is_scalar:
853                check_constraint(<double>_ib, b_name, b_constraint)
854    else:
855        if narg_int64 > 0:
856            _ia = <int64_t>a
857            if a_constraint != CONS_NONE and is_scalar:
858                check_constraint(<double>_ia, a_name, a_constraint)
859        if narg_int64 > 1:
860            _ib = <int64_t>b
861            if b_constraint != CONS_NONE and is_scalar:
862                check_constraint(<double>_ib, b_name, b_constraint)
863        if narg_int64 > 2:
864            _ic = <int64_t>c
865            if c_constraint != CONS_NONE and is_scalar:
866                check_constraint(<double>_ic, c_name, c_constraint)
867
868    if size is None:
869        with lock:
870            if narg_int64 == 0:
871                if narg_double == 0:
872                    return (<random_uint_0>func)(state)
873                elif narg_double == 1:
874                    return (<random_uint_d>func)(state, _da)
875                elif narg_double == 2:
876                    return (<random_uint_dd>func)(state, _da, _db)
877            elif narg_int64 == 1:
878                if narg_double == 0:
879                    return (<random_uint_i>func)(state, _ia)
880                if narg_double == 1:
881                    return (<random_uint_di>func)(state, _da, _ib)
882            else:
883                return (<random_uint_iii>func)(state, _ia, _ib, _ic)
884
885    cdef np.npy_intp i, n
886    cdef np.ndarray randoms = <np.ndarray>np.empty(size, np.int64)
887    cdef np.int64_t *randoms_data
888    cdef random_uint_0 f0
889    cdef random_uint_d fd
890    cdef random_uint_dd fdd
891    cdef random_uint_di fdi
892    cdef random_uint_i fi
893    cdef random_uint_iii fiii
894
895    n = np.PyArray_SIZE(randoms)
896    randoms_data = <np.int64_t *>np.PyArray_DATA(randoms)
897
898    with lock, nogil:
899        if narg_int64 == 0:
900            if narg_double == 0:
901                f0 = (<random_uint_0>func)
902                for i in range(n):
903                    randoms_data[i] = f0(state)
904            elif narg_double == 1:
905                fd = (<random_uint_d>func)
906                for i in range(n):
907                    randoms_data[i] = fd(state, _da)
908            elif narg_double == 2:
909                fdd = (<random_uint_dd>func)
910                for i in range(n):
911                    randoms_data[i] = fdd(state, _da, _db)
912        elif narg_int64 == 1:
913            if narg_double == 0:
914                fi = (<random_uint_i>func)
915                for i in range(n):
916                    randoms_data[i] = fi(state, _ia)
917            if narg_double == 1:
918                fdi = (<random_uint_di>func)
919                for i in range(n):
920                    randoms_data[i] = fdi(state, _da, _ib)
921        else:
922            fiii = (<random_uint_iii>func)
923            for i in range(n):
924                randoms_data[i] = fiii(state, _ia, _ib, _ic)
925
926    return randoms
927
928
929cdef object cont_broadcast_1_f(void *func, bitgen_t *state, object size, object lock,
930                               np.ndarray a_arr, object a_name, constraint_type a_constraint,
931                               object out):
932
933    cdef np.ndarray randoms
934    cdef float a_val
935    cdef float *randoms_data
936    cdef np.broadcast it
937    cdef random_float_1 f = (<random_float_1>func)
938    cdef np.npy_intp i, n
939
940    if a_constraint != CONS_NONE:
941        check_array_constraint(a_arr, a_name, a_constraint)
942
943    if size is not None and out is None:
944        randoms = <np.ndarray>np.empty(size, np.float32)
945    elif out is None:
946        randoms = np.PyArray_SimpleNew(np.PyArray_NDIM(a_arr),
947                                       np.PyArray_DIMS(a_arr),
948                                       np.NPY_FLOAT32)
949    else:
950        randoms = <np.ndarray>out
951
952    randoms_data = <float *>np.PyArray_DATA(randoms)
953    n = np.PyArray_SIZE(randoms)
954    it = np.PyArray_MultiIterNew2(randoms, a_arr)
955    validate_output_shape(it.shape, randoms)
956
957    with lock, nogil:
958        for i in range(n):
959            a_val = (<float*>np.PyArray_MultiIter_DATA(it, 1))[0]
960            randoms_data[i] = f(state, a_val)
961
962            np.PyArray_MultiIter_NEXT(it)
963
964    return randoms
965
966cdef object cont_f(void *func, bitgen_t *state, object size, object lock,
967                   object a, object a_name, constraint_type a_constraint,
968                   object out):
969
970    cdef np.ndarray a_arr, b_arr, c_arr
971    cdef float _a
972    cdef bint is_scalar = True
973    cdef int requirements = np.NPY_ALIGNED | np.NPY_FORCECAST
974    check_output(out, np.float32, size)
975    a_arr = <np.ndarray>np.PyArray_FROMANY(a, np.NPY_FLOAT32, 0, 0, requirements)
976    is_scalar = np.PyArray_NDIM(a_arr) == 0
977
978    if not is_scalar:
979        return cont_broadcast_1_f(func, state, size, lock, a_arr, a_name, a_constraint, out)
980
981    _a = <float>PyFloat_AsDouble(a)
982    if a_constraint != CONS_NONE:
983        check_constraint(_a, a_name, a_constraint)
984
985    if size is None and out is None:
986        with lock:
987            return (<random_float_1>func)(state, _a)
988
989    cdef np.npy_intp i, n
990    cdef np.ndarray randoms
991    if out is None:
992        randoms = <np.ndarray>np.empty(size, np.float32)
993    else:
994        randoms = <np.ndarray>out
995    n = np.PyArray_SIZE(randoms)
996
997    cdef float *randoms_data = <float *>np.PyArray_DATA(randoms)
998    cdef random_float_1 f1 = <random_float_1>func
999
1000    with lock, nogil:
1001        for i in range(n):
1002            randoms_data[i] = f1(state, _a)
1003
1004    if out is None:
1005        return randoms
1006    else:
1007        return out
1008