shape_base.py 37 KB

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