/
sort.py
217 lines (159 loc) · 6.52 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
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
import warnings
import cupy
import numpy
from cupy.cuda import thrust
def sort(a, axis=-1, kind=None):
"""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.
kind: Default is `None`, which is equivalent to 'stable'. Unlike in
NumPy any other options are not accepted here.
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 kind is not None and kind != 'stable':
raise ValueError("kind can only be None or 'stable'")
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.available:
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 cupy.array(0, dtype=numpy.intp)
# 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')
# thrust.lexsort() assumes a C-contiguous array
if not keys.flags.c_contiguous:
keys = keys.copy('C')
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, kind=None):
"""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.
kind: Default is `None`, which is equivalent to 'stable'. Unlike in
NumPy any other options are not accepted here.
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`
"""
if kind is not None and kind != 'stable':
raise ValueError("kind can only be None or 'stable'")
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`
"""
warnings.warn(
'msort is deprecated, use cupy.sort(a, axis=0) instead',
DeprecationWarning)
return sort(a, axis=0)
def sort_complex(a):
"""Sort a complex array using the real part first,
then the imaginary part.
Args:
a (cupy.ndarray): Array to be sorted.
Returns:
cupy.ndarray: sorted complex array.
.. seealso:: :func:`numpy.sort_complex`
"""
if a.dtype.char in 'bhBHF':
a = a.astype('F')
else:
a = a.astype('D')
a.sort()
return a
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)