Merge pull request #14 from bashtage/small-sync

ENH: Add fast path for randint broadcasting

Author: mattip

numpy/random/bounded_integers.pyx.in

@@ -125,19 +125,29 @@ cdef object _rand_{{nptype}}_broadcast(object low, object high, object size,
 
     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)
-    cnt = np.PyArray_SIZE(high_arr)
-    flat = high_arr.flat
-    for i in range(cnt):
-        # Subtract 1 since generator produces values on the closed int [off, off+rng]
-        closed_upper = int(flat[i]) - 1
-        if closed_upper > {{ub}}:
-            raise ValueError('high is out of bounds for {{nptype}}')
-        if closed_upper < {{lb}}:
+    dt = high_arr.dtype
+    if np.issubdtype(dt, np.integer):
+        # Avoid object dtype path if already an integer
+        if np.any(np.less_equal(high_arr, {{lb}})):
             raise ValueError('low >= high')
-        highm1_data[i] = <{{nptype}}_t>closed_upper
+        high_m1 = high_arr - dt.type(1)
+        if np.any(np.greater(high_m1, {{ub}})):
+            raise ValueError('high is out of bounds for {{nptype}}')
+        highm1_arr = <np.ndarray>np.PyArray_FROM_OTF(high_m1, np.{{npctype}}, np.NPY_ALIGNED | np.NPY_FORCECAST)
+    else:
+        # If input is object or a floating type
+        highm1_arr = <np.ndarray>np.empty_like(high_arr, dtype=np.{{nptype}})
+        highm1_data = <{{nptype}}_t *>np.PyArray_DATA(highm1_arr)
+        cnt = np.PyArray_SIZE(high_arr)
+        flat = high_arr.flat
+        for i in range(cnt):
+            # Subtract 1 since generator produces values on the closed int [off, off+rng]
+            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')

numpy/random/generator.pyx

@@ -349,7 +349,7 @@ cdef class RandomGenerator:
         --------
         randint : Uniform sampling over a given half-open interval of integers.
         random_integers : Uniform sampling over a given closed interval of
-            integers.
+                          integers.
 
         Examples
         --------
@@ -413,12 +413,17 @@ cdef class RandomGenerator:
             `size`-shaped array of random integers from the appropriate
             distribution, or a single such random int if `size` not provided.
 
+        Notes
+        -----
+        When using broadcasting with uint64 dtypes, the maximum value (2**64)
+        cannot be represented as a standard integer type. The high array (or
+        low if high is None) must have object dtype, e.g., array([2**64]).
+
         See Also
         --------
-        random_integers : similar to `randint`, only for the closed
-            interval [`low`, `high`], and 1 is the lowest value if `high` is
-            omitted. In particular, this other one is the one to use to generate
-            uniformly distributed discrete non-integers.
+        random_integers : similar to `randint`, only for the closed interval
+                          [`low`, `high`], where 1 is the lowest value if
+                          `high` is omitted.
 
         Examples
         --------

numpy/random/tests/test_direct.py

@@ -51,6 +51,7 @@ def uniform32_from_uint64(x):
     out = (joined >> np.uint32(9)) * (1.0 / 2 ** 23)
     return out.astype(np.float32)
 
+
 def uniform32_from_uint53(x):
     x = np.uint64(x) >> np.uint64(16)
     x = np.uint32(x & np.uint64(0xffffffff))
@@ -92,6 +93,7 @@ def uniform_from_uint32(x):
         out[i // 2] = (a * 67108864.0 + b) / 9007199254740992.0
     return out
 
+
 def uniform_from_dsfmt(x):
     return x.view(np.double) - 1.0
 
@@ -414,7 +416,8 @@ def test_seed_float_array(self):
         rs = RandomGenerator(self.brng(*self.data1['seed']))
         assert_raises(self.seed_error_type, rs.brng.seed, np.array([np.pi]))
         assert_raises(self.seed_error_type, rs.brng.seed, np.array([-np.pi]))
-        assert_raises(self.seed_error_type, rs.brng.seed, np.array([np.pi, -np.pi]))
+        assert_raises(self.seed_error_type, rs.brng.seed,
+                      np.array([np.pi, -np.pi]))
         assert_raises(self.seed_error_type, rs.brng.seed, np.array([0, np.pi]))
         assert_raises(self.seed_error_type, rs.brng.seed, [np.pi])
         assert_raises(self.seed_error_type, rs.brng.seed, [0, np.pi])

numpy/random/tests/test_generator_mt19937.py

@@ -335,6 +335,33 @@ def test_repeatability_broadcasting(self):
 
             assert_array_equal(val, val_bc)
 
+    def test_int64_uint64_broadcast_exceptions(self):
+        configs = {np.uint64: ((0, 2**65), (-1, 2**62), (10, 9), (0, 0)),
+                   np.int64: ((0, 2**64), (-(2**64), 2**62), (10, 9), (0, 0),
+                              (-2**63-1, -2**63-1))}
+        for dtype in configs:
+            for config in configs[dtype]:
+                low, high = config
+                low_a = np.array([[low]*10])
+                high_a = np.array([high] * 10)
+                assert_raises(ValueError, random.randint, low, high,
+                              dtype=dtype)
+                assert_raises(ValueError, random.randint, low_a, high,
+                              dtype=dtype)
+                assert_raises(ValueError, random.randint, low, high_a,
+                              dtype=dtype)
+                assert_raises(ValueError, random.randint, low_a, high_a,
+                              dtype=dtype)
+
+                low_o = np.array([[low]*10], dtype=np.object)
+                high_o = np.array([high] * 10, dtype=np.object)
+                assert_raises(ValueError, random.randint, low_o, high,
+                              dtype=dtype)
+                assert_raises(ValueError, random.randint, low, high_o,
+                              dtype=dtype)
+                assert_raises(ValueError, random.randint, low_o, high_o,
+                              dtype=dtype)
+
     def test_int64_uint64_corner_case(self):
         # When stored in Numpy arrays, `lbnd` is casted
         # as np.int64, and `ubnd` is casted as np.uint64.