_shape_base_impl.py 38 KB

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