codekingpro/portable-devtools
115k
1import functools
2import warnings
3
4import numpy as np
5import numpy._core.numeric as _nx
6from numpy._core import atleast_3d, overrides, vstack
7from numpy._core._multiarray_umath import _array_converter
8from numpy._core.fromnumeric import reshape, transpose
9from numpy._core.multiarray import normalize_axis_index
10from numpy._core.numeric import (
11 array,
12 asanyarray,
13 asarray,
14 normalize_axis_tuple,
15 zeros,
16 zeros_like,
17)
18from numpy._core.overrides import set_module
19from numpy._core.shape_base import _arrays_for_stack_dispatcher
20from numpy.lib._index_tricks_impl import ndindex
21from numpy.matrixlib.defmatrix import matrix # this raises all the right alarm bells
22
23__all__ = [
24 'column_stack', 'row_stack', 'dstack', 'array_split', 'split',
25 'hsplit', 'vsplit', 'dsplit', 'apply_over_axes', 'expand_dims',
26 'apply_along_axis', 'kron', 'tile', 'take_along_axis',
27 'put_along_axis'
28 ]
29
30
31array_function_dispatch = functools.partial(
32 overrides.array_function_dispatch, module='numpy')
33
34
35def _make_along_axis_idx(arr_shape, indices, axis):
36 # compute dimensions to iterate over
37 if not _nx.issubdtype(indices.dtype, _nx.integer):
38 raise IndexError('`indices` must be an integer array')
39 if len(arr_shape) != indices.ndim:
40 raise ValueError(
41 "`indices` and `arr` must have the same number of dimensions")
42 shape_ones = (1,) * indices.ndim
43 dest_dims = list(range(axis)) + [None] + list(range(axis + 1, indices.ndim))
44
45 # build a fancy index, consisting of orthogonal aranges, with the
46 # requested index inserted at the right location
47 fancy_index = []
48 for dim, n in zip(dest_dims, arr_shape):
49 if dim is None:
50 fancy_index.append(indices)
51 else:
52 ind_shape = shape_ones[:dim] + (-1,) + shape_ones[dim + 1:]
53 fancy_index.append(_nx.arange(n).reshape(ind_shape))
54
55 return tuple(fancy_index)
56
57
58def _take_along_axis_dispatcher(arr, indices, axis=None):
59 return (arr, indices)
60
61
62@array_function_dispatch(_take_along_axis_dispatcher)
63def take_along_axis(arr, indices, axis=-1):
64 """
65 Take values from the input array by matching 1d index and data slices.
66
67 This iterates over matching 1d slices oriented along the specified axis in
68 the index and data arrays, and uses the former to look up values in the
69 latter. These slices can be different lengths.
70
71 Functions returning an index along an axis, like `argsort` and
72 `argpartition`, produce suitable indices for this function.
73
74 Parameters
75 ----------
76 arr : ndarray (Ni..., M, Nk...)
77 Source array
78 indices : ndarray (Ni..., J, Nk...)
79 Indices to take along each 1d slice of ``arr``. This must match the
80 dimension of ``arr``, but dimensions Ni and Nj only need to broadcast
81 against ``arr``.
82 axis : int or None, optional
83 The axis to take 1d slices along. If axis is None, the input array is
84 treated as if it had first been flattened to 1d, for consistency with
85 `sort` and `argsort`.
86
87 .. versionchanged:: 2.3
88 The default value is now ``-1``.
89
90 Returns
91 -------
92 out: ndarray (Ni..., J, Nk...)
93 The indexed result.
94
95 Notes
96 -----
97 This is equivalent to (but faster than) the following use of `ndindex` and
98 `s_`, which sets each of ``ii`` and ``kk`` to a tuple of indices::
99
100 Ni, M, Nk = a.shape[:axis], a.shape[axis], a.shape[axis+1:]
101 J = indices.shape[axis] # Need not equal M
102 out = np.empty(Ni + (J,) + Nk)
103
104 for ii in ndindex(Ni):
105 for kk in ndindex(Nk):
106 a_1d = a [ii + s_[:,] + kk]
107 indices_1d = indices[ii + s_[:,] + kk]
108 out_1d = out [ii + s_[:,] + kk]
109 for j in range(J):
110 out_1d[j] = a_1d[indices_1d[j]]
111
112 Equivalently, eliminating the inner loop, the last two lines would be::
113
114 out_1d[:] = a_1d[indices_1d]
115
116 See Also
117 --------
118 take : Take along an axis, using the same indices for every 1d slice
119 put_along_axis :
120 Put values into the destination array by matching 1d index and data slices
121
122 Examples
123 --------
124 >>> import numpy as np
125
126 For this sample array
127
128 >>> a = np.array([[10, 30, 20], [60, 40, 50]])
129
130 We can sort either by using sort directly, or argsort and this function
131
132 >>> np.sort(a, axis=1)
133 array([[10, 20, 30],
134 [40, 50, 60]])
135 >>> ai = np.argsort(a, axis=1)
136 >>> ai
137 array([[0, 2, 1],
138 [1, 2, 0]])
139 >>> np.take_along_axis(a, ai, axis=1)
140 array([[10, 20, 30],
141 [40, 50, 60]])
142
143 The same works for max and min, if you maintain the trivial dimension
144 with ``keepdims``:
145
146 >>> np.max(a, axis=1, keepdims=True)
147 array([[30],
148 [60]])
149 >>> ai = np.argmax(a, axis=1, keepdims=True)
150 >>> ai
151 array([[1],
152 [0]])
153 >>> np.take_along_axis(a, ai, axis=1)
154 array([[30],
155 [60]])
156
157 If we want to get the max and min at the same time, we can stack the
158 indices first
159
160 >>> ai_min = np.argmin(a, axis=1, keepdims=True)
161 >>> ai_max = np.argmax(a, axis=1, keepdims=True)
162 >>> ai = np.concatenate([ai_min, ai_max], axis=1)
163 >>> ai
164 array([[0, 1],
165 [1, 0]])
166 >>> np.take_along_axis(a, ai, axis=1)
167 array([[10, 30],
168 [40, 60]])
169 """
170 # normalize inputs
171 if axis is None:
172 if indices.ndim != 1:
173 raise ValueError(
174 'when axis=None, `indices` must have a single dimension.')
175 arr = np.array(arr.flat)
176 axis = 0
177 else:
178 axis = normalize_axis_index(axis, arr.ndim)
179
180 # use the fancy index
181 return arr[_make_along_axis_idx(arr.shape, indices, axis)]
182
183
184def _put_along_axis_dispatcher(arr, indices, values, axis):
185 return (arr, indices, values)
186
187
188@array_function_dispatch(_put_along_axis_dispatcher)
189def put_along_axis(arr, indices, values, axis):
190 """
191 Put values into the destination array by matching 1d index and data slices.
192
193 This iterates over matching 1d slices oriented along the specified axis in
194 the index and data arrays, and uses the former to place values into the
195 latter. These slices can be different lengths.
196
197 Functions returning an index along an axis, like `argsort` and
198 `argpartition`, produce suitable indices for this function.
199
200 Parameters
201 ----------
202 arr : ndarray (Ni..., M, Nk...)
203 Destination array.
204 indices : ndarray (Ni..., J, Nk...)
205 Indices to change along each 1d slice of `arr`. This must match the
206 dimension of arr, but dimensions in Ni and Nj may be 1 to broadcast
207 against `arr`.
208 values : array_like (Ni..., J, Nk...)
209 values to insert at those indices. Its shape and dimension are
210 broadcast to match that of `indices`.
211 axis : int
212 The axis to take 1d slices along. If axis is None, the destination
213 array is treated as if a flattened 1d view had been created of it.
214
215 Notes
216 -----
217 This is equivalent to (but faster than) the following use of `ndindex` and
218 `s_`, which sets each of ``ii`` and ``kk`` to a tuple of indices::
219
220 Ni, M, Nk = a.shape[:axis], a.shape[axis], a.shape[axis+1:]
221 J = indices.shape[axis] # Need not equal M
222
223 for ii in ndindex(Ni):
224 for kk in ndindex(Nk):
225 a_1d = a [ii + s_[:,] + kk]
226 indices_1d = indices[ii + s_[:,] + kk]
227 values_1d = values [ii + s_[:,] + kk]
228 for j in range(J):
229 a_1d[indices_1d[j]] = values_1d[j]
230
231 Equivalently, eliminating the inner loop, the last two lines would be::
232
233 a_1d[indices_1d] = values_1d
234
235 See Also
236 --------
237 take_along_axis :
238 Take values from the input array by matching 1d index and data slices
239
240 Examples
241 --------
242 >>> import numpy as np
243
244 For this sample array
245
246 >>> a = np.array([[10, 30, 20], [60, 40, 50]])
247
248 We can replace the maximum values with:
249
250 >>> ai = np.argmax(a, axis=1, keepdims=True)
251 >>> ai
252 array([[1],
253 [0]])
254 >>> np.put_along_axis(a, ai, 99, axis=1)
255 >>> a
256 array([[10, 99, 20],
257 [99, 40, 50]])
258
259 """
260 # normalize inputs
261 if axis is None:
262 if indices.ndim != 1:
263 raise ValueError(
264 'when axis=None, `indices` must have a single dimension.')
265 arr = np.array(arr.flat)
266 axis = 0
267 else:
268 axis = normalize_axis_index(axis, arr.ndim)
269
270 # use the fancy index
271 arr[_make_along_axis_idx(arr.shape, indices, axis)] = values
272
273
274def _apply_along_axis_dispatcher(func1d, axis, arr, *args, **kwargs):
275 return (arr,)
276
277
278@array_function_dispatch(_apply_along_axis_dispatcher)
279def apply_along_axis(func1d, axis, arr, *args, **kwargs):
280 """
281 Apply a function to 1-D slices along the given axis.
282
283 Execute `func1d(a, *args, **kwargs)` where `func1d` operates on 1-D arrays
284 and `a` is a 1-D slice of `arr` along `axis`.
285
286 This is equivalent to (but faster than) the following use of `ndindex` and
287 `s_`, which sets each of ``ii``, ``jj``, and ``kk`` to a tuple of indices::
288
289 Ni, Nk = a.shape[:axis], a.shape[axis+1:]
290 for ii in ndindex(Ni):
291 for kk in ndindex(Nk):
292 f = func1d(arr[ii + s_[:,] + kk])
293 Nj = f.shape
294 for jj in ndindex(Nj):
295 out[ii + jj + kk] = f[jj]
296
297 Equivalently, eliminating the inner loop, this can be expressed as::
298
299 Ni, Nk = a.shape[:axis], a.shape[axis+1:]
300 for ii in ndindex(Ni):
301 for kk in ndindex(Nk):
302 out[ii + s_[...,] + kk] = func1d(arr[ii + s_[:,] + kk])
303
304 Parameters
305 ----------
306 func1d : function (M,) -> (Nj...)
307 This function should accept 1-D arrays. It is applied to 1-D
308 slices of `arr` along the specified axis.
309 axis : integer
310 Axis along which `arr` is sliced.
311 arr : ndarray (Ni..., M, Nk...)
312 Input array.
313 args : any
314 Additional arguments to `func1d`.
315 kwargs : any
316 Additional named arguments to `func1d`.
317
318 Returns
319 -------
320 out : ndarray (Ni..., Nj..., Nk...)
321 The output array. The shape of `out` is identical to the shape of
322 `arr`, except along the `axis` dimension. This axis is removed, and
323 replaced with new dimensions equal to the shape of the return value
324 of `func1d`. So if `func1d` returns a scalar `out` will have one
325 fewer dimensions than `arr`.
326
327 See Also
328 --------
329 apply_over_axes : Apply a function repeatedly over multiple axes.
330
331 Examples
332 --------
333 >>> import numpy as np
334 >>> def my_func(a):
335 ... \"\"\"Average first and last element of a 1-D array\"\"\"
336 ... return (a[0] + a[-1]) * 0.5
337 >>> b = np.array([[1,2,3], [4,5,6], [7,8,9]])
338 >>> np.apply_along_axis(my_func, 0, b)
339 array([4., 5., 6.])
340 >>> np.apply_along_axis(my_func, 1, b)
341 array([2., 5., 8.])
342
343 For a function that returns a 1D array, the number of dimensions in
344 `outarr` is the same as `arr`.
345
346 >>> b = np.array([[8,1,7], [4,3,9], [5,2,6]])
347 >>> np.apply_along_axis(sorted, 1, b)
348 array([[1, 7, 8],
349 [3, 4, 9],
350 [2, 5, 6]])
351
352 For a function that returns a higher dimensional array, those dimensions
353 are inserted in place of the `axis` dimension.
354
355 >>> b = np.array([[1,2,3], [4,5,6], [7,8,9]])
356 >>> np.apply_along_axis(np.diag, -1, b)
357 array([[[1, 0, 0],
358 [0, 2, 0],
359 [0, 0, 3]],
360 [[4, 0, 0],
361 [0, 5, 0],
362 [0, 0, 6]],
363 [[7, 0, 0],
364 [0, 8, 0],
365 [0, 0, 9]]])
366 """
367 # handle negative axes
368 conv = _array_converter(arr)
369 arr = conv[0]
370
371 nd = arr.ndim
372 axis = normalize_axis_index(axis, nd)
373
374 # arr, with the iteration axis at the end
375 in_dims = list(range(nd))
376 inarr_view = transpose(arr, in_dims[:axis] + in_dims[axis + 1:] + [axis])
377
378 # compute indices for the iteration axes, and append a trailing ellipsis to
379 # prevent 0d arrays decaying to scalars, which fixes gh-8642
380 inds = ndindex(inarr_view.shape[:-1])
381 inds = (ind + (Ellipsis,) for ind in inds)
382
383 # invoke the function on the first item
384 try:
385 ind0 = next(inds)
386 except StopIteration:
387 raise ValueError(
388 'Cannot apply_along_axis when any iteration dimensions are 0'
389 ) from None
390 res = asanyarray(func1d(inarr_view[ind0], *args, **kwargs))
391
392 # build a buffer for storing evaluations of func1d.
393 # remove the requested axis, and add the new ones on the end.
394 # laid out so that each write is contiguous.
395 # for a tuple index inds, buff[inds] = func1d(inarr_view[inds])
396 if not isinstance(res, matrix):
397 buff = zeros_like(res, shape=inarr_view.shape[:-1] + res.shape)
398 else:
399 # Matrices are nasty with reshaping, so do not preserve them here.
400 buff = zeros(inarr_view.shape[:-1] + res.shape, dtype=res.dtype)
401
402 # permutation of axes such that out = buff.transpose(buff_permute)
403 buff_dims = list(range(buff.ndim))
404 buff_permute = (
405 buff_dims[0 : axis] +
406 buff_dims[buff.ndim - res.ndim : buff.ndim] +
407 buff_dims[axis : buff.ndim - res.ndim]
408 )
409
410 # save the first result, then compute and save all remaining results
411 buff[ind0] = res
412 for ind in inds:
413 buff[ind] = asanyarray(func1d(inarr_view[ind], *args, **kwargs))
414
415 res = transpose(buff, buff_permute)
416 return conv.wrap(res)
417
418
419def _apply_over_axes_dispatcher(func, a, axes):
420 return (a,)
421
422
423@array_function_dispatch(_apply_over_axes_dispatcher)
424def apply_over_axes(func, a, axes):
425 """
426 Apply a function repeatedly over multiple axes.
427
428 `func` is called as `res = func(a, axis)`, where `axis` is the first
429 element of `axes`. The result `res` of the function call must have
430 either the same dimensions as `a` or one less dimension. If `res`
431 has one less dimension than `a`, a dimension is inserted before
432 `axis`. The call to `func` is then repeated for each axis in `axes`,
433 with `res` as the first argument.
434
435 Parameters
436 ----------
437 func : function
438 This function must take two arguments, `func(a, axis)`.
439 a : array_like
440 Input array.
441 axes : array_like
442 Axes over which `func` is applied; the elements must be integers.
443
444 Returns
445 -------
446 apply_over_axis : ndarray
447 The output array. The number of dimensions is the same as `a`,
448 but the shape can be different. This depends on whether `func`
449 changes the shape of its output with respect to its input.
450
451 See Also
452 --------
453 apply_along_axis :
454 Apply a function to 1-D slices of an array along the given axis.
455
456 Notes
457 -----
458 This function is equivalent to tuple axis arguments to reorderable ufuncs
459 with keepdims=True. Tuple axis arguments to ufuncs have been available since
460 version 1.7.0.
461
462 Examples
463 --------
464 >>> import numpy as np
465 >>> a = np.arange(24).reshape(2,3,4)
466 >>> a
467 array([[[ 0, 1, 2, 3],
468 [ 4, 5, 6, 7],
469 [ 8, 9, 10, 11]],
470 [[12, 13, 14, 15],
471 [16, 17, 18, 19],
472 [20, 21, 22, 23]]])
473
474 Sum over axes 0 and 2. The result has same number of dimensions
475 as the original array:
476
477 >>> np.apply_over_axes(np.sum, a, [0,2])
478 array([[[ 60],
479 [ 92],
480 [124]]])
481
482 Tuple axis arguments to ufuncs are equivalent:
483
484 >>> np.sum(a, axis=(0,2), keepdims=True)
485 array([[[ 60],
486 [ 92],
487 [124]]])
488
489 """
490 val = asarray(a)
491 N = a.ndim
492 if array(axes).ndim == 0:
493 axes = (axes,)
494 for axis in axes:
495 if axis < 0:
496 axis = N + axis
497 args = (val, axis)
498 res = func(*args)
499 if res.ndim == val.ndim:
500 val = res
501 else:
502 res = expand_dims(res, axis)
503 if res.ndim == val.ndim:
504 val = res
505 else:
506 raise ValueError("function is not returning "
507 "an array of the correct shape")
508 return val
509
510
511def _expand_dims_dispatcher(a, axis):
512 return (a,)
513
514
515@array_function_dispatch(_expand_dims_dispatcher)
516def expand_dims(a, axis):
517 """
518 Expand the shape of an array.
519
520 Insert a new axis that will appear at the `axis` position in the expanded
521 array shape.
522
523 Parameters
524 ----------
525 a : array_like
526 Input array.
527 axis : int or tuple of ints
528 Position in the expanded axes where the new axis (or axes) is placed.
529
530 .. deprecated:: 1.13.0
531 Passing an axis where ``axis > a.ndim`` will be treated as
532 ``axis == a.ndim``, and passing ``axis < -a.ndim - 1`` will
533 be treated as ``axis == 0``. This behavior is deprecated.
534
535 Returns
536 -------
537 result : ndarray
538 View of `a` with the number of dimensions increased.
539
540 See Also
541 --------
542 squeeze : The inverse operation, removing singleton dimensions
543 reshape : Insert, remove, and combine dimensions, and resize existing ones
544 atleast_1d, atleast_2d, atleast_3d
545
546 Examples
547 --------
548 >>> import numpy as np
549 >>> x = np.array([1, 2])
550 >>> x.shape
551 (2,)
552
553 The following is equivalent to ``x[np.newaxis, :]`` or ``x[np.newaxis]``:
554
555 >>> y = np.expand_dims(x, axis=0)
556 >>> y
557 array([[1, 2]])
558 >>> y.shape
559 (1, 2)
560
561 The following is equivalent to ``x[:, np.newaxis]``:
562
563 >>> y = np.expand_dims(x, axis=1)
564 >>> y
565 array([[1],
566 [2]])
567 >>> y.shape
568 (2, 1)
569
570 ``axis`` may also be a tuple:
571
572 >>> y = np.expand_dims(x, axis=(0, 1))
573 >>> y
574 array([[[1, 2]]])
575
576 >>> y = np.expand_dims(x, axis=(2, 0))
577 >>> y
578 array([[[1],
579 [2]]])
580
581 Note that some examples may use ``None`` instead of ``np.newaxis``. These
582 are the same objects:
583
584 >>> np.newaxis is None
585 True
586
587 """
588 if isinstance(a, matrix):
589 a = asarray(a)
590 else:
591 a = asanyarray(a)
592
593 if not isinstance(axis, (tuple, list)):
594 axis = (axis,)
595
596 out_ndim = len(axis) + a.ndim
597 axis = normalize_axis_tuple(axis, out_ndim)
598
599 shape_it = iter(a.shape)
600 shape = [1 if ax in axis else next(shape_it) for ax in range(out_ndim)]
601
602 return a.reshape(shape)
603
604
605# NOTE: Remove once deprecation period passes
606@set_module("numpy")
607def row_stack(tup, *, dtype=None, casting="same_kind"):
608 # Deprecated in NumPy 2.0, 2023-08-18
609 warnings.warn(
610 "`row_stack` alias is deprecated. "
611 "Use `np.vstack` directly.",
612 DeprecationWarning,
613 stacklevel=2
614 )
615 return vstack(tup, dtype=dtype, casting=casting)
616
617
618row_stack.__doc__ = vstack.__doc__
619
620
621def _column_stack_dispatcher(tup):
622 return _arrays_for_stack_dispatcher(tup)
623
624
625@array_function_dispatch(_column_stack_dispatcher)
626def column_stack(tup):
627 """
628 Stack 1-D arrays as columns into a 2-D array.
629
630 Take a sequence of 1-D arrays and stack them as columns
631 to make a single 2-D array. 2-D arrays are stacked as-is,
632 just like with `hstack`. 1-D arrays are turned into 2-D columns
633 first.
634
635 Parameters
636 ----------
637 tup : sequence of 1-D or 2-D arrays.
638 Arrays to stack. All of them must have the same first dimension.
639
640 Returns
641 -------
642 stacked : 2-D array
643 The array formed by stacking the given arrays.
644
645 See Also
646 --------
647 stack, hstack, vstack, concatenate
648
649 Examples
650 --------
651 >>> import numpy as np
652 >>> a = np.array((1,2,3))
653 >>> b = np.array((4,5,6))
654 >>> np.column_stack((a,b))
655 array([[1, 4],
656 [2, 5],
657 [3, 6]])
658
659 """
660 arrays = []
661 for v in tup:
662 arr = asanyarray(v)
663 if arr.ndim < 2:
664 arr = array(arr, copy=None, subok=True, ndmin=2).T
665 arrays.append(arr)
666 return _nx.concatenate(arrays, 1)
667
668
669def _dstack_dispatcher(tup):
670 return _arrays_for_stack_dispatcher(tup)
671
672
673@array_function_dispatch(_dstack_dispatcher)
674def dstack(tup):
675 """
676 Stack arrays in sequence depth wise (along third axis).
677
678 This is equivalent to concatenation along the third axis after 2-D arrays
679 of shape `(M,N)` have been reshaped to `(M,N,1)` and 1-D arrays of shape
680 `(N,)` have been reshaped to `(1,N,1)`. Rebuilds arrays divided by
681 `dsplit`.
682
683 This function makes most sense for arrays with up to 3 dimensions. For
684 instance, for pixel-data with a height (first axis), width (second axis),
685 and r/g/b channels (third axis). The functions `concatenate`, `stack` and
686 `block` provide more general stacking and concatenation operations.
687
688 Parameters
689 ----------
690 tup : sequence of arrays
691 The arrays must have the same shape along all but the third axis.
692 1-D or 2-D arrays must have the same shape.
693
694 Returns
695 -------
696 stacked : ndarray
697 The array formed by stacking the given arrays, will be at least 3-D.
698
699 See Also
700 --------
701 concatenate : Join a sequence of arrays along an existing axis.
702 stack : Join a sequence of arrays along a new axis.
703 block : Assemble an nd-array from nested lists of blocks.
704 vstack : Stack arrays in sequence vertically (row wise).
705 hstack : Stack arrays in sequence horizontally (column wise).
706 column_stack : Stack 1-D arrays as columns into a 2-D array.
707 dsplit : Split array along third axis.
708
709 Examples
710 --------
711 >>> import numpy as np
712 >>> a = np.array((1,2,3))
713 >>> b = np.array((4,5,6))
714 >>> np.dstack((a,b))
715 array([[[1, 4],
716 [2, 5],
717 [3, 6]]])
718
719 >>> a = np.array([[1],[2],[3]])
720 >>> b = np.array([[4],[5],[6]])
721 >>> np.dstack((a,b))
722 array([[[1, 4]],
723 [[2, 5]],
724 [[3, 6]]])
725
726 """
727 arrs = atleast_3d(*tup)
728 if not isinstance(arrs, tuple):
729 arrs = (arrs,)
730 return _nx.concatenate(arrs, 2)
731
732
733def _array_split_dispatcher(ary, indices_or_sections, axis=None):
734 return (ary, indices_or_sections)
735
736
737@array_function_dispatch(_array_split_dispatcher)
738def array_split(ary, indices_or_sections, axis=0):
739 """
740 Split an array into multiple sub-arrays.
741
742 Please refer to the ``split`` documentation. The only difference
743 between these functions is that ``array_split`` allows
744 `indices_or_sections` to be an integer that does *not* equally
745 divide the axis. For an array of length l that should be split
746 into n sections, it returns l % n sub-arrays of size l//n + 1
747 and the rest of size l//n.
748
749 See Also
750 --------
751 split : Split array into multiple sub-arrays of equal size.
752
753 Examples
754 --------
755 >>> import numpy as np
756 >>> x = np.arange(8.0)
757 >>> np.array_split(x, 3)
758 [array([0., 1., 2.]), array([3., 4., 5.]), array([6., 7.])]
759
760 >>> x = np.arange(9)
761 >>> np.array_split(x, 4)
762 [array([0, 1, 2]), array([3, 4]), array([5, 6]), array([7, 8])]
763
764 """
765 try:
766 Ntotal = ary.shape[axis]
767 except AttributeError:
768 Ntotal = len(ary)
769 try:
770 # handle array case.
771 Nsections = len(indices_or_sections) + 1
772 div_points = [0] + list(indices_or_sections) + [Ntotal]
773 except TypeError:
774 # indices_or_sections is a scalar, not an array.
775 Nsections = int(indices_or_sections)
776 if Nsections <= 0:
777 raise ValueError('number sections must be larger than 0.') from None
778 Neach_section, extras = divmod(Ntotal, Nsections)
779 section_sizes = ([0] +
780 extras * [Neach_section + 1] +
781 (Nsections - extras) * [Neach_section])
782 div_points = _nx.array(section_sizes, dtype=_nx.intp).cumsum()
783
784 sub_arys = []
785 sary = _nx.swapaxes(ary, axis, 0)
786 for i in range(Nsections):
787 st = div_points[i]
788 end = div_points[i + 1]
789 sub_arys.append(_nx.swapaxes(sary[st:end], axis, 0))
790
791 return sub_arys
792
793
794def _split_dispatcher(ary, indices_or_sections, axis=None):
795 return (ary, indices_or_sections)
796
797
798@array_function_dispatch(_split_dispatcher)
799def split(ary, indices_or_sections, axis=0):
800 """
801 Split an array into multiple sub-arrays as views into `ary`.
802
803 Parameters
804 ----------
805 ary : ndarray
806 Array to be divided into sub-arrays.
807 indices_or_sections : int or 1-D array
808 If `indices_or_sections` is an integer, N, the array will be divided
809 into N equal arrays along `axis`. If such a split is not possible,
810 an error is raised.
811
812 If `indices_or_sections` is a 1-D array of sorted integers, the entries
813 indicate where along `axis` the array is split. For example,
814 ``[2, 3]`` would, for ``axis=0``, result in
815
816 - ary[:2]
817 - ary[2:3]
818 - ary[3:]
819
820 If an index exceeds the dimension of the array along `axis`,
821 an empty sub-array is returned correspondingly.
822 axis : int, optional
823 The axis along which to split, default is 0.
824
825 Returns
826 -------
827 sub-arrays : list of ndarrays
828 A list of sub-arrays as views into `ary`.
829
830 Raises
831 ------
832 ValueError
833 If `indices_or_sections` is given as an integer, but
834 a split does not result in equal division.
835
836 See Also
837 --------
838 array_split : Split an array into multiple sub-arrays of equal or
839 near-equal size. Does not raise an exception if
840 an equal division cannot be made.
841 hsplit : Split array into multiple sub-arrays horizontally (column-wise).
842 vsplit : Split array into multiple sub-arrays vertically (row wise).
843 dsplit : Split array into multiple sub-arrays along the 3rd axis (depth).
844 concatenate : Join a sequence of arrays along an existing axis.
845 stack : Join a sequence of arrays along a new axis.
846 hstack : Stack arrays in sequence horizontally (column wise).
847 vstack : Stack arrays in sequence vertically (row wise).
848 dstack : Stack arrays in sequence depth wise (along third dimension).
849
850 Examples
851 --------
852 >>> import numpy as np
853 >>> x = np.arange(9.0)
854 >>> np.split(x, 3)
855 [array([0., 1., 2.]), array([3., 4., 5.]), array([6., 7., 8.])]
856
857 >>> x = np.arange(8.0)
858 >>> np.split(x, [3, 5, 6, 10])
859 [array([0., 1., 2.]),
860 array([3., 4.]),
861 array([5.]),
862 array([6., 7.]),
863 array([], dtype=float64)]
864
865 """
866 try:
867 len(indices_or_sections)
868 except TypeError:
869 sections = indices_or_sections
870 N = ary.shape[axis]
871 if N % sections:
872 raise ValueError(
873 'array split does not result in an equal division') from None
874 return array_split(ary, indices_or_sections, axis)
875
876
877def _hvdsplit_dispatcher(ary, indices_or_sections):
878 return (ary, indices_or_sections)
879
880
881@array_function_dispatch(_hvdsplit_dispatcher)
882def hsplit(ary, indices_or_sections):
883 """
884 Split an array into multiple sub-arrays horizontally (column-wise).
885
886 Please refer to the `split` documentation. `hsplit` is equivalent
887 to `split` with ``axis=1``, the array is always split along the second
888 axis except for 1-D arrays, where it is split at ``axis=0``.
889
890 See Also
891 --------
892 split : Split an array into multiple sub-arrays of equal size.
893
894 Examples
895 --------
896 >>> import numpy as np
897 >>> x = np.arange(16.0).reshape(4, 4)
898 >>> x
899 array([[ 0., 1., 2., 3.],
900 [ 4., 5., 6., 7.],
901 [ 8., 9., 10., 11.],
902 [12., 13., 14., 15.]])
903 >>> np.hsplit(x, 2)
904 [array([[ 0., 1.],
905 [ 4., 5.],
906 [ 8., 9.],
907 [12., 13.]]),
908 array([[ 2., 3.],
909 [ 6., 7.],
910 [10., 11.],
911 [14., 15.]])]
912 >>> np.hsplit(x, np.array([3, 6]))
913 [array([[ 0., 1., 2.],
914 [ 4., 5., 6.],
915 [ 8., 9., 10.],
916 [12., 13., 14.]]),
917 array([[ 3.],
918 [ 7.],
919 [11.],
920 [15.]]),
921 array([], shape=(4, 0), dtype=float64)]
922
923 With a higher dimensional array the split is still along the second axis.
924
925 >>> x = np.arange(8.0).reshape(2, 2, 2)
926 >>> x
927 array([[[0., 1.],
928 [2., 3.]],
929 [[4., 5.],
930 [6., 7.]]])
931 >>> np.hsplit(x, 2)
932 [array([[[0., 1.]],
933 [[4., 5.]]]),
934 array([[[2., 3.]],
935 [[6., 7.]]])]
936
937 With a 1-D array, the split is along axis 0.
938
939 >>> x = np.array([0, 1, 2, 3, 4, 5])
940 >>> np.hsplit(x, 2)
941 [array([0, 1, 2]), array([3, 4, 5])]
942
943 """
944 if _nx.ndim(ary) == 0:
945 raise ValueError('hsplit only works on arrays of 1 or more dimensions')
946 if ary.ndim > 1:
947 return split(ary, indices_or_sections, 1)
948 else:
949 return split(ary, indices_or_sections, 0)
950
951
952@array_function_dispatch(_hvdsplit_dispatcher)
953def vsplit(ary, indices_or_sections):
954 """
955 Split an array into multiple sub-arrays vertically (row-wise).
956
957 Please refer to the ``split`` documentation. ``vsplit`` is equivalent
958 to ``split`` with `axis=0` (default), the array is always split along the
959 first axis regardless of the array dimension.
960
961 See Also
962 --------
963 split : Split an array into multiple sub-arrays of equal size.
964
965 Examples
966 --------
967 >>> import numpy as np
968 >>> x = np.arange(16.0).reshape(4, 4)
969 >>> x
970 array([[ 0., 1., 2., 3.],
971 [ 4., 5., 6., 7.],
972 [ 8., 9., 10., 11.],
973 [12., 13., 14., 15.]])
974 >>> np.vsplit(x, 2)
975 [array([[0., 1., 2., 3.],
976 [4., 5., 6., 7.]]),
977 array([[ 8., 9., 10., 11.],
978 [12., 13., 14., 15.]])]
979 >>> np.vsplit(x, np.array([3, 6]))
980 [array([[ 0., 1., 2., 3.],
981 [ 4., 5., 6., 7.],
982 [ 8., 9., 10., 11.]]),
983 array([[12., 13., 14., 15.]]),
984 array([], shape=(0, 4), dtype=float64)]
985
986 With a higher dimensional array the split is still along the first axis.
987
988 >>> x = np.arange(8.0).reshape(2, 2, 2)
989 >>> x
990 array([[[0., 1.],
991 [2., 3.]],
992 [[4., 5.],
993 [6., 7.]]])
994 >>> np.vsplit(x, 2)
995 [array([[[0., 1.],
996 [2., 3.]]]),
997 array([[[4., 5.],
998 [6., 7.]]])]
999
1000 """
1001 if _nx.ndim(ary) < 2:
1002 raise ValueError('vsplit only works on arrays of 2 or more dimensions')
1003 return split(ary, indices_or_sections, 0)
1004
1005
1006@array_function_dispatch(_hvdsplit_dispatcher)
1007def dsplit(ary, indices_or_sections):
1008 """
1009 Split array into multiple sub-arrays along the 3rd axis (depth).
1010
1011 Please refer to the `split` documentation. `dsplit` is equivalent
1012 to `split` with ``axis=2``, the array is always split along the third
1013 axis provided the array dimension is greater than or equal to 3.
1014
1015 See Also
1016 --------
1017 split : Split an array into multiple sub-arrays of equal size.
1018
1019 Examples
1020 --------
1021 >>> import numpy as np
1022 >>> x = np.arange(16.0).reshape(2, 2, 4)
1023 >>> x
1024 array([[[ 0., 1., 2., 3.],
1025 [ 4., 5., 6., 7.]],
1026 [[ 8., 9., 10., 11.],
1027 [12., 13., 14., 15.]]])
1028 >>> np.dsplit(x, 2)
1029 [array([[[ 0., 1.],
1030 [ 4., 5.]],
1031 [[ 8., 9.],
1032 [12., 13.]]]), array([[[ 2., 3.],
1033 [ 6., 7.]],
1034 [[10., 11.],
1035 [14., 15.]]])]
1036 >>> np.dsplit(x, np.array([3, 6]))
1037 [array([[[ 0., 1., 2.],
1038 [ 4., 5., 6.]],
1039 [[ 8., 9., 10.],
1040 [12., 13., 14.]]]),
1041 array([[[ 3.],
1042 [ 7.]],
1043 [[11.],
1044 [15.]]]),
1045 array([], shape=(2, 2, 0), dtype=float64)]
1046 """
1047 if _nx.ndim(ary) < 3:
1048 raise ValueError('dsplit only works on arrays of 3 or more dimensions')
1049 return split(ary, indices_or_sections, 2)
1050
1051
1052def get_array_wrap(*args):
1053 """Find the wrapper for the array with the highest priority.
1054
1055 In case of ties, leftmost wins. If no wrapper is found, return None.
1056
1057 .. deprecated:: 2.0
1058 """
1059
1060 # Deprecated in NumPy 2.0, 2023-07-11
1061 warnings.warn(
1062 "`get_array_wrap` is deprecated. "
1063 "(deprecated in NumPy 2.0)",
1064 DeprecationWarning,
1065 stacklevel=2
1066 )
1067
1068 wrappers = sorted((getattr(x, '__array_priority__', 0), -i,
1069 x.__array_wrap__) for i, x in enumerate(args)
1070 if hasattr(x, '__array_wrap__'))
1071 if wrappers:
1072 return wrappers[-1][-1]
1073 return None
1074
1075
1076def _kron_dispatcher(a, b):
1077 return (a, b)
1078
1079
1080@array_function_dispatch(_kron_dispatcher)
1081def kron(a, b):
1082 """
1083 Kronecker product of two arrays.
1084
1085 Computes the Kronecker product, a composite array made of blocks of the
1086 second array scaled by the first.
1087
1088 Parameters
1089 ----------
1090 a, b : array_like
1091
1092 Returns
1093 -------
1094 out : ndarray
1095
1096 See Also
1097 --------
1098 outer : The outer product
1099
1100 Notes
1101 -----
1102 The function assumes that the number of dimensions of `a` and `b`
1103 are the same, if necessary prepending the smallest with ones.
1104 If ``a.shape = (r0,r1,...,rN)`` and ``b.shape = (s0,s1,...,sN)``,
1105 the Kronecker product has shape ``(r0*s0, r1*s1, ..., rN*SN)``.
1106 The elements are products of elements from `a` and `b`, organized
1107 explicitly by::
1108
1109 kron(a,b)[k0,k1,...,kN] = a[i0,i1,...,iN] * b[j0,j1,...,jN]
1110
1111 where::
1112
1113 kt = it * st + jt, t = 0,...,N
1114
1115 In the common 2-D case (N=1), the block structure can be visualized::
1116
1117 [[ a[0,0]*b, a[0,1]*b, ... , a[0,-1]*b ],
1118 [ ... ... ],
1119 [ a[-1,0]*b, a[-1,1]*b, ... , a[-1,-1]*b ]]
1120
1121
1122 Examples
1123 --------
1124 >>> import numpy as np
1125 >>> np.kron([1,10,100], [5,6,7])
1126 array([ 5, 6, 7, ..., 500, 600, 700])
1127 >>> np.kron([5,6,7], [1,10,100])
1128 array([ 5, 50, 500, ..., 7, 70, 700])
1129
1130 >>> np.kron(np.eye(2), np.ones((2,2)))
1131 array([[1., 1., 0., 0.],
1132 [1., 1., 0., 0.],
1133 [0., 0., 1., 1.],
1134 [0., 0., 1., 1.]])
1135
1136 >>> a = np.arange(100).reshape((2,5,2,5))
1137 >>> b = np.arange(24).reshape((2,3,4))
1138 >>> c = np.kron(a,b)
1139 >>> c.shape
1140 (2, 10, 6, 20)
1141 >>> I = (1,3,0,2)
1142 >>> J = (0,2,1)
1143 >>> J1 = (0,) + J # extend to ndim=4
1144 >>> S1 = (1,) + b.shape
1145 >>> K = tuple(np.array(I) * np.array(S1) + np.array(J1))
1146 >>> c[K] == a[I]*b[J]
1147 True
1148
1149 """
1150 # Working:
1151 # 1. Equalise the shapes by prepending smaller array with 1s
1152 # 2. Expand shapes of both the arrays by adding new axes at
1153 # odd positions for 1st array and even positions for 2nd
1154 # 3. Compute the product of the modified array
1155 # 4. The inner most array elements now contain the rows of
1156 # the Kronecker product
1157 # 5. Reshape the result to kron's shape, which is same as
1158 # product of shapes of the two arrays.
1159 b = asanyarray(b)
1160 a = array(a, copy=None, subok=True, ndmin=b.ndim)
1161 is_any_mat = isinstance(a, matrix) or isinstance(b, matrix)
1162 ndb, nda = b.ndim, a.ndim
1163 nd = max(ndb, nda)
1164
1165 if (nda == 0 or ndb == 0):
1166 return _nx.multiply(a, b)
1167
1168 as_ = a.shape
1169 bs = b.shape
1170 if not a.flags.contiguous:
1171 a = reshape(a, as_)
1172 if not b.flags.contiguous:
1173 b = reshape(b, bs)
1174
1175 # Equalise the shapes by prepending smaller one with 1s
1176 as_ = (1,) * max(0, ndb - nda) + as_
1177 bs = (1,) * max(0, nda - ndb) + bs
1178
1179 # Insert empty dimensions
1180 a_arr = expand_dims(a, axis=tuple(range(ndb - nda)))
1181 b_arr = expand_dims(b, axis=tuple(range(nda - ndb)))
1182
1183 # Compute the product
1184 a_arr = expand_dims(a_arr, axis=tuple(range(1, nd * 2, 2)))
1185 b_arr = expand_dims(b_arr, axis=tuple(range(0, nd * 2, 2)))
1186 # In case of `mat`, convert result to `array`
1187 result = _nx.multiply(a_arr, b_arr, subok=(not is_any_mat))
1188
1189 # Reshape back
1190 result = result.reshape(_nx.multiply(as_, bs))
1191
1192 return result if not is_any_mat else matrix(result, copy=False)
1193
1194
1195def _tile_dispatcher(A, reps):
1196 return (A, reps)
1197
1198
1199@array_function_dispatch(_tile_dispatcher)
1200def tile(A, reps):