/
sort.py
185 lines (132 loc) · 5.48 KB
/
sort.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
import cupy
import numpy
if cupy.cuda.thrust_enabled:
from cupy.cuda import thrust
def sort(a, axis=-1):
"""Returns a sorted copy of an array with a stable sorting algorithm.
Args:
a (cupy.ndarray): Array to be sorted.
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, ``cupy.sort`` currently does not support
``kind`` and ``order`` parameters that ``numpy.sort`` does
support.
.. seealso:: :func:`numpy.sort`
"""
if axis is None:
ret = a.flatten()
axis = -1
else:
ret = a.copy()
ret.sort(axis=axis)
return ret
def lexsort(keys):
"""Perform an indirect sort using an array of keys.
Args:
keys (cupy.ndarray): ``(k, N)`` array containing ``k`` ``(N,)``-shaped
arrays. The ``k`` different "rows" to be sorted. The last row is
the primary sort key.
Returns:
cupy.ndarray: Array of indices that sort the keys.
.. note::
For its implementation reason, ``cupy.lexsort`` currently supports only
keys with their rank of one or two and does not support ``axis``
parameter that ``numpy.lexsort`` supports.
.. seealso:: :func:`numpy.lexsort`
"""
# TODO(takagi): Support axis argument.
if not cupy.cuda.thrust_enabled:
raise RuntimeError('Thrust is needed to use cupy.lexsort. Please '
'install CUDA Toolkit with Thrust then reinstall '
'CuPy after uninstalling it.')
if keys.ndim == ():
# as numpy.lexsort() raises
raise TypeError('need sequence of keys with len > 0 in lexsort')
if keys.ndim == 1:
return 0
# TODO(takagi): Support ranks of three or more.
if keys.ndim > 2:
raise NotImplementedError('Keys with the rank of three or more is not '
'supported in lexsort')
idx_array = cupy.ndarray(keys._shape[1:], dtype=numpy.intp)
k = keys._shape[0]
n = keys._shape[1]
thrust.lexsort(keys.dtype, idx_array.data.ptr, keys.data.ptr, k, n)
return idx_array
def argsort(a, axis=-1):
"""Returns the indices that would sort an array with a stable sorting.
Args:
a (cupy.ndarray): Array to sort.
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 indices that sort ``a``.
.. note::
For its implementation reason, ``cupy.argsort`` does not support
``kind`` and ``order`` parameters.
.. seealso:: :func:`numpy.argsort`
"""
return a.argsort(axis=axis)
def msort(a):
"""Returns a copy of an array sorted along the first axis.
Args:
a (cupy.ndarray): Array to be sorted.
Returns:
cupy.ndarray: Array of the same type and shape as ``a``.
.. note:
``cupy.msort(a)``, the CuPy counterpart of ``numpy.msort(a)``, is
equivalent to ``cupy.sort(a, axis=0)``.
.. seealso:: :func:`numpy.msort`
"""
# TODO(takagi): Support float16 and bool.
return sort(a, axis=0)
# TODO(okuta): Implement sort_complex
def partition(a, kth, axis=-1):
"""Returns a partitioned 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``.
.. seealso:: :func:`numpy.partition`
"""
if axis is None:
ret = a.flatten()
axis = -1
else:
ret = a.copy()
ret.partition(kth, axis=axis)
return ret
def argpartition(a, kth, axis=-1):
"""Returns the indices that would partially sort an array.
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, `cupy.argpartition` fully sorts the
given array as `cupy.argsort` does. It also does not support ``kind``
and ``order`` parameters that ``numpy.argpartition`` supports.
.. seealso:: :func:`numpy.argpartition`
"""
return a.argpartition(kth, axis=axis)