diff --git a/numpy/random/mtrand/mtrand.pyx b/numpy/random/mtrand/mtrand.pyx index 110b60a9bc27..5956c7333875 100644 --- a/numpy/random/mtrand/mtrand.pyx +++ b/numpy/random/mtrand/mtrand.pyx @@ -24,6 +24,14 @@ include "Python.pxi" include "numpy.pxd" +from libc cimport stdint + +ctypedef fused stdsized: + stdint.int8_t + stdint.int16_t + stdint.int32_t + stdint.int64_t + cdef extern from "math.h": double exp(double x) double log(double x) @@ -4971,33 +4979,53 @@ cdef class RandomState: [0, 1, 2]]) """ - cdef npy_intp i, j - - i = len(x) - 1 - - # Logic adapted from random.shuffle() - if isinstance(x, np.ndarray) and \ - (x.ndim > 1 or x.dtype.fields is not None): - # For a multi-dimensional ndarray, indexing returns a view onto - # each row. So we can't just use ordinary assignment to swap the - # rows; we need a bounce buffer. + cdef: + npy_intp i, j, n = len(x) + stdint.int8_t[:] int8_buf + stdint.int16_t[:] int16_buf + stdint.int32_t[:] int32_buf + stdint.int64_t[:] int64_buf + + # Fast, statically typed path: shuffle the underlying buffer. + # We exclude subclasses as this approach fails e.g. with MaskedArrays. + if (type(x) is np.ndarray and x.ndim == 1 and + x.dtype.itemsize in [1, 2, 4, 8]): + if x.dtype.itemsize == 1: + int8_buf = x.view(np.int8) + self._shuffle_stdsize(int8_buf) + elif x.dtype.itemsize == 2: + int16_buf = x.view(np.int16) + self._shuffle_stdsize(int16_buf) + elif x.dtype.itemsize == 4: + int32_buf = x.view(np.int32) + self._shuffle_stdsize(int32_buf) + elif x.dtype.itemsize == 8: + int64_buf = x.view(np.int64) + self._shuffle_stdsize(int64_buf) + # Untyped path 1: multidimensional arrays require an bounce buffer + # because indexing returns views. + elif isinstance(x, np.ndarray) and x.ndim > 1: buf = np.empty_like(x[0]) with self.lock: - while i > 0: + for i in reversed(range(1, n)): j = rk_interval(i, self.internal_state) - buf[...] = x[j] + buf[:] = x[j] x[j] = x[i] - x[i] = buf - i = i - 1 + x[i] = buf[:] + # Untyped path 2 (including for 1d, non-standard-sized arrays). else: - # For single-dimensional arrays, lists, and any other Python - # sequence types, indexing returns a real object that's - # independent of the array contents, so we can just swap directly. with self.lock: - while i > 0: + for i in reversed(range(1, n)): j = rk_interval(i, self.internal_state) x[i], x[j] = x[j], x[i] - i = i - 1 + + cdef void _shuffle_stdsize(self, stdsized[:] x): + cdef: + npy_intp i, j, n = x.size + with self.lock, cython.boundscheck(False), cython.wraparound(False): + for i in reversed(range(1, n)): + j = rk_interval(i, self.internal_state) + x[i], x[j] = x[j], x[i] def permutation(self, object x): """ diff --git a/numpy/random/tests/test_random.py b/numpy/random/tests/test_random.py index a6783fe8f478..d599a21365cf 100644 --- a/numpy/random/tests/test_random.py +++ b/numpy/random/tests/test_random.py @@ -8,6 +8,7 @@ from numpy.compat import asbytes import sys + class TestSeed(TestCase): def test_scalar(self): s = np.random.RandomState(0) @@ -38,6 +39,7 @@ def test_invalid_array(self): assert_raises(ValueError, np.random.RandomState, [1, 2, 4294967296]) assert_raises(ValueError, np.random.RandomState, [1, -2, 4294967296]) + class TestBinomial(TestCase): def test_n_zero(self): # Tests the corner case of n == 0 for the binomial distribution. @@ -128,6 +130,7 @@ def test_negative_binomial(self): # arguments without truncation. self.prng.negative_binomial(0.5, 0.5) + class TestRandint(TestCase): rfunc = np.random.randint @@ -352,9 +355,12 @@ def test_bytes(self): np.testing.assert_equal(actual, desired) def test_shuffle(self): - # Test lists, arrays, and multidimensional versions of both: + # Test lists, arrays (of various dtypes), and multidimensional versions + # of both: for conv in [lambda x: x, - np.asarray, + lambda x: np.asarray(x).astype(np.int8), + lambda x: np.asarray(x).astype(np.float32), + lambda x: np.asarray(x).astype(np.complex64), lambda x: [(i, i) for i in x], lambda x: np.asarray([(i, i) for i in x])]: np.random.seed(self.seed) @@ -379,13 +385,15 @@ def test_shuffle_masked(self): # gh-3263 a = np.ma.masked_values(np.reshape(range(20), (5,4)) % 3 - 1, -1) b = np.ma.masked_values(np.arange(20) % 3 - 1, -1) - ma = np.ma.count_masked(a) - mb = np.ma.count_masked(b) + a_orig = a.copy() + b_orig = b.copy() for i in range(50): np.random.shuffle(a) - self.assertEqual(ma, np.ma.count_masked(a)) + assert_equal( + sorted(a.data[~a.mask]), sorted(a_orig.data[~a_orig.mask])) np.random.shuffle(b) - self.assertEqual(mb, np.ma.count_masked(b)) + assert_equal( + sorted(b.data[~b.mask]), sorted(b_orig.data[~b_orig.mask])) def test_beta(self): np.random.seed(self.seed)