REF: Refactor to a common path for more data types

Refactor random integers to a commmon path

Author: bashtage

randomstate/bounded_integers.pxi.in

@@ -84,7 +84,6 @@ cdef object _rand_{{nptype}}(low, high, size, aug_state *state, lock):
 
 {{endfor}}
 
-
 cdef inline uint64_t _gen_mask(uint64_t max_val) nogil:
     # Smallest bit mask >= max
     cdef uint64_t mask = max_val
@@ -96,6 +95,137 @@ cdef inline uint64_t _gen_mask(uint64_t max_val) nogil:
     mask |= mask >> 32
     return mask
 
+{{
+py:
+bc_ctypes = (('uint32', 'uint32', 'uint64', 'NPY_UINT64', 0, 0, 0, '0X100000000ULL'),
+          ('uint16', 'uint16', 'uint32', 'NPY_UINT32', 1, 16, 0, '0X10000UL'),
+          ('uint8', 'uint8', 'uint16', 'NPY_UINT16', 3, 8, 0, '0X100UL'),
+          ('bool','uint8', 'uint8', 'NPY_UINT8', 31, 1, 0, '0x2UL'),
+          ('int32', 'uint32', 'uint64', 'NPY_INT64', 0, 0, '-0x80000000LL', '0x80000000LL'),
+          ('int16', 'uint16', 'uint32', 'NPY_INT32', 1, 16, '-0x8000LL', '0x8000LL' ),
+          ('int8', 'uint8', 'uint16', 'NPY_INT16', 3, 8, '-0x80LL', '0x80LL' ),
+)}}
+
+{{for  nptype, utype, nptype_up, npctype, remaining, bitshift, lb, ub in bc_ctypes}}
+
+{{ py: otype = nptype + '_' if nptype == 'bool' else nptype }}
+
+cdef object _rand_{{nptype}}_combined(object low, object high, object size, aug_state *state, object lock):
+    """
+    _rand_{{nptype}}_combined(low, high, size, *state, lock)
+
+    Return random np.{{nptype}} integers between `low` and `high`, inclusive.
+
+    Return random integers from the "discrete uniform" distribution in the
+    closed interval [`low`, `high`).  If `high` is None (the default),
+    then results are from [0, `low`). On entry the arguments are presumed
+    to have been validated for size and order for the np.{{nptype}} type.
+
+    Parameters
+    ----------
+    low : int
+        Lowest (signed) integer to be drawn from the distribution (unless
+        ``high=None``, in which case this parameter is the *highest* such
+        integer).
+    high : int
+        If provided, the largest (signed) integer to be drawn from the
+        distribution (see above for behavior if ``high=None``).
+    size : int or tuple of ints
+        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.
+    rngstate : encapsulated pointer to rk_state
+        The specific type depends on the python version. In Python 2 it is
+        a PyCObject, in Python 3 a PyCapsule object.
+
+    Returns
+    -------
+    out : python scalar or ndarray of np.{{nptype}}
+          `size`-shaped array of random integers from the appropriate
+          distribution, or a single such random int if `size` not provided.
+
+    """
+    cdef {{utype}}_t off, val, mask
+    cdef uint32_t buf
+    cdef {{utype}}_t *out
+    cdef {{utype}}_t *out_data
+    cdef {{nptype_up}}_t rng, last_rng, low_v, high_v
+    cdef np.ndarray low_arr, high_arr, out_arr
+    cdef np.npy_intp i, cnt
+    cdef np.broadcast it
+    cdef int buf_rem = 0
+
+
+    low = np.asarray(low)
+    high = np.asarray(high)
+    if low.shape == high.shape == ():
+        low = int(low) # TODO: Cast appropriately?
+        high = int(high) # TODO: Cast appropriately?
+
+        if low < {{lb}}:
+            raise ValueError("low is out of bounds for {{nptype}}")
+        if high > {{ub}}:
+            raise ValueError("high is out of bounds for {{nptype}}")
+        if low >= high:
+            raise ValueError("low >= high")
+
+        high -= 1
+        rng = <{{utype}}_t>(high - low)
+        off = <{{utype}}_t>(<{{nptype}}_t>low)
+        if size is None:
+            with lock:
+                random_bounded_{{utype}}_fill(state, off, rng, 1, out)
+            return np.{{otype}}(<{{nptype}}_t>out[0])
+        else:
+            out_arr = <np.ndarray>np.empty(size, np.{{nptype}})
+            cnt = np.PyArray_SIZE(out_arr)
+            out = <{{utype}}_t *>np.PyArray_DATA(out_arr)
+            with lock, nogil:
+                random_bounded_{{utype}}_fill(state, off, rng, cnt, out)
+            return out_arr
+
+    # Array path
+    low_arr = <np.ndarray>low
+    high_arr = <np.ndarray>high
+    if np.any(np.less(low_arr, {{lb}})):
+        raise ValueError('low is out of bounds for {{nptype}}')
+    if np.any(np.greater(high_arr, {{ub}})):
+        raise ValueError('high is out of bounds for {{nptype}}')
+    if np.any(np.greater_equal(low_arr, high_arr)):
+        raise ValueError('low >= high')
+
+    low_arr = <np.ndarray>np.PyArray_FROM_OTF(low, np.{{npctype}}, np.NPY_ALIGNED | np.NPY_FORCECAST)
+    high_arr = <np.ndarray>np.PyArray_FROM_OTF(high, np.{{npctype}}, np.NPY_ALIGNED | np.NPY_FORCECAST)
+
+    if size is not None:
+        out_arr = <np.ndarray>np.empty(size, np.{{otype}})
+    else:
+        it = np.PyArray_MultiIterNew2(low_arr, high_arr)
+        out_arr = <np.ndarray>np.empty(it.shape, np.{{otype}})
+
+    it = np.PyArray_MultiIterNew3(low_arr, high_arr, out_arr)
+    out_data = <{{utype}}_t *>np.PyArray_DATA(out_arr)
+    n = np.PyArray_SIZE(out_arr)
+    mask = last_rng = 0
+    with lock, nogil:
+        for i in range(n):
+            low_v = (<{{nptype_up}}_t*>np.PyArray_MultiIter_DATA(it, 0))[0]
+            high_v = (<{{nptype_up}}_t*>np.PyArray_MultiIter_DATA(it, 1))[0]
+            rng = (high_v - 1) - low_v
+            off = <{{utype}}_t>(<{{nptype_up}}_t>low_v)
+
+            if rng != last_rng:
+                # Smallest bit mask >= max
+                mask = <{{utype}}_t>_gen_mask(rng)
+
+            out[i] = random_buffered_bounded_{{utype}}(state, off, rng, mask, &buf_rem, &buf)
+
+            np.PyArray_MultiIter_NEXT(it)
+
+    return out_arr
+{{endfor}}
+
+
 {{
 py:
 bc_ctypes = (('uint32', 'uint32', 'uint64', 'NPY_UINT64', 0, 0, 0, '0X100000000ULL'),

randomstate/distributions.c

@@ -1419,8 +1419,14 @@ inline uint64_t random_bounded_uint64(aug_state *state, uint64_t off, uint64_t r
 }
 
 
-inline uint32_t random_bounded_uint32(aug_state *state, uint32_t off, uint32_t rng, uint32_t mask)
+inline uint32_t random_buffered_bounded_uint32(aug_state *state, uint32_t off, uint32_t rng, uint32_t mask, int *bcnt, uint32_t *buf)
 {
+    /*
+     * The buffer and buffer count are not used here but are included to allow
+     * this function to be templated with the similar uint8 and uint16
+     * functions
+     */
+
     uint32_t val;
     if (rng == 0)
         return off;
@@ -1489,11 +1495,13 @@ void random_bounded_uint32_fill(aug_state *state, uint32_t off, uint32_t rng, np
 {
     uint32_t val, mask;
     npy_intp i;
+    uint32_t buf = 0;
+    int bcnt = 0;
 
     /* Smallest bit mask >= max */
     mask = (uint32_t)gen_mask(rng);
     for (i = 0; i < cnt; i++) {
-        out[i] =  random_bounded_uint32(state, off, rng, mask);
+        out[i] =  random_buffered_bounded_uint32(state, off, rng, mask, &bcnt, &buf);
     }
 }
 

randomstate/distributions.h

@@ -175,3 +175,9 @@ extern void random_standard_exponential_zig_float_fill(aug_state* state, npy_int
 extern double random_standard_gamma_zig_double(aug_state* state, double shape);
 
 extern float random_standard_gamma_zig_float(aug_state* state, float shape);
+
+inline uint32_t random_buffered_bounded_uint32(aug_state *state, uint32_t off, uint32_t rng, uint32_t mask, int *bcnt, uint32_t *buf);
+
+inline uint16_t random_buffered_bounded_uint16(aug_state *state, uint16_t off, uint16_t rng, uint16_t mask, int *bcnt, uint32_t *buf);
+
+inline uint8_t random_buffered_bounded_uint8(aug_state *state, uint8_t off, uint8_t rng, uint8_t mask, int *bcnt, uint32_t *buf);

randomstate/randomstate.pyx

@@ -64,6 +64,10 @@ cdef extern from "distributions.h":
     cdef uint64_t random_uint64(aug_state* state) nogil
     cdef uint32_t random_uint32(aug_state* state) nogil
     cdef uint64_t random_raw_values(aug_state* state) nogil
+    
+    cdef uint32_t random_buffered_bounded_uint32(aug_state *state, uint32_t off, uint32_t rng, uint32_t mask, int *bcnt, uint32_t *buf) nogil
+    cdef uint16_t random_buffered_bounded_uint16(aug_state *state, uint16_t off, uint16_t rng, uint16_t mask, int *bcnt, uint32_t *buf) nogil
+    cdef uint8_t random_buffered_bounded_uint8(aug_state *state, uint8_t off, uint8_t rng, uint8_t mask, int *bcnt, uint32_t *buf) nogil
 
     cdef long random_positive_int(aug_state* state) nogil
     cdef unsigned long random_uint(aug_state* state) nogil

setup.py

@@ -34,7 +34,7 @@
         'RNG_PCG32', 'RNG_PCG64', 'RNG_XORSHIFT128', 'RNG_XOROSHIRO128PLUS',
         'RNG_XORSHIFT1024', 'RNG_SFMT']
 
-compile_rngs = rngs[:]
+compile_rngs = rngs[:1]
 
 extra_defs = [('_CRT_SECURE_NO_WARNINGS', '1')] if os.name == 'nt' else []
 extra_link_args = ['/LTCG', 'Advapi32.lib', 'Kernel32.lib'] if os.name == 'nt' else []