From 7377e44d319b877da29c9b1239c44fd7da47c3d3 Mon Sep 17 00:00:00 2001 From: Antony Lee Date: Sat, 5 Dec 2015 20:46:09 -0800 Subject: [PATCH 1/4] Faster random.shuffle via static typing. This patch modifies random.shuffle so that (when working on a ndarray) an array of indices is shuffled and then elements are take()n from that array in that order. This allows the inner loop to be statically typed (it turns out this is not so easy to write a generic shuffling code using Cython fused types) and thus much faster (~6x for me), at the expense of a threefold increase in memory use (I guess take() needs to create a copy, and an additional array of indices is created.). See #5514. --- numpy/random/mtrand/mtrand.pyx | 31 +++++++++++-------------------- 1 file changed, 11 insertions(+), 20 deletions(-) diff --git a/numpy/random/mtrand/mtrand.pyx b/numpy/random/mtrand/mtrand.pyx index 110b60a9bc27..10e5043c4a60 100644 --- a/numpy/random/mtrand/mtrand.pyx +++ b/numpy/random/mtrand/mtrand.pyx @@ -4971,33 +4971,24 @@ cdef class RandomState: [0, 1, 2]]) """ - cdef npy_intp i, j - - i = len(x) - 1 + cdef: + npy_intp[::1] idxs + npy_intp i, j, n = len(x) # 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. - buf = np.empty_like(x[0]) - with self.lock: - while i > 0: + if isinstance(x, np.ndarray): + # Take from a shuffled range to benefit from static typing. + idxs = np.arange(n, dtype=np.intp) + with self.lock, cython.boundscheck(False), cython.wraparound(False): + for i in reversed(range(1, n)): j = rk_interval(i, self.internal_state) - buf[...] = x[j] - x[j] = x[i] - x[i] = buf - i = i - 1 + idxs[i], idxs[j] = idxs[j], idxs[i] + x.take(idxs, 0, out=x) 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 def permutation(self, object x): """ From 2d9ff3c1091edbee86e185985d5efc0a9ea4a139 Mon Sep 17 00:00:00 2001 From: Antony Lee Date: Sat, 5 Dec 2015 22:05:18 -0800 Subject: [PATCH 2/4] Specialize shuffling for intp-sized 1d arrays. This avoids the memory overhead of allocating an index array and take()ing from the original array when possible. --- numpy/random/mtrand/mtrand.pyx | 23 ++++++++++++++++------- 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/numpy/random/mtrand/mtrand.pyx b/numpy/random/mtrand/mtrand.pyx index 10e5043c4a60..e25ba3f3720f 100644 --- a/numpy/random/mtrand/mtrand.pyx +++ b/numpy/random/mtrand/mtrand.pyx @@ -4977,19 +4977,28 @@ cdef class RandomState: # Logic adapted from random.shuffle() if isinstance(x, np.ndarray): - # Take from a shuffled range to benefit from static typing. - idxs = np.arange(n, dtype=np.intp) - with self.lock, cython.boundscheck(False), cython.wraparound(False): - for i in reversed(range(1, n)): - j = rk_interval(i, self.internal_state) - idxs[i], idxs[j] = idxs[j], idxs[i] - x.take(idxs, 0, out=x) + if x.ndim == 1 and x.dtype.itemsize == np.dtype(np.intp).itemsize: + # Directly shuffle the array if possible. + self._shuffle_intpsized(x.view(np.intp)) + else: + # Take from a shuffled range to benefit from static typing. + idxs = np.arange(n, dtype=np.intp) + self._shuffle_intpsized(idxs) + x.take(idxs, 0, out=x) else: with self.lock: for i in reversed(range(1, n)): j = rk_interval(i, self.internal_state) x[i], x[j] = x[j], x[i] + cdef void _shuffle_intpsized(self, npy_intp[:] 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): """ permutation(x) From a535de864416f65390997597d967eba039c91bc4 Mon Sep 17 00:00:00 2001 From: Antony Lee Date: Wed, 9 Dec 2015 10:41:15 -0800 Subject: [PATCH 3/4] Fix shuffling of masked arrays. Also ensure that the masked data is preserved upon shuffling, which was not the case before. --- numpy/random/mtrand/mtrand.pyx | 4 +++- numpy/random/tests/test_random.py | 17 +++++++++++++---- 2 files changed, 16 insertions(+), 5 deletions(-) diff --git a/numpy/random/mtrand/mtrand.pyx b/numpy/random/mtrand/mtrand.pyx index e25ba3f3720f..8cb1949aba3c 100644 --- a/numpy/random/mtrand/mtrand.pyx +++ b/numpy/random/mtrand/mtrand.pyx @@ -4977,7 +4977,9 @@ cdef class RandomState: # Logic adapted from random.shuffle() if isinstance(x, np.ndarray): - if x.ndim == 1 and x.dtype.itemsize == np.dtype(np.intp).itemsize: + if (x.ndim == 1 and + x.dtype.itemsize == np.dtype(np.intp).itemsize and + not isinstance(x, np.ma.MaskedArray)): # Directly shuffle the array if possible. self._shuffle_intpsized(x.view(np.intp)) else: diff --git a/numpy/random/tests/test_random.py b/numpy/random/tests/test_random.py index a6783fe8f478..785a37eb30a3 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 @@ -379,13 +382,19 @@ 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])) + 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])) + assert_equal( + sorted(b.data[~b.mask]), sorted(b_orig.data[~b_orig.mask])) def test_beta(self): np.random.seed(self.seed) From 1c4cc5d6a2c7bc53e94b63a8322131c5930e6f57 Mon Sep 17 00:00:00 2001 From: Antony Lee Date: Thu, 10 Dec 2015 00:50:30 -0800 Subject: [PATCH 4/4] shuffle: specialize std sizes; keep memory perf. Do not rely on take(), which for non-standard sized arrays, thus ensuing the previous memory performance at the expense of speed. Also get rid of the guarantee that shuffling masked arrays maintains the masked values as well, which should probably be handled on np.ma's side anyways. --- numpy/random/mtrand/mtrand.pyx | 56 ++++++++++++++++++++++--------- numpy/random/tests/test_random.py | 11 +++--- 2 files changed, 46 insertions(+), 21 deletions(-) diff --git a/numpy/random/mtrand/mtrand.pyx b/numpy/random/mtrand/mtrand.pyx index 8cb1949aba3c..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) @@ -4972,28 +4980,46 @@ cdef class RandomState: """ cdef: - npy_intp[::1] idxs npy_intp i, j, n = len(x) - - # Logic adapted from random.shuffle() - if isinstance(x, np.ndarray): - if (x.ndim == 1 and - x.dtype.itemsize == np.dtype(np.intp).itemsize and - not isinstance(x, np.ma.MaskedArray)): - # Directly shuffle the array if possible. - self._shuffle_intpsized(x.view(np.intp)) - else: - # Take from a shuffled range to benefit from static typing. - idxs = np.arange(n, dtype=np.intp) - self._shuffle_intpsized(idxs) - x.take(idxs, 0, out=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: + for i in reversed(range(1, n)): + j = rk_interval(i, self.internal_state) + buf[:] = x[j] + x[j] = x[i] + x[i] = buf[:] + # Untyped path 2 (including for 1d, non-standard-sized arrays). else: with self.lock: for i in reversed(range(1, n)): j = rk_interval(i, self.internal_state) x[i], x[j] = x[j], x[i] - cdef void _shuffle_intpsized(self, npy_intp[:] x): + cdef void _shuffle_stdsize(self, stdsized[:] x): cdef: npy_intp i, j, n = x.size with self.lock, cython.boundscheck(False), cython.wraparound(False): diff --git a/numpy/random/tests/test_random.py b/numpy/random/tests/test_random.py index 785a37eb30a3..d599a21365cf 100644 --- a/numpy/random/tests/test_random.py +++ b/numpy/random/tests/test_random.py @@ -355,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) @@ -386,13 +389,9 @@ def test_shuffle_masked(self): b_orig = b.copy() for i in range(50): np.random.shuffle(a) - assert_equal( - sorted(a.data[a.mask]), sorted(a_orig.data[a_orig.mask])) assert_equal( sorted(a.data[~a.mask]), sorted(a_orig.data[~a_orig.mask])) np.random.shuffle(b) - assert_equal( - sorted(b.data[b.mask]), sorted(b_orig.data[b_orig.mask])) assert_equal( sorted(b.data[~b.mask]), sorted(b_orig.data[~b_orig.mask]))