numeric.py 80 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954195519561957195819591960196119621963196419651966196719681969197019711972197319741975197619771978197919801981198219831984198519861987198819891990199119921993199419951996199719981999200020012002200320042005200620072008200920102011201220132014201520162017201820192020202120222023202420252026202720282029203020312032203320342035203620372038203920402041204220432044204520462047204820492050205120522053205420552056205720582059206020612062206320642065206620672068206920702071207220732074207520762077207820792080208120822083208420852086208720882089209020912092209320942095209620972098209921002101210221032104210521062107210821092110211121122113211421152116211721182119212021212122212321242125212621272128212921302131213221332134213521362137213821392140214121422143214421452146214721482149215021512152215321542155215621572158215921602161216221632164216521662167216821692170217121722173217421752176217721782179218021812182218321842185218621872188218921902191219221932194219521962197219821992200220122022203220422052206220722082209221022112212221322142215221622172218221922202221222222232224222522262227222822292230223122322233223422352236223722382239224022412242224322442245224622472248224922502251225222532254225522562257225822592260226122622263226422652266226722682269227022712272227322742275227622772278227922802281228222832284228522862287228822892290229122922293229422952296229722982299230023012302230323042305230623072308230923102311231223132314231523162317231823192320232123222323232423252326232723282329233023312332233323342335233623372338233923402341234223432344234523462347234823492350235123522353235423552356235723582359236023612362236323642365236623672368236923702371237223732374237523762377237823792380238123822383238423852386238723882389239023912392239323942395239623972398239924002401240224032404240524062407240824092410241124122413241424152416241724182419242024212422242324242425242624272428242924302431243224332434243524362437243824392440244124422443244424452446244724482449245024512452245324542455245624572458245924602461246224632464246524662467246824692470247124722473247424752476247724782479248024812482248324842485248624872488248924902491249224932494249524962497249824992500250125022503250425052506250725082509251025112512251325142515251625172518251925202521252225232524252525262527252825292530253125322533253425352536253725382539254025412542254325442545254625472548254925502551255225532554255525562557255825592560256125622563256425652566256725682569257025712572257325742575257625772578257925802581258225832584258525862587258825892590259125922593259425952596259725982599260026012602260326042605260626072608260926102611261226132614261526162617261826192620262126222623262426252626262726282629263026312632263326342635263626372638263926402641264226432644264526462647264826492650265126522653265426552656265726582659266026612662266326642665266626672668266926702671267226732674267526762677267826792680268126822683268426852686268726882689269026912692269326942695269626972698269927002701270227032704270527062707270827092710271127122713271427152716271727182719272027212722272327242725272627272728272927302731273227332734273527362737273827392740274127422743274427452746274727482749275027512752275327542755275627572758
  1. import builtins
  2. import functools
  3. import itertools
  4. import math
  5. import numbers
  6. import operator
  7. import sys
  8. import warnings
  9. import numpy as np
  10. from numpy.exceptions import AxisError
  11. from . import multiarray, numerictypes, numerictypes as nt, overrides, shape_base, umath
  12. from ._ufunc_config import errstate
  13. from .multiarray import ( # noqa: F401
  14. ALLOW_THREADS,
  15. BUFSIZE,
  16. CLIP,
  17. MAXDIMS,
  18. MAY_SHARE_BOUNDS,
  19. MAY_SHARE_EXACT,
  20. RAISE,
  21. WRAP,
  22. arange,
  23. array,
  24. asanyarray,
  25. asarray,
  26. ascontiguousarray,
  27. asfortranarray,
  28. broadcast,
  29. can_cast,
  30. concatenate,
  31. copyto,
  32. dot,
  33. dtype,
  34. empty,
  35. empty_like,
  36. flatiter,
  37. from_dlpack,
  38. frombuffer,
  39. fromfile,
  40. fromiter,
  41. fromstring,
  42. inner,
  43. lexsort,
  44. matmul,
  45. may_share_memory,
  46. min_scalar_type,
  47. ndarray,
  48. nditer,
  49. nested_iters,
  50. normalize_axis_index,
  51. promote_types,
  52. putmask,
  53. result_type,
  54. shares_memory,
  55. vdot,
  56. vecdot,
  57. where,
  58. zeros,
  59. )
  60. from .overrides import finalize_array_function_like, set_module
  61. from .umath import NAN, PINF, invert, multiply, sin
  62. bitwise_not = invert
  63. ufunc = type(sin)
  64. newaxis = None
  65. array_function_dispatch = functools.partial(
  66. overrides.array_function_dispatch, module='numpy')
  67. __all__ = [
  68. 'newaxis', 'ndarray', 'flatiter', 'nditer', 'nested_iters', 'ufunc',
  69. 'arange', 'array', 'asarray', 'asanyarray', 'ascontiguousarray',
  70. 'asfortranarray', 'zeros', 'count_nonzero', 'empty', 'broadcast', 'dtype',
  71. 'fromstring', 'fromfile', 'frombuffer', 'from_dlpack', 'where',
  72. 'argwhere', 'copyto', 'concatenate', 'lexsort', 'astype',
  73. 'can_cast', 'promote_types', 'min_scalar_type',
  74. 'result_type', 'isfortran', 'empty_like', 'zeros_like', 'ones_like',
  75. 'correlate', 'convolve', 'inner', 'dot', 'outer', 'vdot', 'roll',
  76. 'rollaxis', 'moveaxis', 'cross', 'tensordot', 'little_endian',
  77. 'fromiter', 'array_equal', 'array_equiv', 'indices', 'fromfunction',
  78. 'isclose', 'isscalar', 'binary_repr', 'base_repr', 'ones',
  79. 'identity', 'allclose', 'putmask',
  80. 'flatnonzero', 'inf', 'nan', 'False_', 'True_', 'bitwise_not',
  81. 'full', 'full_like', 'matmul', 'vecdot', 'shares_memory',
  82. 'may_share_memory']
  83. def _zeros_like_dispatcher(
  84. a, dtype=None, order=None, subok=None, shape=None, *, device=None
  85. ):
  86. return (a,)
  87. @array_function_dispatch(_zeros_like_dispatcher)
  88. def zeros_like(
  89. a, dtype=None, order='K', subok=True, shape=None, *, device=None
  90. ):
  91. """
  92. Return an array of zeros with the same shape and type as a given array.
  93. Parameters
  94. ----------
  95. a : array_like
  96. The shape and data-type of `a` define these same attributes of
  97. the returned array.
  98. dtype : data-type, optional
  99. Overrides the data type of the result.
  100. order : {'C', 'F', 'A', or 'K'}, optional
  101. Overrides the memory layout of the result. 'C' means C-order,
  102. 'F' means F-order, 'A' means 'F' if `a` is Fortran contiguous,
  103. 'C' otherwise. 'K' means match the layout of `a` as closely
  104. as possible.
  105. subok : bool, optional.
  106. If True, then the newly created array will use the sub-class
  107. type of `a`, otherwise it will be a base-class array. Defaults
  108. to True.
  109. shape : int or sequence of ints, optional.
  110. Overrides the shape of the result. If order='K' and the number of
  111. dimensions is unchanged, will try to keep order, otherwise,
  112. order='C' is implied.
  113. device : str, optional
  114. The device on which to place the created array. Default: None.
  115. For Array-API interoperability only, so must be ``"cpu"`` if passed.
  116. .. versionadded:: 2.0.0
  117. Returns
  118. -------
  119. out : ndarray
  120. Array of zeros with the same shape and type as `a`.
  121. See Also
  122. --------
  123. empty_like : Return an empty array with shape and type of input.
  124. ones_like : Return an array of ones with shape and type of input.
  125. full_like : Return a new array with shape of input filled with value.
  126. zeros : Return a new array setting values to zero.
  127. Examples
  128. --------
  129. >>> import numpy as np
  130. >>> x = np.arange(6)
  131. >>> x = x.reshape((2, 3))
  132. >>> x
  133. array([[0, 1, 2],
  134. [3, 4, 5]])
  135. >>> np.zeros_like(x)
  136. array([[0, 0, 0],
  137. [0, 0, 0]])
  138. >>> y = np.arange(3, dtype=float)
  139. >>> y
  140. array([0., 1., 2.])
  141. >>> np.zeros_like(y)
  142. array([0., 0., 0.])
  143. """
  144. res = empty_like(
  145. a, dtype=dtype, order=order, subok=subok, shape=shape, device=device
  146. )
  147. # needed instead of a 0 to get same result as zeros for string dtypes
  148. z = zeros(1, dtype=res.dtype)
  149. multiarray.copyto(res, z, casting='unsafe')
  150. return res
  151. @finalize_array_function_like
  152. @set_module('numpy')
  153. def ones(shape, dtype=None, order='C', *, device=None, like=None):
  154. """
  155. Return a new array of given shape and type, filled with ones.
  156. Parameters
  157. ----------
  158. shape : int or sequence of ints
  159. Shape of the new array, e.g., ``(2, 3)`` or ``2``.
  160. dtype : data-type, optional
  161. The desired data-type for the array, e.g., `numpy.int8`. Default is
  162. `numpy.float64`.
  163. order : {'C', 'F'}, optional, default: C
  164. Whether to store multi-dimensional data in row-major
  165. (C-style) or column-major (Fortran-style) order in
  166. memory.
  167. device : str, optional
  168. The device on which to place the created array. Default: None.
  169. For Array-API interoperability only, so must be ``"cpu"`` if passed.
  170. .. versionadded:: 2.0.0
  171. ${ARRAY_FUNCTION_LIKE}
  172. .. versionadded:: 1.20.0
  173. Returns
  174. -------
  175. out : ndarray
  176. Array of ones with the given shape, dtype, and order.
  177. See Also
  178. --------
  179. ones_like : Return an array of ones with shape and type of input.
  180. empty : Return a new uninitialized array.
  181. zeros : Return a new array setting values to zero.
  182. full : Return a new array of given shape filled with value.
  183. Examples
  184. --------
  185. >>> import numpy as np
  186. >>> np.ones(5)
  187. array([1., 1., 1., 1., 1.])
  188. >>> np.ones((5,), dtype=int)
  189. array([1, 1, 1, 1, 1])
  190. >>> np.ones((2, 1))
  191. array([[1.],
  192. [1.]])
  193. >>> s = (2,2)
  194. >>> np.ones(s)
  195. array([[1., 1.],
  196. [1., 1.]])
  197. """
  198. if like is not None:
  199. return _ones_with_like(
  200. like, shape, dtype=dtype, order=order, device=device
  201. )
  202. a = empty(shape, dtype, order, device=device)
  203. multiarray.copyto(a, 1, casting='unsafe')
  204. return a
  205. _ones_with_like = array_function_dispatch()(ones)
  206. def _ones_like_dispatcher(
  207. a, dtype=None, order=None, subok=None, shape=None, *, device=None
  208. ):
  209. return (a,)
  210. @array_function_dispatch(_ones_like_dispatcher)
  211. def ones_like(
  212. a, dtype=None, order='K', subok=True, shape=None, *, device=None
  213. ):
  214. """
  215. Return an array of ones with the same shape and type as a given array.
  216. Parameters
  217. ----------
  218. a : array_like
  219. The shape and data-type of `a` define these same attributes of
  220. the returned array.
  221. dtype : data-type, optional
  222. Overrides the data type of the result.
  223. order : {'C', 'F', 'A', or 'K'}, optional
  224. Overrides the memory layout of the result. 'C' means C-order,
  225. 'F' means F-order, 'A' means 'F' if `a` is Fortran contiguous,
  226. 'C' otherwise. 'K' means match the layout of `a` as closely
  227. as possible.
  228. subok : bool, optional.
  229. If True, then the newly created array will use the sub-class
  230. type of `a`, otherwise it will be a base-class array. Defaults
  231. to True.
  232. shape : int or sequence of ints, optional.
  233. Overrides the shape of the result. If order='K' and the number of
  234. dimensions is unchanged, will try to keep order, otherwise,
  235. order='C' is implied.
  236. device : str, optional
  237. The device on which to place the created array. Default: None.
  238. For Array-API interoperability only, so must be ``"cpu"`` if passed.
  239. .. versionadded:: 2.0.0
  240. Returns
  241. -------
  242. out : ndarray
  243. Array of ones with the same shape and type as `a`.
  244. See Also
  245. --------
  246. empty_like : Return an empty array with shape and type of input.
  247. zeros_like : Return an array of zeros with shape and type of input.
  248. full_like : Return a new array with shape of input filled with value.
  249. ones : Return a new array setting values to one.
  250. Examples
  251. --------
  252. >>> import numpy as np
  253. >>> x = np.arange(6)
  254. >>> x = x.reshape((2, 3))
  255. >>> x
  256. array([[0, 1, 2],
  257. [3, 4, 5]])
  258. >>> np.ones_like(x)
  259. array([[1, 1, 1],
  260. [1, 1, 1]])
  261. >>> y = np.arange(3, dtype=float)
  262. >>> y
  263. array([0., 1., 2.])
  264. >>> np.ones_like(y)
  265. array([1., 1., 1.])
  266. """
  267. res = empty_like(
  268. a, dtype=dtype, order=order, subok=subok, shape=shape, device=device
  269. )
  270. multiarray.copyto(res, 1, casting='unsafe')
  271. return res
  272. def _full_dispatcher(
  273. shape, fill_value, dtype=None, order=None, *, device=None, like=None
  274. ):
  275. return (like,)
  276. @finalize_array_function_like
  277. @set_module('numpy')
  278. def full(shape, fill_value, dtype=None, order='C', *, device=None, like=None):
  279. """
  280. Return a new array of given shape and type, filled with `fill_value`.
  281. Parameters
  282. ----------
  283. shape : int or sequence of ints
  284. Shape of the new array, e.g., ``(2, 3)`` or ``2``.
  285. fill_value : scalar or array_like
  286. Fill value.
  287. dtype : data-type, optional
  288. The desired data-type for the array The default, None, means
  289. ``np.array(fill_value).dtype``.
  290. order : {'C', 'F'}, optional
  291. Whether to store multidimensional data in C- or Fortran-contiguous
  292. (row- or column-wise) order in memory.
  293. device : str, optional
  294. The device on which to place the created array. Default: None.
  295. For Array-API interoperability only, so must be ``"cpu"`` if passed.
  296. .. versionadded:: 2.0.0
  297. ${ARRAY_FUNCTION_LIKE}
  298. .. versionadded:: 1.20.0
  299. Returns
  300. -------
  301. out : ndarray
  302. Array of `fill_value` with the given shape, dtype, and order.
  303. See Also
  304. --------
  305. full_like : Return a new array with shape of input filled with value.
  306. empty : Return a new uninitialized array.
  307. ones : Return a new array setting values to one.
  308. zeros : Return a new array setting values to zero.
  309. Examples
  310. --------
  311. >>> import numpy as np
  312. >>> np.full((2, 2), np.inf)
  313. array([[inf, inf],
  314. [inf, inf]])
  315. >>> np.full((2, 2), 10)
  316. array([[10, 10],
  317. [10, 10]])
  318. >>> np.full((2, 2), [1, 2])
  319. array([[1, 2],
  320. [1, 2]])
  321. """
  322. if like is not None:
  323. return _full_with_like(
  324. like, shape, fill_value, dtype=dtype, order=order, device=device
  325. )
  326. if dtype is None:
  327. fill_value = asarray(fill_value)
  328. dtype = fill_value.dtype
  329. a = empty(shape, dtype, order, device=device)
  330. multiarray.copyto(a, fill_value, casting='unsafe')
  331. return a
  332. _full_with_like = array_function_dispatch()(full)
  333. def _full_like_dispatcher(
  334. a, fill_value, dtype=None, order=None, subok=None, shape=None,
  335. *, device=None
  336. ):
  337. return (a,)
  338. @array_function_dispatch(_full_like_dispatcher)
  339. def full_like(
  340. a, fill_value, dtype=None, order='K', subok=True, shape=None,
  341. *, device=None
  342. ):
  343. """
  344. Return a full array with the same shape and type as a given array.
  345. Parameters
  346. ----------
  347. a : array_like
  348. The shape and data-type of `a` define these same attributes of
  349. the returned array.
  350. fill_value : array_like
  351. Fill value.
  352. dtype : data-type, optional
  353. Overrides the data type of the result.
  354. order : {'C', 'F', 'A', or 'K'}, optional
  355. Overrides the memory layout of the result. 'C' means C-order,
  356. 'F' means F-order, 'A' means 'F' if `a` is Fortran contiguous,
  357. 'C' otherwise. 'K' means match the layout of `a` as closely
  358. as possible.
  359. subok : bool, optional.
  360. If True, then the newly created array will use the sub-class
  361. type of `a`, otherwise it will be a base-class array. Defaults
  362. to True.
  363. shape : int or sequence of ints, optional.
  364. Overrides the shape of the result. If order='K' and the number of
  365. dimensions is unchanged, will try to keep order, otherwise,
  366. order='C' is implied.
  367. device : str, optional
  368. The device on which to place the created array. Default: None.
  369. For Array-API interoperability only, so must be ``"cpu"`` if passed.
  370. .. versionadded:: 2.0.0
  371. Returns
  372. -------
  373. out : ndarray
  374. Array of `fill_value` with the same shape and type as `a`.
  375. See Also
  376. --------
  377. empty_like : Return an empty array with shape and type of input.
  378. ones_like : Return an array of ones with shape and type of input.
  379. zeros_like : Return an array of zeros with shape and type of input.
  380. full : Return a new array of given shape filled with value.
  381. Examples
  382. --------
  383. >>> import numpy as np
  384. >>> x = np.arange(6, dtype=int)
  385. >>> np.full_like(x, 1)
  386. array([1, 1, 1, 1, 1, 1])
  387. >>> np.full_like(x, 0.1)
  388. array([0, 0, 0, 0, 0, 0])
  389. >>> np.full_like(x, 0.1, dtype=np.double)
  390. array([0.1, 0.1, 0.1, 0.1, 0.1, 0.1])
  391. >>> np.full_like(x, np.nan, dtype=np.double)
  392. array([nan, nan, nan, nan, nan, nan])
  393. >>> y = np.arange(6, dtype=np.double)
  394. >>> np.full_like(y, 0.1)
  395. array([0.1, 0.1, 0.1, 0.1, 0.1, 0.1])
  396. >>> y = np.zeros([2, 2, 3], dtype=int)
  397. >>> np.full_like(y, [0, 0, 255])
  398. array([[[ 0, 0, 255],
  399. [ 0, 0, 255]],
  400. [[ 0, 0, 255],
  401. [ 0, 0, 255]]])
  402. """
  403. res = empty_like(
  404. a, dtype=dtype, order=order, subok=subok, shape=shape, device=device
  405. )
  406. multiarray.copyto(res, fill_value, casting='unsafe')
  407. return res
  408. def _count_nonzero_dispatcher(a, axis=None, *, keepdims=None):
  409. return (a,)
  410. @array_function_dispatch(_count_nonzero_dispatcher)
  411. def count_nonzero(a, axis=None, *, keepdims=False):
  412. """
  413. Counts the number of non-zero values in the array ``a``.
  414. The word "non-zero" is in reference to the Python 2.x
  415. built-in method ``__nonzero__()`` (renamed ``__bool__()``
  416. in Python 3.x) of Python objects that tests an object's
  417. "truthfulness". For example, any number is considered
  418. truthful if it is nonzero, whereas any string is considered
  419. truthful if it is not the empty string. Thus, this function
  420. (recursively) counts how many elements in ``a`` (and in
  421. sub-arrays thereof) have their ``__nonzero__()`` or ``__bool__()``
  422. method evaluated to ``True``.
  423. Parameters
  424. ----------
  425. a : array_like
  426. The array for which to count non-zeros.
  427. axis : int or tuple, optional
  428. Axis or tuple of axes along which to count non-zeros.
  429. Default is None, meaning that non-zeros will be counted
  430. along a flattened version of ``a``.
  431. keepdims : bool, optional
  432. If this is set to True, the axes that are counted are left
  433. in the result as dimensions with size one. With this option,
  434. the result will broadcast correctly against the input array.
  435. Returns
  436. -------
  437. count : int or array of int
  438. Number of non-zero values in the array along a given axis.
  439. Otherwise, the total number of non-zero values in the array
  440. is returned.
  441. See Also
  442. --------
  443. nonzero : Return the coordinates of all the non-zero values.
  444. Examples
  445. --------
  446. >>> import numpy as np
  447. >>> np.count_nonzero(np.eye(4))
  448. np.int64(4)
  449. >>> a = np.array([[0, 1, 7, 0],
  450. ... [3, 0, 2, 19]])
  451. >>> np.count_nonzero(a)
  452. np.int64(5)
  453. >>> np.count_nonzero(a, axis=0)
  454. array([1, 1, 2, 1])
  455. >>> np.count_nonzero(a, axis=1)
  456. array([2, 3])
  457. >>> np.count_nonzero(a, axis=1, keepdims=True)
  458. array([[2],
  459. [3]])
  460. """
  461. if axis is None and not keepdims:
  462. return multiarray.count_nonzero(a)
  463. a = asanyarray(a)
  464. # TODO: this works around .astype(bool) not working properly (gh-9847)
  465. if np.issubdtype(a.dtype, np.character):
  466. a_bool = a != a.dtype.type()
  467. else:
  468. a_bool = a.astype(np.bool, copy=False)
  469. return a_bool.sum(axis=axis, dtype=np.intp, keepdims=keepdims)
  470. @set_module('numpy')
  471. def isfortran(a):
  472. """
  473. Check if the array is Fortran contiguous but *not* C contiguous.
  474. This function is obsolete. If you only want to check if an array is Fortran
  475. contiguous use ``a.flags.f_contiguous`` instead.
  476. Parameters
  477. ----------
  478. a : ndarray
  479. Input array.
  480. Returns
  481. -------
  482. isfortran : bool
  483. Returns True if the array is Fortran contiguous but *not* C contiguous.
  484. Examples
  485. --------
  486. np.array allows to specify whether the array is written in C-contiguous
  487. order (last index varies the fastest), or FORTRAN-contiguous order in
  488. memory (first index varies the fastest).
  489. >>> import numpy as np
  490. >>> a = np.array([[1, 2, 3], [4, 5, 6]], order='C')
  491. >>> a
  492. array([[1, 2, 3],
  493. [4, 5, 6]])
  494. >>> np.isfortran(a)
  495. False
  496. >>> b = np.array([[1, 2, 3], [4, 5, 6]], order='F')
  497. >>> b
  498. array([[1, 2, 3],
  499. [4, 5, 6]])
  500. >>> np.isfortran(b)
  501. True
  502. The transpose of a C-ordered array is a FORTRAN-ordered array.
  503. >>> a = np.array([[1, 2, 3], [4, 5, 6]], order='C')
  504. >>> a
  505. array([[1, 2, 3],
  506. [4, 5, 6]])
  507. >>> np.isfortran(a)
  508. False
  509. >>> b = a.T
  510. >>> b
  511. array([[1, 4],
  512. [2, 5],
  513. [3, 6]])
  514. >>> np.isfortran(b)
  515. True
  516. C-ordered arrays evaluate as False even if they are also FORTRAN-ordered.
  517. >>> np.isfortran(np.array([1, 2], order='F'))
  518. False
  519. """
  520. return a.flags.fnc
  521. def _argwhere_dispatcher(a):
  522. return (a,)
  523. @array_function_dispatch(_argwhere_dispatcher)
  524. def argwhere(a):
  525. """
  526. Find the indices of array elements that are non-zero, grouped by element.
  527. Parameters
  528. ----------
  529. a : array_like
  530. Input data.
  531. Returns
  532. -------
  533. index_array : (N, a.ndim) ndarray
  534. Indices of elements that are non-zero. Indices are grouped by element.
  535. This array will have shape ``(N, a.ndim)`` where ``N`` is the number of
  536. non-zero items.
  537. See Also
  538. --------
  539. where, nonzero
  540. Notes
  541. -----
  542. ``np.argwhere(a)`` is almost the same as ``np.transpose(np.nonzero(a))``,
  543. but produces a result of the correct shape for a 0D array.
  544. The output of ``argwhere`` is not suitable for indexing arrays.
  545. For this purpose use ``nonzero(a)`` instead.
  546. Examples
  547. --------
  548. >>> import numpy as np
  549. >>> x = np.arange(6).reshape(2,3)
  550. >>> x
  551. array([[0, 1, 2],
  552. [3, 4, 5]])
  553. >>> np.argwhere(x>1)
  554. array([[0, 2],
  555. [1, 0],
  556. [1, 1],
  557. [1, 2]])
  558. """
  559. # nonzero does not behave well on 0d, so promote to 1d
  560. if np.ndim(a) == 0:
  561. a = shape_base.atleast_1d(a)
  562. # then remove the added dimension
  563. return argwhere(a)[:, :0]
  564. return transpose(nonzero(a))
  565. def _flatnonzero_dispatcher(a):
  566. return (a,)
  567. @array_function_dispatch(_flatnonzero_dispatcher)
  568. def flatnonzero(a):
  569. """
  570. Return indices that are non-zero in the flattened version of a.
  571. This is equivalent to ``np.nonzero(np.ravel(a))[0]``.
  572. Parameters
  573. ----------
  574. a : array_like
  575. Input data.
  576. Returns
  577. -------
  578. res : ndarray
  579. Output array, containing the indices of the elements of ``a.ravel()``
  580. that are non-zero.
  581. See Also
  582. --------
  583. nonzero : Return the indices of the non-zero elements of the input array.
  584. ravel : Return a 1-D array containing the elements of the input array.
  585. Examples
  586. --------
  587. >>> import numpy as np
  588. >>> x = np.arange(-2, 3)
  589. >>> x
  590. array([-2, -1, 0, 1, 2])
  591. >>> np.flatnonzero(x)
  592. array([0, 1, 3, 4])
  593. Use the indices of the non-zero elements as an index array to extract
  594. these elements:
  595. >>> x.ravel()[np.flatnonzero(x)]
  596. array([-2, -1, 1, 2])
  597. """
  598. return np.nonzero(np.ravel(a))[0]
  599. def _correlate_dispatcher(a, v, mode=None):
  600. return (a, v)
  601. @array_function_dispatch(_correlate_dispatcher)
  602. def correlate(a, v, mode='valid'):
  603. r"""
  604. Cross-correlation of two 1-dimensional sequences.
  605. This function computes the correlation as generally defined in signal
  606. processing texts [1]_:
  607. .. math:: c_k = \sum_n a_{n+k} \cdot \overline{v}_n
  608. with a and v sequences being zero-padded where necessary and
  609. :math:`\overline v` denoting complex conjugation.
  610. Parameters
  611. ----------
  612. a, v : array_like
  613. Input sequences.
  614. mode : {'valid', 'same', 'full'}, optional
  615. Refer to the `convolve` docstring. Note that the default
  616. is 'valid', unlike `convolve`, which uses 'full'.
  617. Returns
  618. -------
  619. out : ndarray
  620. Discrete cross-correlation of `a` and `v`.
  621. See Also
  622. --------
  623. convolve : Discrete, linear convolution of two one-dimensional sequences.
  624. scipy.signal.correlate : uses FFT which has superior performance
  625. on large arrays.
  626. Notes
  627. -----
  628. The definition of correlation above is not unique and sometimes
  629. correlation may be defined differently. Another common definition is [1]_:
  630. .. math:: c'_k = \sum_n a_{n} \cdot \overline{v_{n+k}}
  631. which is related to :math:`c_k` by :math:`c'_k = c_{-k}`.
  632. `numpy.correlate` may perform slowly in large arrays (i.e. n = 1e5)
  633. because it does not use the FFT to compute the convolution; in that case,
  634. `scipy.signal.correlate` might be preferable.
  635. References
  636. ----------
  637. .. [1] Wikipedia, "Cross-correlation",
  638. https://en.wikipedia.org/wiki/Cross-correlation
  639. Examples
  640. --------
  641. >>> import numpy as np
  642. >>> np.correlate([1, 2, 3], [0, 1, 0.5])
  643. array([3.5])
  644. >>> np.correlate([1, 2, 3], [0, 1, 0.5], "same")
  645. array([2. , 3.5, 3. ])
  646. >>> np.correlate([1, 2, 3], [0, 1, 0.5], "full")
  647. array([0.5, 2. , 3.5, 3. , 0. ])
  648. Using complex sequences:
  649. >>> np.correlate([1+1j, 2, 3-1j], [0, 1, 0.5j], 'full')
  650. array([ 0.5-0.5j, 1.0+0.j , 1.5-1.5j, 3.0-1.j , 0.0+0.j ])
  651. Note that you get the time reversed, complex conjugated result
  652. (:math:`\overline{c_{-k}}`) when the two input sequences a and v change
  653. places:
  654. >>> np.correlate([0, 1, 0.5j], [1+1j, 2, 3-1j], 'full')
  655. array([ 0.0+0.j , 3.0+1.j , 1.5+1.5j, 1.0+0.j , 0.5+0.5j])
  656. """
  657. return multiarray.correlate2(a, v, mode)
  658. def _convolve_dispatcher(a, v, mode=None):
  659. return (a, v)
  660. @array_function_dispatch(_convolve_dispatcher)
  661. def convolve(a, v, mode='full'):
  662. """
  663. Returns the discrete, linear convolution of two one-dimensional sequences.
  664. The convolution operator is often seen in signal processing, where it
  665. models the effect of a linear time-invariant system on a signal [1]_. In
  666. probability theory, the sum of two independent random variables is
  667. distributed according to the convolution of their individual
  668. distributions.
  669. If `v` is longer than `a`, the arrays are swapped before computation.
  670. Parameters
  671. ----------
  672. a : (N,) array_like
  673. First one-dimensional input array.
  674. v : (M,) array_like
  675. Second one-dimensional input array.
  676. mode : {'full', 'valid', 'same'}, optional
  677. 'full':
  678. By default, mode is 'full'. This returns the convolution
  679. at each point of overlap, with an output shape of (N+M-1,). At
  680. the end-points of the convolution, the signals do not overlap
  681. completely, and boundary effects may be seen.
  682. 'same':
  683. Mode 'same' returns output of length ``max(M, N)``. Boundary
  684. effects are still visible.
  685. 'valid':
  686. Mode 'valid' returns output of length
  687. ``max(M, N) - min(M, N) + 1``. The convolution product is only given
  688. for points where the signals overlap completely. Values outside
  689. the signal boundary have no effect.
  690. Returns
  691. -------
  692. out : ndarray
  693. Discrete, linear convolution of `a` and `v`.
  694. See Also
  695. --------
  696. scipy.signal.fftconvolve : Convolve two arrays using the Fast Fourier
  697. Transform.
  698. scipy.linalg.toeplitz : Used to construct the convolution operator.
  699. polymul : Polynomial multiplication. Same output as convolve, but also
  700. accepts poly1d objects as input.
  701. Notes
  702. -----
  703. The discrete convolution operation is defined as
  704. .. math:: (a * v)_n = \\sum_{m = -\\infty}^{\\infty} a_m v_{n - m}
  705. It can be shown that a convolution :math:`x(t) * y(t)` in time/space
  706. is equivalent to the multiplication :math:`X(f) Y(f)` in the Fourier
  707. domain, after appropriate padding (padding is necessary to prevent
  708. circular convolution). Since multiplication is more efficient (faster)
  709. than convolution, the function `scipy.signal.fftconvolve` exploits the
  710. FFT to calculate the convolution of large data-sets.
  711. References
  712. ----------
  713. .. [1] Wikipedia, "Convolution",
  714. https://en.wikipedia.org/wiki/Convolution
  715. Examples
  716. --------
  717. Note how the convolution operator flips the second array
  718. before "sliding" the two across one another:
  719. >>> import numpy as np
  720. >>> np.convolve([1, 2, 3], [0, 1, 0.5])
  721. array([0. , 1. , 2.5, 4. , 1.5])
  722. Only return the middle values of the convolution.
  723. Contains boundary effects, where zeros are taken
  724. into account:
  725. >>> np.convolve([1,2,3],[0,1,0.5], 'same')
  726. array([1. , 2.5, 4. ])
  727. The two arrays are of the same length, so there
  728. is only one position where they completely overlap:
  729. >>> np.convolve([1,2,3],[0,1,0.5], 'valid')
  730. array([2.5])
  731. """
  732. a, v = array(a, copy=None, ndmin=1), array(v, copy=None, ndmin=1)
  733. if len(a) == 0:
  734. raise ValueError('a cannot be empty')
  735. if len(v) == 0:
  736. raise ValueError('v cannot be empty')
  737. if len(v) > len(a):
  738. a, v = v, a
  739. return multiarray.correlate(a, v[::-1], mode)
  740. def _outer_dispatcher(a, b, out=None):
  741. return (a, b, out)
  742. @array_function_dispatch(_outer_dispatcher)
  743. def outer(a, b, out=None):
  744. """
  745. Compute the outer product of two vectors.
  746. Given two vectors `a` and `b` of length ``M`` and ``N``, respectively,
  747. the outer product [1]_ is::
  748. [[a_0*b_0 a_0*b_1 ... a_0*b_{N-1} ]
  749. [a_1*b_0 .
  750. [ ... .
  751. [a_{M-1}*b_0 a_{M-1}*b_{N-1} ]]
  752. Parameters
  753. ----------
  754. a : (M,) array_like
  755. First input vector. Input is flattened if
  756. not already 1-dimensional.
  757. b : (N,) array_like
  758. Second input vector. Input is flattened if
  759. not already 1-dimensional.
  760. out : (M, N) ndarray, optional
  761. A location where the result is stored
  762. Returns
  763. -------
  764. out : (M, N) ndarray
  765. ``out[i, j] = a[i] * b[j]``
  766. See also
  767. --------
  768. inner
  769. einsum : ``einsum('i,j->ij', a.ravel(), b.ravel())`` is the equivalent.
  770. ufunc.outer : A generalization to dimensions other than 1D and other
  771. operations. ``np.multiply.outer(a.ravel(), b.ravel())``
  772. is the equivalent.
  773. linalg.outer : An Array API compatible variation of ``np.outer``,
  774. which accepts 1-dimensional inputs only.
  775. tensordot : ``np.tensordot(a.ravel(), b.ravel(), axes=((), ()))``
  776. is the equivalent.
  777. References
  778. ----------
  779. .. [1] G. H. Golub and C. F. Van Loan, *Matrix Computations*, 3rd
  780. ed., Baltimore, MD, Johns Hopkins University Press, 1996,
  781. pg. 8.
  782. Examples
  783. --------
  784. Make a (*very* coarse) grid for computing a Mandelbrot set:
  785. >>> import numpy as np
  786. >>> rl = np.outer(np.ones((5,)), np.linspace(-2, 2, 5))
  787. >>> rl
  788. array([[-2., -1., 0., 1., 2.],
  789. [-2., -1., 0., 1., 2.],
  790. [-2., -1., 0., 1., 2.],
  791. [-2., -1., 0., 1., 2.],
  792. [-2., -1., 0., 1., 2.]])
  793. >>> im = np.outer(1j*np.linspace(2, -2, 5), np.ones((5,)))
  794. >>> im
  795. array([[0.+2.j, 0.+2.j, 0.+2.j, 0.+2.j, 0.+2.j],
  796. [0.+1.j, 0.+1.j, 0.+1.j, 0.+1.j, 0.+1.j],
  797. [0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j],
  798. [0.-1.j, 0.-1.j, 0.-1.j, 0.-1.j, 0.-1.j],
  799. [0.-2.j, 0.-2.j, 0.-2.j, 0.-2.j, 0.-2.j]])
  800. >>> grid = rl + im
  801. >>> grid
  802. array([[-2.+2.j, -1.+2.j, 0.+2.j, 1.+2.j, 2.+2.j],
  803. [-2.+1.j, -1.+1.j, 0.+1.j, 1.+1.j, 2.+1.j],
  804. [-2.+0.j, -1.+0.j, 0.+0.j, 1.+0.j, 2.+0.j],
  805. [-2.-1.j, -1.-1.j, 0.-1.j, 1.-1.j, 2.-1.j],
  806. [-2.-2.j, -1.-2.j, 0.-2.j, 1.-2.j, 2.-2.j]])
  807. An example using a "vector" of letters:
  808. >>> x = np.array(['a', 'b', 'c'], dtype=object)
  809. >>> np.outer(x, [1, 2, 3])
  810. array([['a', 'aa', 'aaa'],
  811. ['b', 'bb', 'bbb'],
  812. ['c', 'cc', 'ccc']], dtype=object)
  813. """
  814. a = asarray(a)
  815. b = asarray(b)
  816. return multiply(a.ravel()[:, newaxis], b.ravel()[newaxis, :], out)
  817. def _tensordot_dispatcher(a, b, axes=None):
  818. return (a, b)
  819. @array_function_dispatch(_tensordot_dispatcher)
  820. def tensordot(a, b, axes=2):
  821. """
  822. Compute tensor dot product along specified axes.
  823. Given two tensors, `a` and `b`, and an array_like object containing
  824. two array_like objects, ``(a_axes, b_axes)``, sum the products of
  825. `a`'s and `b`'s elements (components) over the axes specified by
  826. ``a_axes`` and ``b_axes``. The third argument can be a single non-negative
  827. integer_like scalar, ``N``; if it is such, then the last ``N`` dimensions
  828. of `a` and the first ``N`` dimensions of `b` are summed over.
  829. Parameters
  830. ----------
  831. a, b : array_like
  832. Tensors to "dot".
  833. axes : int or (2,) array_like
  834. * integer_like
  835. If an int N, sum over the last N axes of `a` and the first N axes
  836. of `b` in order. The sizes of the corresponding axes must match.
  837. * (2,) array_like
  838. Or, a list of axes to be summed over, first sequence applying to `a`,
  839. second to `b`. Both elements array_like must be of the same length.
  840. Returns
  841. -------
  842. output : ndarray
  843. The tensor dot product of the input.
  844. See Also
  845. --------
  846. dot, einsum
  847. Notes
  848. -----
  849. Three common use cases are:
  850. * ``axes = 0`` : tensor product :math:`a\\otimes b`
  851. * ``axes = 1`` : tensor dot product :math:`a\\cdot b`
  852. * ``axes = 2`` : (default) tensor double contraction :math:`a:b`
  853. When `axes` is integer_like, the sequence of axes for evaluation
  854. will be: from the -Nth axis to the -1th axis in `a`,
  855. and from the 0th axis to (N-1)th axis in `b`.
  856. For example, ``axes = 2`` is the equal to
  857. ``axes = [[-2, -1], [0, 1]]``.
  858. When N-1 is smaller than 0, or when -N is larger than -1,
  859. the element of `a` and `b` are defined as the `axes`.
  860. When there is more than one axis to sum over - and they are not the last
  861. (first) axes of `a` (`b`) - the argument `axes` should consist of
  862. two sequences of the same length, with the first axis to sum over given
  863. first in both sequences, the second axis second, and so forth.
  864. The calculation can be referred to ``numpy.einsum``.
  865. The shape of the result consists of the non-contracted axes of the
  866. first tensor, followed by the non-contracted axes of the second.
  867. Examples
  868. --------
  869. An example on integer_like:
  870. >>> a_0 = np.array([[1, 2], [3, 4]])
  871. >>> b_0 = np.array([[5, 6], [7, 8]])
  872. >>> c_0 = np.tensordot(a_0, b_0, axes=0)
  873. >>> c_0.shape
  874. (2, 2, 2, 2)
  875. >>> c_0
  876. array([[[[ 5, 6],
  877. [ 7, 8]],
  878. [[10, 12],
  879. [14, 16]]],
  880. [[[15, 18],
  881. [21, 24]],
  882. [[20, 24],
  883. [28, 32]]]])
  884. An example on array_like:
  885. >>> a = np.arange(60.).reshape(3,4,5)
  886. >>> b = np.arange(24.).reshape(4,3,2)
  887. >>> c = np.tensordot(a,b, axes=([1,0],[0,1]))
  888. >>> c.shape
  889. (5, 2)
  890. >>> c
  891. array([[4400., 4730.],
  892. [4532., 4874.],
  893. [4664., 5018.],
  894. [4796., 5162.],
  895. [4928., 5306.]])
  896. A slower but equivalent way of computing the same...
  897. >>> d = np.zeros((5,2))
  898. >>> for i in range(5):
  899. ... for j in range(2):
  900. ... for k in range(3):
  901. ... for n in range(4):
  902. ... d[i,j] += a[k,n,i] * b[n,k,j]
  903. >>> c == d
  904. array([[ True, True],
  905. [ True, True],
  906. [ True, True],
  907. [ True, True],
  908. [ True, True]])
  909. An extended example taking advantage of the overloading of + and \\*:
  910. >>> a = np.array(range(1, 9)).reshape((2, 2, 2))
  911. >>> A = np.array(('a', 'b', 'c', 'd'), dtype=object)
  912. >>> A = A.reshape((2, 2))
  913. >>> a; A
  914. array([[[1, 2],
  915. [3, 4]],
  916. [[5, 6],
  917. [7, 8]]])
  918. array([['a', 'b'],
  919. ['c', 'd']], dtype=object)
  920. >>> np.tensordot(a, A) # third argument default is 2 for double-contraction
  921. array(['abbcccdddd', 'aaaaabbbbbbcccccccdddddddd'], dtype=object)
  922. >>> np.tensordot(a, A, 1)
  923. array([[['acc', 'bdd'],
  924. ['aaacccc', 'bbbdddd']],
  925. [['aaaaacccccc', 'bbbbbdddddd'],
  926. ['aaaaaaacccccccc', 'bbbbbbbdddddddd']]], dtype=object)
  927. >>> np.tensordot(a, A, 0) # tensor product (result too long to incl.)
  928. array([[[[['a', 'b'],
  929. ['c', 'd']],
  930. ...
  931. >>> np.tensordot(a, A, (0, 1))
  932. array([[['abbbbb', 'cddddd'],
  933. ['aabbbbbb', 'ccdddddd']],
  934. [['aaabbbbbbb', 'cccddddddd'],
  935. ['aaaabbbbbbbb', 'ccccdddddddd']]], dtype=object)
  936. >>> np.tensordot(a, A, (2, 1))
  937. array([[['abb', 'cdd'],
  938. ['aaabbbb', 'cccdddd']],
  939. [['aaaaabbbbbb', 'cccccdddddd'],
  940. ['aaaaaaabbbbbbbb', 'cccccccdddddddd']]], dtype=object)
  941. >>> np.tensordot(a, A, ((0, 1), (0, 1)))
  942. array(['abbbcccccddddddd', 'aabbbbccccccdddddddd'], dtype=object)
  943. >>> np.tensordot(a, A, ((2, 1), (1, 0)))
  944. array(['acccbbdddd', 'aaaaacccccccbbbbbbdddddddd'], dtype=object)
  945. """
  946. try:
  947. iter(axes)
  948. except Exception:
  949. axes_a = list(range(-axes, 0))
  950. axes_b = list(range(axes))
  951. else:
  952. axes_a, axes_b = axes
  953. try:
  954. na = len(axes_a)
  955. axes_a = list(axes_a)
  956. except TypeError:
  957. axes_a = [axes_a]
  958. na = 1
  959. try:
  960. nb = len(axes_b)
  961. axes_b = list(axes_b)
  962. except TypeError:
  963. axes_b = [axes_b]
  964. nb = 1
  965. a, b = asarray(a), asarray(b)
  966. as_ = a.shape
  967. nda = a.ndim
  968. bs = b.shape
  969. ndb = b.ndim
  970. equal = True
  971. if na != nb:
  972. equal = False
  973. else:
  974. for k in range(na):
  975. if as_[axes_a[k]] != bs[axes_b[k]]:
  976. equal = False
  977. break
  978. if axes_a[k] < 0:
  979. axes_a[k] += nda
  980. if axes_b[k] < 0:
  981. axes_b[k] += ndb
  982. if not equal:
  983. raise ValueError("shape-mismatch for sum")
  984. # Move the axes to sum over to the end of "a"
  985. # and to the front of "b"
  986. notin = [k for k in range(nda) if k not in axes_a]
  987. newaxes_a = notin + axes_a
  988. N2 = math.prod(as_[axis] for axis in axes_a)
  989. newshape_a = (math.prod(as_[ax] for ax in notin), N2)
  990. olda = [as_[axis] for axis in notin]
  991. notin = [k for k in range(ndb) if k not in axes_b]
  992. newaxes_b = axes_b + notin
  993. N2 = math.prod(bs[axis] for axis in axes_b)
  994. newshape_b = (N2, math.prod(bs[ax] for ax in notin))
  995. oldb = [bs[axis] for axis in notin]
  996. at = a.transpose(newaxes_a).reshape(newshape_a)
  997. bt = b.transpose(newaxes_b).reshape(newshape_b)
  998. res = dot(at, bt)
  999. return res.reshape(olda + oldb)
  1000. def _roll_dispatcher(a, shift, axis=None):
  1001. return (a,)
  1002. @array_function_dispatch(_roll_dispatcher)
  1003. def roll(a, shift, axis=None):
  1004. """
  1005. Roll array elements along a given axis.
  1006. Elements that roll beyond the last position are re-introduced at
  1007. the first.
  1008. Parameters
  1009. ----------
  1010. a : array_like
  1011. Input array.
  1012. shift : int or tuple of ints
  1013. The number of places by which elements are shifted. If a tuple,
  1014. then `axis` must be a tuple of the same size, and each of the
  1015. given axes is shifted by the corresponding number. If an int
  1016. while `axis` is a tuple of ints, then the same value is used for
  1017. all given axes.
  1018. axis : int or tuple of ints, optional
  1019. Axis or axes along which elements are shifted. By default, the
  1020. array is flattened before shifting, after which the original
  1021. shape is restored.
  1022. Returns
  1023. -------
  1024. res : ndarray
  1025. Output array, with the same shape as `a`.
  1026. See Also
  1027. --------
  1028. rollaxis : Roll the specified axis backwards, until it lies in a
  1029. given position.
  1030. Notes
  1031. -----
  1032. Supports rolling over multiple dimensions simultaneously.
  1033. Examples
  1034. --------
  1035. >>> import numpy as np
  1036. >>> x = np.arange(10)
  1037. >>> np.roll(x, 2)
  1038. array([8, 9, 0, 1, 2, 3, 4, 5, 6, 7])
  1039. >>> np.roll(x, -2)
  1040. array([2, 3, 4, 5, 6, 7, 8, 9, 0, 1])
  1041. >>> x2 = np.reshape(x, (2, 5))
  1042. >>> x2
  1043. array([[0, 1, 2, 3, 4],
  1044. [5, 6, 7, 8, 9]])
  1045. >>> np.roll(x2, 1)
  1046. array([[9, 0, 1, 2, 3],
  1047. [4, 5, 6, 7, 8]])
  1048. >>> np.roll(x2, -1)
  1049. array([[1, 2, 3, 4, 5],
  1050. [6, 7, 8, 9, 0]])
  1051. >>> np.roll(x2, 1, axis=0)
  1052. array([[5, 6, 7, 8, 9],
  1053. [0, 1, 2, 3, 4]])
  1054. >>> np.roll(x2, -1, axis=0)
  1055. array([[5, 6, 7, 8, 9],
  1056. [0, 1, 2, 3, 4]])
  1057. >>> np.roll(x2, 1, axis=1)
  1058. array([[4, 0, 1, 2, 3],
  1059. [9, 5, 6, 7, 8]])
  1060. >>> np.roll(x2, -1, axis=1)
  1061. array([[1, 2, 3, 4, 0],
  1062. [6, 7, 8, 9, 5]])
  1063. >>> np.roll(x2, (1, 1), axis=(1, 0))
  1064. array([[9, 5, 6, 7, 8],
  1065. [4, 0, 1, 2, 3]])
  1066. >>> np.roll(x2, (2, 1), axis=(1, 0))
  1067. array([[8, 9, 5, 6, 7],
  1068. [3, 4, 0, 1, 2]])
  1069. """
  1070. a = asanyarray(a)
  1071. if axis is None:
  1072. return roll(a.ravel(), shift, 0).reshape(a.shape)
  1073. else:
  1074. axis = normalize_axis_tuple(axis, a.ndim, allow_duplicate=True)
  1075. broadcasted = broadcast(shift, axis)
  1076. if broadcasted.ndim > 1:
  1077. raise ValueError(
  1078. "'shift' and 'axis' should be scalars or 1D sequences")
  1079. shifts = dict.fromkeys(range(a.ndim), 0)
  1080. for sh, ax in broadcasted:
  1081. shifts[ax] += int(sh)
  1082. rolls = [((slice(None), slice(None)),)] * a.ndim
  1083. for ax, offset in shifts.items():
  1084. offset %= a.shape[ax] or 1 # If `a` is empty, nothing matters.
  1085. if offset:
  1086. # (original, result), (original, result)
  1087. rolls[ax] = ((slice(None, -offset), slice(offset, None)),
  1088. (slice(-offset, None), slice(None, offset)))
  1089. result = empty_like(a)
  1090. for indices in itertools.product(*rolls):
  1091. arr_index, res_index = zip(*indices)
  1092. result[res_index] = a[arr_index]
  1093. return result
  1094. def _rollaxis_dispatcher(a, axis, start=None):
  1095. return (a,)
  1096. @array_function_dispatch(_rollaxis_dispatcher)
  1097. def rollaxis(a, axis, start=0):
  1098. """
  1099. Roll the specified axis backwards, until it lies in a given position.
  1100. This function continues to be supported for backward compatibility, but you
  1101. should prefer `moveaxis`. The `moveaxis` function was added in NumPy
  1102. 1.11.
  1103. Parameters
  1104. ----------
  1105. a : ndarray
  1106. Input array.
  1107. axis : int
  1108. The axis to be rolled. The positions of the other axes do not
  1109. change relative to one another.
  1110. start : int, optional
  1111. When ``start <= axis``, the axis is rolled back until it lies in
  1112. this position. When ``start > axis``, the axis is rolled until it
  1113. lies before this position. The default, 0, results in a "complete"
  1114. roll. The following table describes how negative values of ``start``
  1115. are interpreted:
  1116. .. table::
  1117. :align: left
  1118. +-------------------+----------------------+
  1119. | ``start`` | Normalized ``start`` |
  1120. +===================+======================+
  1121. | ``-(arr.ndim+1)`` | raise ``AxisError`` |
  1122. +-------------------+----------------------+
  1123. | ``-arr.ndim`` | 0 |
  1124. +-------------------+----------------------+
  1125. | |vdots| | |vdots| |
  1126. +-------------------+----------------------+
  1127. | ``-1`` | ``arr.ndim-1`` |
  1128. +-------------------+----------------------+
  1129. | ``0`` | ``0`` |
  1130. +-------------------+----------------------+
  1131. | |vdots| | |vdots| |
  1132. +-------------------+----------------------+
  1133. | ``arr.ndim`` | ``arr.ndim`` |
  1134. +-------------------+----------------------+
  1135. | ``arr.ndim + 1`` | raise ``AxisError`` |
  1136. +-------------------+----------------------+
  1137. .. |vdots| unicode:: U+22EE .. Vertical Ellipsis
  1138. Returns
  1139. -------
  1140. res : ndarray
  1141. For NumPy >= 1.10.0 a view of `a` is always returned. For earlier
  1142. NumPy versions a view of `a` is returned only if the order of the
  1143. axes is changed, otherwise the input array is returned.
  1144. See Also
  1145. --------
  1146. moveaxis : Move array axes to new positions.
  1147. roll : Roll the elements of an array by a number of positions along a
  1148. given axis.
  1149. Examples
  1150. --------
  1151. >>> import numpy as np
  1152. >>> a = np.ones((3,4,5,6))
  1153. >>> np.rollaxis(a, 3, 1).shape
  1154. (3, 6, 4, 5)
  1155. >>> np.rollaxis(a, 2).shape
  1156. (5, 3, 4, 6)
  1157. >>> np.rollaxis(a, 1, 4).shape
  1158. (3, 5, 6, 4)
  1159. """
  1160. n = a.ndim
  1161. axis = normalize_axis_index(axis, n)
  1162. if start < 0:
  1163. start += n
  1164. msg = "'%s' arg requires %d <= %s < %d, but %d was passed in"
  1165. if not (0 <= start < n + 1):
  1166. raise AxisError(msg % ('start', -n, 'start', n + 1, start))
  1167. if axis < start:
  1168. # it's been removed
  1169. start -= 1
  1170. if axis == start:
  1171. return a[...]
  1172. axes = list(range(n))
  1173. axes.remove(axis)
  1174. axes.insert(start, axis)
  1175. return a.transpose(axes)
  1176. @set_module("numpy.lib.array_utils")
  1177. def normalize_axis_tuple(axis, ndim, argname=None, allow_duplicate=False):
  1178. """
  1179. Normalizes an axis argument into a tuple of non-negative integer axes.
  1180. This handles shorthands such as ``1`` and converts them to ``(1,)``,
  1181. as well as performing the handling of negative indices covered by
  1182. `normalize_axis_index`.
  1183. By default, this forbids axes from being specified multiple times.
  1184. Used internally by multi-axis-checking logic.
  1185. Parameters
  1186. ----------
  1187. axis : int, iterable of int
  1188. The un-normalized index or indices of the axis.
  1189. ndim : int
  1190. The number of dimensions of the array that `axis` should be normalized
  1191. against.
  1192. argname : str, optional
  1193. A prefix to put before the error message, typically the name of the
  1194. argument.
  1195. allow_duplicate : bool, optional
  1196. If False, the default, disallow an axis from being specified twice.
  1197. Returns
  1198. -------
  1199. normalized_axes : tuple of int
  1200. The normalized axis index, such that `0 <= normalized_axis < ndim`
  1201. Raises
  1202. ------
  1203. AxisError
  1204. If any axis provided is out of range
  1205. ValueError
  1206. If an axis is repeated
  1207. See also
  1208. --------
  1209. normalize_axis_index : normalizing a single scalar axis
  1210. """
  1211. # Optimization to speed-up the most common cases.
  1212. if not isinstance(axis, (tuple, list)):
  1213. try:
  1214. axis = [operator.index(axis)]
  1215. except TypeError:
  1216. pass
  1217. # Going via an iterator directly is slower than via list comprehension.
  1218. axis = tuple(normalize_axis_index(ax, ndim, argname) for ax in axis)
  1219. if not allow_duplicate and len(set(axis)) != len(axis):
  1220. if argname:
  1221. raise ValueError(f'repeated axis in `{argname}` argument')
  1222. else:
  1223. raise ValueError('repeated axis')
  1224. return axis
  1225. def _moveaxis_dispatcher(a, source, destination):
  1226. return (a,)
  1227. @array_function_dispatch(_moveaxis_dispatcher)
  1228. def moveaxis(a, source, destination):
  1229. """
  1230. Move axes of an array to new positions.
  1231. Other axes remain in their original order.
  1232. Parameters
  1233. ----------
  1234. a : np.ndarray
  1235. The array whose axes should be reordered.
  1236. source : int or sequence of int
  1237. Original positions of the axes to move. These must be unique.
  1238. destination : int or sequence of int
  1239. Destination positions for each of the original axes. These must also be
  1240. unique.
  1241. Returns
  1242. -------
  1243. result : np.ndarray
  1244. Array with moved axes. This array is a view of the input array.
  1245. See Also
  1246. --------
  1247. transpose : Permute the dimensions of an array.
  1248. swapaxes : Interchange two axes of an array.
  1249. Examples
  1250. --------
  1251. >>> import numpy as np
  1252. >>> x = np.zeros((3, 4, 5))
  1253. >>> np.moveaxis(x, 0, -1).shape
  1254. (4, 5, 3)
  1255. >>> np.moveaxis(x, -1, 0).shape
  1256. (5, 3, 4)
  1257. These all achieve the same result:
  1258. >>> np.transpose(x).shape
  1259. (5, 4, 3)
  1260. >>> np.swapaxes(x, 0, -1).shape
  1261. (5, 4, 3)
  1262. >>> np.moveaxis(x, [0, 1], [-1, -2]).shape
  1263. (5, 4, 3)
  1264. >>> np.moveaxis(x, [0, 1, 2], [-1, -2, -3]).shape
  1265. (5, 4, 3)
  1266. """
  1267. try:
  1268. # allow duck-array types if they define transpose
  1269. transpose = a.transpose
  1270. except AttributeError:
  1271. a = asarray(a)
  1272. transpose = a.transpose
  1273. source = normalize_axis_tuple(source, a.ndim, 'source')
  1274. destination = normalize_axis_tuple(destination, a.ndim, 'destination')
  1275. if len(source) != len(destination):
  1276. raise ValueError('`source` and `destination` arguments must have '
  1277. 'the same number of elements')
  1278. order = [n for n in range(a.ndim) if n not in source]
  1279. for dest, src in sorted(zip(destination, source)):
  1280. order.insert(dest, src)
  1281. result = transpose(order)
  1282. return result
  1283. def _cross_dispatcher(a, b, axisa=None, axisb=None, axisc=None, axis=None):
  1284. return (a, b)
  1285. @array_function_dispatch(_cross_dispatcher)
  1286. def cross(a, b, axisa=-1, axisb=-1, axisc=-1, axis=None):
  1287. """
  1288. Return the cross product of two (arrays of) vectors.
  1289. The cross product of `a` and `b` in :math:`R^3` is a vector perpendicular
  1290. to both `a` and `b`. If `a` and `b` are arrays of vectors, the vectors
  1291. are defined by the last axis of `a` and `b` by default, and these axes
  1292. can have dimensions 2 or 3. Where the dimension of either `a` or `b` is
  1293. 2, the third component of the input vector is assumed to be zero and the
  1294. cross product calculated accordingly. In cases where both input vectors
  1295. have dimension 2, the z-component of the cross product is returned.
  1296. Parameters
  1297. ----------
  1298. a : array_like
  1299. Components of the first vector(s).
  1300. b : array_like
  1301. Components of the second vector(s).
  1302. axisa : int, optional
  1303. Axis of `a` that defines the vector(s). By default, the last axis.
  1304. axisb : int, optional
  1305. Axis of `b` that defines the vector(s). By default, the last axis.
  1306. axisc : int, optional
  1307. Axis of `c` containing the cross product vector(s). Ignored if
  1308. both input vectors have dimension 2, as the return is scalar.
  1309. By default, the last axis.
  1310. axis : int, optional
  1311. If defined, the axis of `a`, `b` and `c` that defines the vector(s)
  1312. and cross product(s). Overrides `axisa`, `axisb` and `axisc`.
  1313. Returns
  1314. -------
  1315. c : ndarray
  1316. Vector cross product(s).
  1317. Raises
  1318. ------
  1319. ValueError
  1320. When the dimension of the vector(s) in `a` and/or `b` does not
  1321. equal 2 or 3.
  1322. See Also
  1323. --------
  1324. inner : Inner product
  1325. outer : Outer product.
  1326. linalg.cross : An Array API compatible variation of ``np.cross``,
  1327. which accepts (arrays of) 3-element vectors only.
  1328. ix_ : Construct index arrays.
  1329. Notes
  1330. -----
  1331. Supports full broadcasting of the inputs.
  1332. Dimension-2 input arrays were deprecated in 2.0.0. If you do need this
  1333. functionality, you can use::
  1334. def cross2d(x, y):
  1335. return x[..., 0] * y[..., 1] - x[..., 1] * y[..., 0]
  1336. Examples
  1337. --------
  1338. Vector cross-product.
  1339. >>> import numpy as np
  1340. >>> x = [1, 2, 3]
  1341. >>> y = [4, 5, 6]
  1342. >>> np.cross(x, y)
  1343. array([-3, 6, -3])
  1344. One vector with dimension 2.
  1345. >>> x = [1, 2]
  1346. >>> y = [4, 5, 6]
  1347. >>> np.cross(x, y)
  1348. array([12, -6, -3])
  1349. Equivalently:
  1350. >>> x = [1, 2, 0]
  1351. >>> y = [4, 5, 6]
  1352. >>> np.cross(x, y)
  1353. array([12, -6, -3])
  1354. Both vectors with dimension 2.
  1355. >>> x = [1,2]
  1356. >>> y = [4,5]
  1357. >>> np.cross(x, y)
  1358. array(-3)
  1359. Multiple vector cross-products. Note that the direction of the cross
  1360. product vector is defined by the *right-hand rule*.
  1361. >>> x = np.array([[1,2,3], [4,5,6]])
  1362. >>> y = np.array([[4,5,6], [1,2,3]])
  1363. >>> np.cross(x, y)
  1364. array([[-3, 6, -3],
  1365. [ 3, -6, 3]])
  1366. The orientation of `c` can be changed using the `axisc` keyword.
  1367. >>> np.cross(x, y, axisc=0)
  1368. array([[-3, 3],
  1369. [ 6, -6],
  1370. [-3, 3]])
  1371. Change the vector definition of `x` and `y` using `axisa` and `axisb`.
  1372. >>> x = np.array([[1,2,3], [4,5,6], [7, 8, 9]])
  1373. >>> y = np.array([[7, 8, 9], [4,5,6], [1,2,3]])
  1374. >>> np.cross(x, y)
  1375. array([[ -6, 12, -6],
  1376. [ 0, 0, 0],
  1377. [ 6, -12, 6]])
  1378. >>> np.cross(x, y, axisa=0, axisb=0)
  1379. array([[-24, 48, -24],
  1380. [-30, 60, -30],
  1381. [-36, 72, -36]])
  1382. """
  1383. if axis is not None:
  1384. axisa, axisb, axisc = (axis,) * 3
  1385. a = asarray(a)
  1386. b = asarray(b)
  1387. if (a.ndim < 1) or (b.ndim < 1):
  1388. raise ValueError("At least one array has zero dimension")
  1389. # Check axisa and axisb are within bounds
  1390. axisa = normalize_axis_index(axisa, a.ndim, msg_prefix='axisa')
  1391. axisb = normalize_axis_index(axisb, b.ndim, msg_prefix='axisb')
  1392. # Move working axis to the end of the shape
  1393. a = moveaxis(a, axisa, -1)
  1394. b = moveaxis(b, axisb, -1)
  1395. msg = ("incompatible dimensions for cross product\n"
  1396. "(dimension must be 2 or 3)")
  1397. if a.shape[-1] not in (2, 3) or b.shape[-1] not in (2, 3):
  1398. raise ValueError(msg)
  1399. if a.shape[-1] == 2 or b.shape[-1] == 2:
  1400. # Deprecated in NumPy 2.0, 2023-09-26
  1401. warnings.warn(
  1402. "Arrays of 2-dimensional vectors are deprecated. Use arrays of "
  1403. "3-dimensional vectors instead. (deprecated in NumPy 2.0)",
  1404. DeprecationWarning, stacklevel=2
  1405. )
  1406. # Create the output array
  1407. shape = broadcast(a[..., 0], b[..., 0]).shape
  1408. if a.shape[-1] == 3 or b.shape[-1] == 3:
  1409. shape += (3,)
  1410. # Check axisc is within bounds
  1411. axisc = normalize_axis_index(axisc, len(shape), msg_prefix='axisc')
  1412. dtype = promote_types(a.dtype, b.dtype)
  1413. cp = empty(shape, dtype)
  1414. # recast arrays as dtype
  1415. a = a.astype(dtype)
  1416. b = b.astype(dtype)
  1417. # create local aliases for readability
  1418. a0 = a[..., 0]
  1419. a1 = a[..., 1]
  1420. if a.shape[-1] == 3:
  1421. a2 = a[..., 2]
  1422. b0 = b[..., 0]
  1423. b1 = b[..., 1]
  1424. if b.shape[-1] == 3:
  1425. b2 = b[..., 2]
  1426. if cp.ndim != 0 and cp.shape[-1] == 3:
  1427. cp0 = cp[..., 0]
  1428. cp1 = cp[..., 1]
  1429. cp2 = cp[..., 2]
  1430. if a.shape[-1] == 2:
  1431. if b.shape[-1] == 2:
  1432. # a0 * b1 - a1 * b0
  1433. multiply(a0, b1, out=cp)
  1434. cp -= a1 * b0
  1435. return cp
  1436. else:
  1437. assert b.shape[-1] == 3
  1438. # cp0 = a1 * b2 - 0 (a2 = 0)
  1439. # cp1 = 0 - a0 * b2 (a2 = 0)
  1440. # cp2 = a0 * b1 - a1 * b0
  1441. multiply(a1, b2, out=cp0)
  1442. multiply(a0, b2, out=cp1)
  1443. negative(cp1, out=cp1)
  1444. multiply(a0, b1, out=cp2)
  1445. cp2 -= a1 * b0
  1446. else:
  1447. assert a.shape[-1] == 3
  1448. if b.shape[-1] == 3:
  1449. # cp0 = a1 * b2 - a2 * b1
  1450. # cp1 = a2 * b0 - a0 * b2
  1451. # cp2 = a0 * b1 - a1 * b0
  1452. multiply(a1, b2, out=cp0)
  1453. tmp = np.multiply(a2, b1, out=...)
  1454. cp0 -= tmp
  1455. multiply(a2, b0, out=cp1)
  1456. multiply(a0, b2, out=tmp)
  1457. cp1 -= tmp
  1458. multiply(a0, b1, out=cp2)
  1459. multiply(a1, b0, out=tmp)
  1460. cp2 -= tmp
  1461. else:
  1462. assert b.shape[-1] == 2
  1463. # cp0 = 0 - a2 * b1 (b2 = 0)
  1464. # cp1 = a2 * b0 - 0 (b2 = 0)
  1465. # cp2 = a0 * b1 - a1 * b0
  1466. multiply(a2, b1, out=cp0)
  1467. negative(cp0, out=cp0)
  1468. multiply(a2, b0, out=cp1)
  1469. multiply(a0, b1, out=cp2)
  1470. cp2 -= a1 * b0
  1471. return moveaxis(cp, -1, axisc)
  1472. little_endian = (sys.byteorder == 'little')
  1473. @set_module('numpy')
  1474. def indices(dimensions, dtype=int, sparse=False):
  1475. """
  1476. Return an array representing the indices of a grid.
  1477. Compute an array where the subarrays contain index values 0, 1, ...
  1478. varying only along the corresponding axis.
  1479. Parameters
  1480. ----------
  1481. dimensions : sequence of ints
  1482. The shape of the grid.
  1483. dtype : dtype, optional
  1484. Data type of the result.
  1485. sparse : boolean, optional
  1486. Return a sparse representation of the grid instead of a dense
  1487. representation. Default is False.
  1488. Returns
  1489. -------
  1490. grid : one ndarray or tuple of ndarrays
  1491. If sparse is False:
  1492. Returns one array of grid indices,
  1493. ``grid.shape = (len(dimensions),) + tuple(dimensions)``.
  1494. If sparse is True:
  1495. Returns a tuple of arrays, with
  1496. ``grid[i].shape = (1, ..., 1, dimensions[i], 1, ..., 1)`` with
  1497. dimensions[i] in the ith place
  1498. See Also
  1499. --------
  1500. mgrid, ogrid, meshgrid
  1501. Notes
  1502. -----
  1503. The output shape in the dense case is obtained by prepending the number
  1504. of dimensions in front of the tuple of dimensions, i.e. if `dimensions`
  1505. is a tuple ``(r0, ..., rN-1)`` of length ``N``, the output shape is
  1506. ``(N, r0, ..., rN-1)``.
  1507. The subarrays ``grid[k]`` contains the N-D array of indices along the
  1508. ``k-th`` axis. Explicitly::
  1509. grid[k, i0, i1, ..., iN-1] = ik
  1510. Examples
  1511. --------
  1512. >>> import numpy as np
  1513. >>> grid = np.indices((2, 3))
  1514. >>> grid.shape
  1515. (2, 2, 3)
  1516. >>> grid[0] # row indices
  1517. array([[0, 0, 0],
  1518. [1, 1, 1]])
  1519. >>> grid[1] # column indices
  1520. array([[0, 1, 2],
  1521. [0, 1, 2]])
  1522. The indices can be used as an index into an array.
  1523. >>> x = np.arange(20).reshape(5, 4)
  1524. >>> row, col = np.indices((2, 3))
  1525. >>> x[row, col]
  1526. array([[0, 1, 2],
  1527. [4, 5, 6]])
  1528. Note that it would be more straightforward in the above example to
  1529. extract the required elements directly with ``x[:2, :3]``.
  1530. If sparse is set to true, the grid will be returned in a sparse
  1531. representation.
  1532. >>> i, j = np.indices((2, 3), sparse=True)
  1533. >>> i.shape
  1534. (2, 1)
  1535. >>> j.shape
  1536. (1, 3)
  1537. >>> i # row indices
  1538. array([[0],
  1539. [1]])
  1540. >>> j # column indices
  1541. array([[0, 1, 2]])
  1542. """
  1543. dimensions = tuple(dimensions)
  1544. N = len(dimensions)
  1545. shape = (1,) * N
  1546. if sparse:
  1547. res = ()
  1548. else:
  1549. res = empty((N,) + dimensions, dtype=dtype)
  1550. for i, dim in enumerate(dimensions):
  1551. idx = arange(dim, dtype=dtype).reshape(
  1552. shape[:i] + (dim,) + shape[i + 1:]
  1553. )
  1554. if sparse:
  1555. res = res + (idx,)
  1556. else:
  1557. res[i] = idx
  1558. return res
  1559. @finalize_array_function_like
  1560. @set_module('numpy')
  1561. def fromfunction(function, shape, *, dtype=float, like=None, **kwargs):
  1562. """
  1563. Construct an array by executing a function over each coordinate.
  1564. The resulting array therefore has a value ``fn(x, y, z)`` at
  1565. coordinate ``(x, y, z)``.
  1566. Parameters
  1567. ----------
  1568. function : callable
  1569. The function is called with N parameters, where N is the rank of
  1570. `shape`. Each parameter represents the coordinates of the array
  1571. varying along a specific axis. For example, if `shape`
  1572. were ``(2, 2)``, then the parameters would be
  1573. ``array([[0, 0], [1, 1]])`` and ``array([[0, 1], [0, 1]])``
  1574. shape : (N,) tuple of ints
  1575. Shape of the output array, which also determines the shape of
  1576. the coordinate arrays passed to `function`.
  1577. dtype : data-type, optional
  1578. Data-type of the coordinate arrays passed to `function`.
  1579. By default, `dtype` is float.
  1580. ${ARRAY_FUNCTION_LIKE}
  1581. .. versionadded:: 1.20.0
  1582. Returns
  1583. -------
  1584. fromfunction : any
  1585. The result of the call to `function` is passed back directly.
  1586. Therefore the shape of `fromfunction` is completely determined by
  1587. `function`. If `function` returns a scalar value, the shape of
  1588. `fromfunction` would not match the `shape` parameter.
  1589. See Also
  1590. --------
  1591. indices, meshgrid
  1592. Notes
  1593. -----
  1594. Keywords other than `dtype` and `like` are passed to `function`.
  1595. Examples
  1596. --------
  1597. >>> import numpy as np
  1598. >>> np.fromfunction(lambda i, j: i, (2, 2), dtype=float)
  1599. array([[0., 0.],
  1600. [1., 1.]])
  1601. >>> np.fromfunction(lambda i, j: j, (2, 2), dtype=float)
  1602. array([[0., 1.],
  1603. [0., 1.]])
  1604. >>> np.fromfunction(lambda i, j: i == j, (3, 3), dtype=int)
  1605. array([[ True, False, False],
  1606. [False, True, False],
  1607. [False, False, True]])
  1608. >>> np.fromfunction(lambda i, j: i + j, (3, 3), dtype=int)
  1609. array([[0, 1, 2],
  1610. [1, 2, 3],
  1611. [2, 3, 4]])
  1612. """
  1613. if like is not None:
  1614. return _fromfunction_with_like(
  1615. like, function, shape, dtype=dtype, **kwargs)
  1616. args = indices(shape, dtype=dtype)
  1617. return function(*args, **kwargs)
  1618. _fromfunction_with_like = array_function_dispatch()(fromfunction)
  1619. def _frombuffer(buf, dtype, shape, order, axis_order=None):
  1620. array = frombuffer(buf, dtype=dtype)
  1621. if order == 'K' and axis_order is not None:
  1622. return array.reshape(shape, order='C').transpose(axis_order)
  1623. return array.reshape(shape, order=order)
  1624. @set_module('numpy')
  1625. def isscalar(element):
  1626. """
  1627. Returns True if the type of `element` is a scalar type.
  1628. Parameters
  1629. ----------
  1630. element : any
  1631. Input argument, can be of any type and shape.
  1632. Returns
  1633. -------
  1634. val : bool
  1635. True if `element` is a scalar type, False if it is not.
  1636. See Also
  1637. --------
  1638. ndim : Get the number of dimensions of an array
  1639. Notes
  1640. -----
  1641. If you need a stricter way to identify a *numerical* scalar, use
  1642. ``isinstance(x, numbers.Number)``, as that returns ``False`` for most
  1643. non-numerical elements such as strings.
  1644. In most cases ``np.ndim(x) == 0`` should be used instead of this function,
  1645. as that will also return true for 0d arrays. This is how numpy overloads
  1646. functions in the style of the ``dx`` arguments to `gradient` and
  1647. the ``bins`` argument to `histogram`. Some key differences:
  1648. +------------------------------------+---------------+-------------------+
  1649. | x |``isscalar(x)``|``np.ndim(x) == 0``|
  1650. +====================================+===============+===================+
  1651. | PEP 3141 numeric objects | ``True`` | ``True`` |
  1652. | (including builtins) | | |
  1653. +------------------------------------+---------------+-------------------+
  1654. | builtin string and buffer objects | ``True`` | ``True`` |
  1655. +------------------------------------+---------------+-------------------+
  1656. | other builtin objects, like | ``False`` | ``True`` |
  1657. | `pathlib.Path`, `Exception`, | | |
  1658. | the result of `re.compile` | | |
  1659. +------------------------------------+---------------+-------------------+
  1660. | third-party objects like | ``False`` | ``True`` |
  1661. | `matplotlib.figure.Figure` | | |
  1662. +------------------------------------+---------------+-------------------+
  1663. | zero-dimensional numpy arrays | ``False`` | ``True`` |
  1664. +------------------------------------+---------------+-------------------+
  1665. | other numpy arrays | ``False`` | ``False`` |
  1666. +------------------------------------+---------------+-------------------+
  1667. | `list`, `tuple`, and other | ``False`` | ``False`` |
  1668. | sequence objects | | |
  1669. +------------------------------------+---------------+-------------------+
  1670. Examples
  1671. --------
  1672. >>> import numpy as np
  1673. >>> np.isscalar(3.1)
  1674. True
  1675. >>> np.isscalar(np.array(3.1))
  1676. False
  1677. >>> np.isscalar([3.1])
  1678. False
  1679. >>> np.isscalar(False)
  1680. True
  1681. >>> np.isscalar('numpy')
  1682. True
  1683. NumPy supports PEP 3141 numbers:
  1684. >>> from fractions import Fraction
  1685. >>> np.isscalar(Fraction(5, 17))
  1686. True
  1687. >>> from numbers import Number
  1688. >>> np.isscalar(Number())
  1689. True
  1690. """
  1691. return (isinstance(element, generic)
  1692. or type(element) in ScalarType
  1693. or isinstance(element, numbers.Number))
  1694. @set_module('numpy')
  1695. def binary_repr(num, width=None):
  1696. """
  1697. Return the binary representation of the input number as a string.
  1698. For negative numbers, if width is not given, a minus sign is added to the
  1699. front. If width is given, the two's complement of the number is
  1700. returned, with respect to that width.
  1701. In a two's-complement system negative numbers are represented by the two's
  1702. complement of the absolute value. This is the most common method of
  1703. representing signed integers on computers [1]_. A N-bit two's-complement
  1704. system can represent every integer in the range
  1705. :math:`-2^{N-1}` to :math:`+2^{N-1}-1`.
  1706. Parameters
  1707. ----------
  1708. num : int
  1709. Only an integer decimal number can be used.
  1710. width : int, optional
  1711. The length of the returned string if `num` is positive, or the length
  1712. of the two's complement if `num` is negative, provided that `width` is
  1713. at least a sufficient number of bits for `num` to be represented in
  1714. the designated form. If the `width` value is insufficient, an error is
  1715. raised.
  1716. Returns
  1717. -------
  1718. bin : str
  1719. Binary representation of `num` or two's complement of `num`.
  1720. See Also
  1721. --------
  1722. base_repr: Return a string representation of a number in the given base
  1723. system.
  1724. bin: Python's built-in binary representation generator of an integer.
  1725. Notes
  1726. -----
  1727. `binary_repr` is equivalent to using `base_repr` with base 2, but about 25x
  1728. faster.
  1729. References
  1730. ----------
  1731. .. [1] Wikipedia, "Two's complement",
  1732. https://en.wikipedia.org/wiki/Two's_complement
  1733. Examples
  1734. --------
  1735. >>> import numpy as np
  1736. >>> np.binary_repr(3)
  1737. '11'
  1738. >>> np.binary_repr(-3)
  1739. '-11'
  1740. >>> np.binary_repr(3, width=4)
  1741. '0011'
  1742. The two's complement is returned when the input number is negative and
  1743. width is specified:
  1744. >>> np.binary_repr(-3, width=3)
  1745. '101'
  1746. >>> np.binary_repr(-3, width=5)
  1747. '11101'
  1748. """
  1749. def err_if_insufficient(width, binwidth):
  1750. if width is not None and width < binwidth:
  1751. raise ValueError(
  1752. f"Insufficient bit {width=} provided for {binwidth=}"
  1753. )
  1754. # Ensure that num is a Python integer to avoid overflow or unwanted
  1755. # casts to floating point.
  1756. num = operator.index(num)
  1757. if num == 0:
  1758. return '0' * (width or 1)
  1759. elif num > 0:
  1760. binary = f'{num:b}'
  1761. binwidth = len(binary)
  1762. outwidth = (binwidth if width is None
  1763. else builtins.max(binwidth, width))
  1764. err_if_insufficient(width, binwidth)
  1765. return binary.zfill(outwidth)
  1766. elif width is None:
  1767. return f'-{-num:b}'
  1768. else:
  1769. poswidth = len(f'{-num:b}')
  1770. # See gh-8679: remove extra digit
  1771. # for numbers at boundaries.
  1772. if 2**(poswidth - 1) == -num:
  1773. poswidth -= 1
  1774. twocomp = 2**(poswidth + 1) + num
  1775. binary = f'{twocomp:b}'
  1776. binwidth = len(binary)
  1777. outwidth = builtins.max(binwidth, width)
  1778. err_if_insufficient(width, binwidth)
  1779. return '1' * (outwidth - binwidth) + binary
  1780. @set_module('numpy')
  1781. def base_repr(number, base=2, padding=0):
  1782. """
  1783. Return a string representation of a number in the given base system.
  1784. Parameters
  1785. ----------
  1786. number : int
  1787. The value to convert. Positive and negative values are handled.
  1788. base : int, optional
  1789. Convert `number` to the `base` number system. The valid range is 2-36,
  1790. the default value is 2.
  1791. padding : int, optional
  1792. Number of zeros padded on the left. Default is 0 (no padding).
  1793. Returns
  1794. -------
  1795. out : str
  1796. String representation of `number` in `base` system.
  1797. See Also
  1798. --------
  1799. binary_repr : Faster version of `base_repr` for base 2.
  1800. Examples
  1801. --------
  1802. >>> import numpy as np
  1803. >>> np.base_repr(5)
  1804. '101'
  1805. >>> np.base_repr(6, 5)
  1806. '11'
  1807. >>> np.base_repr(7, base=5, padding=3)
  1808. '00012'
  1809. >>> np.base_repr(10, base=16)
  1810. 'A'
  1811. >>> np.base_repr(32, base=16)
  1812. '20'
  1813. """
  1814. digits = '0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ'
  1815. if base > len(digits):
  1816. raise ValueError("Bases greater than 36 not handled in base_repr.")
  1817. elif base < 2:
  1818. raise ValueError("Bases less than 2 not handled in base_repr.")
  1819. num = abs(int(number))
  1820. res = []
  1821. while num:
  1822. res.append(digits[num % base])
  1823. num //= base
  1824. if padding:
  1825. res.append('0' * padding)
  1826. if number < 0:
  1827. res.append('-')
  1828. return ''.join(reversed(res or '0'))
  1829. # These are all essentially abbreviations
  1830. # These might wind up in a special abbreviations module
  1831. def _maketup(descr, val):
  1832. dt = dtype(descr)
  1833. # Place val in all scalar tuples:
  1834. fields = dt.fields
  1835. if fields is None:
  1836. return val
  1837. else:
  1838. res = [_maketup(fields[name][0], val) for name in dt.names]
  1839. return tuple(res)
  1840. @finalize_array_function_like
  1841. @set_module('numpy')
  1842. def identity(n, dtype=None, *, like=None):
  1843. """
  1844. Return the identity array.
  1845. The identity array is a square array with ones on
  1846. the main diagonal.
  1847. Parameters
  1848. ----------
  1849. n : int
  1850. Number of rows (and columns) in `n` x `n` output.
  1851. dtype : data-type, optional
  1852. Data-type of the output. Defaults to ``float``.
  1853. ${ARRAY_FUNCTION_LIKE}
  1854. .. versionadded:: 1.20.0
  1855. Returns
  1856. -------
  1857. out : ndarray
  1858. `n` x `n` array with its main diagonal set to one,
  1859. and all other elements 0.
  1860. Examples
  1861. --------
  1862. >>> import numpy as np
  1863. >>> np.identity(3)
  1864. array([[1., 0., 0.],
  1865. [0., 1., 0.],
  1866. [0., 0., 1.]])
  1867. """
  1868. if like is not None:
  1869. return _identity_with_like(like, n, dtype=dtype)
  1870. from numpy import eye
  1871. return eye(n, dtype=dtype, like=like)
  1872. _identity_with_like = array_function_dispatch()(identity)
  1873. def _allclose_dispatcher(a, b, rtol=None, atol=None, equal_nan=None):
  1874. return (a, b, rtol, atol)
  1875. @array_function_dispatch(_allclose_dispatcher)
  1876. def allclose(a, b, rtol=1.e-5, atol=1.e-8, equal_nan=False):
  1877. """
  1878. Returns True if two arrays are element-wise equal within a tolerance.
  1879. The tolerance values are positive, typically very small numbers. The
  1880. relative difference (`rtol` * abs(`b`)) and the absolute difference
  1881. `atol` are added together to compare against the absolute difference
  1882. between `a` and `b`.
  1883. .. warning:: The default `atol` is not appropriate for comparing numbers
  1884. with magnitudes much smaller than one (see Notes).
  1885. NaNs are treated as equal if they are in the same place and if
  1886. ``equal_nan=True``. Infs are treated as equal if they are in the same
  1887. place and of the same sign in both arrays.
  1888. Parameters
  1889. ----------
  1890. a, b : array_like
  1891. Input arrays to compare.
  1892. rtol : array_like
  1893. The relative tolerance parameter (see Notes).
  1894. atol : array_like
  1895. The absolute tolerance parameter (see Notes).
  1896. equal_nan : bool
  1897. Whether to compare NaN's as equal. If True, NaN's in `a` will be
  1898. considered equal to NaN's in `b` in the output array.
  1899. Returns
  1900. -------
  1901. allclose : bool
  1902. Returns True if the two arrays are equal within the given
  1903. tolerance; False otherwise.
  1904. See Also
  1905. --------
  1906. isclose, all, any, equal
  1907. Notes
  1908. -----
  1909. If the following equation is element-wise True, then allclose returns
  1910. True.::
  1911. absolute(a - b) <= (atol + rtol * absolute(b))
  1912. The above equation is not symmetric in `a` and `b`, so that
  1913. ``allclose(a, b)`` might be different from ``allclose(b, a)`` in
  1914. some rare cases.
  1915. The default value of `atol` is not appropriate when the reference value
  1916. `b` has magnitude smaller than one. For example, it is unlikely that
  1917. ``a = 1e-9`` and ``b = 2e-9`` should be considered "close", yet
  1918. ``allclose(1e-9, 2e-9)`` is ``True`` with default settings. Be sure
  1919. to select `atol` for the use case at hand, especially for defining the
  1920. threshold below which a non-zero value in `a` will be considered "close"
  1921. to a very small or zero value in `b`.
  1922. The comparison of `a` and `b` uses standard broadcasting, which
  1923. means that `a` and `b` need not have the same shape in order for
  1924. ``allclose(a, b)`` to evaluate to True. The same is true for
  1925. `equal` but not `array_equal`.
  1926. `allclose` is not defined for non-numeric data types.
  1927. `bool` is considered a numeric data-type for this purpose.
  1928. Examples
  1929. --------
  1930. >>> import numpy as np
  1931. >>> np.allclose([1e10,1e-7], [1.00001e10,1e-8])
  1932. False
  1933. >>> np.allclose([1e10,1e-8], [1.00001e10,1e-9])
  1934. True
  1935. >>> np.allclose([1e10,1e-8], [1.0001e10,1e-9])
  1936. False
  1937. >>> np.allclose([1.0, np.nan], [1.0, np.nan])
  1938. False
  1939. >>> np.allclose([1.0, np.nan], [1.0, np.nan], equal_nan=True)
  1940. True
  1941. """
  1942. res = all(isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan))
  1943. return builtins.bool(res)
  1944. def _isclose_dispatcher(a, b, rtol=None, atol=None, equal_nan=None):
  1945. return (a, b, rtol, atol)
  1946. @array_function_dispatch(_isclose_dispatcher)
  1947. def isclose(a, b, rtol=1.e-5, atol=1.e-8, equal_nan=False):
  1948. """
  1949. Returns a boolean array where two arrays are element-wise equal within a
  1950. tolerance.
  1951. The tolerance values are positive, typically very small numbers. The
  1952. relative difference (`rtol` * abs(`b`)) and the absolute difference
  1953. `atol` are added together to compare against the absolute difference
  1954. between `a` and `b`.
  1955. .. warning:: The default `atol` is not appropriate for comparing numbers
  1956. with magnitudes much smaller than one (see Notes).
  1957. Parameters
  1958. ----------
  1959. a, b : array_like
  1960. Input arrays to compare.
  1961. rtol : array_like
  1962. The relative tolerance parameter (see Notes).
  1963. atol : array_like
  1964. The absolute tolerance parameter (see Notes).
  1965. equal_nan : bool
  1966. Whether to compare NaN's as equal. If True, NaN's in `a` will be
  1967. considered equal to NaN's in `b` in the output array.
  1968. Returns
  1969. -------
  1970. y : array_like
  1971. Returns a boolean array of where `a` and `b` are equal within the
  1972. given tolerance. If both `a` and `b` are scalars, returns a single
  1973. boolean value.
  1974. See Also
  1975. --------
  1976. allclose
  1977. math.isclose
  1978. Notes
  1979. -----
  1980. For finite values, isclose uses the following equation to test whether
  1981. two floating point values are equivalent.::
  1982. absolute(a - b) <= (atol + rtol * absolute(b))
  1983. Unlike the built-in `math.isclose`, the above equation is not symmetric
  1984. in `a` and `b` -- it assumes `b` is the reference value -- so that
  1985. `isclose(a, b)` might be different from `isclose(b, a)`.
  1986. The default value of `atol` is not appropriate when the reference value
  1987. `b` has magnitude smaller than one. For example, it is unlikely that
  1988. ``a = 1e-9`` and ``b = 2e-9`` should be considered "close", yet
  1989. ``isclose(1e-9, 2e-9)`` is ``True`` with default settings. Be sure
  1990. to select `atol` for the use case at hand, especially for defining the
  1991. threshold below which a non-zero value in `a` will be considered "close"
  1992. to a very small or zero value in `b`.
  1993. `isclose` is not defined for non-numeric data types.
  1994. :class:`bool` is considered a numeric data-type for this purpose.
  1995. Examples
  1996. --------
  1997. >>> import numpy as np
  1998. >>> np.isclose([1e10,1e-7], [1.00001e10,1e-8])
  1999. array([ True, False])
  2000. >>> np.isclose([1e10,1e-8], [1.00001e10,1e-9])
  2001. array([ True, True])
  2002. >>> np.isclose([1e10,1e-8], [1.0001e10,1e-9])
  2003. array([False, True])
  2004. >>> np.isclose([1.0, np.nan], [1.0, np.nan])
  2005. array([ True, False])
  2006. >>> np.isclose([1.0, np.nan], [1.0, np.nan], equal_nan=True)
  2007. array([ True, True])
  2008. >>> np.isclose([1e-8, 1e-7], [0.0, 0.0])
  2009. array([ True, False])
  2010. >>> np.isclose([1e-100, 1e-7], [0.0, 0.0], atol=0.0)
  2011. array([False, False])
  2012. >>> np.isclose([1e-10, 1e-10], [1e-20, 0.0])
  2013. array([ True, True])
  2014. >>> np.isclose([1e-10, 1e-10], [1e-20, 0.999999e-10], atol=0.0)
  2015. array([False, True])
  2016. """
  2017. # Turn all but python scalars into arrays.
  2018. x, y, atol, rtol = (
  2019. a if isinstance(a, (int, float, complex)) else asanyarray(a)
  2020. for a in (a, b, atol, rtol))
  2021. # Make sure y is an inexact type to avoid bad behavior on abs(MIN_INT).
  2022. # This will cause casting of x later. Also, make sure to allow subclasses
  2023. # (e.g., for numpy.ma).
  2024. # NOTE: We explicitly allow timedelta, which used to work. This could
  2025. # possibly be deprecated. See also gh-18286.
  2026. # timedelta works if `atol` is an integer or also a timedelta.
  2027. # Although, the default tolerances are unlikely to be useful
  2028. if (dtype := getattr(y, "dtype", None)) is not None and dtype.kind != "m":
  2029. dt = multiarray.result_type(y, 1.)
  2030. y = asanyarray(y, dtype=dt)
  2031. elif isinstance(y, int):
  2032. y = float(y)
  2033. # atol and rtol can be arrays
  2034. if not (np.all(np.isfinite(atol)) and np.all(np.isfinite(rtol))):
  2035. err_s = np.geterr()["invalid"]
  2036. err_msg = f"One of rtol or atol is not valid, atol: {atol}, rtol: {rtol}"
  2037. if err_s == "warn":
  2038. warnings.warn(err_msg, RuntimeWarning, stacklevel=2)
  2039. elif err_s == "raise":
  2040. raise FloatingPointError(err_msg)
  2041. elif err_s == "print":
  2042. print(err_msg)
  2043. with errstate(invalid='ignore'):
  2044. result = (less_equal(abs(x - y), atol + rtol * abs(y))
  2045. & isfinite(y)
  2046. | (x == y))
  2047. if equal_nan:
  2048. result |= isnan(x) & isnan(y)
  2049. return result[()] # Flatten 0d arrays to scalars
  2050. def _array_equal_dispatcher(a1, a2, equal_nan=None):
  2051. return (a1, a2)
  2052. _no_nan_types = {
  2053. # should use np.dtype.BoolDType, but as of writing
  2054. # that fails the reloading test.
  2055. type(dtype(nt.bool)),
  2056. type(dtype(nt.int8)),
  2057. type(dtype(nt.int16)),
  2058. type(dtype(nt.int32)),
  2059. type(dtype(nt.int64)),
  2060. }
  2061. def _dtype_cannot_hold_nan(dtype):
  2062. return type(dtype) in _no_nan_types
  2063. @array_function_dispatch(_array_equal_dispatcher)
  2064. def array_equal(a1, a2, equal_nan=False):
  2065. """
  2066. True if two arrays have the same shape and elements, False otherwise.
  2067. Parameters
  2068. ----------
  2069. a1, a2 : array_like
  2070. Input arrays.
  2071. equal_nan : bool
  2072. Whether to compare NaN's as equal. If the dtype of a1 and a2 is
  2073. complex, values will be considered equal if either the real or the
  2074. imaginary component of a given value is ``nan``.
  2075. Returns
  2076. -------
  2077. b : bool
  2078. Returns True if the arrays are equal.
  2079. See Also
  2080. --------
  2081. allclose: Returns True if two arrays are element-wise equal within a
  2082. tolerance.
  2083. array_equiv: Returns True if input arrays are shape consistent and all
  2084. elements equal.
  2085. Examples
  2086. --------
  2087. >>> import numpy as np
  2088. >>> np.array_equal([1, 2], [1, 2])
  2089. True
  2090. >>> np.array_equal(np.array([1, 2]), np.array([1, 2]))
  2091. True
  2092. >>> np.array_equal([1, 2], [1, 2, 3])
  2093. False
  2094. >>> np.array_equal([1, 2], [1, 4])
  2095. False
  2096. >>> a = np.array([1, np.nan])
  2097. >>> np.array_equal(a, a)
  2098. False
  2099. >>> np.array_equal(a, a, equal_nan=True)
  2100. True
  2101. When ``equal_nan`` is True, complex values with nan components are
  2102. considered equal if either the real *or* the imaginary components are nan.
  2103. >>> a = np.array([1 + 1j])
  2104. >>> b = a.copy()
  2105. >>> a.real = np.nan
  2106. >>> b.imag = np.nan
  2107. >>> np.array_equal(a, b, equal_nan=True)
  2108. True
  2109. """
  2110. try:
  2111. a1, a2 = asarray(a1), asarray(a2)
  2112. except Exception:
  2113. return False
  2114. if a1.shape != a2.shape:
  2115. return False
  2116. if not equal_nan:
  2117. return builtins.bool((asanyarray(a1 == a2)).all())
  2118. if a1 is a2:
  2119. # nan will compare equal so an array will compare equal to itself.
  2120. return True
  2121. cannot_have_nan = (_dtype_cannot_hold_nan(a1.dtype)
  2122. and _dtype_cannot_hold_nan(a2.dtype))
  2123. if cannot_have_nan:
  2124. return builtins.bool(asarray(a1 == a2).all())
  2125. # Handling NaN values if equal_nan is True
  2126. a1nan, a2nan = isnan(a1), isnan(a2)
  2127. # NaN's occur at different locations
  2128. if not (a1nan == a2nan).all():
  2129. return False
  2130. # Shapes of a1, a2 and masks are guaranteed to be consistent by this point
  2131. return builtins.bool((a1[~a1nan] == a2[~a1nan]).all())
  2132. def _array_equiv_dispatcher(a1, a2):
  2133. return (a1, a2)
  2134. @array_function_dispatch(_array_equiv_dispatcher)
  2135. def array_equiv(a1, a2):
  2136. """
  2137. Returns True if input arrays are shape consistent and all elements equal.
  2138. Shape consistent means they are either the same shape, or one input array
  2139. can be broadcasted to create the same shape as the other one.
  2140. Parameters
  2141. ----------
  2142. a1, a2 : array_like
  2143. Input arrays.
  2144. Returns
  2145. -------
  2146. out : bool
  2147. True if equivalent, False otherwise.
  2148. Examples
  2149. --------
  2150. >>> import numpy as np
  2151. >>> np.array_equiv([1, 2], [1, 2])
  2152. True
  2153. >>> np.array_equiv([1, 2], [1, 3])
  2154. False
  2155. Showing the shape equivalence:
  2156. >>> np.array_equiv([1, 2], [[1, 2], [1, 2]])
  2157. True
  2158. >>> np.array_equiv([1, 2], [[1, 2, 1, 2], [1, 2, 1, 2]])
  2159. False
  2160. >>> np.array_equiv([1, 2], [[1, 2], [1, 3]])
  2161. False
  2162. """
  2163. try:
  2164. a1, a2 = asarray(a1), asarray(a2)
  2165. except Exception:
  2166. return False
  2167. try:
  2168. multiarray.broadcast(a1, a2)
  2169. except Exception:
  2170. return False
  2171. return builtins.bool(asanyarray(a1 == a2).all())
  2172. def _astype_dispatcher(x, dtype, /, *, copy=None, device=None):
  2173. return (x, dtype)
  2174. @array_function_dispatch(_astype_dispatcher)
  2175. def astype(x, dtype, /, *, copy=True, device=None):
  2176. """
  2177. Copies an array to a specified data type.
  2178. This function is an Array API compatible alternative to
  2179. `numpy.ndarray.astype`.
  2180. Parameters
  2181. ----------
  2182. x : ndarray
  2183. Input NumPy array to cast. ``array_likes`` are explicitly not
  2184. supported here.
  2185. dtype : dtype
  2186. Data type of the result.
  2187. copy : bool, optional
  2188. Specifies whether to copy an array when the specified dtype matches
  2189. the data type of the input array ``x``. If ``True``, a newly allocated
  2190. array must always be returned. If ``False`` and the specified dtype
  2191. matches the data type of the input array, the input array must be
  2192. returned; otherwise, a newly allocated array must be returned.
  2193. Defaults to ``True``.
  2194. device : str, optional
  2195. The device on which to place the returned array. Default: None.
  2196. For Array-API interoperability only, so must be ``"cpu"`` if passed.
  2197. .. versionadded:: 2.1.0
  2198. Returns
  2199. -------
  2200. out : ndarray
  2201. An array having the specified data type.
  2202. See Also
  2203. --------
  2204. ndarray.astype
  2205. Examples
  2206. --------
  2207. >>> import numpy as np
  2208. >>> arr = np.array([1, 2, 3]); arr
  2209. array([1, 2, 3])
  2210. >>> np.astype(arr, np.float64)
  2211. array([1., 2., 3.])
  2212. Non-copy case:
  2213. >>> arr = np.array([1, 2, 3])
  2214. >>> arr_noncpy = np.astype(arr, arr.dtype, copy=False)
  2215. >>> np.shares_memory(arr, arr_noncpy)
  2216. True
  2217. """
  2218. if not (isinstance(x, np.ndarray) or isscalar(x)):
  2219. raise TypeError(
  2220. "Input should be a NumPy array or scalar. "
  2221. f"It is a {type(x)} instead."
  2222. )
  2223. if device is not None and device != "cpu":
  2224. raise ValueError(
  2225. 'Device not understood. Only "cpu" is allowed, but received:'
  2226. f' {device}'
  2227. )
  2228. return x.astype(dtype, copy=copy)
  2229. inf = PINF
  2230. nan = NAN
  2231. False_ = nt.bool(False)
  2232. True_ = nt.bool(True)
  2233. def extend_all(module):
  2234. existing = set(__all__)
  2235. mall = module.__all__
  2236. for a in mall:
  2237. if a not in existing:
  2238. __all__.append(a)
  2239. from . import _asarray, _ufunc_config, arrayprint, fromnumeric
  2240. from ._asarray import *
  2241. from ._ufunc_config import *
  2242. from .arrayprint import *
  2243. from .fromnumeric import *
  2244. from .numerictypes import *
  2245. from .umath import *
  2246. extend_all(fromnumeric)
  2247. extend_all(umath)
  2248. extend_all(numerictypes)
  2249. extend_all(arrayprint)
  2250. extend_all(_asarray)
  2251. extend_all(_ufunc_config)