ENH: Allow broadcasting in randint

Allow array broadcasting for random integers

Author: bashtage

appveyor.yml

@@ -21,13 +21,12 @@ build_script:
   - SET PATH=C:\Py;C:\Py\Scripts;C:\Py\Library\bin;%PATH%
   - conda config --set always_yes yes
   - conda update conda --quiet
-  - conda install numpy cython nose pandas --quiet
+  - conda install numpy cython nose pandas pytest --quiet
   - python setup.py develop
   - set "GIT_DIR=%cd%"
 
 test_script:
-  - cd ..
-  - nosetests randomstate
+  - pytest randomstate
 
 on_success:
   - cd %GIT_DIR%\randomstate

doc/source/change-log.rst

@@ -2,6 +2,11 @@
 
 Change Log
 ==========
+Since 1.13.2
+------------
+* Allow (:meth:`~randomstate.prng.mt19937.randint`) to broadcast inputs
+* Sync with upstream NumPy changes
+* Add protection against negative inputs in (:meth:`~randomstate.prng.mt19937.dirichlet`)
 
 Version 1.13.2
 --------------

randomstate/bounded_integers.pxi.in

@@ -84,3 +84,241 @@ 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
+    mask |= mask >> 1
+    mask |= mask >> 2
+    mask |= mask >> 4
+    mask |= mask >> 8
+    mask |= mask >> 16
+    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}}_broadcast(object low, object high, object size, aug_state *state, object lock):
+    """
+    Generate bounded random {{nptype}} values using broadcasting
+
+    Parameters
+    ----------
+    low : int or array-like
+        Array containing the lowest (signed) integers to be drawn from the
+        distribution.
+    high : int or array-like
+        Array containing the the open interval bound for the distribution.
+    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
+        the output shape is determined by the broadcast shapes of low and
+        high
+    state : augmented random state
+        State to use in the core random number generators
+    lock : threading.Lock
+        Lock to prevent multiple using a single RandomState simultaneously
+
+    Returns
+    -------
+    out : ndarray of np.{{nptype}}
+          array of random integers from the appropriate distribution where the
+          size is determined by size if provided or the broadcast shape of low
+          and high
+    """
+    cdef np.ndarray low_arr, high_arr, out_arr
+    cdef np.npy_intp i
+    cdef np.broadcast it
+    cdef int buf_rem = 0
+
+    cdef {{utype}}_t *out_data
+    cdef {{utype}}_t val, mask, off
+    cdef {{nptype_up}}_t rng, last_rng, low_v, high_v
+    cdef uint32_t buf
+
+    # TODO: Direct error check? Probably not
+    # TODO: Make constant?
+    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 == 0:
+                out_data[i] = 0
+                continue
+
+            if rng != last_rng:
+                # TODO: Is this too much of an optimization?  is it worth it?
+                # Smallest bit mask >= max
+                mask = <{{utype}}_t>_gen_mask(rng)
+
+            while True:
+                if not buf_rem:
+                    buf = random_uint32(state)
+                    buf_rem = {{remaining}}
+                else:
+                    buf >>= {{bitshift}}
+                    buf_rem -= 1
+                val = <{{utype}}_t>buf & mask
+                if val <= rng:
+                    break
+            out_data[i] = off + val
+
+            np.PyArray_MultiIter_NEXT(it)
+
+    return out_arr
+
+{{endfor}}
+
+
+{{
+py:
+big_bc_ctypes = (('uint64', 'uint64', 'NPY_UINT64', '0x0ULL', '0xFFFFFFFFFFFFFFFFULL'),
+                 ('int64', 'uint64', 'NPY_INT64', '-0x8000000000000000LL', '0x7FFFFFFFFFFFFFFFLL' )
+)}}
+
+{{for  nptype, utype, npctype, lb, ub in big_bc_ctypes}}
+
+{{ py: otype = nptype}}
+
+cdef object _rand_{{nptype}}_broadcast(object low, object high, object size, aug_state *state, object lock):
+    """
+    Generate bounded random {{nptype}} values using broadcasting
+
+    Parameters
+    ----------
+    low : int or array-like
+        Array containing the lowest (signed) integers to be drawn from the
+        distribution.
+    high : int or array-like
+        Array containing the the open interval bound for the distribution.
+    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
+        the output shape is determined by the broadcast shapes of low and
+        high
+    state : augmented random state
+        State to use in the core random number generators
+    lock : threading.Lock
+        Lock to prevent multiple using a single RandomState simultaneously
+
+    Returns
+    -------
+    out : ndarray of np.{{nptype}}
+          array of random integers from the appropriate distribution where the
+          size is determined by size if provided or the broadcast shape of low
+          and high
+    """
+    cdef np.ndarray low_arr, high_arr, out_arr, highm1_arr
+    cdef np.npy_intp i
+    cdef np.broadcast it
+    cdef int buf_rem = 0
+    cdef object closed_upper
+
+    cdef uint64_t *out_data
+    cdef {{nptype}}_t *highm1_data
+    cdef {{nptype}}_t low_v, high_v
+    cdef uint64_t rng, last_rng, val, mask, off
+
+    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}}')
+
+    highm1_arr = <np.ndarray>np.empty_like(high_arr, dtype=np.{{nptype}})
+    highm1_data = <{{nptype}}_t *>np.PyArray_DATA(highm1_arr)
+    n = np.PyArray_SIZE(high_arr)
+    flat = high_arr.flat
+    for i in range(n):
+        closed_upper = int(flat[i]) - 1
+        if closed_upper > {{ub}}:
+            raise ValueError('high is out of bounds for {{nptype}}')
+        if closed_upper < {{lb}}:
+            raise ValueError('low >= high')
+
+        highm1_data[i] = <{{nptype}}_t>closed_upper
+
+    if np.any(np.greater(low_arr, highm1_arr)):
+        raise ValueError('low >= high')
+
+    high_arr = highm1_arr
+    low_arr = <np.ndarray>np.PyArray_FROM_OTF(low, np.{{npctype}}, np.NPY_ALIGNED | np.NPY_FORCECAST)
+
+    if size is not None:
+        out_arr = <np.ndarray>np.empty(size, np.{{nptype}})
+    else:
+        it = np.PyArray_MultiIterNew2(low_arr, high_arr)
+        out_arr = <np.ndarray>np.empty(it.shape, np.{{nptype}})
+
+    it = np.PyArray_MultiIterNew3(low_arr, high_arr, out_arr)
+    out_data = <uint64_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}}_t*>np.PyArray_MultiIter_DATA(it, 0))[0]
+            high_v = (<{{nptype}}_t*>np.PyArray_MultiIter_DATA(it, 1))[0]
+            rng = <{{nptype}}_t>(high_v - low_v) # No -1 here since implemented above
+            off = low_v
+            if rng == 0:
+                out_data[i] = 0
+                continue
+
+            if rng != last_rng:
+                mask = <uint64_t>_gen_mask(rng)
+
+            if rng <= 0xFFFFFFFFULL:
+                while True:
+                    val = random_uint32(state) & mask
+                    if val <= rng:
+                        break
+            else:
+                while True:
+                    val = random_uint64(state) & mask
+                    if val <= rng:
+                        break
+
+            out_data[i] = off + val
+
+            np.PyArray_MultiIter_NEXT(it)
+
+    return out_arr
+
+{{endfor}}