REF: Complete bounded interger refactor

Complete refactor by implementing uint64 and int64 refactorings

Author: bashtage

randomstate/bounded_integers.pxi.in

@@ -147,7 +147,6 @@ cdef object _rand_{{nptype}}_combined(object low, object high, object size, aug_
     """
     cdef {{utype}}_t rng, last_rng, off, val, mask, out_val
     cdef uint32_t buf
-    cdef {{utype}}_t *out
     cdef {{utype}}_t *out_data
     cdef {{nptype_up}}_t low_v, high_v
     cdef np.ndarray low_arr, high_arr, out_arr
@@ -179,9 +178,9 @@ cdef object _rand_{{nptype}}_combined(object low, object high, object size, aug_
         else:
             out_arr = <np.ndarray>np.empty(size, np.{{nptype}})
             cnt = np.PyArray_SIZE(out_arr)
-            out = <{{utype}}_t *>np.PyArray_DATA(out_arr)
+            out_data = <{{utype}}_t *>np.PyArray_DATA(out_arr)
             with lock, nogil:
-                random_bounded_{{utype}}_fill(state, off, rng, cnt, out)
+                random_bounded_{{utype}}_fill(state, off, rng, cnt, out_data)
             return out_arr
 
     # Array path
@@ -432,20 +431,140 @@ cdef object _rand_{{nptype}}_broadcast(object low, object high, object size, aug
                 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
+                mask = _gen_mask(rng)
+            out_data[i] = random_bounded_uint64(state, off, rng, mask)
 
-            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}}_combined(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, cnt
+    cdef np.broadcast it
+    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, out_val
+
+    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?
+        high -= 1  # Use a closed interval
+
+        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")
+
+        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_val)
+            return np.{{otype}}(<{{nptype}}_t>out_val)
+        else:
+            out_arr = <np.ndarray>np.empty(size, np.{{nptype}})
+            cnt = np.PyArray_SIZE(out_arr)
+            out_data = <{{utype}}_t *>np.PyArray_DATA(out_arr)
+            with lock, nogil:
+                random_bounded_{{utype}}_fill(state, off, rng, cnt, out_data)
+            return out_arr
+
+    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 = _gen_mask(rng)
+            out_data[i] = random_bounded_uint64(state, off, rng, mask)
 
             np.PyArray_MultiIter_NEXT(it)
 

randomstate/randomstate.pyx

@@ -65,6 +65,7 @@ cdef extern from "distributions.h":
     cdef uint32_t random_uint32(aug_state* state) nogil
     cdef uint64_t random_raw_values(aug_state* state) nogil
 
+    cdef uint64_t random_bounded_uint64(aug_state *state, uint64_t off, uint64_t rng, uint64_t mask) 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
@@ -1024,10 +1025,14 @@ cdef class RandomState:
 
         if key == 'int32':
             ret =  _rand_int32_combined(low, high, size, &self.rng_state, self.lock)
+        elif key == 'int64':
+            ret =  _rand_int64_combined(low, high, size, &self.rng_state, self.lock)
         elif key == 'int16':
             ret =  _rand_int16_combined(low, high, size, &self.rng_state, self.lock)
         elif key == 'int8':
             ret =  _rand_int8_combined(low, high, size, &self.rng_state, self.lock)
+        elif key == 'uint64':
+            ret =  _rand_uint64_combined(low, high, size, &self.rng_state, self.lock)
         elif key == 'uint32':
             ret =  _rand_uint32_combined(low, high, size, &self.rng_state, self.lock)
         elif key == 'uint16':
@@ -1037,52 +1042,10 @@ cdef class RandomState:
         elif key == 'bool':
             ret =  _rand_bool_combined(low, high, size, &self.rng_state, self.lock)
 
-        if key != 'int64' and key != 'uint64':
-            if size is None and dtype in (np.bool, np.int, np.long):
-                    if np.array(ret).shape == ():
-                        return dtype(ret)
-            return ret
-
-        lowbnd, highbnd = _randint_type[key]
-
-        low = np.asarray(low)
-        high = np.asarray(high)
-
-        if low.shape == high.shape == ():
-            # TODO: Do not cast these inputs to Python int
-            #
-            # This is a workaround until gh-8851 is resolved (bug in NumPy
-            # integer comparison and subtraction involving uint64 and non-
-            # uint64). Afterwards, remove these two lines.
-            ilow = int(low)
-            ihigh = int(high)
-
-            if ilow < lowbnd:
-                raise ValueError("low is out of bounds for %s" % (key,))
-            if ihigh > highbnd:
-                raise ValueError("high is out of bounds for %s" % (key,))
-            if ilow >= ihigh:
-                raise ValueError("low >= high")
-
-            if key == 'int64':
-                ret = _rand_int64(ilow, ihigh - 1, size, &self.rng_state, self.lock)
-            elif key == 'uint64':
-                ret = _rand_uint64(ilow, ihigh - 1, size, &self.rng_state, self.lock)
-            elif key == 'bool':
-                ret = _rand_bool(ilow, ihigh - 1, size, &self.rng_state, self.lock)
 
-            if size is None:
-                if dtype in (np.bool, np.int, np.long):
+        if size is None and dtype in (np.bool, np.int, np.long):
+                if np.array(ret).shape == ():
                     return dtype(ret)
-            return ret
-
-        if key == 'int64':
-            ret = _rand_int64_broadcast(low, high, size, &self.rng_state, self.lock)
-        elif key == 'uint64':
-            ret = _rand_uint64_broadcast(low, high, size, &self.rng_state, self.lock)
-        elif key == 'bool':
-            ret = _rand_bool_broadcast(low, high, size, &self.rng_state, self.lock)
-
         return ret
 
     def bytes(self, np.npy_intp length):