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