ENH: Add fast past for randint

Add fast path for 1d inputs
Add tests to ensure all paths are tested

Author: bashtage

randomstate/bounded_integers.pxi.in

@@ -80,7 +80,10 @@ cdef object _rand_{{nptype}}(object low, object high, object size, aug_state *st
 
     low = np.asarray(low)
     high = np.asarray(high)
-    if low.shape == high.shape == ():
+    low_ndim = np.PyArray_NDIM(<np.ndarray>low)
+    high_ndim = np.PyArray_NDIM(<np.ndarray>high)
+    if ((low_ndim == 0 or (low_ndim==1 and low.size==1 and size is not None)) and
+            (high_ndim == 0 or (high_ndim==1 and high.size==1 and size is not None))):
         low = int(low) # TODO: Cast appropriately?
         high = int(high) # TODO: Cast appropriately?
 
@@ -199,7 +202,10 @@ cdef object _rand_{{nptype}}(object low, object high, object size, aug_state *st
 
     low = np.asarray(low)
     high = np.asarray(high)
-    if low.shape == high.shape == ():
+    low_ndim = np.PyArray_NDIM(<np.ndarray>low)
+    high_ndim = np.PyArray_NDIM(<np.ndarray>high)
+    if ((low_ndim == 0 or (low_ndim==1 and low.size==1 and size is not None)) and
+            (high_ndim == 0 or (high_ndim==1 and high.size==1 and size is not None))):
         low = int(low) # TODO: Cast appropriately?
         high = int(high) # TODO: Cast appropriately?
         high -= 1  # Use a closed interval

randomstate/tests/test_numpy_mt19937.py

@@ -160,10 +160,10 @@ def test_bounds_checking_array(self):
         for dt in self.itype:
             lbnd = 0 if dt is bool else np.iinfo(dt).min
             ubnd = 2 if dt is bool else np.iinfo(dt).max + 1
-            assert_raises(ValueError, self.rfunc, [lbnd - 1], [ubnd], dtype=dt)
-            assert_raises(ValueError, self.rfunc, [lbnd], [ubnd + 1], dtype=dt)
-            assert_raises(ValueError, self.rfunc, ubnd, [lbnd], dtype=dt)
-            assert_raises(ValueError, self.rfunc, [1], 0, dtype=dt)
+            assert_raises(ValueError, self.rfunc, [lbnd - 1] * 2, [ubnd] * 2, dtype=dt)
+            assert_raises(ValueError, self.rfunc, [lbnd] * 2, [ubnd + 1] * 2, dtype=dt)
+            assert_raises(ValueError, self.rfunc, ubnd, [lbnd] * 2, dtype=dt)
+            assert_raises(ValueError, self.rfunc, [1] * 2, 0, dtype=dt)
 
     def test_rng_zero_and_extremes(self):
         for dt in self.itype:
@@ -180,19 +180,25 @@ def test_rng_zero_and_extremes(self):
             assert_equal(self.rfunc(tgt, tgt + 1, size=1000, dtype=dt), tgt)
 
     def test_rng_zero_and_extremes_array(self):
+        size = 1000
         for dt in self.itype:
             lbnd = 0 if dt is bool else np.iinfo(dt).min
             ubnd = 2 if dt is bool else np.iinfo(dt).max + 1
 
             tgt = ubnd - 1
-            print(dt)
-            assert_equal(self.rfunc([tgt], [tgt + 1], size=1000, dtype=dt), tgt)
+            assert_equal(self.rfunc([tgt], [tgt + 1], size=size, dtype=dt), tgt)
+            assert_equal(self.rfunc([tgt] * size, [tgt + 1]  * size, dtype=dt), tgt)
+            assert_equal(self.rfunc([tgt] * size, [tgt + 1]  * size, size=size, dtype=dt), tgt)
 
             tgt = lbnd
-            assert_equal(self.rfunc([tgt], [tgt + 1], size=1000, dtype=dt), tgt)
+            assert_equal(self.rfunc([tgt], [tgt + 1], size=size, dtype=dt), tgt)
+            assert_equal(self.rfunc([tgt] * size, [tgt + 1] * size, dtype=dt), tgt)
+            assert_equal(self.rfunc([tgt] * size, [tgt + 1] * size, size=size, dtype=dt), tgt)
 
             tgt = (lbnd + ubnd) // 2
-            assert_equal(self.rfunc([tgt], [tgt + 1], size=1000, dtype=dt), tgt)
+            assert_equal(self.rfunc([tgt], [tgt + 1], size=size, dtype=dt), tgt)
+            assert_equal(self.rfunc([tgt] * size, [tgt + 1] * size, dtype=dt), tgt)
+            assert_equal(self.rfunc([tgt] * size, [tgt + 1] * size, size=size, dtype=dt), tgt)
 
     def test_full_range(self):
         # Test for ticket #1690
@@ -222,6 +228,24 @@ def test_full_range_array(self):
                                      "but one was with the following "
                                      "message:\n\n%s" % str(e))
 
+    def test_scalar_array_equiv(self):
+        for dt in self.itype:
+            lbnd = 0 if dt is bool else np.iinfo(dt).min
+            ubnd = 2 if dt is bool else np.iinfo(dt).max + 1
+
+            size = 1000
+            mt19937.seed(1234)
+            scalar = self.rfunc(lbnd, ubnd, size=size, dtype=dt)
+
+            mt19937.seed(1234)
+            scalar_array = self.rfunc(lbnd, ubnd, size=size, dtype=dt)
+
+            mt19937.seed(1234)
+            array = self.rfunc([lbnd] * size, [ubnd] * size, size=size, dtype=dt)
+            assert_array_equal(scalar, scalar_array)
+            assert_array_equal(scalar, array)
+
+
     def test_in_bounds_fuzz(self):
         # Don't use fixed seed
         mt19937.seed()
@@ -270,6 +294,27 @@ def test_repeatability(self):
         res = hashlib.md5(val).hexdigest()
         assert_(tgt[np.dtype(bool).name] == res)
 
+    def test_repeatability_broadcasting(self):
+
+        for dt in self.itype:
+
+            lbnd = 0 if dt in (np.bool, bool, np.bool_) else np.iinfo(dt).min
+            ubnd = 2 if dt in (np.bool, bool, np.bool_) else np.iinfo(dt).max + 1
+
+            # view as little endian for hash
+            mt19937.seed(1234)
+            val = self.rfunc(lbnd, ubnd, size=1000, dtype=dt)
+
+            mt19937.seed(1234)
+            val_bc = self.rfunc([lbnd] * 1000, ubnd, dtype=dt)
+
+            assert_array_equal(val, val_bc)
+
+            mt19937.seed(1234)
+            val_bc = self.rfunc([lbnd] * 1000, [ubnd] * 1000, dtype=dt)
+
+            assert_array_equal(val, val_bc)
+
     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.
@@ -322,6 +367,8 @@ def test_respect_dtype_array(self):
 
             sample = self.rfunc([lbnd], [ubnd], dtype=dt)
             self.assertEqual(sample.dtype, dt)
+            sample = self.rfunc([lbnd] * 2, [ubnd] * 2, dtype=dt)
+            self.assertEqual(sample.dtype, dt)
 
 
 class TestRandomDist(TestCase):