codekingpro/portable-devtools
115k
1"""Module containing non-deprecated functions borrowed from Numeric.
2
3"""
4import functools
5import math
6import types
7
8import numpy as np
9from numpy._utils import set_module
10
11from . import _methods, multiarray as mu, numerictypes as nt, overrides, umath as um
12from ._multiarray_umath import _array_converter
13from .multiarray import asanyarray, asarray, concatenate
14
15_dt_ = nt.sctype2char
16
17# functions that are methods
18__all__ = [
19 'all', 'amax', 'amin', 'any', 'argmax',
20 'argmin', 'argpartition', 'argsort', 'around', 'choose', 'clip',
21 'compress', 'cumprod', 'cumsum', 'cumulative_prod', 'cumulative_sum',
22 'diagonal', 'mean', 'max', 'min', 'matrix_transpose',
23 'ndim', 'nonzero', 'partition', 'prod', 'ptp', 'put',
24 'ravel', 'repeat', 'reshape', 'resize', 'round',
25 'searchsorted', 'shape', 'size', 'sort', 'squeeze',
26 'std', 'sum', 'swapaxes', 'take', 'trace', 'transpose', 'var',
27]
28
29_gentype = types.GeneratorType
30# save away Python sum
31_sum_ = sum
32
33array_function_dispatch = functools.partial(
34 overrides.array_function_dispatch, module='numpy')
35
36
37# functions that are now methods
38def _wrapit(obj, method, *args, **kwds):
39 conv = _array_converter(obj)
40 # As this already tried the method, subok is maybe quite reasonable here
41 # but this follows what was done before. TODO: revisit this.
42 arr, = conv.as_arrays(subok=False)
43 result = getattr(arr, method)(*args, **kwds)
44
45 return conv.wrap(result, to_scalar=False)
46
47
48def _wrapfunc(obj, method, *args, **kwds):
49 bound = getattr(obj, method, None)
50 if bound is None:
51 return _wrapit(obj, method, *args, **kwds)
52
53 try:
54 return bound(*args, **kwds)
55 except TypeError:
56 # A TypeError occurs if the object does have such a method in its
57 # class, but its signature is not identical to that of NumPy's. This
58 # situation has occurred in the case of a downstream library like
59 # 'pandas'.
60 #
61 # Call _wrapit from within the except clause to ensure a potential
62 # exception has a traceback chain.
63 return _wrapit(obj, method, *args, **kwds)
64
65
66def _wrapreduction(obj, ufunc, method, axis, dtype, out, **kwargs):
67 passkwargs = {k: v for k, v in kwargs.items()
68 if v is not np._NoValue}
69
70 if type(obj) is not mu.ndarray:
71 try:
72 reduction = getattr(obj, method)
73 except AttributeError:
74 pass
75 else:
76 # This branch is needed for reductions like any which don't
77 # support a dtype.
78 if dtype is not None:
79 return reduction(axis=axis, dtype=dtype, out=out, **passkwargs)
80 else:
81 return reduction(axis=axis, out=out, **passkwargs)
82
83 return ufunc.reduce(obj, axis, dtype, out, **passkwargs)
84
85
86def _wrapreduction_any_all(obj, ufunc, method, axis, out, **kwargs):
87 # Same as above function, but dtype is always bool (but never passed on)
88 passkwargs = {k: v for k, v in kwargs.items()
89 if v is not np._NoValue}
90
91 if type(obj) is not mu.ndarray:
92 try:
93 reduction = getattr(obj, method)
94 except AttributeError:
95 pass
96 else:
97 return reduction(axis=axis, out=out, **passkwargs)
98
99 return ufunc.reduce(obj, axis, bool, out, **passkwargs)
100
101
102def _take_dispatcher(a, indices, axis=None, out=None, mode=None):
103 return (a, out)
104
105
106@array_function_dispatch(_take_dispatcher)
107def take(a, indices, axis=None, out=None, mode='raise'):
108 """
109 Take elements from an array along an axis.
110
111 When axis is not None, this function does the same thing as "fancy"
112 indexing (indexing arrays using arrays); however, it can be easier to use
113 if you need elements along a given axis. A call such as
114 ``np.take(arr, indices, axis=3)`` is equivalent to
115 ``arr[:,:,:,indices,...]``.
116
117 Explained without fancy indexing, this is equivalent to the following use
118 of `ndindex`, which sets each of ``ii``, ``jj``, and ``kk`` to a tuple of
119 indices::
120
121 Ni, Nk = a.shape[:axis], a.shape[axis+1:]
122 Nj = indices.shape
123 for ii in ndindex(Ni):
124 for jj in ndindex(Nj):
125 for kk in ndindex(Nk):
126 out[ii + jj + kk] = a[ii + (indices[jj],) + kk]
127
128 Parameters
129 ----------
130 a : array_like (Ni..., M, Nk...)
131 The source array.
132 indices : array_like (Nj...)
133 The indices of the values to extract.
134 Also allow scalars for indices.
135 axis : int, optional
136 The axis over which to select values. By default, the flattened
137 input array is used.
138 out : ndarray, optional (Ni..., Nj..., Nk...)
139 If provided, the result will be placed in this array. It should
140 be of the appropriate shape and dtype. Note that `out` is always
141 buffered if `mode='raise'`; use other modes for better performance.
142 mode : {'raise', 'wrap', 'clip'}, optional
143 Specifies how out-of-bounds indices will behave.
144
145 * 'raise' -- raise an error (default)
146 * 'wrap' -- wrap around
147 * 'clip' -- clip to the range
148
149 'clip' mode means that all indices that are too large are replaced
150 by the index that addresses the last element along that axis. Note
151 that this disables indexing with negative numbers.
152
153 Returns
154 -------
155 out : ndarray (Ni..., Nj..., Nk...)
156 The returned array has the same type as `a`.
157
158 See Also
159 --------
160 compress : Take elements using a boolean mask
161 ndarray.take : equivalent method
162 take_along_axis : Take elements by matching the array and the index arrays
163
164 Notes
165 -----
166 By eliminating the inner loop in the description above, and using `s_` to
167 build simple slice objects, `take` can be expressed in terms of applying
168 fancy indexing to each 1-d slice::
169
170 Ni, Nk = a.shape[:axis], a.shape[axis+1:]
171 for ii in ndindex(Ni):
172 for kk in ndindex(Nk):
173 out[ii + s_[...,] + kk] = a[ii + s_[:,] + kk][indices]
174
175 For this reason, it is equivalent to (but faster than) the following use
176 of `apply_along_axis`::
177
178 out = np.apply_along_axis(lambda a_1d: a_1d[indices], axis, a)
179
180 Examples
181 --------
182 >>> import numpy as np
183 >>> a = [4, 3, 5, 7, 6, 8]
184 >>> indices = [0, 1, 4]
185 >>> np.take(a, indices)
186 array([4, 3, 6])
187
188 In this example if `a` is an ndarray, "fancy" indexing can be used.
189
190 >>> a = np.array(a)
191 >>> a[indices]
192 array([4, 3, 6])
193
194 If `indices` is not one dimensional, the output also has these dimensions.
195
196 >>> np.take(a, [[0, 1], [2, 3]])
197 array([[4, 3],
198 [5, 7]])
199 """
200 return _wrapfunc(a, 'take', indices, axis=axis, out=out, mode=mode)
201
202
203def _reshape_dispatcher(a, /, shape, order=None, *, copy=None):
204 return (a,)
205
206
207@array_function_dispatch(_reshape_dispatcher)
208def reshape(a, /, shape, order='C', *, copy=None):
209 """
210 Gives a new shape to an array without changing its data.
211
212 Parameters
213 ----------
214 a : array_like
215 Array to be reshaped.
216 shape : int or tuple of ints
217 The new shape should be compatible with the original shape. If
218 an integer, then the result will be a 1-D array of that length.
219 One shape dimension can be -1. In this case, the value is
220 inferred from the length of the array and remaining dimensions.
221 order : {'C', 'F', 'A'}, optional
222 Read the elements of ``a`` using this index order, and place the
223 elements into the reshaped array using this index order. 'C'
224 means to read / write the elements using C-like index order,
225 with the last axis index changing fastest, back to the first
226 axis index changing slowest. 'F' means to read / write the
227 elements using Fortran-like index order, with the first index
228 changing fastest, and the last index changing slowest. Note that
229 the 'C' and 'F' options take no account of the memory layout of
230 the underlying array, and only refer to the order of indexing.
231 'A' means to read / write the elements in Fortran-like index
232 order if ``a`` is Fortran *contiguous* in memory, C-like order
233 otherwise.
234 copy : bool, optional
235 If ``True``, then the array data is copied. If ``None``, a copy will
236 only be made if it's required by ``order``. For ``False`` it raises
237 a ``ValueError`` if a copy cannot be avoided. Default: ``None``.
238
239 Returns
240 -------
241 reshaped_array : ndarray
242 This will be a new view object if possible; otherwise, it will
243 be a copy. Note there is no guarantee of the *memory layout* (C- or
244 Fortran- contiguous) of the returned array.
245
246 See Also
247 --------
248 ndarray.reshape : Equivalent method.
249
250 Notes
251 -----
252 It is not always possible to change the shape of an array without copying
253 the data.
254
255 The ``order`` keyword gives the index ordering both for *fetching*
256 the values from ``a``, and then *placing* the values into the output
257 array. For example, let's say you have an array:
258
259 >>> a = np.arange(6).reshape((3, 2))
260 >>> a
261 array([[0, 1],
262 [2, 3],
263 [4, 5]])
264
265 You can think of reshaping as first raveling the array (using the given
266 index order), then inserting the elements from the raveled array into the
267 new array using the same kind of index ordering as was used for the
268 raveling.
269
270 >>> np.reshape(a, (2, 3)) # C-like index ordering
271 array([[0, 1, 2],
272 [3, 4, 5]])
273 >>> np.reshape(np.ravel(a), (2, 3)) # equivalent to C ravel then C reshape
274 array([[0, 1, 2],
275 [3, 4, 5]])
276 >>> np.reshape(a, (2, 3), order='F') # Fortran-like index ordering
277 array([[0, 4, 3],
278 [2, 1, 5]])
279 >>> np.reshape(np.ravel(a, order='F'), (2, 3), order='F')
280 array([[0, 4, 3],
281 [2, 1, 5]])
282
283 Examples
284 --------
285 >>> import numpy as np
286 >>> a = np.array([[1,2,3], [4,5,6]])
287 >>> np.reshape(a, 6)
288 array([1, 2, 3, 4, 5, 6])
289 >>> np.reshape(a, 6, order='F')
290 array([1, 4, 2, 5, 3, 6])
291
292 >>> np.reshape(a, (3,-1)) # the unspecified value is inferred to be 2
293 array([[1, 2],
294 [3, 4],
295 [5, 6]])
296 """
297 if copy is not None:
298 return _wrapfunc(a, 'reshape', shape, order=order, copy=copy)
299 return _wrapfunc(a, 'reshape', shape, order=order)
300
301
302def _choose_dispatcher(a, choices, out=None, mode=None):
303 yield a
304 yield from choices
305 yield out
306
307
308@array_function_dispatch(_choose_dispatcher)
309def choose(a, choices, out=None, mode='raise'):
310 """
311 Construct an array from an index array and a list of arrays to choose from.
312
313 First of all, if confused or uncertain, definitely look at the Examples -
314 in its full generality, this function is less simple than it might
315 seem from the following code description::
316
317 np.choose(a,c) == np.array([c[a[I]][I] for I in np.ndindex(a.shape)])
318
319 But this omits some subtleties. Here is a fully general summary:
320
321 Given an "index" array (`a`) of integers and a sequence of ``n`` arrays
322 (`choices`), `a` and each choice array are first broadcast, as necessary,
323 to arrays of a common shape; calling these *Ba* and *Bchoices[i], i =
324 0,...,n-1* we have that, necessarily, ``Ba.shape == Bchoices[i].shape``
325 for each ``i``. Then, a new array with shape ``Ba.shape`` is created as
326 follows:
327
328 * if ``mode='raise'`` (the default), then, first of all, each element of
329 ``a`` (and thus ``Ba``) must be in the range ``[0, n-1]``; now, suppose
330 that ``i`` (in that range) is the value at the ``(j0, j1, ..., jm)``
331 position in ``Ba`` - then the value at the same position in the new array
332 is the value in ``Bchoices[i]`` at that same position;
333
334 * if ``mode='wrap'``, values in `a` (and thus `Ba`) may be any (signed)
335 integer; modular arithmetic is used to map integers outside the range
336 `[0, n-1]` back into that range; and then the new array is constructed
337 as above;
338
339 * if ``mode='clip'``, values in `a` (and thus ``Ba``) may be any (signed)
340 integer; negative integers are mapped to 0; values greater than ``n-1``
341 are mapped to ``n-1``; and then the new array is constructed as above.
342
343 Parameters
344 ----------
345 a : int array
346 This array must contain integers in ``[0, n-1]``, where ``n`` is the
347 number of choices, unless ``mode=wrap`` or ``mode=clip``, in which
348 cases any integers are permissible.
349 choices : sequence of arrays
350 Choice arrays. `a` and all of the choices must be broadcastable to the
351 same shape. If `choices` is itself an array (not recommended), then
352 its outermost dimension (i.e., the one corresponding to
353 ``choices.shape[0]``) is taken as defining the "sequence".
354 out : array, optional
355 If provided, the result will be inserted into this array. It should
356 be of the appropriate shape and dtype. Note that `out` is always
357 buffered if ``mode='raise'``; use other modes for better performance.
358 mode : {'raise' (default), 'wrap', 'clip'}, optional
359 Specifies how indices outside ``[0, n-1]`` will be treated:
360
361 * 'raise' : an exception is raised
362 * 'wrap' : value becomes value mod ``n``
363 * 'clip' : values < 0 are mapped to 0, values > n-1 are mapped to n-1
364
365 Returns
366 -------
367 merged_array : array
368 The merged result.
369
370 Raises
371 ------
372 ValueError: shape mismatch
373 If `a` and each choice array are not all broadcastable to the same
374 shape.
375
376 See Also
377 --------
378 ndarray.choose : equivalent method
379 numpy.take_along_axis : Preferable if `choices` is an array
380
381 Notes
382 -----
383 To reduce the chance of misinterpretation, even though the following
384 "abuse" is nominally supported, `choices` should neither be, nor be
385 thought of as, a single array, i.e., the outermost sequence-like container
386 should be either a list or a tuple.
387
388 Examples
389 --------
390
391 >>> import numpy as np
392 >>> choices = [[0, 1, 2, 3], [10, 11, 12, 13],
393 ... [20, 21, 22, 23], [30, 31, 32, 33]]
394 >>> np.choose([2, 3, 1, 0], choices
395 ... # the first element of the result will be the first element of the
396 ... # third (2+1) "array" in choices, namely, 20; the second element
397 ... # will be the second element of the fourth (3+1) choice array, i.e.,
398 ... # 31, etc.
399 ... )
400 array([20, 31, 12, 3])
401 >>> np.choose([2, 4, 1, 0], choices, mode='clip') # 4 goes to 3 (4-1)
402 array([20, 31, 12, 3])
403 >>> # because there are 4 choice arrays
404 >>> np.choose([2, 4, 1, 0], choices, mode='wrap') # 4 goes to (4 mod 4)
405 array([20, 1, 12, 3])
406 >>> # i.e., 0
407
408 A couple examples illustrating how choose broadcasts:
409
410 >>> a = [[1, 0, 1], [0, 1, 0], [1, 0, 1]]
411 >>> choices = [-10, 10]
412 >>> np.choose(a, choices)
413 array([[ 10, -10, 10],
414 [-10, 10, -10],
415 [ 10, -10, 10]])
416
417 >>> # With thanks to Anne Archibald
418 >>> a = np.array([0, 1]).reshape((2,1,1))
419 >>> c1 = np.array([1, 2, 3]).reshape((1,3,1))
420 >>> c2 = np.array([-1, -2, -3, -4, -5]).reshape((1,1,5))
421 >>> np.choose(a, (c1, c2)) # result is 2x3x5, res[0,:,:]=c1, res[1,:,:]=c2
422 array([[[ 1, 1, 1, 1, 1],
423 [ 2, 2, 2, 2, 2],
424 [ 3, 3, 3, 3, 3]],
425 [[-1, -2, -3, -4, -5],
426 [-1, -2, -3, -4, -5],
427 [-1, -2, -3, -4, -5]]])
428
429 """
430 return _wrapfunc(a, 'choose', choices, out=out, mode=mode)
431
432
433def _repeat_dispatcher(a, repeats, axis=None):
434 return (a,)
435
436
437@array_function_dispatch(_repeat_dispatcher)
438def repeat(a, repeats, axis=None):
439 """
440 Repeat each element of an array after themselves
441
442 Parameters
443 ----------
444 a : array_like
445 Input array.
446 repeats : int or array of ints
447 The number of repetitions for each element. `repeats` is broadcasted
448 to fit the shape of the given axis.
449 axis : int, optional
450 The axis along which to repeat values. By default, use the
451 flattened input array, and return a flat output array.
452
453 Returns
454 -------
455 repeated_array : ndarray
456 Output array which has the same shape as `a`, except along
457 the given axis.
458
459 See Also
460 --------
461 tile : Tile an array.
462 unique : Find the unique elements of an array.
463
464 Examples
465 --------
466 >>> import numpy as np
467 >>> np.repeat(3, 4)
468 array([3, 3, 3, 3])
469 >>> x = np.array([[1,2],[3,4]])
470 >>> np.repeat(x, 2)
471 array([1, 1, 2, 2, 3, 3, 4, 4])
472 >>> np.repeat(x, 3, axis=1)
473 array([[1, 1, 1, 2, 2, 2],
474 [3, 3, 3, 4, 4, 4]])
475 >>> np.repeat(x, [1, 2], axis=0)
476 array([[1, 2],
477 [3, 4],
478 [3, 4]])
479
480 """
481 return _wrapfunc(a, 'repeat', repeats, axis=axis)
482
483
484def _put_dispatcher(a, ind, v, mode=None):
485 return (a, ind, v)
486
487
488@array_function_dispatch(_put_dispatcher)
489def put(a, ind, v, mode='raise'):
490 """
491 Replaces specified elements of an array with given values.
492
493 The indexing works on the flattened target array. `put` is roughly
494 equivalent to:
495
496 ::
497
498 a.flat[ind] = v
499
500 Parameters
501 ----------
502 a : ndarray
503 Target array.
504 ind : array_like
505 Target indices, interpreted as integers.
506 v : array_like
507 Values to place in `a` at target indices. If `v` is shorter than
508 `ind` it will be repeated as necessary.
509 mode : {'raise', 'wrap', 'clip'}, optional
510 Specifies how out-of-bounds indices will behave.
511
512 * 'raise' -- raise an error (default)
513 * 'wrap' -- wrap around
514 * 'clip' -- clip to the range
515
516 'clip' mode means that all indices that are too large are replaced
517 by the index that addresses the last element along that axis. Note
518 that this disables indexing with negative numbers. In 'raise' mode,
519 if an exception occurs the target array may still be modified.
520
521 See Also
522 --------
523 putmask, place
524 put_along_axis : Put elements by matching the array and the index arrays
525
526 Examples
527 --------
528 >>> import numpy as np
529 >>> a = np.arange(5)
530 >>> np.put(a, [0, 2], [-44, -55])
531 >>> a
532 array([-44, 1, -55, 3, 4])
533
534 >>> a = np.arange(5)
535 >>> np.put(a, 22, -5, mode='clip')
536 >>> a
537 array([ 0, 1, 2, 3, -5])
538
539 """
540 try:
541 put = a.put
542 except AttributeError as e:
543 raise TypeError(f"argument 1 must be numpy.ndarray, not {type(a)}") from e
544
545 return put(ind, v, mode=mode)
546
547
548def _swapaxes_dispatcher(a, axis1, axis2):
549 return (a,)
550
551
552@array_function_dispatch(_swapaxes_dispatcher)
553def swapaxes(a, axis1, axis2):
554 """
555 Interchange two axes of an array.
556
557 Parameters
558 ----------
559 a : array_like
560 Input array.
561 axis1 : int
562 First axis.
563 axis2 : int
564 Second axis.
565
566 Returns
567 -------
568 a_swapped : ndarray
569 For NumPy >= 1.10.0, if `a` is an ndarray, then a view of `a` is
570 returned; otherwise a new array is created. For earlier NumPy
571 versions a view of `a` is returned only if the order of the
572 axes is changed, otherwise the input array is returned.
573
574 Examples
575 --------
576 >>> import numpy as np
577 >>> x = np.array([[1,2,3]])
578 >>> np.swapaxes(x,0,1)
579 array([[1],
580 [2],
581 [3]])
582
583 >>> x = np.array([[[0,1],[2,3]],[[4,5],[6,7]]])
584 >>> x
585 array([[[0, 1],
586 [2, 3]],
587 [[4, 5],
588 [6, 7]]])
589
590 >>> np.swapaxes(x,0,2)
591 array([[[0, 4],
592 [2, 6]],
593 [[1, 5],
594 [3, 7]]])
595
596 """
597 return _wrapfunc(a, 'swapaxes', axis1, axis2)
598
599
600def _transpose_dispatcher(a, axes=None):
601 return (a,)
602
603
604@array_function_dispatch(_transpose_dispatcher)
605def transpose(a, axes=None):
606 """
607 Returns an array with axes transposed.
608
609 For a 1-D array, this returns an unchanged view of the original array, as a
610 transposed vector is simply the same vector.
611 To convert a 1-D array into a 2-D column vector, an additional dimension
612 must be added, e.g., ``np.atleast_2d(a).T`` achieves this, as does
613 ``a[:, np.newaxis]``.
614 For a 2-D array, this is the standard matrix transpose.
615 For an n-D array, if axes are given, their order indicates how the
616 axes are permuted (see Examples). If axes are not provided, then
617 ``transpose(a).shape == a.shape[::-1]``.
618
619 Parameters
620 ----------
621 a : array_like
622 Input array.
623 axes : tuple or list of ints, optional
624 If specified, it must be a tuple or list which contains a permutation
625 of [0, 1, ..., N-1] where N is the number of axes of `a`. Negative
626 indices can also be used to specify axes. The i-th axis of the returned
627 array will correspond to the axis numbered ``axes[i]`` of the input.
628 If not specified, defaults to ``range(a.ndim)[::-1]``, which reverses
629 the order of the axes.
630
631 Returns
632 -------
633 p : ndarray
634 `a` with its axes permuted. A view is returned whenever possible.
635
636 See Also
637 --------
638 ndarray.transpose : Equivalent method.
639 moveaxis : Move axes of an array to new positions.
640 argsort : Return the indices that would sort an array.
641
642 Notes
643 -----
644 Use ``transpose(a, argsort(axes))`` to invert the transposition of tensors
645 when using the `axes` keyword argument.
646
647 Examples
648 --------
649 >>> import numpy as np
650 >>> a = np.array([[1, 2], [3, 4]])
651 >>> a
652 array([[1, 2],
653 [3, 4]])
654 >>> np.transpose(a)
655 array([[1, 3],
656 [2, 4]])
657
658 >>> a = np.array([1, 2, 3, 4])
659 >>> a
660 array([1, 2, 3, 4])
661 >>> np.transpose(a)
662 array([1, 2, 3, 4])
663
664 >>> a = np.ones((1, 2, 3))
665 >>> np.transpose(a, (1, 0, 2)).shape
666 (2, 1, 3)
667
668 >>> a = np.ones((2, 3, 4, 5))
669 >>> np.transpose(a).shape
670 (5, 4, 3, 2)
671
672 >>> a = np.arange(3*4*5).reshape((3, 4, 5))
673 >>> np.transpose(a, (-1, 0, -2)).shape
674 (5, 3, 4)
675
676 """
677 return _wrapfunc(a, 'transpose', axes)
678
679
680def _matrix_transpose_dispatcher(x):
681 return (x,)
682
683@array_function_dispatch(_matrix_transpose_dispatcher)
684def matrix_transpose(x, /):
685 """
686 Transposes a matrix (or a stack of matrices) ``x``.
687
688 This function is Array API compatible.
689
690 Parameters
691 ----------
692 x : array_like
693 Input array having shape (..., M, N) and whose two innermost
694 dimensions form ``MxN`` matrices.
695
696 Returns
697 -------
698 out : ndarray
699 An array containing the transpose for each matrix and having shape
700 (..., N, M).
701
702 See Also
703 --------
704 transpose : Generic transpose method.
705
706 Examples
707 --------
708 >>> import numpy as np
709 >>> np.matrix_transpose([[1, 2], [3, 4]])
710 array([[1, 3],
711 [2, 4]])
712
713 >>> np.matrix_transpose([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
714 array([[[1, 3],
715 [2, 4]],
716 [[5, 7],
717 [6, 8]]])
718
719 """
720 x = asanyarray(x)
721 if x.ndim < 2:
722 raise ValueError(
723 f"Input array must be at least 2-dimensional, but it is {x.ndim}"
724 )
725 return swapaxes(x, -1, -2)
726
727
728def _partition_dispatcher(a, kth, axis=None, kind=None, order=None):
729 return (a,)
730
731
732@array_function_dispatch(_partition_dispatcher)
733def partition(a, kth, axis=-1, kind='introselect', order=None):
734 """
735 Return a partitioned copy of an array.
736
737 Creates a copy of the array and partially sorts it in such a way that
738 the value of the element in k-th position is in the position it would be
739 in a sorted array. In the output array, all elements smaller than the k-th
740 element are located to the left of this element and all equal or greater
741 are located to its right. The ordering of the elements in the two
742 partitions on the either side of the k-th element in the output array is
743 undefined.
744
745 Parameters
746 ----------
747 a : array_like
748 Array to be sorted.
749 kth : int or sequence of ints
750 Element index to partition by. The k-th value of the element
751 will be in its final sorted position and all smaller elements
752 will be moved before it and all equal or greater elements behind
753 it. The order of all elements in the partitions is undefined. If
754 provided with a sequence of k-th it will partition all elements
755 indexed by k-th of them into their sorted position at once.
756
757 axis : int or None, optional
758 Axis along which to sort. If None, the array is flattened before
759 sorting. The default is -1, which sorts along the last axis.
760 kind : {'introselect'}, optional
761 Selection algorithm. Default is 'introselect'.
762 order : str or list of str, optional
763 When `a` is an array with fields defined, this argument
764 specifies which fields to compare first, second, etc. A single
765 field can be specified as a string. Not all fields need be
766 specified, but unspecified fields will still be used, in the
767 order in which they come up in the dtype, to break ties.
768
769 Returns
770 -------
771 partitioned_array : ndarray
772 Array of the same type and shape as `a`.
773
774 See Also
775 --------
776 ndarray.partition : Method to sort an array in-place.
777 argpartition : Indirect partition.
778 sort : Full sorting
779
780 Notes
781 -----
782 The various selection algorithms are characterized by their average
783 speed, worst case performance, work space size, and whether they are
784 stable. A stable sort keeps items with the same key in the same
785 relative order. The available algorithms have the following
786 properties:
787
788 ================= ======= ============= ============ =======
789 kind speed worst case work space stable
790 ================= ======= ============= ============ =======
791 'introselect' 1 O(n) 0 no
792 ================= ======= ============= ============ =======
793
794 All the partition algorithms make temporary copies of the data when
795 partitioning along any but the last axis. Consequently,
796 partitioning along the last axis is faster and uses less space than
797 partitioning along any other axis.
798
799 The sort order for complex numbers is lexicographic. If both the
800 real and imaginary parts are non-nan then the order is determined by
801 the real parts except when they are equal, in which case the order
802 is determined by the imaginary parts.
803
804 The sort order of ``np.nan`` is bigger than ``np.inf``.
805
806 Examples
807 --------
808 >>> import numpy as np
809 >>> a = np.array([7, 1, 7, 7, 1, 5, 7, 2, 3, 2, 6, 2, 3, 0])
810 >>> p = np.partition(a, 4)
811 >>> p
812 array([0, 1, 2, 1, 2, 5, 2, 3, 3, 6, 7, 7, 7, 7]) # may vary
813
814 ``p[4]`` is 2; all elements in ``p[:4]`` are less than or equal
815 to ``p[4]``, and all elements in ``p[5:]`` are greater than or
816 equal to ``p[4]``. The partition is::
817
818 [0, 1, 2, 1], [2], [5, 2, 3, 3, 6, 7, 7, 7, 7]
819
820 The next example shows the use of multiple values passed to `kth`.
821
822 >>> p2 = np.partition(a, (4, 8))
823 >>> p2
824 array([0, 1, 2, 1, 2, 3, 3, 2, 5, 6, 7, 7, 7, 7])
825
826 ``p2[4]`` is 2 and ``p2[8]`` is 5. All elements in ``p2[:4]``
827 are less than or equal to ``p2[4]``, all elements in ``p2[5:8]``
828 are greater than or equal to ``p2[4]`` and less than or equal to
829 ``p2[8]``, and all elements in ``p2[9:]`` are greater than or
830 equal to ``p2[8]``. The partition is::
831
832 [0, 1, 2, 1], [2], [3, 3, 2], [5], [6, 7, 7, 7, 7]
833 """
834 if axis is None:
835 # flatten returns (1, N) for np.matrix, so always use the last axis
836 a = asanyarray(a).flatten()
837 axis = -1
838 else:
839 a = asanyarray(a).copy(order="K")
840 a.partition(kth, axis=axis, kind=kind, order=order)
841 return a
842
843
844def _argpartition_dispatcher(a, kth, axis=None, kind=None, order=None):
845 return (a,)
846
847
848@array_function_dispatch(_argpartition_dispatcher)
849def argpartition(a, kth, axis=-1, kind='introselect', order=None):
850 """
851 Perform an indirect partition along the given axis using the
852 algorithm specified by the `kind` keyword. It returns an array of
853 indices of the same shape as `a` that index data along the given
854 axis in partitioned order.
855
856 Parameters
857 ----------
858 a : array_like
859 Array to sort.
860 kth : int or sequence of ints
861 Element index to partition by. The k-th element will be in its
862 final sorted position and all smaller elements will be moved
863 before it and all larger elements behind it. The order of all
864 elements in the partitions is undefined. If provided with a
865 sequence of k-th it will partition all of them into their sorted
866 position at once.
867
868 axis : int or None, optional
869 Axis along which to sort. The default is -1 (the last axis). If
870 None, the flattened array is used.
871 kind : {'introselect'}, optional
872 Selection algorithm. Default is 'introselect'
873 order : str or list of str, optional
874 When `a` is an array with fields defined, this argument
875 specifies which fields to compare first, second, etc. A single
876 field can be specified as a string, and not all fields need be
877 specified, but unspecified fields will still be used, in the
878 order in which they come up in the dtype, to break ties.
879
880 Returns
881 -------
882 index_array : ndarray, int
883 Array of indices that partition `a` along the specified axis.
884 If `a` is one-dimensional, ``a[index_array]`` yields a partitioned `a`.
885 More generally, ``np.take_along_axis(a, index_array, axis=axis)``
886 always yields the partitioned `a`, irrespective of dimensionality.
887
888 See Also
889 --------
890 partition : Describes partition algorithms used.
891 ndarray.partition : Inplace partition.
892 argsort : Full indirect sort.
893 take_along_axis : Apply ``index_array`` from argpartition
894 to an array as if by calling partition.
895
896 Notes
897 -----
898 The returned indices are not guaranteed to be sorted according to
899 the values. Furthermore, the default selection algorithm ``introselect``
900 is unstable, and hence the returned indices are not guaranteed
901 to be the earliest/latest occurrence of the element.
902
903 `argpartition` works for real/complex inputs with nan values,
904 see `partition` for notes on the enhanced sort order and
905 different selection algorithms.
906
907 Examples
908 --------
909 One dimensional array:
910
911 >>> import numpy as np
912 >>> x = np.array([3, 4, 2, 1])
913 >>> x[np.argpartition(x, 3)]
914 array([2, 1, 3, 4]) # may vary
915 >>> x[np.argpartition(x, (1, 3))]
916 array([1, 2, 3, 4]) # may vary
917
918 >>> x = [3, 4, 2, 1]
919 >>> np.array(x)[np.argpartition(x, 3)]
920 array([2, 1, 3, 4]) # may vary
921
922 Multi-dimensional array:
923
924 >>> x = np.array([[3, 4, 2], [1, 3, 1]])
925 >>> index_array = np.argpartition(x, kth=1, axis=-1)
926 >>> # below is the same as np.partition(x, kth=1)
927 >>> np.take_along_axis(x, index_array, axis=-1)
928 array([[2, 3, 4],
929 [1, 1, 3]])
930
931 """
932 return _wrapfunc(a, 'argpartition', kth, axis=axis, kind=kind, order=order)
933
934
935def _sort_dispatcher(a, axis=None, kind=None, order=None, *, stable=None):
936 return (a,)
937
938
939@array_function_dispatch(_sort_dispatcher)
940def sort(a, axis=-1, kind=None, order=None, *, stable=None):
941 """
942 Return a sorted copy of an array.
943
944 Parameters
945 ----------
946 a : array_like
947 Array to be sorted.
948 axis : int or None, optional
949 Axis along which to sort. If None, the array is flattened before
950 sorting. The default is -1, which sorts along the last axis.
951 kind : {'quicksort', 'mergesort', 'heapsort', 'stable'}, optional
952 Sorting algorithm. The default is 'quicksort'. Note that both 'stable'
953 and 'mergesort' use timsort or radix sort under the covers and,
954 in general, the actual implementation will vary with data type.
955 The 'mergesort' option is retained for backwards compatibility.
956 order : str or list of str, optional
957 When `a` is an array with fields defined, this argument specifies
958 which fields to compare first, second, etc. A single field can
959 be specified as a string, and not all fields need be specified,
960 but unspecified fields will still be used, in the order in which
961 they come up in the dtype, to break ties.
962 stable : bool, optional
963 Sort stability. If ``True``, the returned array will maintain
964 the relative order of ``a`` values which compare as equal.
965 If ``False`` or ``None``, this is not guaranteed. Internally,
966 this option selects ``kind='stable'``. Default: ``None``.
967
968 .. versionadded:: 2.0.0
969
970 Returns
971 -------
972 sorted_array : ndarray
973 Array of the same type and shape as `a`.
974
975 See Also
976 --------
977 ndarray.sort : Method to sort an array in-place.
978 argsort : Indirect sort.
979 lexsort : Indirect stable sort on multiple keys.
980 searchsorted : Find elements in a sorted array.
981 partition : Partial sort.
982
983 Notes
984 -----
985 The various sorting algorithms are characterized by their average speed,
986 worst case performance, work space size, and whether they are stable. A
987 stable sort keeps items with the same key in the same relative
988 order. The four algorithms implemented in NumPy have the following
989 properties:
990
991 =========== ======= ============= ============ ========
992 kind speed worst case work space stable
993 =========== ======= ============= ============ ========
994 'quicksort' 1 O(n^2) 0 no
995 'heapsort' 3 O(n*log(n)) 0 no
996 'mergesort' 2 O(n*log(n)) ~n/2 yes
997 'timsort' 2 O(n*log(n)) ~n/2 yes
998 =========== ======= ============= ============ ========
999
1000 .. note:: The datatype determines which of 'mergesort' or 'timsort'
1001 is actually used, even if 'mergesort' is specified. User selection
1002 at a finer scale is not currently available.
1003
1004 For performance, ``sort`` makes a temporary copy if needed to make the data
1005 `contiguous <https://numpy.org/doc/stable/glossary.html#term-contiguous>`_
1006 in memory along the sort axis. For even better performance and reduced
1007 memory consumption, ensure that the array is already contiguous along the
1008 sort axis.
1009
1010 The sort order for complex numbers is lexicographic. If both the real
1011 and imaginary parts are non-nan then the order is determined by the
1012 real parts except when they are equal, in which case the order is
1013 determined by the imaginary parts.
1014
1015 Previous to numpy 1.4.0 sorting real and complex arrays containing nan
1016 values led to undefined behaviour. In numpy versions >= 1.4.0 nan
1017 values are sorted to the end. The extended sort order is:
1018
1019 * Real: [R, nan]
1020 * Complex: [R + Rj, R + nanj, nan + Rj, nan + nanj]
1021
1022 where R is a non-nan real value. Complex values with the same nan
1023 placements are sorted according to the non-nan part if it exists.
1024 Non-nan values are sorted as before.
1025
1026 quicksort has been changed to:
1027 `introsort <https://en.wikipedia.org/wiki/Introsort>`_.
1028 When sorting does not make enough progress it switches to
1029 `heapsort <https://en.wikipedia.org/wiki/Heapsort>`_.
1030 This implementation makes quicksort O(n*log(n)) in the worst case.
1031
1032 'stable' automatically chooses the best stable sorting algorithm
1033 for the data type being sorted.
1034 It, along with 'mergesort' is currently mapped to
1035 `timsort <https://en.wikipedia.org/wiki/Timsort>`_
1036 or `radix sort <https://en.wikipedia.org/wiki/Radix_sort>`_
1037 depending on the data type.
1038 API forward compatibility currently limits the
1039 ability to select the implementation and it is hardwired for the different
1040 data types.
1041
1042 Timsort is added for better performance on already or nearly
1043 sorted data. On random data timsort is almost identical to
1044 mergesort. It is now used for stable sort while quicksort is still the
1045 default sort if none is chosen. For timsort details, refer to
1046 `CPython listsort.txt
1047 <https://github.com/python/cpython/blob/3.7/Objects/listsort.txt>`_
1048 'mergesort' and 'stable' are mapped to radix sort for integer data types.
1049 Radix sort is an O(n) sort instead of O(n log n).
1050
1051 NaT now sorts to the end of arrays for consistency with NaN.
1052
1053 Examples
1054 --------
1055 >>> import numpy as np
1056 >>> a = np.array([[1,4],[3,1]])
1057 >>> np.sort(a) # sort along the last axis
1058 array([[1, 4],
1059 [1, 3]])
1060 >>> np.sort(a, axis=None) # sort the flattened array
1061 array([1, 1, 3, 4])
1062 >>> np.sort(a, axis=0) # sort along the first axis
1063 array([[1, 1],
1064 [3, 4]])
1065
1066 Use the `order` keyword to specify a field to use when sorting a
1067 structured array:
1068
1069 >>> dtype = [('name', 'S10'), ('height', float), ('age', int)]
1070 >>> values = [('Arthur', 1.8, 41), ('Lancelot', 1.9, 38),
1071 ... ('Galahad', 1.7, 38)]
1072 >>> a = np.array(values, dtype=dtype) # create a structured array
1073 >>> np.sort(a, order='height') # doctest: +SKIP
1074 array([('Galahad', 1.7, 38), ('Arthur', 1.8, 41),
1075 ('Lancelot', 1.8999999999999999, 38)],
1076 dtype=[('name', '|S10'), ('height', '<f8'), ('age', '<i4')])
1077
1078 Sort by age, then height if ages are equal:
1079
1080 >>> np.sort(a, order=['age', 'height']) # doctest: +SKIP
1081 array([('Galahad', 1.7, 38), ('Lancelot', 1.8999999999999999, 38),
1082 ('Arthur', 1.8, 41)],
1083 dtype=[('name', '|S10'), ('height', '<f8'), ('age', '<i4')])
1084
1085 """
1086 if axis is None:
1087 # flatten returns (1, N) for np.matrix, so always use the last axis
1088 a = asanyarray(a).flatten()
1089 axis = -1
1090 else:
1091 a = asanyarray(a).copy(order="K")
1092 a.sort(axis=axis, kind=kind, order=order, stable=stable)
1093 return a
1094
1095
1096def _argsort_dispatcher(a, axis=None, kind=None, order=None, *, stable=None):
1097 return (a,)
1098
1099
1100@array_function_dispatch(_argsort_dispatcher)
1101def argsort(a, axis=-1, kind=None, order=None, *, stable=None):
1102 """
1103 Returns the indices that would sort an array.
1104
1105 Perform an indirect sort along the given axis using the algorithm specified
1106 by the `kind` keyword. It returns an array of indices of the same shape as
1107 `a` that index data along the given axis in sorted order.
1108
1109 Parameters
1110 ----------
1111 a : array_like
1112 Array to sort.
1113 axis : int or None, optional
1114 Axis along which to sort. The default is -1 (the last axis). If None,
1115 the flattened array is used.
1116 kind : {'quicksort', 'mergesort', 'heapsort', 'stable'}, optional
1117 Sorting algorithm. The default is 'quicksort'. Note that both 'stable'
1118 and 'mergesort' use timsort under the covers and, in general, the
1119 actual implementation will vary with data type. The 'mergesort' option
1120 is retained for backwards compatibility.
1121 order : str or list of str, optional
1122 When `a` is an array with fields defined, this argument specifies
1123 which fields to compare first, second, etc. A single field can
1124 be specified as a string, and not all fields need be specified,
1125 but unspecified fields will still be used, in the order in which
1126 they come up in the dtype, to break ties.
1127 stable : bool, optional
1128 Sort stability. If ``True``, the returned array will maintain
1129 the relative order of ``a`` values which compare as equal.
1130 If ``False`` or ``None``, this is not guaranteed. Internally,
1131 this option selects ``kind='stable'``. Default: ``None``.
1132
1133 .. versionadded:: 2.0.0
1134
1135 Returns
1136 -------
1137 index_array : ndarray, int
1138 Array of indices that sort `a` along the specified `axis`.
1139 If `a` is one-dimensional, ``a[index_array]`` yields a sorted `a`.
1140 More generally, ``np.take_along_axis(a, index_array, axis=axis)``
1141 always yields the sorted `a`, irrespective of dimensionality.
1142
1143 See Also
1144 --------
1145 sort : Describes sorting algorithms used.
1146 lexsort : Indirect stable sort with multiple keys.
1147 ndarray.sort : Inplace sort.
1148 argpartition : Indirect partial sort.
1149 take_along_axis : Apply ``index_array`` from argsort
1150 to an array as if by calling sort.
1151
1152 Notes
1153 -----
1154 See `sort` for notes on the different sorting algorithms.
1155
1156 As of NumPy 1.4.0 `argsort` works with real/complex arrays containing
1157 nan values. The enhanced sort order is documented in `sort`.
1158
1159 Examples
1160 --------
1161 One dimensional array:
1162
1163 >>> import numpy as np
1164 >>> x = np.array([3, 1, 2])
1165 >>> np.argsort(x)
1166 array([1, 2, 0])
1167
1168 Two-dimensional array:
1169
1170 >>> x = np.array([[0, 3], [2, 2]])
1171 >>> x
1172 array([[0, 3],
1173 [2, 2]])
1174
1175 >>> ind = np.argsort(x, axis=0) # sorts along first axis (down)
1176 >>> ind
1177 array([[0, 1],
1178 [1, 0]])
1179 >>> np.take_along_axis(x, ind, axis=0) # same as np.sort(x, axis=0)
1180 array([[0, 2],
1181 [2, 3]])
1182
1183 >>> ind = np.argsort(x, axis=1) # sorts along last axis (across)
1184 >>> ind
1185 array([[0, 1],
1186 [0, 1]])
1187 >>> np.take_along_axis(x, ind, axis=1) # same as np.sort(x, axis=1)
1188 array([[0, 3],
1189 [2, 2]])
1190
1191 Indices of the sorted elements of a N-dimensional array:
1192
1193 >>> ind = np.unravel_index(np.argsort(x, axis=None), x.shape)
1194 >>> ind
1195 (array([0, 1, 1, 0]), array([0, 0, 1, 1]))
1196 >>> x[ind] # same as np.sort(x, axis=None)
1197 array([0, 2, 2, 3])
1198
1199 Sorting with keys:
1200