Repository: cupy/cupy

Issue:
Implement cupy.partition
`cupy.partition`, the counterpart of `numpy.partition`, creates a copy of an array with its elements rearranged in such a way that the value of the element in k-th position is in the position it would be in a sorted array. This function is equivalent to `std::nth_element` of C++ STL.

numpy.partition
https://docs.scipy.org/doc/numpy/reference/generated/numpy.partition.html

std::nth_element
http://en.cppreference.com/w/cpp/algorithm/nth_element

Unfortunately, neither of the following parallel algorithm libraries implement this kind of parallel selection algorithm:

- Thrust (https://github.com/thrust/thrust)
- Modern GPU (https://github.com/moderngpu/moderngpu)
- CUB (https://nvlabs.github.io/cub/)

For now, the most likely is to implement radix select based on CUB's radix sort implementation.



Discussion:
`cupy.partiotion` as well as `numpy.partition` belong to a problem domain called unordered partial sort, which is a very similar problem with selection of kth element.

https://en.wikipedia.org/wiki/Selection_algorithm#Unordered_partial_sorting

There are the following algorithms to solve such kind of problems:

- Quickselect
- Median of median
- Heap select
- Introselect
- Radixselect

`numpy.partition` uses introselect, which is a hybrid algorithm of quickselect and median of medians and does not paralellize well. For GPU environment with CuPy, radixselect would be the most likely option.

- https://docs.scipy.org/doc/numpy/reference/generated/numpy.partition.html
- https://github.com/moderngpu/moderngpu/issues/10#issuecomment-210596229

The following papers mention on radixselect:

- Fast K-selection Algorithms for Graphics Processing Units (http://www.math.grin.edu/~blanchaj/Research/ABGS_KSelection.pdf)

- Radix Selection Algorithm for the kth Order Statistics (http://csis.pace.edu/~ctappert/srd2012/a4.pdf)

The first paper has its public source code and Tensorflow guys have some discussion on it recently.

- https://code.google.com/archive/p/ggks/downloads
- https://github.com/tensorflow/tensorflow/issues/5719

I compared the performance of:

- radixselect on CPU
- std::sort & choose on CPU
- thrust::sort & choose on GPU

(Celeron G3900 @ 2.80GHz and GeForce GTX 750 Ti)

Just `thrust::sort` & choose may be enough at first because its performance fit well to `O(n)` in time complexity.

https://gist.github.com/sonots/7e11231af2f2cc6b1f519ad832632566

```
$ nvcc radixselect.cu -O2 -std=c++11 -Wno-deprecated-gpu-targets && ./a.out

----- n = 1000, k = 100
radixselect : 59233 Time elapsed 0.000006 sec
std::sort   : 59233 Time elapsed 0.000043 sec
thrust::sort: 59233 Time elapsed 0.000583 sec

----- n = 10000, k = 1000
radixselect : 59017 Time elapsed 0.000032 sec
std::sort   : 59017 Time elapsed 0.000597 sec
thrust::sort: 59017 Time elapsed 0.000586 sec

----- n = 100000, k = 10000
radixselect : 58992 Time elapsed 0.000347 sec
std::sort   : 58992 Time elapsed 0.006836 sec
thrust::sort: 58992 Time elapsed 0.000436 sec

----- n = 1000000, k = 100000
radixselect : 58953 Time elapsed 0.004255 sec
std::sort   : 58953 Time elapsed 0.074907 sec
thrust::sort: 58953 Time elapsed 0.002840 sec

----- n = 10000000, k = 1000000
radixselect : 58986 Time elapsed 0.039464 sec
std::sort   : 58986 Time elapsed 0.745077 sec
thrust::sort: 58986 Time elapsed 0.023981 sec

----- n = 100000000, k = 10000000
radixselect : 58983 Time elapsed 0.489075 sec
std::sort   : 58983 Time elapsed 7.569387 sec
thrust::sort: 58983 Time elapsed 0.235514 sec
```
According to Fig. 7 in the paper,

- Fast K-selection Algorithms for Graphics Processing Units (http://www.math.grin.edu/~blanchaj/Research/ABGS_KSelection.pdf)

sort & choose has O(N) complexity in time for vector length `N` as theoretically shown, and is at least just four times slower than radixselect in case vector length is 2^27.

It seems reasonable to first implement `cupy.partition` with naive sort & choose.



Resolving patch:
diff --git a/cupy/__init__.py b/cupy/__init__.py
--- a/cupy/__init__.py
+++ b/cupy/__init__.py
@@ -410,6 +410,7 @@ def is_available():
 from cupy.sorting.sort import argsort  # NOQA
 from cupy.sorting.sort import lexsort  # NOQA
 from cupy.sorting.sort import msort  # NOQA
+from cupy.sorting.sort import partition  # NOQA
 from cupy.sorting.sort import sort  # NOQA
 
 # -----------------------------------------------------------------------------
diff --git a/cupy/sorting/sort.py b/cupy/sorting/sort.py
--- a/cupy/sorting/sort.py
+++ b/cupy/sorting/sort.py
@@ -127,7 +127,41 @@ def msort(a):
 # TODO(okuta): Implement sort_complex
 
 
-# TODO(okuta): Implement partition
+def partition(a, kth, axis=-1):
+    """Returns a partially sorted copy of an array.
+
+    Creates a copy of the array whose elements are rearranged such that the
+    value of the element in k-th position would occur in that position in a
+    sorted array. All of the elements before the new k-th element are less
+    than or equal to the elements after the new k-th element.
+
+    Args:
+        a (cupy.ndarray): Array to be sorted.
+        kth (int or sequence of ints): Element index to partition by. If
+            supplied with a sequence of k-th it will partition all elements
+            indexed by k-th of them into their sorted position at once.
+        axis (int or None): Axis along which to sort. Default is -1, which
+            means sort along the last axis. If None is supplied, the array is
+            flattened before sorting.
+
+    Returns:
+        cupy.ndarray: Array of the same type and shape as ``a``.
+
+    .. note::
+       For its implementation reason, :func:`cupy.partition` fully sorts the
+       given array as :func:`cupy.sort` does. It also does not support
+       ``kind`` and ``order`` parameters that :func:`numpy.partition` supports.
+
+    .. seealso:: :func:`numpy.partition`
+
+    """
+    if axis is None:
+        ret = a.flatten()
+        axis = -1
+    else:
+        ret = a.copy()
+    ret.partition(kth, axis=axis)
+    return ret
 
 
 # TODO(okuta): Implement argpartition


Test patch:
diff --git a/tests/cupy_tests/sorting_tests/test_sort.py b/tests/cupy_tests/sorting_tests/test_sort.py
--- a/tests/cupy_tests/sorting_tests/test_sort.py
+++ b/tests/cupy_tests/sorting_tests/test_sort.py
@@ -298,3 +298,142 @@ def test_msort_unsupported_dtype(self, dtype):
         a = testing.shaped_random((10,), cupy, dtype)
         with self.assertRaises(NotImplementedError):
             return cupy.msort(a)
+
+
+@testing.parameterize(*testing.product({
+    'external': [False, True],
+}))
+@testing.gpu
+class TestPartition(unittest.TestCase):
+
+    _multiprocess_can_split_ = True
+
+    def partition(self, a, kth, axis=-1):
+        if self.external:
+            xp = cupy.get_array_module(a)
+            return xp.partition(a, kth, axis=axis)
+        else:
+            a.partition(kth, axis=axis)
+            return a
+
+    # Test base cases
+
+    @testing.numpy_cupy_raises()
+    def test_partition_zero_dim(self, xp):
+        a = testing.shaped_random((), xp)
+        kth = 2
+        return self.partition(a, kth)
+
+    @testing.for_all_dtypes(no_float16=True, no_bool=True)
+    @testing.numpy_cupy_equal()
+    def test_partition_one_dim(self, xp, dtype):
+        a = testing.shaped_random((10,), xp, dtype)
+        kth = 2
+        x = self.partition(a, kth)
+        self.assertTrue(xp.all(x[0:kth] <= x[kth:kth + 1]))
+        self.assertTrue(xp.all(x[kth:kth + 1] <= x[kth + 1:]))
+        return x[kth]
+
+    @testing.for_all_dtypes(no_float16=True, no_bool=True)
+    @testing.numpy_cupy_array_equal()
+    def test_partition_multi_dim(self, xp, dtype):
+        a = testing.shaped_random((10, 10, 10), xp, dtype)
+        kth = 2
+        x = self.partition(a, kth)
+        self.assertTrue(xp.all(x[:, :, 0:kth] <= x[:, :, kth:kth + 1]))
+        self.assertTrue(xp.all(x[:, :, kth:kth + 1] <= x[:, :, kth + 1:]))
+        return x[:, :, kth:kth + 1]
+
+    # Test unsupported dtype
+
+    @testing.for_dtypes([numpy.float16, numpy.bool_])
+    def test_partition_unsupported_dtype(self, dtype):
+        a = testing.shaped_random((10,), cupy, dtype)
+        kth = 2
+        with self.assertRaises(NotImplementedError):
+            return self.partition(a, kth)
+
+    # Test non-contiguous array
+
+    @testing.numpy_cupy_array_equal()
+    def test_partition_non_contiguous(self, xp):
+        a = testing.shaped_random((10,), xp)[::2]
+        kth = 2
+        if not self.external:
+            if xp is cupy:
+                with self.assertRaises(NotImplementedError):
+                    return self.partition(a, kth)
+            return xp.array([])  # dummy
+        else:
+            return self.partition(a, kth)
+
+    # Test kth
+
+    @testing.numpy_cupy_equal()
+    def test_partition_sequence_kth(self, xp):
+        a = testing.shaped_random((10,), xp)
+        kth = (2, 4)
+        x = self.partition(a, kth)
+        return x[kth[0]], x[kth[1]]
+
+    @testing.numpy_cupy_equal()
+    def test_partition_negative_kth(self, xp):
+        a = testing.shaped_random((10,), xp)
+        kth = -3
+        x = self.partition(a, kth)
+        return x[kth]
+
+    @testing.numpy_cupy_raises()
+    def test_partition_invalid_kth(self, xp):
+        a = testing.shaped_random((10,), xp)
+        kth = 10
+        return self.partition(a, kth)
+
+    @testing.numpy_cupy_raises()
+    def test_partition_invalid_negative_kth(self, xp):
+        a = testing.shaped_random((10,), xp)
+        kth = -11
+        return self.partition(a, kth)
+
+    # Test axis
+
+    @testing.numpy_cupy_array_equal()
+    def test_partition_axis(self, xp):
+        a = testing.shaped_random((10, 10, 10), xp)
+        kth = 2
+        axis = 0
+        x = self.partition(a, kth, axis=axis)
+        return x[kth, :, :]
+
+    @testing.numpy_cupy_array_equal()
+    def test_partition_negative_axis(self, xp):
+        a = testing.shaped_random((10, 10, 10), xp)
+        kth = 2
+        axis = -1
+        x = self.partition(a, kth, axis=axis)
+        return x[:, :, kth]
+
+    @testing.numpy_cupy_equal()
+    def test_partition_none_axis(self, xp):
+        if self.external:
+            a = testing.shaped_random((2, 2), xp)
+            kth = 2
+            axis = None
+            x = self.partition(a, kth, axis=axis)
+            return x[kth]
+        else:
+            return None
+
+    @testing.numpy_cupy_raises()
+    def test_partition_invalid_axis(self, xp):
+        a = testing.shaped_random((2, 2, 2), xp)
+        kth = 2
+        axis = 3
+        return self.partition(a, kth, axis=axis)
+
+    @testing.numpy_cupy_raises()
+    def test_partition_invalid_negative_axis(self, xp):
+        a = testing.shaped_random((2, 2, 2), xp)
+        kth = 2
+        axis = -4
+        return self.partition(a, kth, axis=axis)