test_nditer.py 127 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954195519561957195819591960196119621963196419651966196719681969197019711972197319741975197619771978197919801981198219831984198519861987198819891990199119921993199419951996199719981999200020012002200320042005200620072008200920102011201220132014201520162017201820192020202120222023202420252026202720282029203020312032203320342035203620372038203920402041204220432044204520462047204820492050205120522053205420552056205720582059206020612062206320642065206620672068206920702071207220732074207520762077207820792080208120822083208420852086208720882089209020912092209320942095209620972098209921002101210221032104210521062107210821092110211121122113211421152116211721182119212021212122212321242125212621272128212921302131213221332134213521362137213821392140214121422143214421452146214721482149215021512152215321542155215621572158215921602161216221632164216521662167216821692170217121722173217421752176217721782179218021812182218321842185218621872188218921902191219221932194219521962197219821992200220122022203220422052206220722082209221022112212221322142215221622172218221922202221222222232224222522262227222822292230223122322233223422352236223722382239224022412242224322442245224622472248224922502251225222532254225522562257225822592260226122622263226422652266226722682269227022712272227322742275227622772278227922802281228222832284228522862287228822892290229122922293229422952296229722982299230023012302230323042305230623072308230923102311231223132314231523162317231823192320232123222323232423252326232723282329233023312332233323342335233623372338233923402341234223432344234523462347234823492350235123522353235423552356235723582359236023612362236323642365236623672368236923702371237223732374237523762377237823792380238123822383238423852386238723882389239023912392239323942395239623972398239924002401240224032404240524062407240824092410241124122413241424152416241724182419242024212422242324242425242624272428242924302431243224332434243524362437243824392440244124422443244424452446244724482449245024512452245324542455245624572458245924602461246224632464246524662467246824692470247124722473247424752476247724782479248024812482248324842485248624872488248924902491249224932494249524962497249824992500250125022503250425052506250725082509251025112512251325142515251625172518251925202521252225232524252525262527252825292530253125322533253425352536253725382539254025412542254325442545254625472548254925502551255225532554255525562557255825592560256125622563256425652566256725682569257025712572257325742575257625772578257925802581258225832584258525862587258825892590259125922593259425952596259725982599260026012602260326042605260626072608260926102611261226132614261526162617261826192620262126222623262426252626262726282629263026312632263326342635263626372638263926402641264226432644264526462647264826492650265126522653265426552656265726582659266026612662266326642665266626672668266926702671267226732674267526762677267826792680268126822683268426852686268726882689269026912692269326942695269626972698269927002701270227032704270527062707270827092710271127122713271427152716271727182719272027212722272327242725272627272728272927302731273227332734273527362737273827392740274127422743274427452746274727482749275027512752275327542755275627572758275927602761276227632764276527662767276827692770277127722773277427752776277727782779278027812782278327842785278627872788278927902791279227932794279527962797279827992800280128022803280428052806280728082809281028112812281328142815281628172818281928202821282228232824282528262827282828292830283128322833283428352836283728382839284028412842284328442845284628472848284928502851285228532854285528562857285828592860286128622863286428652866286728682869287028712872287328742875287628772878287928802881288228832884288528862887288828892890289128922893289428952896289728982899290029012902290329042905290629072908290929102911291229132914291529162917291829192920292129222923292429252926292729282929293029312932293329342935293629372938293929402941294229432944294529462947294829492950295129522953295429552956295729582959296029612962296329642965296629672968296929702971297229732974297529762977297829792980298129822983298429852986298729882989299029912992299329942995299629972998299930003001300230033004300530063007300830093010301130123013301430153016301730183019302030213022302330243025302630273028302930303031303230333034303530363037303830393040304130423043304430453046304730483049305030513052305330543055305630573058305930603061306230633064306530663067306830693070307130723073307430753076307730783079308030813082308330843085308630873088308930903091309230933094309530963097309830993100310131023103310431053106310731083109311031113112311331143115311631173118311931203121312231233124312531263127312831293130313131323133313431353136313731383139314031413142314331443145314631473148314931503151315231533154315531563157315831593160316131623163316431653166316731683169317031713172317331743175317631773178317931803181318231833184318531863187318831893190319131923193319431953196319731983199320032013202320332043205320632073208320932103211321232133214321532163217321832193220322132223223322432253226322732283229323032313232323332343235323632373238323932403241324232433244324532463247324832493250325132523253325432553256325732583259326032613262326332643265326632673268326932703271327232733274327532763277327832793280328132823283328432853286328732883289329032913292329332943295329632973298329933003301330233033304330533063307330833093310331133123313331433153316
  1. import sys
  2. import pytest
  3. import textwrap
  4. import subprocess
  5. import numpy as np
  6. import numpy.core._multiarray_tests as _multiarray_tests
  7. from numpy import array, arange, nditer, all
  8. from numpy.testing import (
  9. assert_, assert_equal, assert_array_equal, assert_raises,
  10. IS_WASM, HAS_REFCOUNT, suppress_warnings, break_cycles
  11. )
  12. def iter_multi_index(i):
  13. ret = []
  14. while not i.finished:
  15. ret.append(i.multi_index)
  16. i.iternext()
  17. return ret
  18. def iter_indices(i):
  19. ret = []
  20. while not i.finished:
  21. ret.append(i.index)
  22. i.iternext()
  23. return ret
  24. def iter_iterindices(i):
  25. ret = []
  26. while not i.finished:
  27. ret.append(i.iterindex)
  28. i.iternext()
  29. return ret
  30. @pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts")
  31. def test_iter_refcount():
  32. # Make sure the iterator doesn't leak
  33. # Basic
  34. a = arange(6)
  35. dt = np.dtype('f4').newbyteorder()
  36. rc_a = sys.getrefcount(a)
  37. rc_dt = sys.getrefcount(dt)
  38. with nditer(a, [],
  39. [['readwrite', 'updateifcopy']],
  40. casting='unsafe',
  41. op_dtypes=[dt]) as it:
  42. assert_(not it.iterationneedsapi)
  43. assert_(sys.getrefcount(a) > rc_a)
  44. assert_(sys.getrefcount(dt) > rc_dt)
  45. # del 'it'
  46. it = None
  47. assert_equal(sys.getrefcount(a), rc_a)
  48. assert_equal(sys.getrefcount(dt), rc_dt)
  49. # With a copy
  50. a = arange(6, dtype='f4')
  51. dt = np.dtype('f4')
  52. rc_a = sys.getrefcount(a)
  53. rc_dt = sys.getrefcount(dt)
  54. it = nditer(a, [],
  55. [['readwrite']],
  56. op_dtypes=[dt])
  57. rc2_a = sys.getrefcount(a)
  58. rc2_dt = sys.getrefcount(dt)
  59. it2 = it.copy()
  60. assert_(sys.getrefcount(a) > rc2_a)
  61. assert_(sys.getrefcount(dt) > rc2_dt)
  62. it = None
  63. assert_equal(sys.getrefcount(a), rc2_a)
  64. assert_equal(sys.getrefcount(dt), rc2_dt)
  65. it2 = None
  66. assert_equal(sys.getrefcount(a), rc_a)
  67. assert_equal(sys.getrefcount(dt), rc_dt)
  68. del it2 # avoid pyflakes unused variable warning
  69. def test_iter_best_order():
  70. # The iterator should always find the iteration order
  71. # with increasing memory addresses
  72. # Test the ordering for 1-D to 5-D shapes
  73. for shape in [(5,), (3, 4), (2, 3, 4), (2, 3, 4, 3), (2, 3, 2, 2, 3)]:
  74. a = arange(np.prod(shape))
  75. # Test each combination of positive and negative strides
  76. for dirs in range(2**len(shape)):
  77. dirs_index = [slice(None)]*len(shape)
  78. for bit in range(len(shape)):
  79. if ((2**bit) & dirs):
  80. dirs_index[bit] = slice(None, None, -1)
  81. dirs_index = tuple(dirs_index)
  82. aview = a.reshape(shape)[dirs_index]
  83. # C-order
  84. i = nditer(aview, [], [['readonly']])
  85. assert_equal([x for x in i], a)
  86. # Fortran-order
  87. i = nditer(aview.T, [], [['readonly']])
  88. assert_equal([x for x in i], a)
  89. # Other order
  90. if len(shape) > 2:
  91. i = nditer(aview.swapaxes(0, 1), [], [['readonly']])
  92. assert_equal([x for x in i], a)
  93. def test_iter_c_order():
  94. # Test forcing C order
  95. # Test the ordering for 1-D to 5-D shapes
  96. for shape in [(5,), (3, 4), (2, 3, 4), (2, 3, 4, 3), (2, 3, 2, 2, 3)]:
  97. a = arange(np.prod(shape))
  98. # Test each combination of positive and negative strides
  99. for dirs in range(2**len(shape)):
  100. dirs_index = [slice(None)]*len(shape)
  101. for bit in range(len(shape)):
  102. if ((2**bit) & dirs):
  103. dirs_index[bit] = slice(None, None, -1)
  104. dirs_index = tuple(dirs_index)
  105. aview = a.reshape(shape)[dirs_index]
  106. # C-order
  107. i = nditer(aview, order='C')
  108. assert_equal([x for x in i], aview.ravel(order='C'))
  109. # Fortran-order
  110. i = nditer(aview.T, order='C')
  111. assert_equal([x for x in i], aview.T.ravel(order='C'))
  112. # Other order
  113. if len(shape) > 2:
  114. i = nditer(aview.swapaxes(0, 1), order='C')
  115. assert_equal([x for x in i],
  116. aview.swapaxes(0, 1).ravel(order='C'))
  117. def test_iter_f_order():
  118. # Test forcing F order
  119. # Test the ordering for 1-D to 5-D shapes
  120. for shape in [(5,), (3, 4), (2, 3, 4), (2, 3, 4, 3), (2, 3, 2, 2, 3)]:
  121. a = arange(np.prod(shape))
  122. # Test each combination of positive and negative strides
  123. for dirs in range(2**len(shape)):
  124. dirs_index = [slice(None)]*len(shape)
  125. for bit in range(len(shape)):
  126. if ((2**bit) & dirs):
  127. dirs_index[bit] = slice(None, None, -1)
  128. dirs_index = tuple(dirs_index)
  129. aview = a.reshape(shape)[dirs_index]
  130. # C-order
  131. i = nditer(aview, order='F')
  132. assert_equal([x for x in i], aview.ravel(order='F'))
  133. # Fortran-order
  134. i = nditer(aview.T, order='F')
  135. assert_equal([x for x in i], aview.T.ravel(order='F'))
  136. # Other order
  137. if len(shape) > 2:
  138. i = nditer(aview.swapaxes(0, 1), order='F')
  139. assert_equal([x for x in i],
  140. aview.swapaxes(0, 1).ravel(order='F'))
  141. def test_iter_c_or_f_order():
  142. # Test forcing any contiguous (C or F) order
  143. # Test the ordering for 1-D to 5-D shapes
  144. for shape in [(5,), (3, 4), (2, 3, 4), (2, 3, 4, 3), (2, 3, 2, 2, 3)]:
  145. a = arange(np.prod(shape))
  146. # Test each combination of positive and negative strides
  147. for dirs in range(2**len(shape)):
  148. dirs_index = [slice(None)]*len(shape)
  149. for bit in range(len(shape)):
  150. if ((2**bit) & dirs):
  151. dirs_index[bit] = slice(None, None, -1)
  152. dirs_index = tuple(dirs_index)
  153. aview = a.reshape(shape)[dirs_index]
  154. # C-order
  155. i = nditer(aview, order='A')
  156. assert_equal([x for x in i], aview.ravel(order='A'))
  157. # Fortran-order
  158. i = nditer(aview.T, order='A')
  159. assert_equal([x for x in i], aview.T.ravel(order='A'))
  160. # Other order
  161. if len(shape) > 2:
  162. i = nditer(aview.swapaxes(0, 1), order='A')
  163. assert_equal([x for x in i],
  164. aview.swapaxes(0, 1).ravel(order='A'))
  165. def test_nditer_multi_index_set():
  166. # Test the multi_index set
  167. a = np.arange(6).reshape(2, 3)
  168. it = np.nditer(a, flags=['multi_index'])
  169. # Removes the iteration on two first elements of a[0]
  170. it.multi_index = (0, 2,)
  171. assert_equal([i for i in it], [2, 3, 4, 5])
  172. @pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts")
  173. def test_nditer_multi_index_set_refcount():
  174. # Test if the reference count on index variable is decreased
  175. index = 0
  176. i = np.nditer(np.array([111, 222, 333, 444]), flags=['multi_index'])
  177. start_count = sys.getrefcount(index)
  178. i.multi_index = (index,)
  179. end_count = sys.getrefcount(index)
  180. assert_equal(start_count, end_count)
  181. def test_iter_best_order_multi_index_1d():
  182. # The multi-indices should be correct with any reordering
  183. a = arange(4)
  184. # 1D order
  185. i = nditer(a, ['multi_index'], [['readonly']])
  186. assert_equal(iter_multi_index(i), [(0,), (1,), (2,), (3,)])
  187. # 1D reversed order
  188. i = nditer(a[::-1], ['multi_index'], [['readonly']])
  189. assert_equal(iter_multi_index(i), [(3,), (2,), (1,), (0,)])
  190. def test_iter_best_order_multi_index_2d():
  191. # The multi-indices should be correct with any reordering
  192. a = arange(6)
  193. # 2D C-order
  194. i = nditer(a.reshape(2, 3), ['multi_index'], [['readonly']])
  195. assert_equal(iter_multi_index(i), [(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2)])
  196. # 2D Fortran-order
  197. i = nditer(a.reshape(2, 3).copy(order='F'), ['multi_index'], [['readonly']])
  198. assert_equal(iter_multi_index(i), [(0, 0), (1, 0), (0, 1), (1, 1), (0, 2), (1, 2)])
  199. # 2D reversed C-order
  200. i = nditer(a.reshape(2, 3)[::-1], ['multi_index'], [['readonly']])
  201. assert_equal(iter_multi_index(i), [(1, 0), (1, 1), (1, 2), (0, 0), (0, 1), (0, 2)])
  202. i = nditer(a.reshape(2, 3)[:, ::-1], ['multi_index'], [['readonly']])
  203. assert_equal(iter_multi_index(i), [(0, 2), (0, 1), (0, 0), (1, 2), (1, 1), (1, 0)])
  204. i = nditer(a.reshape(2, 3)[::-1, ::-1], ['multi_index'], [['readonly']])
  205. assert_equal(iter_multi_index(i), [(1, 2), (1, 1), (1, 0), (0, 2), (0, 1), (0, 0)])
  206. # 2D reversed Fortran-order
  207. i = nditer(a.reshape(2, 3).copy(order='F')[::-1], ['multi_index'], [['readonly']])
  208. assert_equal(iter_multi_index(i), [(1, 0), (0, 0), (1, 1), (0, 1), (1, 2), (0, 2)])
  209. i = nditer(a.reshape(2, 3).copy(order='F')[:, ::-1],
  210. ['multi_index'], [['readonly']])
  211. assert_equal(iter_multi_index(i), [(0, 2), (1, 2), (0, 1), (1, 1), (0, 0), (1, 0)])
  212. i = nditer(a.reshape(2, 3).copy(order='F')[::-1, ::-1],
  213. ['multi_index'], [['readonly']])
  214. assert_equal(iter_multi_index(i), [(1, 2), (0, 2), (1, 1), (0, 1), (1, 0), (0, 0)])
  215. def test_iter_best_order_multi_index_3d():
  216. # The multi-indices should be correct with any reordering
  217. a = arange(12)
  218. # 3D C-order
  219. i = nditer(a.reshape(2, 3, 2), ['multi_index'], [['readonly']])
  220. assert_equal(iter_multi_index(i),
  221. [(0, 0, 0), (0, 0, 1), (0, 1, 0), (0, 1, 1), (0, 2, 0), (0, 2, 1),
  222. (1, 0, 0), (1, 0, 1), (1, 1, 0), (1, 1, 1), (1, 2, 0), (1, 2, 1)])
  223. # 3D Fortran-order
  224. i = nditer(a.reshape(2, 3, 2).copy(order='F'), ['multi_index'], [['readonly']])
  225. assert_equal(iter_multi_index(i),
  226. [(0, 0, 0), (1, 0, 0), (0, 1, 0), (1, 1, 0), (0, 2, 0), (1, 2, 0),
  227. (0, 0, 1), (1, 0, 1), (0, 1, 1), (1, 1, 1), (0, 2, 1), (1, 2, 1)])
  228. # 3D reversed C-order
  229. i = nditer(a.reshape(2, 3, 2)[::-1], ['multi_index'], [['readonly']])
  230. assert_equal(iter_multi_index(i),
  231. [(1, 0, 0), (1, 0, 1), (1, 1, 0), (1, 1, 1), (1, 2, 0), (1, 2, 1),
  232. (0, 0, 0), (0, 0, 1), (0, 1, 0), (0, 1, 1), (0, 2, 0), (0, 2, 1)])
  233. i = nditer(a.reshape(2, 3, 2)[:, ::-1], ['multi_index'], [['readonly']])
  234. assert_equal(iter_multi_index(i),
  235. [(0, 2, 0), (0, 2, 1), (0, 1, 0), (0, 1, 1), (0, 0, 0), (0, 0, 1),
  236. (1, 2, 0), (1, 2, 1), (1, 1, 0), (1, 1, 1), (1, 0, 0), (1, 0, 1)])
  237. i = nditer(a.reshape(2, 3, 2)[:,:, ::-1], ['multi_index'], [['readonly']])
  238. assert_equal(iter_multi_index(i),
  239. [(0, 0, 1), (0, 0, 0), (0, 1, 1), (0, 1, 0), (0, 2, 1), (0, 2, 0),
  240. (1, 0, 1), (1, 0, 0), (1, 1, 1), (1, 1, 0), (1, 2, 1), (1, 2, 0)])
  241. # 3D reversed Fortran-order
  242. i = nditer(a.reshape(2, 3, 2).copy(order='F')[::-1],
  243. ['multi_index'], [['readonly']])
  244. assert_equal(iter_multi_index(i),
  245. [(1, 0, 0), (0, 0, 0), (1, 1, 0), (0, 1, 0), (1, 2, 0), (0, 2, 0),
  246. (1, 0, 1), (0, 0, 1), (1, 1, 1), (0, 1, 1), (1, 2, 1), (0, 2, 1)])
  247. i = nditer(a.reshape(2, 3, 2).copy(order='F')[:, ::-1],
  248. ['multi_index'], [['readonly']])
  249. assert_equal(iter_multi_index(i),
  250. [(0, 2, 0), (1, 2, 0), (0, 1, 0), (1, 1, 0), (0, 0, 0), (1, 0, 0),
  251. (0, 2, 1), (1, 2, 1), (0, 1, 1), (1, 1, 1), (0, 0, 1), (1, 0, 1)])
  252. i = nditer(a.reshape(2, 3, 2).copy(order='F')[:,:, ::-1],
  253. ['multi_index'], [['readonly']])
  254. assert_equal(iter_multi_index(i),
  255. [(0, 0, 1), (1, 0, 1), (0, 1, 1), (1, 1, 1), (0, 2, 1), (1, 2, 1),
  256. (0, 0, 0), (1, 0, 0), (0, 1, 0), (1, 1, 0), (0, 2, 0), (1, 2, 0)])
  257. def test_iter_best_order_c_index_1d():
  258. # The C index should be correct with any reordering
  259. a = arange(4)
  260. # 1D order
  261. i = nditer(a, ['c_index'], [['readonly']])
  262. assert_equal(iter_indices(i), [0, 1, 2, 3])
  263. # 1D reversed order
  264. i = nditer(a[::-1], ['c_index'], [['readonly']])
  265. assert_equal(iter_indices(i), [3, 2, 1, 0])
  266. def test_iter_best_order_c_index_2d():
  267. # The C index should be correct with any reordering
  268. a = arange(6)
  269. # 2D C-order
  270. i = nditer(a.reshape(2, 3), ['c_index'], [['readonly']])
  271. assert_equal(iter_indices(i), [0, 1, 2, 3, 4, 5])
  272. # 2D Fortran-order
  273. i = nditer(a.reshape(2, 3).copy(order='F'),
  274. ['c_index'], [['readonly']])
  275. assert_equal(iter_indices(i), [0, 3, 1, 4, 2, 5])
  276. # 2D reversed C-order
  277. i = nditer(a.reshape(2, 3)[::-1], ['c_index'], [['readonly']])
  278. assert_equal(iter_indices(i), [3, 4, 5, 0, 1, 2])
  279. i = nditer(a.reshape(2, 3)[:, ::-1], ['c_index'], [['readonly']])
  280. assert_equal(iter_indices(i), [2, 1, 0, 5, 4, 3])
  281. i = nditer(a.reshape(2, 3)[::-1, ::-1], ['c_index'], [['readonly']])
  282. assert_equal(iter_indices(i), [5, 4, 3, 2, 1, 0])
  283. # 2D reversed Fortran-order
  284. i = nditer(a.reshape(2, 3).copy(order='F')[::-1],
  285. ['c_index'], [['readonly']])
  286. assert_equal(iter_indices(i), [3, 0, 4, 1, 5, 2])
  287. i = nditer(a.reshape(2, 3).copy(order='F')[:, ::-1],
  288. ['c_index'], [['readonly']])
  289. assert_equal(iter_indices(i), [2, 5, 1, 4, 0, 3])
  290. i = nditer(a.reshape(2, 3).copy(order='F')[::-1, ::-1],
  291. ['c_index'], [['readonly']])
  292. assert_equal(iter_indices(i), [5, 2, 4, 1, 3, 0])
  293. def test_iter_best_order_c_index_3d():
  294. # The C index should be correct with any reordering
  295. a = arange(12)
  296. # 3D C-order
  297. i = nditer(a.reshape(2, 3, 2), ['c_index'], [['readonly']])
  298. assert_equal(iter_indices(i),
  299. [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11])
  300. # 3D Fortran-order
  301. i = nditer(a.reshape(2, 3, 2).copy(order='F'),
  302. ['c_index'], [['readonly']])
  303. assert_equal(iter_indices(i),
  304. [0, 6, 2, 8, 4, 10, 1, 7, 3, 9, 5, 11])
  305. # 3D reversed C-order
  306. i = nditer(a.reshape(2, 3, 2)[::-1], ['c_index'], [['readonly']])
  307. assert_equal(iter_indices(i),
  308. [6, 7, 8, 9, 10, 11, 0, 1, 2, 3, 4, 5])
  309. i = nditer(a.reshape(2, 3, 2)[:, ::-1], ['c_index'], [['readonly']])
  310. assert_equal(iter_indices(i),
  311. [4, 5, 2, 3, 0, 1, 10, 11, 8, 9, 6, 7])
  312. i = nditer(a.reshape(2, 3, 2)[:,:, ::-1], ['c_index'], [['readonly']])
  313. assert_equal(iter_indices(i),
  314. [1, 0, 3, 2, 5, 4, 7, 6, 9, 8, 11, 10])
  315. # 3D reversed Fortran-order
  316. i = nditer(a.reshape(2, 3, 2).copy(order='F')[::-1],
  317. ['c_index'], [['readonly']])
  318. assert_equal(iter_indices(i),
  319. [6, 0, 8, 2, 10, 4, 7, 1, 9, 3, 11, 5])
  320. i = nditer(a.reshape(2, 3, 2).copy(order='F')[:, ::-1],
  321. ['c_index'], [['readonly']])
  322. assert_equal(iter_indices(i),
  323. [4, 10, 2, 8, 0, 6, 5, 11, 3, 9, 1, 7])
  324. i = nditer(a.reshape(2, 3, 2).copy(order='F')[:,:, ::-1],
  325. ['c_index'], [['readonly']])
  326. assert_equal(iter_indices(i),
  327. [1, 7, 3, 9, 5, 11, 0, 6, 2, 8, 4, 10])
  328. def test_iter_best_order_f_index_1d():
  329. # The Fortran index should be correct with any reordering
  330. a = arange(4)
  331. # 1D order
  332. i = nditer(a, ['f_index'], [['readonly']])
  333. assert_equal(iter_indices(i), [0, 1, 2, 3])
  334. # 1D reversed order
  335. i = nditer(a[::-1], ['f_index'], [['readonly']])
  336. assert_equal(iter_indices(i), [3, 2, 1, 0])
  337. def test_iter_best_order_f_index_2d():
  338. # The Fortran index should be correct with any reordering
  339. a = arange(6)
  340. # 2D C-order
  341. i = nditer(a.reshape(2, 3), ['f_index'], [['readonly']])
  342. assert_equal(iter_indices(i), [0, 2, 4, 1, 3, 5])
  343. # 2D Fortran-order
  344. i = nditer(a.reshape(2, 3).copy(order='F'),
  345. ['f_index'], [['readonly']])
  346. assert_equal(iter_indices(i), [0, 1, 2, 3, 4, 5])
  347. # 2D reversed C-order
  348. i = nditer(a.reshape(2, 3)[::-1], ['f_index'], [['readonly']])
  349. assert_equal(iter_indices(i), [1, 3, 5, 0, 2, 4])
  350. i = nditer(a.reshape(2, 3)[:, ::-1], ['f_index'], [['readonly']])
  351. assert_equal(iter_indices(i), [4, 2, 0, 5, 3, 1])
  352. i = nditer(a.reshape(2, 3)[::-1, ::-1], ['f_index'], [['readonly']])
  353. assert_equal(iter_indices(i), [5, 3, 1, 4, 2, 0])
  354. # 2D reversed Fortran-order
  355. i = nditer(a.reshape(2, 3).copy(order='F')[::-1],
  356. ['f_index'], [['readonly']])
  357. assert_equal(iter_indices(i), [1, 0, 3, 2, 5, 4])
  358. i = nditer(a.reshape(2, 3).copy(order='F')[:, ::-1],
  359. ['f_index'], [['readonly']])
  360. assert_equal(iter_indices(i), [4, 5, 2, 3, 0, 1])
  361. i = nditer(a.reshape(2, 3).copy(order='F')[::-1, ::-1],
  362. ['f_index'], [['readonly']])
  363. assert_equal(iter_indices(i), [5, 4, 3, 2, 1, 0])
  364. def test_iter_best_order_f_index_3d():
  365. # The Fortran index should be correct with any reordering
  366. a = arange(12)
  367. # 3D C-order
  368. i = nditer(a.reshape(2, 3, 2), ['f_index'], [['readonly']])
  369. assert_equal(iter_indices(i),
  370. [0, 6, 2, 8, 4, 10, 1, 7, 3, 9, 5, 11])
  371. # 3D Fortran-order
  372. i = nditer(a.reshape(2, 3, 2).copy(order='F'),
  373. ['f_index'], [['readonly']])
  374. assert_equal(iter_indices(i),
  375. [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11])
  376. # 3D reversed C-order
  377. i = nditer(a.reshape(2, 3, 2)[::-1], ['f_index'], [['readonly']])
  378. assert_equal(iter_indices(i),
  379. [1, 7, 3, 9, 5, 11, 0, 6, 2, 8, 4, 10])
  380. i = nditer(a.reshape(2, 3, 2)[:, ::-1], ['f_index'], [['readonly']])
  381. assert_equal(iter_indices(i),
  382. [4, 10, 2, 8, 0, 6, 5, 11, 3, 9, 1, 7])
  383. i = nditer(a.reshape(2, 3, 2)[:,:, ::-1], ['f_index'], [['readonly']])
  384. assert_equal(iter_indices(i),
  385. [6, 0, 8, 2, 10, 4, 7, 1, 9, 3, 11, 5])
  386. # 3D reversed Fortran-order
  387. i = nditer(a.reshape(2, 3, 2).copy(order='F')[::-1],
  388. ['f_index'], [['readonly']])
  389. assert_equal(iter_indices(i),
  390. [1, 0, 3, 2, 5, 4, 7, 6, 9, 8, 11, 10])
  391. i = nditer(a.reshape(2, 3, 2).copy(order='F')[:, ::-1],
  392. ['f_index'], [['readonly']])
  393. assert_equal(iter_indices(i),
  394. [4, 5, 2, 3, 0, 1, 10, 11, 8, 9, 6, 7])
  395. i = nditer(a.reshape(2, 3, 2).copy(order='F')[:,:, ::-1],
  396. ['f_index'], [['readonly']])
  397. assert_equal(iter_indices(i),
  398. [6, 7, 8, 9, 10, 11, 0, 1, 2, 3, 4, 5])
  399. def test_iter_no_inner_full_coalesce():
  400. # Check no_inner iterators which coalesce into a single inner loop
  401. for shape in [(5,), (3, 4), (2, 3, 4), (2, 3, 4, 3), (2, 3, 2, 2, 3)]:
  402. size = np.prod(shape)
  403. a = arange(size)
  404. # Test each combination of forward and backwards indexing
  405. for dirs in range(2**len(shape)):
  406. dirs_index = [slice(None)]*len(shape)
  407. for bit in range(len(shape)):
  408. if ((2**bit) & dirs):
  409. dirs_index[bit] = slice(None, None, -1)
  410. dirs_index = tuple(dirs_index)
  411. aview = a.reshape(shape)[dirs_index]
  412. # C-order
  413. i = nditer(aview, ['external_loop'], [['readonly']])
  414. assert_equal(i.ndim, 1)
  415. assert_equal(i[0].shape, (size,))
  416. # Fortran-order
  417. i = nditer(aview.T, ['external_loop'], [['readonly']])
  418. assert_equal(i.ndim, 1)
  419. assert_equal(i[0].shape, (size,))
  420. # Other order
  421. if len(shape) > 2:
  422. i = nditer(aview.swapaxes(0, 1),
  423. ['external_loop'], [['readonly']])
  424. assert_equal(i.ndim, 1)
  425. assert_equal(i[0].shape, (size,))
  426. def test_iter_no_inner_dim_coalescing():
  427. # Check no_inner iterators whose dimensions may not coalesce completely
  428. # Skipping the last element in a dimension prevents coalescing
  429. # with the next-bigger dimension
  430. a = arange(24).reshape(2, 3, 4)[:,:, :-1]
  431. i = nditer(a, ['external_loop'], [['readonly']])
  432. assert_equal(i.ndim, 2)
  433. assert_equal(i[0].shape, (3,))
  434. a = arange(24).reshape(2, 3, 4)[:, :-1,:]
  435. i = nditer(a, ['external_loop'], [['readonly']])
  436. assert_equal(i.ndim, 2)
  437. assert_equal(i[0].shape, (8,))
  438. a = arange(24).reshape(2, 3, 4)[:-1,:,:]
  439. i = nditer(a, ['external_loop'], [['readonly']])
  440. assert_equal(i.ndim, 1)
  441. assert_equal(i[0].shape, (12,))
  442. # Even with lots of 1-sized dimensions, should still coalesce
  443. a = arange(24).reshape(1, 1, 2, 1, 1, 3, 1, 1, 4, 1, 1)
  444. i = nditer(a, ['external_loop'], [['readonly']])
  445. assert_equal(i.ndim, 1)
  446. assert_equal(i[0].shape, (24,))
  447. def test_iter_dim_coalescing():
  448. # Check that the correct number of dimensions are coalesced
  449. # Tracking a multi-index disables coalescing
  450. a = arange(24).reshape(2, 3, 4)
  451. i = nditer(a, ['multi_index'], [['readonly']])
  452. assert_equal(i.ndim, 3)
  453. # A tracked index can allow coalescing if it's compatible with the array
  454. a3d = arange(24).reshape(2, 3, 4)
  455. i = nditer(a3d, ['c_index'], [['readonly']])
  456. assert_equal(i.ndim, 1)
  457. i = nditer(a3d.swapaxes(0, 1), ['c_index'], [['readonly']])
  458. assert_equal(i.ndim, 3)
  459. i = nditer(a3d.T, ['c_index'], [['readonly']])
  460. assert_equal(i.ndim, 3)
  461. i = nditer(a3d.T, ['f_index'], [['readonly']])
  462. assert_equal(i.ndim, 1)
  463. i = nditer(a3d.T.swapaxes(0, 1), ['f_index'], [['readonly']])
  464. assert_equal(i.ndim, 3)
  465. # When C or F order is forced, coalescing may still occur
  466. a3d = arange(24).reshape(2, 3, 4)
  467. i = nditer(a3d, order='C')
  468. assert_equal(i.ndim, 1)
  469. i = nditer(a3d.T, order='C')
  470. assert_equal(i.ndim, 3)
  471. i = nditer(a3d, order='F')
  472. assert_equal(i.ndim, 3)
  473. i = nditer(a3d.T, order='F')
  474. assert_equal(i.ndim, 1)
  475. i = nditer(a3d, order='A')
  476. assert_equal(i.ndim, 1)
  477. i = nditer(a3d.T, order='A')
  478. assert_equal(i.ndim, 1)
  479. def test_iter_broadcasting():
  480. # Standard NumPy broadcasting rules
  481. # 1D with scalar
  482. i = nditer([arange(6), np.int32(2)], ['multi_index'], [['readonly']]*2)
  483. assert_equal(i.itersize, 6)
  484. assert_equal(i.shape, (6,))
  485. # 2D with scalar
  486. i = nditer([arange(6).reshape(2, 3), np.int32(2)],
  487. ['multi_index'], [['readonly']]*2)
  488. assert_equal(i.itersize, 6)
  489. assert_equal(i.shape, (2, 3))
  490. # 2D with 1D
  491. i = nditer([arange(6).reshape(2, 3), arange(3)],
  492. ['multi_index'], [['readonly']]*2)
  493. assert_equal(i.itersize, 6)
  494. assert_equal(i.shape, (2, 3))
  495. i = nditer([arange(2).reshape(2, 1), arange(3)],
  496. ['multi_index'], [['readonly']]*2)
  497. assert_equal(i.itersize, 6)
  498. assert_equal(i.shape, (2, 3))
  499. # 2D with 2D
  500. i = nditer([arange(2).reshape(2, 1), arange(3).reshape(1, 3)],
  501. ['multi_index'], [['readonly']]*2)
  502. assert_equal(i.itersize, 6)
  503. assert_equal(i.shape, (2, 3))
  504. # 3D with scalar
  505. i = nditer([np.int32(2), arange(24).reshape(4, 2, 3)],
  506. ['multi_index'], [['readonly']]*2)
  507. assert_equal(i.itersize, 24)
  508. assert_equal(i.shape, (4, 2, 3))
  509. # 3D with 1D
  510. i = nditer([arange(3), arange(24).reshape(4, 2, 3)],
  511. ['multi_index'], [['readonly']]*2)
  512. assert_equal(i.itersize, 24)
  513. assert_equal(i.shape, (4, 2, 3))
  514. i = nditer([arange(3), arange(8).reshape(4, 2, 1)],
  515. ['multi_index'], [['readonly']]*2)
  516. assert_equal(i.itersize, 24)
  517. assert_equal(i.shape, (4, 2, 3))
  518. # 3D with 2D
  519. i = nditer([arange(6).reshape(2, 3), arange(24).reshape(4, 2, 3)],
  520. ['multi_index'], [['readonly']]*2)
  521. assert_equal(i.itersize, 24)
  522. assert_equal(i.shape, (4, 2, 3))
  523. i = nditer([arange(2).reshape(2, 1), arange(24).reshape(4, 2, 3)],
  524. ['multi_index'], [['readonly']]*2)
  525. assert_equal(i.itersize, 24)
  526. assert_equal(i.shape, (4, 2, 3))
  527. i = nditer([arange(3).reshape(1, 3), arange(8).reshape(4, 2, 1)],
  528. ['multi_index'], [['readonly']]*2)
  529. assert_equal(i.itersize, 24)
  530. assert_equal(i.shape, (4, 2, 3))
  531. # 3D with 3D
  532. i = nditer([arange(2).reshape(1, 2, 1), arange(3).reshape(1, 1, 3),
  533. arange(4).reshape(4, 1, 1)],
  534. ['multi_index'], [['readonly']]*3)
  535. assert_equal(i.itersize, 24)
  536. assert_equal(i.shape, (4, 2, 3))
  537. i = nditer([arange(6).reshape(1, 2, 3), arange(4).reshape(4, 1, 1)],
  538. ['multi_index'], [['readonly']]*2)
  539. assert_equal(i.itersize, 24)
  540. assert_equal(i.shape, (4, 2, 3))
  541. i = nditer([arange(24).reshape(4, 2, 3), arange(12).reshape(4, 1, 3)],
  542. ['multi_index'], [['readonly']]*2)
  543. assert_equal(i.itersize, 24)
  544. assert_equal(i.shape, (4, 2, 3))
  545. def test_iter_itershape():
  546. # Check that allocated outputs work with a specified shape
  547. a = np.arange(6, dtype='i2').reshape(2, 3)
  548. i = nditer([a, None], [], [['readonly'], ['writeonly', 'allocate']],
  549. op_axes=[[0, 1, None], None],
  550. itershape=(-1, -1, 4))
  551. assert_equal(i.operands[1].shape, (2, 3, 4))
  552. assert_equal(i.operands[1].strides, (24, 8, 2))
  553. i = nditer([a.T, None], [], [['readonly'], ['writeonly', 'allocate']],
  554. op_axes=[[0, 1, None], None],
  555. itershape=(-1, -1, 4))
  556. assert_equal(i.operands[1].shape, (3, 2, 4))
  557. assert_equal(i.operands[1].strides, (8, 24, 2))
  558. i = nditer([a.T, None], [], [['readonly'], ['writeonly', 'allocate']],
  559. order='F',
  560. op_axes=[[0, 1, None], None],
  561. itershape=(-1, -1, 4))
  562. assert_equal(i.operands[1].shape, (3, 2, 4))
  563. assert_equal(i.operands[1].strides, (2, 6, 12))
  564. # If we specify 1 in the itershape, it shouldn't allow broadcasting
  565. # of that dimension to a bigger value
  566. assert_raises(ValueError, nditer, [a, None], [],
  567. [['readonly'], ['writeonly', 'allocate']],
  568. op_axes=[[0, 1, None], None],
  569. itershape=(-1, 1, 4))
  570. # Test bug that for no op_axes but itershape, they are NULLed correctly
  571. i = np.nditer([np.ones(2), None, None], itershape=(2,))
  572. def test_iter_broadcasting_errors():
  573. # Check that errors are thrown for bad broadcasting shapes
  574. # 1D with 1D
  575. assert_raises(ValueError, nditer, [arange(2), arange(3)],
  576. [], [['readonly']]*2)
  577. # 2D with 1D
  578. assert_raises(ValueError, nditer,
  579. [arange(6).reshape(2, 3), arange(2)],
  580. [], [['readonly']]*2)
  581. # 2D with 2D
  582. assert_raises(ValueError, nditer,
  583. [arange(6).reshape(2, 3), arange(9).reshape(3, 3)],
  584. [], [['readonly']]*2)
  585. assert_raises(ValueError, nditer,
  586. [arange(6).reshape(2, 3), arange(4).reshape(2, 2)],
  587. [], [['readonly']]*2)
  588. # 3D with 3D
  589. assert_raises(ValueError, nditer,
  590. [arange(36).reshape(3, 3, 4), arange(24).reshape(2, 3, 4)],
  591. [], [['readonly']]*2)
  592. assert_raises(ValueError, nditer,
  593. [arange(8).reshape(2, 4, 1), arange(24).reshape(2, 3, 4)],
  594. [], [['readonly']]*2)
  595. # Verify that the error message mentions the right shapes
  596. try:
  597. nditer([arange(2).reshape(1, 2, 1),
  598. arange(3).reshape(1, 3),
  599. arange(6).reshape(2, 3)],
  600. [],
  601. [['readonly'], ['readonly'], ['writeonly', 'no_broadcast']])
  602. raise AssertionError('Should have raised a broadcast error')
  603. except ValueError as e:
  604. msg = str(e)
  605. # The message should contain the shape of the 3rd operand
  606. assert_(msg.find('(2,3)') >= 0,
  607. 'Message "%s" doesn\'t contain operand shape (2,3)' % msg)
  608. # The message should contain the broadcast shape
  609. assert_(msg.find('(1,2,3)') >= 0,
  610. 'Message "%s" doesn\'t contain broadcast shape (1,2,3)' % msg)
  611. try:
  612. nditer([arange(6).reshape(2, 3), arange(2)],
  613. [],
  614. [['readonly'], ['readonly']],
  615. op_axes=[[0, 1], [0, np.newaxis]],
  616. itershape=(4, 3))
  617. raise AssertionError('Should have raised a broadcast error')
  618. except ValueError as e:
  619. msg = str(e)
  620. # The message should contain "shape->remappedshape" for each operand
  621. assert_(msg.find('(2,3)->(2,3)') >= 0,
  622. 'Message "%s" doesn\'t contain operand shape (2,3)->(2,3)' % msg)
  623. assert_(msg.find('(2,)->(2,newaxis)') >= 0,
  624. ('Message "%s" doesn\'t contain remapped operand shape' +
  625. '(2,)->(2,newaxis)') % msg)
  626. # The message should contain the itershape parameter
  627. assert_(msg.find('(4,3)') >= 0,
  628. 'Message "%s" doesn\'t contain itershape parameter (4,3)' % msg)
  629. try:
  630. nditer([np.zeros((2, 1, 1)), np.zeros((2,))],
  631. [],
  632. [['writeonly', 'no_broadcast'], ['readonly']])
  633. raise AssertionError('Should have raised a broadcast error')
  634. except ValueError as e:
  635. msg = str(e)
  636. # The message should contain the shape of the bad operand
  637. assert_(msg.find('(2,1,1)') >= 0,
  638. 'Message "%s" doesn\'t contain operand shape (2,1,1)' % msg)
  639. # The message should contain the broadcast shape
  640. assert_(msg.find('(2,1,2)') >= 0,
  641. 'Message "%s" doesn\'t contain the broadcast shape (2,1,2)' % msg)
  642. def test_iter_flags_errors():
  643. # Check that bad combinations of flags produce errors
  644. a = arange(6)
  645. # Not enough operands
  646. assert_raises(ValueError, nditer, [], [], [])
  647. # Too many operands
  648. assert_raises(ValueError, nditer, [a]*100, [], [['readonly']]*100)
  649. # Bad global flag
  650. assert_raises(ValueError, nditer, [a], ['bad flag'], [['readonly']])
  651. # Bad op flag
  652. assert_raises(ValueError, nditer, [a], [], [['readonly', 'bad flag']])
  653. # Bad order parameter
  654. assert_raises(ValueError, nditer, [a], [], [['readonly']], order='G')
  655. # Bad casting parameter
  656. assert_raises(ValueError, nditer, [a], [], [['readonly']], casting='noon')
  657. # op_flags must match ops
  658. assert_raises(ValueError, nditer, [a]*3, [], [['readonly']]*2)
  659. # Cannot track both a C and an F index
  660. assert_raises(ValueError, nditer, a,
  661. ['c_index', 'f_index'], [['readonly']])
  662. # Inner iteration and multi-indices/indices are incompatible
  663. assert_raises(ValueError, nditer, a,
  664. ['external_loop', 'multi_index'], [['readonly']])
  665. assert_raises(ValueError, nditer, a,
  666. ['external_loop', 'c_index'], [['readonly']])
  667. assert_raises(ValueError, nditer, a,
  668. ['external_loop', 'f_index'], [['readonly']])
  669. # Must specify exactly one of readwrite/readonly/writeonly per operand
  670. assert_raises(ValueError, nditer, a, [], [[]])
  671. assert_raises(ValueError, nditer, a, [], [['readonly', 'writeonly']])
  672. assert_raises(ValueError, nditer, a, [], [['readonly', 'readwrite']])
  673. assert_raises(ValueError, nditer, a, [], [['writeonly', 'readwrite']])
  674. assert_raises(ValueError, nditer, a,
  675. [], [['readonly', 'writeonly', 'readwrite']])
  676. # Python scalars are always readonly
  677. assert_raises(TypeError, nditer, 1.5, [], [['writeonly']])
  678. assert_raises(TypeError, nditer, 1.5, [], [['readwrite']])
  679. # Array scalars are always readonly
  680. assert_raises(TypeError, nditer, np.int32(1), [], [['writeonly']])
  681. assert_raises(TypeError, nditer, np.int32(1), [], [['readwrite']])
  682. # Check readonly array
  683. a.flags.writeable = False
  684. assert_raises(ValueError, nditer, a, [], [['writeonly']])
  685. assert_raises(ValueError, nditer, a, [], [['readwrite']])
  686. a.flags.writeable = True
  687. # Multi-indices available only with the multi_index flag
  688. i = nditer(arange(6), [], [['readonly']])
  689. assert_raises(ValueError, lambda i:i.multi_index, i)
  690. # Index available only with an index flag
  691. assert_raises(ValueError, lambda i:i.index, i)
  692. # GotoCoords and GotoIndex incompatible with buffering or no_inner
  693. def assign_multi_index(i):
  694. i.multi_index = (0,)
  695. def assign_index(i):
  696. i.index = 0
  697. def assign_iterindex(i):
  698. i.iterindex = 0
  699. def assign_iterrange(i):
  700. i.iterrange = (0, 1)
  701. i = nditer(arange(6), ['external_loop'])
  702. assert_raises(ValueError, assign_multi_index, i)
  703. assert_raises(ValueError, assign_index, i)
  704. assert_raises(ValueError, assign_iterindex, i)
  705. assert_raises(ValueError, assign_iterrange, i)
  706. i = nditer(arange(6), ['buffered'])
  707. assert_raises(ValueError, assign_multi_index, i)
  708. assert_raises(ValueError, assign_index, i)
  709. assert_raises(ValueError, assign_iterrange, i)
  710. # Can't iterate if size is zero
  711. assert_raises(ValueError, nditer, np.array([]))
  712. def test_iter_slice():
  713. a, b, c = np.arange(3), np.arange(3), np.arange(3.)
  714. i = nditer([a, b, c], [], ['readwrite'])
  715. with i:
  716. i[0:2] = (3, 3)
  717. assert_equal(a, [3, 1, 2])
  718. assert_equal(b, [3, 1, 2])
  719. assert_equal(c, [0, 1, 2])
  720. i[1] = 12
  721. assert_equal(i[0:2], [3, 12])
  722. def test_iter_assign_mapping():
  723. a = np.arange(24, dtype='f8').reshape(2, 3, 4).T
  724. it = np.nditer(a, [], [['readwrite', 'updateifcopy']],
  725. casting='same_kind', op_dtypes=[np.dtype('f4')])
  726. with it:
  727. it.operands[0][...] = 3
  728. it.operands[0][...] = 14
  729. assert_equal(a, 14)
  730. it = np.nditer(a, [], [['readwrite', 'updateifcopy']],
  731. casting='same_kind', op_dtypes=[np.dtype('f4')])
  732. with it:
  733. x = it.operands[0][-1:1]
  734. x[...] = 14
  735. it.operands[0][...] = -1234
  736. assert_equal(a, -1234)
  737. # check for no warnings on dealloc
  738. x = None
  739. it = None
  740. def test_iter_nbo_align_contig():
  741. # Check that byte order, alignment, and contig changes work
  742. # Byte order change by requesting a specific dtype
  743. a = np.arange(6, dtype='f4')
  744. au = a.byteswap().newbyteorder()
  745. assert_(a.dtype.byteorder != au.dtype.byteorder)
  746. i = nditer(au, [], [['readwrite', 'updateifcopy']],
  747. casting='equiv',
  748. op_dtypes=[np.dtype('f4')])
  749. with i:
  750. # context manager triggers WRITEBACKIFCOPY on i at exit
  751. assert_equal(i.dtypes[0].byteorder, a.dtype.byteorder)
  752. assert_equal(i.operands[0].dtype.byteorder, a.dtype.byteorder)
  753. assert_equal(i.operands[0], a)
  754. i.operands[0][:] = 2
  755. assert_equal(au, [2]*6)
  756. del i # should not raise a warning
  757. # Byte order change by requesting NBO
  758. a = np.arange(6, dtype='f4')
  759. au = a.byteswap().newbyteorder()
  760. assert_(a.dtype.byteorder != au.dtype.byteorder)
  761. with nditer(au, [], [['readwrite', 'updateifcopy', 'nbo']],
  762. casting='equiv') as i:
  763. # context manager triggers UPDATEIFCOPY on i at exit
  764. assert_equal(i.dtypes[0].byteorder, a.dtype.byteorder)
  765. assert_equal(i.operands[0].dtype.byteorder, a.dtype.byteorder)
  766. assert_equal(i.operands[0], a)
  767. i.operands[0][:] = 12345
  768. i.operands[0][:] = 2
  769. assert_equal(au, [2]*6)
  770. # Unaligned input
  771. a = np.zeros((6*4+1,), dtype='i1')[1:]
  772. a.dtype = 'f4'
  773. a[:] = np.arange(6, dtype='f4')
  774. assert_(not a.flags.aligned)
  775. # Without 'aligned', shouldn't copy
  776. i = nditer(a, [], [['readonly']])
  777. assert_(not i.operands[0].flags.aligned)
  778. assert_equal(i.operands[0], a)
  779. # With 'aligned', should make a copy
  780. with nditer(a, [], [['readwrite', 'updateifcopy', 'aligned']]) as i:
  781. assert_(i.operands[0].flags.aligned)
  782. # context manager triggers UPDATEIFCOPY on i at exit
  783. assert_equal(i.operands[0], a)
  784. i.operands[0][:] = 3
  785. assert_equal(a, [3]*6)
  786. # Discontiguous input
  787. a = arange(12)
  788. # If it is contiguous, shouldn't copy
  789. i = nditer(a[:6], [], [['readonly']])
  790. assert_(i.operands[0].flags.contiguous)
  791. assert_equal(i.operands[0], a[:6])
  792. # If it isn't contiguous, should buffer
  793. i = nditer(a[::2], ['buffered', 'external_loop'],
  794. [['readonly', 'contig']],
  795. buffersize=10)
  796. assert_(i[0].flags.contiguous)
  797. assert_equal(i[0], a[::2])
  798. def test_iter_array_cast():
  799. # Check that arrays are cast as requested
  800. # No cast 'f4' -> 'f4'
  801. a = np.arange(6, dtype='f4').reshape(2, 3)
  802. i = nditer(a, [], [['readwrite']], op_dtypes=[np.dtype('f4')])
  803. with i:
  804. assert_equal(i.operands[0], a)
  805. assert_equal(i.operands[0].dtype, np.dtype('f4'))
  806. # Byte-order cast '<f4' -> '>f4'
  807. a = np.arange(6, dtype='<f4').reshape(2, 3)
  808. with nditer(a, [], [['readwrite', 'updateifcopy']],
  809. casting='equiv',
  810. op_dtypes=[np.dtype('>f4')]) as i:
  811. assert_equal(i.operands[0], a)
  812. assert_equal(i.operands[0].dtype, np.dtype('>f4'))
  813. # Safe case 'f4' -> 'f8'
  814. a = np.arange(24, dtype='f4').reshape(2, 3, 4).swapaxes(1, 2)
  815. i = nditer(a, [], [['readonly', 'copy']],
  816. casting='safe',
  817. op_dtypes=[np.dtype('f8')])
  818. assert_equal(i.operands[0], a)
  819. assert_equal(i.operands[0].dtype, np.dtype('f8'))
  820. # The memory layout of the temporary should match a (a is (48,4,16))
  821. # except negative strides get flipped to positive strides.
  822. assert_equal(i.operands[0].strides, (96, 8, 32))
  823. a = a[::-1,:, ::-1]
  824. i = nditer(a, [], [['readonly', 'copy']],
  825. casting='safe',
  826. op_dtypes=[np.dtype('f8')])
  827. assert_equal(i.operands[0], a)
  828. assert_equal(i.operands[0].dtype, np.dtype('f8'))
  829. assert_equal(i.operands[0].strides, (96, 8, 32))
  830. # Same-kind cast 'f8' -> 'f4' -> 'f8'
  831. a = np.arange(24, dtype='f8').reshape(2, 3, 4).T
  832. with nditer(a, [],
  833. [['readwrite', 'updateifcopy']],
  834. casting='same_kind',
  835. op_dtypes=[np.dtype('f4')]) as i:
  836. assert_equal(i.operands[0], a)
  837. assert_equal(i.operands[0].dtype, np.dtype('f4'))
  838. assert_equal(i.operands[0].strides, (4, 16, 48))
  839. # Check that WRITEBACKIFCOPY is activated at exit
  840. i.operands[0][2, 1, 1] = -12.5
  841. assert_(a[2, 1, 1] != -12.5)
  842. assert_equal(a[2, 1, 1], -12.5)
  843. a = np.arange(6, dtype='i4')[::-2]
  844. with nditer(a, [],
  845. [['writeonly', 'updateifcopy']],
  846. casting='unsafe',
  847. op_dtypes=[np.dtype('f4')]) as i:
  848. assert_equal(i.operands[0].dtype, np.dtype('f4'))
  849. # Even though the stride was negative in 'a', it
  850. # becomes positive in the temporary
  851. assert_equal(i.operands[0].strides, (4,))
  852. i.operands[0][:] = [1, 2, 3]
  853. assert_equal(a, [1, 2, 3])
  854. def test_iter_array_cast_errors():
  855. # Check that invalid casts are caught
  856. # Need to enable copying for casts to occur
  857. assert_raises(TypeError, nditer, arange(2, dtype='f4'), [],
  858. [['readonly']], op_dtypes=[np.dtype('f8')])
  859. # Also need to allow casting for casts to occur
  860. assert_raises(TypeError, nditer, arange(2, dtype='f4'), [],
  861. [['readonly', 'copy']], casting='no',
  862. op_dtypes=[np.dtype('f8')])
  863. assert_raises(TypeError, nditer, arange(2, dtype='f4'), [],
  864. [['readonly', 'copy']], casting='equiv',
  865. op_dtypes=[np.dtype('f8')])
  866. assert_raises(TypeError, nditer, arange(2, dtype='f8'), [],
  867. [['writeonly', 'updateifcopy']],
  868. casting='no',
  869. op_dtypes=[np.dtype('f4')])
  870. assert_raises(TypeError, nditer, arange(2, dtype='f8'), [],
  871. [['writeonly', 'updateifcopy']],
  872. casting='equiv',
  873. op_dtypes=[np.dtype('f4')])
  874. # '<f4' -> '>f4' should not work with casting='no'
  875. assert_raises(TypeError, nditer, arange(2, dtype='<f4'), [],
  876. [['readonly', 'copy']], casting='no',
  877. op_dtypes=[np.dtype('>f4')])
  878. # 'f4' -> 'f8' is a safe cast, but 'f8' -> 'f4' isn't
  879. assert_raises(TypeError, nditer, arange(2, dtype='f4'), [],
  880. [['readwrite', 'updateifcopy']],
  881. casting='safe',
  882. op_dtypes=[np.dtype('f8')])
  883. assert_raises(TypeError, nditer, arange(2, dtype='f8'), [],
  884. [['readwrite', 'updateifcopy']],
  885. casting='safe',
  886. op_dtypes=[np.dtype('f4')])
  887. # 'f4' -> 'i4' is neither a safe nor a same-kind cast
  888. assert_raises(TypeError, nditer, arange(2, dtype='f4'), [],
  889. [['readonly', 'copy']],
  890. casting='same_kind',
  891. op_dtypes=[np.dtype('i4')])
  892. assert_raises(TypeError, nditer, arange(2, dtype='i4'), [],
  893. [['writeonly', 'updateifcopy']],
  894. casting='same_kind',
  895. op_dtypes=[np.dtype('f4')])
  896. def test_iter_scalar_cast():
  897. # Check that scalars are cast as requested
  898. # No cast 'f4' -> 'f4'
  899. i = nditer(np.float32(2.5), [], [['readonly']],
  900. op_dtypes=[np.dtype('f4')])
  901. assert_equal(i.dtypes[0], np.dtype('f4'))
  902. assert_equal(i.value.dtype, np.dtype('f4'))
  903. assert_equal(i.value, 2.5)
  904. # Safe cast 'f4' -> 'f8'
  905. i = nditer(np.float32(2.5), [],
  906. [['readonly', 'copy']],
  907. casting='safe',
  908. op_dtypes=[np.dtype('f8')])
  909. assert_equal(i.dtypes[0], np.dtype('f8'))
  910. assert_equal(i.value.dtype, np.dtype('f8'))
  911. assert_equal(i.value, 2.5)
  912. # Same-kind cast 'f8' -> 'f4'
  913. i = nditer(np.float64(2.5), [],
  914. [['readonly', 'copy']],
  915. casting='same_kind',
  916. op_dtypes=[np.dtype('f4')])
  917. assert_equal(i.dtypes[0], np.dtype('f4'))
  918. assert_equal(i.value.dtype, np.dtype('f4'))
  919. assert_equal(i.value, 2.5)
  920. # Unsafe cast 'f8' -> 'i4'
  921. i = nditer(np.float64(3.0), [],
  922. [['readonly', 'copy']],
  923. casting='unsafe',
  924. op_dtypes=[np.dtype('i4')])
  925. assert_equal(i.dtypes[0], np.dtype('i4'))
  926. assert_equal(i.value.dtype, np.dtype('i4'))
  927. assert_equal(i.value, 3)
  928. # Readonly scalars may be cast even without setting COPY or BUFFERED
  929. i = nditer(3, [], [['readonly']], op_dtypes=[np.dtype('f8')])
  930. assert_equal(i[0].dtype, np.dtype('f8'))
  931. assert_equal(i[0], 3.)
  932. def test_iter_scalar_cast_errors():
  933. # Check that invalid casts are caught
  934. # Need to allow copying/buffering for write casts of scalars to occur
  935. assert_raises(TypeError, nditer, np.float32(2), [],
  936. [['readwrite']], op_dtypes=[np.dtype('f8')])
  937. assert_raises(TypeError, nditer, 2.5, [],
  938. [['readwrite']], op_dtypes=[np.dtype('f4')])
  939. # 'f8' -> 'f4' isn't a safe cast if the value would overflow
  940. assert_raises(TypeError, nditer, np.float64(1e60), [],
  941. [['readonly']],
  942. casting='safe',
  943. op_dtypes=[np.dtype('f4')])
  944. # 'f4' -> 'i4' is neither a safe nor a same-kind cast
  945. assert_raises(TypeError, nditer, np.float32(2), [],
  946. [['readonly']],
  947. casting='same_kind',
  948. op_dtypes=[np.dtype('i4')])
  949. def test_iter_object_arrays_basic():
  950. # Check that object arrays work
  951. obj = {'a':3,'b':'d'}
  952. a = np.array([[1, 2, 3], None, obj, None], dtype='O')
  953. if HAS_REFCOUNT:
  954. rc = sys.getrefcount(obj)
  955. # Need to allow references for object arrays
  956. assert_raises(TypeError, nditer, a)
  957. if HAS_REFCOUNT:
  958. assert_equal(sys.getrefcount(obj), rc)
  959. i = nditer(a, ['refs_ok'], ['readonly'])
  960. vals = [x_[()] for x_ in i]
  961. assert_equal(np.array(vals, dtype='O'), a)
  962. vals, i, x = [None]*3
  963. if HAS_REFCOUNT:
  964. assert_equal(sys.getrefcount(obj), rc)
  965. i = nditer(a.reshape(2, 2).T, ['refs_ok', 'buffered'],
  966. ['readonly'], order='C')
  967. assert_(i.iterationneedsapi)
  968. vals = [x_[()] for x_ in i]
  969. assert_equal(np.array(vals, dtype='O'), a.reshape(2, 2).ravel(order='F'))
  970. vals, i, x = [None]*3
  971. if HAS_REFCOUNT:
  972. assert_equal(sys.getrefcount(obj), rc)
  973. i = nditer(a.reshape(2, 2).T, ['refs_ok', 'buffered'],
  974. ['readwrite'], order='C')
  975. with i:
  976. for x in i:
  977. x[...] = None
  978. vals, i, x = [None]*3
  979. if HAS_REFCOUNT:
  980. assert_(sys.getrefcount(obj) == rc-1)
  981. assert_equal(a, np.array([None]*4, dtype='O'))
  982. def test_iter_object_arrays_conversions():
  983. # Conversions to/from objects
  984. a = np.arange(6, dtype='O')
  985. i = nditer(a, ['refs_ok', 'buffered'], ['readwrite'],
  986. casting='unsafe', op_dtypes='i4')
  987. with i:
  988. for x in i:
  989. x[...] += 1
  990. assert_equal(a, np.arange(6)+1)
  991. a = np.arange(6, dtype='i4')
  992. i = nditer(a, ['refs_ok', 'buffered'], ['readwrite'],
  993. casting='unsafe', op_dtypes='O')
  994. with i:
  995. for x in i:
  996. x[...] += 1
  997. assert_equal(a, np.arange(6)+1)
  998. # Non-contiguous object array
  999. a = np.zeros((6,), dtype=[('p', 'i1'), ('a', 'O')])
  1000. a = a['a']
  1001. a[:] = np.arange(6)
  1002. i = nditer(a, ['refs_ok', 'buffered'], ['readwrite'],
  1003. casting='unsafe', op_dtypes='i4')
  1004. with i:
  1005. for x in i:
  1006. x[...] += 1
  1007. assert_equal(a, np.arange(6)+1)
  1008. #Non-contiguous value array
  1009. a = np.zeros((6,), dtype=[('p', 'i1'), ('a', 'i4')])
  1010. a = a['a']
  1011. a[:] = np.arange(6) + 98172488
  1012. i = nditer(a, ['refs_ok', 'buffered'], ['readwrite'],
  1013. casting='unsafe', op_dtypes='O')
  1014. with i:
  1015. ob = i[0][()]
  1016. if HAS_REFCOUNT:
  1017. rc = sys.getrefcount(ob)
  1018. for x in i:
  1019. x[...] += 1
  1020. if HAS_REFCOUNT:
  1021. assert_(sys.getrefcount(ob) == rc-1)
  1022. assert_equal(a, np.arange(6)+98172489)
  1023. def test_iter_common_dtype():
  1024. # Check that the iterator finds a common data type correctly
  1025. i = nditer([array([3], dtype='f4'), array([0], dtype='f8')],
  1026. ['common_dtype'],
  1027. [['readonly', 'copy']]*2,
  1028. casting='safe')
  1029. assert_equal(i.dtypes[0], np.dtype('f8'))
  1030. assert_equal(i.dtypes[1], np.dtype('f8'))
  1031. i = nditer([array([3], dtype='i4'), array([0], dtype='f4')],
  1032. ['common_dtype'],
  1033. [['readonly', 'copy']]*2,
  1034. casting='safe')
  1035. assert_equal(i.dtypes[0], np.dtype('f8'))
  1036. assert_equal(i.dtypes[1], np.dtype('f8'))
  1037. i = nditer([array([3], dtype='f4'), array(0, dtype='f8')],
  1038. ['common_dtype'],
  1039. [['readonly', 'copy']]*2,
  1040. casting='same_kind')
  1041. assert_equal(i.dtypes[0], np.dtype('f4'))
  1042. assert_equal(i.dtypes[1], np.dtype('f4'))
  1043. i = nditer([array([3], dtype='u4'), array(0, dtype='i4')],
  1044. ['common_dtype'],
  1045. [['readonly', 'copy']]*2,
  1046. casting='safe')
  1047. assert_equal(i.dtypes[0], np.dtype('u4'))
  1048. assert_equal(i.dtypes[1], np.dtype('u4'))
  1049. i = nditer([array([3], dtype='u4'), array(-12, dtype='i4')],
  1050. ['common_dtype'],
  1051. [['readonly', 'copy']]*2,
  1052. casting='safe')
  1053. assert_equal(i.dtypes[0], np.dtype('i8'))
  1054. assert_equal(i.dtypes[1], np.dtype('i8'))
  1055. i = nditer([array([3], dtype='u4'), array(-12, dtype='i4'),
  1056. array([2j], dtype='c8'), array([9], dtype='f8')],
  1057. ['common_dtype'],
  1058. [['readonly', 'copy']]*4,
  1059. casting='safe')
  1060. assert_equal(i.dtypes[0], np.dtype('c16'))
  1061. assert_equal(i.dtypes[1], np.dtype('c16'))
  1062. assert_equal(i.dtypes[2], np.dtype('c16'))
  1063. assert_equal(i.dtypes[3], np.dtype('c16'))
  1064. assert_equal(i.value, (3, -12, 2j, 9))
  1065. # When allocating outputs, other outputs aren't factored in
  1066. i = nditer([array([3], dtype='i4'), None, array([2j], dtype='c16')], [],
  1067. [['readonly', 'copy'],
  1068. ['writeonly', 'allocate'],
  1069. ['writeonly']],
  1070. casting='safe')
  1071. assert_equal(i.dtypes[0], np.dtype('i4'))
  1072. assert_equal(i.dtypes[1], np.dtype('i4'))
  1073. assert_equal(i.dtypes[2], np.dtype('c16'))
  1074. # But, if common data types are requested, they are
  1075. i = nditer([array([3], dtype='i4'), None, array([2j], dtype='c16')],
  1076. ['common_dtype'],
  1077. [['readonly', 'copy'],
  1078. ['writeonly', 'allocate'],
  1079. ['writeonly']],
  1080. casting='safe')
  1081. assert_equal(i.dtypes[0], np.dtype('c16'))
  1082. assert_equal(i.dtypes[1], np.dtype('c16'))
  1083. assert_equal(i.dtypes[2], np.dtype('c16'))
  1084. def test_iter_copy_if_overlap():
  1085. # Ensure the iterator makes copies on read/write overlap, if requested
  1086. # Copy not needed, 1 op
  1087. for flag in ['readonly', 'writeonly', 'readwrite']:
  1088. a = arange(10)
  1089. i = nditer([a], ['copy_if_overlap'], [[flag]])
  1090. with i:
  1091. assert_(i.operands[0] is a)
  1092. # Copy needed, 2 ops, read-write overlap
  1093. x = arange(10)
  1094. a = x[1:]
  1095. b = x[:-1]
  1096. with nditer([a, b], ['copy_if_overlap'], [['readonly'], ['readwrite']]) as i:
  1097. assert_(not np.shares_memory(*i.operands))
  1098. # Copy not needed with elementwise, 2 ops, exactly same arrays
  1099. x = arange(10)
  1100. a = x
  1101. b = x
  1102. i = nditer([a, b], ['copy_if_overlap'], [['readonly', 'overlap_assume_elementwise'],
  1103. ['readwrite', 'overlap_assume_elementwise']])
  1104. with i:
  1105. assert_(i.operands[0] is a and i.operands[1] is b)
  1106. with nditer([a, b], ['copy_if_overlap'], [['readonly'], ['readwrite']]) as i:
  1107. assert_(i.operands[0] is a and not np.shares_memory(i.operands[1], b))
  1108. # Copy not needed, 2 ops, no overlap
  1109. x = arange(10)
  1110. a = x[::2]
  1111. b = x[1::2]
  1112. i = nditer([a, b], ['copy_if_overlap'], [['readonly'], ['writeonly']])
  1113. assert_(i.operands[0] is a and i.operands[1] is b)
  1114. # Copy needed, 2 ops, read-write overlap
  1115. x = arange(4, dtype=np.int8)
  1116. a = x[3:]
  1117. b = x.view(np.int32)[:1]
  1118. with nditer([a, b], ['copy_if_overlap'], [['readonly'], ['writeonly']]) as i:
  1119. assert_(not np.shares_memory(*i.operands))
  1120. # Copy needed, 3 ops, read-write overlap
  1121. for flag in ['writeonly', 'readwrite']:
  1122. x = np.ones([10, 10])
  1123. a = x
  1124. b = x.T
  1125. c = x
  1126. with nditer([a, b, c], ['copy_if_overlap'],
  1127. [['readonly'], ['readonly'], [flag]]) as i:
  1128. a2, b2, c2 = i.operands
  1129. assert_(not np.shares_memory(a2, c2))
  1130. assert_(not np.shares_memory(b2, c2))
  1131. # Copy not needed, 3 ops, read-only overlap
  1132. x = np.ones([10, 10])
  1133. a = x
  1134. b = x.T
  1135. c = x
  1136. i = nditer([a, b, c], ['copy_if_overlap'],
  1137. [['readonly'], ['readonly'], ['readonly']])
  1138. a2, b2, c2 = i.operands
  1139. assert_(a is a2)
  1140. assert_(b is b2)
  1141. assert_(c is c2)
  1142. # Copy not needed, 3 ops, read-only overlap
  1143. x = np.ones([10, 10])
  1144. a = x
  1145. b = np.ones([10, 10])
  1146. c = x.T
  1147. i = nditer([a, b, c], ['copy_if_overlap'],
  1148. [['readonly'], ['writeonly'], ['readonly']])
  1149. a2, b2, c2 = i.operands
  1150. assert_(a is a2)
  1151. assert_(b is b2)
  1152. assert_(c is c2)
  1153. # Copy not needed, 3 ops, write-only overlap
  1154. x = np.arange(7)
  1155. a = x[:3]
  1156. b = x[3:6]
  1157. c = x[4:7]
  1158. i = nditer([a, b, c], ['copy_if_overlap'],
  1159. [['readonly'], ['writeonly'], ['writeonly']])
  1160. a2, b2, c2 = i.operands
  1161. assert_(a is a2)
  1162. assert_(b is b2)
  1163. assert_(c is c2)
  1164. def test_iter_op_axes():
  1165. # Check that custom axes work
  1166. # Reverse the axes
  1167. a = arange(6).reshape(2, 3)
  1168. i = nditer([a, a.T], [], [['readonly']]*2, op_axes=[[0, 1], [1, 0]])
  1169. assert_(all([x == y for (x, y) in i]))
  1170. a = arange(24).reshape(2, 3, 4)
  1171. i = nditer([a.T, a], [], [['readonly']]*2, op_axes=[[2, 1, 0], None])
  1172. assert_(all([x == y for (x, y) in i]))
  1173. # Broadcast 1D to any dimension
  1174. a = arange(1, 31).reshape(2, 3, 5)
  1175. b = arange(1, 3)
  1176. i = nditer([a, b], [], [['readonly']]*2, op_axes=[None, [0, -1, -1]])
  1177. assert_equal([x*y for (x, y) in i], (a*b.reshape(2, 1, 1)).ravel())
  1178. b = arange(1, 4)
  1179. i = nditer([a, b], [], [['readonly']]*2, op_axes=[None, [-1, 0, -1]])
  1180. assert_equal([x*y for (x, y) in i], (a*b.reshape(1, 3, 1)).ravel())
  1181. b = arange(1, 6)
  1182. i = nditer([a, b], [], [['readonly']]*2,
  1183. op_axes=[None, [np.newaxis, np.newaxis, 0]])
  1184. assert_equal([x*y for (x, y) in i], (a*b.reshape(1, 1, 5)).ravel())
  1185. # Inner product-style broadcasting
  1186. a = arange(24).reshape(2, 3, 4)
  1187. b = arange(40).reshape(5, 2, 4)
  1188. i = nditer([a, b], ['multi_index'], [['readonly']]*2,
  1189. op_axes=[[0, 1, -1, -1], [-1, -1, 0, 1]])
  1190. assert_equal(i.shape, (2, 3, 5, 2))
  1191. # Matrix product-style broadcasting
  1192. a = arange(12).reshape(3, 4)
  1193. b = arange(20).reshape(4, 5)
  1194. i = nditer([a, b], ['multi_index'], [['readonly']]*2,
  1195. op_axes=[[0, -1], [-1, 1]])
  1196. assert_equal(i.shape, (3, 5))
  1197. def test_iter_op_axes_errors():
  1198. # Check that custom axes throws errors for bad inputs
  1199. # Wrong number of items in op_axes
  1200. a = arange(6).reshape(2, 3)
  1201. assert_raises(ValueError, nditer, [a, a], [], [['readonly']]*2,
  1202. op_axes=[[0], [1], [0]])
  1203. # Out of bounds items in op_axes
  1204. assert_raises(ValueError, nditer, [a, a], [], [['readonly']]*2,
  1205. op_axes=[[2, 1], [0, 1]])
  1206. assert_raises(ValueError, nditer, [a, a], [], [['readonly']]*2,
  1207. op_axes=[[0, 1], [2, -1]])
  1208. # Duplicate items in op_axes
  1209. assert_raises(ValueError, nditer, [a, a], [], [['readonly']]*2,
  1210. op_axes=[[0, 0], [0, 1]])
  1211. assert_raises(ValueError, nditer, [a, a], [], [['readonly']]*2,
  1212. op_axes=[[0, 1], [1, 1]])
  1213. # Different sized arrays in op_axes
  1214. assert_raises(ValueError, nditer, [a, a], [], [['readonly']]*2,
  1215. op_axes=[[0, 1], [0, 1, 0]])
  1216. # Non-broadcastable dimensions in the result
  1217. assert_raises(ValueError, nditer, [a, a], [], [['readonly']]*2,
  1218. op_axes=[[0, 1], [1, 0]])
  1219. def test_iter_copy():
  1220. # Check that copying the iterator works correctly
  1221. a = arange(24).reshape(2, 3, 4)
  1222. # Simple iterator
  1223. i = nditer(a)
  1224. j = i.copy()
  1225. assert_equal([x[()] for x in i], [x[()] for x in j])
  1226. i.iterindex = 3
  1227. j = i.copy()
  1228. assert_equal([x[()] for x in i], [x[()] for x in j])
  1229. # Buffered iterator
  1230. i = nditer(a, ['buffered', 'ranged'], order='F', buffersize=3)
  1231. j = i.copy()
  1232. assert_equal([x[()] for x in i], [x[()] for x in j])
  1233. i.iterindex = 3
  1234. j = i.copy()
  1235. assert_equal([x[()] for x in i], [x[()] for x in j])
  1236. i.iterrange = (3, 9)
  1237. j = i.copy()
  1238. assert_equal([x[()] for x in i], [x[()] for x in j])
  1239. i.iterrange = (2, 18)
  1240. next(i)
  1241. next(i)
  1242. j = i.copy()
  1243. assert_equal([x[()] for x in i], [x[()] for x in j])
  1244. # Casting iterator
  1245. with nditer(a, ['buffered'], order='F', casting='unsafe',
  1246. op_dtypes='f8', buffersize=5) as i:
  1247. j = i.copy()
  1248. assert_equal([x[()] for x in j], a.ravel(order='F'))
  1249. a = arange(24, dtype='<i4').reshape(2, 3, 4)
  1250. with nditer(a, ['buffered'], order='F', casting='unsafe',
  1251. op_dtypes='>f8', buffersize=5) as i:
  1252. j = i.copy()
  1253. assert_equal([x[()] for x in j], a.ravel(order='F'))
  1254. @pytest.mark.parametrize("dtype", np.typecodes["All"])
  1255. @pytest.mark.parametrize("loop_dtype", np.typecodes["All"])
  1256. @pytest.mark.filterwarnings("ignore::numpy.ComplexWarning")
  1257. def test_iter_copy_casts(dtype, loop_dtype):
  1258. # Ensure the dtype is never flexible:
  1259. if loop_dtype.lower() == "m":
  1260. loop_dtype = loop_dtype + "[ms]"
  1261. elif np.dtype(loop_dtype).itemsize == 0:
  1262. loop_dtype = loop_dtype + "50"
  1263. # Make things a bit more interesting by requiring a byte-swap as well:
  1264. arr = np.ones(1000, dtype=np.dtype(dtype).newbyteorder())
  1265. try:
  1266. expected = arr.astype(loop_dtype)
  1267. except Exception:
  1268. # Some casts are not possible, do not worry about them
  1269. return
  1270. it = np.nditer((arr,), ["buffered", "external_loop", "refs_ok"],
  1271. op_dtypes=[loop_dtype], casting="unsafe")
  1272. if np.issubdtype(np.dtype(loop_dtype), np.number):
  1273. # Casting to strings may be strange, but for simple dtypes do not rely
  1274. # on the cast being correct:
  1275. assert_array_equal(expected, np.ones(1000, dtype=loop_dtype))
  1276. it_copy = it.copy()
  1277. res = next(it)
  1278. del it
  1279. res_copy = next(it_copy)
  1280. del it_copy
  1281. assert_array_equal(res, expected)
  1282. assert_array_equal(res_copy, expected)
  1283. def test_iter_copy_casts_structured():
  1284. # Test a complicated structured dtype for casting, as it requires
  1285. # both multiple steps and a more complex casting setup.
  1286. # Includes a structured -> unstructured (any to object), and many other
  1287. # casts, which cause this to require all steps in the casting machinery
  1288. # one level down as well as the iterator copy (which uses NpyAuxData clone)
  1289. in_dtype = np.dtype([("a", np.dtype("i,")),
  1290. ("b", np.dtype(">i,<i,>d,S17,>d,(3)f,O,i1"))])
  1291. out_dtype = np.dtype([("a", np.dtype("O")),
  1292. ("b", np.dtype(">i,>i,S17,>d,>U3,(3)d,i1,O"))])
  1293. arr = np.ones(1000, dtype=in_dtype)
  1294. it = np.nditer((arr,), ["buffered", "external_loop", "refs_ok"],
  1295. op_dtypes=[out_dtype], casting="unsafe")
  1296. it_copy = it.copy()
  1297. res1 = next(it)
  1298. del it
  1299. res2 = next(it_copy)
  1300. del it_copy
  1301. expected = arr["a"].astype(out_dtype["a"])
  1302. assert_array_equal(res1["a"], expected)
  1303. assert_array_equal(res2["a"], expected)
  1304. for field in in_dtype["b"].names:
  1305. # Note that the .base avoids the subarray field
  1306. expected = arr["b"][field].astype(out_dtype["b"][field].base)
  1307. assert_array_equal(res1["b"][field], expected)
  1308. assert_array_equal(res2["b"][field], expected)
  1309. def test_iter_allocate_output_simple():
  1310. # Check that the iterator will properly allocate outputs
  1311. # Simple case
  1312. a = arange(6)
  1313. i = nditer([a, None], [], [['readonly'], ['writeonly', 'allocate']],
  1314. op_dtypes=[None, np.dtype('f4')])
  1315. assert_equal(i.operands[1].shape, a.shape)
  1316. assert_equal(i.operands[1].dtype, np.dtype('f4'))
  1317. def test_iter_allocate_output_buffered_readwrite():
  1318. # Allocated output with buffering + delay_bufalloc
  1319. a = arange(6)
  1320. i = nditer([a, None], ['buffered', 'delay_bufalloc'],
  1321. [['readonly'], ['allocate', 'readwrite']])
  1322. with i:
  1323. i.operands[1][:] = 1
  1324. i.reset()
  1325. for x in i:
  1326. x[1][...] += x[0][...]
  1327. assert_equal(i.operands[1], a+1)
  1328. def test_iter_allocate_output_itorder():
  1329. # The allocated output should match the iteration order
  1330. # C-order input, best iteration order
  1331. a = arange(6, dtype='i4').reshape(2, 3)
  1332. i = nditer([a, None], [], [['readonly'], ['writeonly', 'allocate']],
  1333. op_dtypes=[None, np.dtype('f4')])
  1334. assert_equal(i.operands[1].shape, a.shape)
  1335. assert_equal(i.operands[1].strides, a.strides)
  1336. assert_equal(i.operands[1].dtype, np.dtype('f4'))
  1337. # F-order input, best iteration order
  1338. a = arange(24, dtype='i4').reshape(2, 3, 4).T
  1339. i = nditer([a, None], [], [['readonly'], ['writeonly', 'allocate']],
  1340. op_dtypes=[None, np.dtype('f4')])
  1341. assert_equal(i.operands[1].shape, a.shape)
  1342. assert_equal(i.operands[1].strides, a.strides)
  1343. assert_equal(i.operands[1].dtype, np.dtype('f4'))
  1344. # Non-contiguous input, C iteration order
  1345. a = arange(24, dtype='i4').reshape(2, 3, 4).swapaxes(0, 1)
  1346. i = nditer([a, None], [],
  1347. [['readonly'], ['writeonly', 'allocate']],
  1348. order='C',
  1349. op_dtypes=[None, np.dtype('f4')])
  1350. assert_equal(i.operands[1].shape, a.shape)
  1351. assert_equal(i.operands[1].strides, (32, 16, 4))
  1352. assert_equal(i.operands[1].dtype, np.dtype('f4'))
  1353. def test_iter_allocate_output_opaxes():
  1354. # Specifying op_axes should work
  1355. a = arange(24, dtype='i4').reshape(2, 3, 4)
  1356. i = nditer([None, a], [], [['writeonly', 'allocate'], ['readonly']],
  1357. op_dtypes=[np.dtype('u4'), None],
  1358. op_axes=[[1, 2, 0], None])
  1359. assert_equal(i.operands[0].shape, (4, 2, 3))
  1360. assert_equal(i.operands[0].strides, (4, 48, 16))
  1361. assert_equal(i.operands[0].dtype, np.dtype('u4'))
  1362. def test_iter_allocate_output_types_promotion():
  1363. # Check type promotion of automatic outputs
  1364. i = nditer([array([3], dtype='f4'), array([0], dtype='f8'), None], [],
  1365. [['readonly']]*2+[['writeonly', 'allocate']])
  1366. assert_equal(i.dtypes[2], np.dtype('f8'))
  1367. i = nditer([array([3], dtype='i4'), array([0], dtype='f4'), None], [],
  1368. [['readonly']]*2+[['writeonly', 'allocate']])
  1369. assert_equal(i.dtypes[2], np.dtype('f8'))
  1370. i = nditer([array([3], dtype='f4'), array(0, dtype='f8'), None], [],
  1371. [['readonly']]*2+[['writeonly', 'allocate']])
  1372. assert_equal(i.dtypes[2], np.dtype('f4'))
  1373. i = nditer([array([3], dtype='u4'), array(0, dtype='i4'), None], [],
  1374. [['readonly']]*2+[['writeonly', 'allocate']])
  1375. assert_equal(i.dtypes[2], np.dtype('u4'))
  1376. i = nditer([array([3], dtype='u4'), array(-12, dtype='i4'), None], [],
  1377. [['readonly']]*2+[['writeonly', 'allocate']])
  1378. assert_equal(i.dtypes[2], np.dtype('i8'))
  1379. def test_iter_allocate_output_types_byte_order():
  1380. # Verify the rules for byte order changes
  1381. # When there's just one input, the output type exactly matches
  1382. a = array([3], dtype='u4').newbyteorder()
  1383. i = nditer([a, None], [],
  1384. [['readonly'], ['writeonly', 'allocate']])
  1385. assert_equal(i.dtypes[0], i.dtypes[1])
  1386. # With two or more inputs, the output type is in native byte order
  1387. i = nditer([a, a, None], [],
  1388. [['readonly'], ['readonly'], ['writeonly', 'allocate']])
  1389. assert_(i.dtypes[0] != i.dtypes[2])
  1390. assert_equal(i.dtypes[0].newbyteorder('='), i.dtypes[2])
  1391. def test_iter_allocate_output_types_scalar():
  1392. # If the inputs are all scalars, the output should be a scalar
  1393. i = nditer([None, 1, 2.3, np.float32(12), np.complex128(3)], [],
  1394. [['writeonly', 'allocate']] + [['readonly']]*4)
  1395. assert_equal(i.operands[0].dtype, np.dtype('complex128'))
  1396. assert_equal(i.operands[0].ndim, 0)
  1397. def test_iter_allocate_output_subtype():
  1398. # Make sure that the subtype with priority wins
  1399. class MyNDArray(np.ndarray):
  1400. __array_priority__ = 15
  1401. # subclass vs ndarray
  1402. a = np.array([[1, 2], [3, 4]]).view(MyNDArray)
  1403. b = np.arange(4).reshape(2, 2).T
  1404. i = nditer([a, b, None], [],
  1405. [['readonly'], ['readonly'], ['writeonly', 'allocate']])
  1406. assert_equal(type(a), type(i.operands[2]))
  1407. assert_(type(b) is not type(i.operands[2]))
  1408. assert_equal(i.operands[2].shape, (2, 2))
  1409. # If subtypes are disabled, we should get back an ndarray.
  1410. i = nditer([a, b, None], [],
  1411. [['readonly'], ['readonly'],
  1412. ['writeonly', 'allocate', 'no_subtype']])
  1413. assert_equal(type(b), type(i.operands[2]))
  1414. assert_(type(a) is not type(i.operands[2]))
  1415. assert_equal(i.operands[2].shape, (2, 2))
  1416. def test_iter_allocate_output_errors():
  1417. # Check that the iterator will throw errors for bad output allocations
  1418. # Need an input if no output data type is specified
  1419. a = arange(6)
  1420. assert_raises(TypeError, nditer, [a, None], [],
  1421. [['writeonly'], ['writeonly', 'allocate']])
  1422. # Allocated output should be flagged for writing
  1423. assert_raises(ValueError, nditer, [a, None], [],
  1424. [['readonly'], ['allocate', 'readonly']])
  1425. # Allocated output can't have buffering without delayed bufalloc
  1426. assert_raises(ValueError, nditer, [a, None], ['buffered'],
  1427. ['allocate', 'readwrite'])
  1428. # Must specify dtype if there are no inputs (cannot promote existing ones;
  1429. # maybe this should use the 'f4' here, but it does not historically.)
  1430. assert_raises(TypeError, nditer, [None, None], [],
  1431. [['writeonly', 'allocate'],
  1432. ['writeonly', 'allocate']],
  1433. op_dtypes=[None, np.dtype('f4')])
  1434. # If using op_axes, must specify all the axes
  1435. a = arange(24, dtype='i4').reshape(2, 3, 4)
  1436. assert_raises(ValueError, nditer, [a, None], [],
  1437. [['readonly'], ['writeonly', 'allocate']],
  1438. op_dtypes=[None, np.dtype('f4')],
  1439. op_axes=[None, [0, np.newaxis, 1]])
  1440. # If using op_axes, the axes must be within bounds
  1441. assert_raises(ValueError, nditer, [a, None], [],
  1442. [['readonly'], ['writeonly', 'allocate']],
  1443. op_dtypes=[None, np.dtype('f4')],
  1444. op_axes=[None, [0, 3, 1]])
  1445. # If using op_axes, there can't be duplicates
  1446. assert_raises(ValueError, nditer, [a, None], [],
  1447. [['readonly'], ['writeonly', 'allocate']],
  1448. op_dtypes=[None, np.dtype('f4')],
  1449. op_axes=[None, [0, 2, 1, 0]])
  1450. # Not all axes may be specified if a reduction. If there is a hole
  1451. # in op_axes, this is an error.
  1452. a = arange(24, dtype='i4').reshape(2, 3, 4)
  1453. assert_raises(ValueError, nditer, [a, None], ["reduce_ok"],
  1454. [['readonly'], ['readwrite', 'allocate']],
  1455. op_dtypes=[None, np.dtype('f4')],
  1456. op_axes=[None, [0, np.newaxis, 2]])
  1457. def test_all_allocated():
  1458. # When no output and no shape is given, `()` is used as shape.
  1459. i = np.nditer([None], op_dtypes=["int64"])
  1460. assert i.operands[0].shape == ()
  1461. assert i.dtypes == (np.dtype("int64"),)
  1462. i = np.nditer([None], op_dtypes=["int64"], itershape=(2, 3, 4))
  1463. assert i.operands[0].shape == (2, 3, 4)
  1464. def test_iter_remove_axis():
  1465. a = arange(24).reshape(2, 3, 4)
  1466. i = nditer(a, ['multi_index'])
  1467. i.remove_axis(1)
  1468. assert_equal([x for x in i], a[:, 0,:].ravel())
  1469. a = a[::-1,:,:]
  1470. i = nditer(a, ['multi_index'])
  1471. i.remove_axis(0)
  1472. assert_equal([x for x in i], a[0,:,:].ravel())
  1473. def test_iter_remove_multi_index_inner_loop():
  1474. # Check that removing multi-index support works
  1475. a = arange(24).reshape(2, 3, 4)
  1476. i = nditer(a, ['multi_index'])
  1477. assert_equal(i.ndim, 3)
  1478. assert_equal(i.shape, (2, 3, 4))
  1479. assert_equal(i.itviews[0].shape, (2, 3, 4))
  1480. # Removing the multi-index tracking causes all dimensions to coalesce
  1481. before = [x for x in i]
  1482. i.remove_multi_index()
  1483. after = [x for x in i]
  1484. assert_equal(before, after)
  1485. assert_equal(i.ndim, 1)
  1486. assert_raises(ValueError, lambda i:i.shape, i)
  1487. assert_equal(i.itviews[0].shape, (24,))
  1488. # Removing the inner loop means there's just one iteration
  1489. i.reset()
  1490. assert_equal(i.itersize, 24)
  1491. assert_equal(i[0].shape, tuple())
  1492. i.enable_external_loop()
  1493. assert_equal(i.itersize, 24)
  1494. assert_equal(i[0].shape, (24,))
  1495. assert_equal(i.value, arange(24))
  1496. def test_iter_iterindex():
  1497. # Make sure iterindex works
  1498. buffersize = 5
  1499. a = arange(24).reshape(4, 3, 2)
  1500. for flags in ([], ['buffered']):
  1501. i = nditer(a, flags, buffersize=buffersize)
  1502. assert_equal(iter_iterindices(i), list(range(24)))
  1503. i.iterindex = 2
  1504. assert_equal(iter_iterindices(i), list(range(2, 24)))
  1505. i = nditer(a, flags, order='F', buffersize=buffersize)
  1506. assert_equal(iter_iterindices(i), list(range(24)))
  1507. i.iterindex = 5
  1508. assert_equal(iter_iterindices(i), list(range(5, 24)))
  1509. i = nditer(a[::-1], flags, order='F', buffersize=buffersize)
  1510. assert_equal(iter_iterindices(i), list(range(24)))
  1511. i.iterindex = 9
  1512. assert_equal(iter_iterindices(i), list(range(9, 24)))
  1513. i = nditer(a[::-1, ::-1], flags, order='C', buffersize=buffersize)
  1514. assert_equal(iter_iterindices(i), list(range(24)))
  1515. i.iterindex = 13
  1516. assert_equal(iter_iterindices(i), list(range(13, 24)))
  1517. i = nditer(a[::1, ::-1], flags, buffersize=buffersize)
  1518. assert_equal(iter_iterindices(i), list(range(24)))
  1519. i.iterindex = 23
  1520. assert_equal(iter_iterindices(i), list(range(23, 24)))
  1521. i.reset()
  1522. i.iterindex = 2
  1523. assert_equal(iter_iterindices(i), list(range(2, 24)))
  1524. def test_iter_iterrange():
  1525. # Make sure getting and resetting the iterrange works
  1526. buffersize = 5
  1527. a = arange(24, dtype='i4').reshape(4, 3, 2)
  1528. a_fort = a.ravel(order='F')
  1529. i = nditer(a, ['ranged'], ['readonly'], order='F',
  1530. buffersize=buffersize)
  1531. assert_equal(i.iterrange, (0, 24))
  1532. assert_equal([x[()] for x in i], a_fort)
  1533. for r in [(0, 24), (1, 2), (3, 24), (5, 5), (0, 20), (23, 24)]:
  1534. i.iterrange = r
  1535. assert_equal(i.iterrange, r)
  1536. assert_equal([x[()] for x in i], a_fort[r[0]:r[1]])
  1537. i = nditer(a, ['ranged', 'buffered'], ['readonly'], order='F',
  1538. op_dtypes='f8', buffersize=buffersize)
  1539. assert_equal(i.iterrange, (0, 24))
  1540. assert_equal([x[()] for x in i], a_fort)
  1541. for r in [(0, 24), (1, 2), (3, 24), (5, 5), (0, 20), (23, 24)]:
  1542. i.iterrange = r
  1543. assert_equal(i.iterrange, r)
  1544. assert_equal([x[()] for x in i], a_fort[r[0]:r[1]])
  1545. def get_array(i):
  1546. val = np.array([], dtype='f8')
  1547. for x in i:
  1548. val = np.concatenate((val, x))
  1549. return val
  1550. i = nditer(a, ['ranged', 'buffered', 'external_loop'],
  1551. ['readonly'], order='F',
  1552. op_dtypes='f8', buffersize=buffersize)
  1553. assert_equal(i.iterrange, (0, 24))
  1554. assert_equal(get_array(i), a_fort)
  1555. for r in [(0, 24), (1, 2), (3, 24), (5, 5), (0, 20), (23, 24)]:
  1556. i.iterrange = r
  1557. assert_equal(i.iterrange, r)
  1558. assert_equal(get_array(i), a_fort[r[0]:r[1]])
  1559. def test_iter_buffering():
  1560. # Test buffering with several buffer sizes and types
  1561. arrays = []
  1562. # F-order swapped array
  1563. arrays.append(np.arange(24,
  1564. dtype='c16').reshape(2, 3, 4).T.newbyteorder().byteswap())
  1565. # Contiguous 1-dimensional array
  1566. arrays.append(np.arange(10, dtype='f4'))
  1567. # Unaligned array
  1568. a = np.zeros((4*16+1,), dtype='i1')[1:]
  1569. a.dtype = 'i4'
  1570. a[:] = np.arange(16, dtype='i4')
  1571. arrays.append(a)
  1572. # 4-D F-order array
  1573. arrays.append(np.arange(120, dtype='i4').reshape(5, 3, 2, 4).T)
  1574. for a in arrays:
  1575. for buffersize in (1, 2, 3, 5, 8, 11, 16, 1024):
  1576. vals = []
  1577. i = nditer(a, ['buffered', 'external_loop'],
  1578. [['readonly', 'nbo', 'aligned']],
  1579. order='C',
  1580. casting='equiv',
  1581. buffersize=buffersize)
  1582. while not i.finished:
  1583. assert_(i[0].size <= buffersize)
  1584. vals.append(i[0].copy())
  1585. i.iternext()
  1586. assert_equal(np.concatenate(vals), a.ravel(order='C'))
  1587. def test_iter_write_buffering():
  1588. # Test that buffering of writes is working
  1589. # F-order swapped array
  1590. a = np.arange(24).reshape(2, 3, 4).T.newbyteorder().byteswap()
  1591. i = nditer(a, ['buffered'],
  1592. [['readwrite', 'nbo', 'aligned']],
  1593. casting='equiv',
  1594. order='C',
  1595. buffersize=16)
  1596. x = 0
  1597. with i:
  1598. while not i.finished:
  1599. i[0] = x
  1600. x += 1
  1601. i.iternext()
  1602. assert_equal(a.ravel(order='C'), np.arange(24))
  1603. def test_iter_buffering_delayed_alloc():
  1604. # Test that delaying buffer allocation works
  1605. a = np.arange(6)
  1606. b = np.arange(1, dtype='f4')
  1607. i = nditer([a, b], ['buffered', 'delay_bufalloc', 'multi_index', 'reduce_ok'],
  1608. ['readwrite'],
  1609. casting='unsafe',
  1610. op_dtypes='f4')
  1611. assert_(i.has_delayed_bufalloc)
  1612. assert_raises(ValueError, lambda i:i.multi_index, i)
  1613. assert_raises(ValueError, lambda i:i[0], i)
  1614. assert_raises(ValueError, lambda i:i[0:2], i)
  1615. def assign_iter(i):
  1616. i[0] = 0
  1617. assert_raises(ValueError, assign_iter, i)
  1618. i.reset()
  1619. assert_(not i.has_delayed_bufalloc)
  1620. assert_equal(i.multi_index, (0,))
  1621. with i:
  1622. assert_equal(i[0], 0)
  1623. i[1] = 1
  1624. assert_equal(i[0:2], [0, 1])
  1625. assert_equal([[x[0][()], x[1][()]] for x in i], list(zip(range(6), [1]*6)))
  1626. def test_iter_buffered_cast_simple():
  1627. # Test that buffering can handle a simple cast
  1628. a = np.arange(10, dtype='f4')
  1629. i = nditer(a, ['buffered', 'external_loop'],
  1630. [['readwrite', 'nbo', 'aligned']],
  1631. casting='same_kind',
  1632. op_dtypes=[np.dtype('f8')],
  1633. buffersize=3)
  1634. with i:
  1635. for v in i:
  1636. v[...] *= 2
  1637. assert_equal(a, 2*np.arange(10, dtype='f4'))
  1638. def test_iter_buffered_cast_byteswapped():
  1639. # Test that buffering can handle a cast which requires swap->cast->swap
  1640. a = np.arange(10, dtype='f4').newbyteorder().byteswap()
  1641. i = nditer(a, ['buffered', 'external_loop'],
  1642. [['readwrite', 'nbo', 'aligned']],
  1643. casting='same_kind',
  1644. op_dtypes=[np.dtype('f8').newbyteorder()],
  1645. buffersize=3)
  1646. with i:
  1647. for v in i:
  1648. v[...] *= 2
  1649. assert_equal(a, 2*np.arange(10, dtype='f4'))
  1650. with suppress_warnings() as sup:
  1651. sup.filter(np.ComplexWarning)
  1652. a = np.arange(10, dtype='f8').newbyteorder().byteswap()
  1653. i = nditer(a, ['buffered', 'external_loop'],
  1654. [['readwrite', 'nbo', 'aligned']],
  1655. casting='unsafe',
  1656. op_dtypes=[np.dtype('c8').newbyteorder()],
  1657. buffersize=3)
  1658. with i:
  1659. for v in i:
  1660. v[...] *= 2
  1661. assert_equal(a, 2*np.arange(10, dtype='f8'))
  1662. def test_iter_buffered_cast_byteswapped_complex():
  1663. # Test that buffering can handle a cast which requires swap->cast->copy
  1664. a = np.arange(10, dtype='c8').newbyteorder().byteswap()
  1665. a += 2j
  1666. i = nditer(a, ['buffered', 'external_loop'],
  1667. [['readwrite', 'nbo', 'aligned']],
  1668. casting='same_kind',
  1669. op_dtypes=[np.dtype('c16')],
  1670. buffersize=3)
  1671. with i:
  1672. for v in i:
  1673. v[...] *= 2
  1674. assert_equal(a, 2*np.arange(10, dtype='c8') + 4j)
  1675. a = np.arange(10, dtype='c8')
  1676. a += 2j
  1677. i = nditer(a, ['buffered', 'external_loop'],
  1678. [['readwrite', 'nbo', 'aligned']],
  1679. casting='same_kind',
  1680. op_dtypes=[np.dtype('c16').newbyteorder()],
  1681. buffersize=3)
  1682. with i:
  1683. for v in i:
  1684. v[...] *= 2
  1685. assert_equal(a, 2*np.arange(10, dtype='c8') + 4j)
  1686. a = np.arange(10, dtype=np.clongdouble).newbyteorder().byteswap()
  1687. a += 2j
  1688. i = nditer(a, ['buffered', 'external_loop'],
  1689. [['readwrite', 'nbo', 'aligned']],
  1690. casting='same_kind',
  1691. op_dtypes=[np.dtype('c16')],
  1692. buffersize=3)
  1693. with i:
  1694. for v in i:
  1695. v[...] *= 2
  1696. assert_equal(a, 2*np.arange(10, dtype=np.clongdouble) + 4j)
  1697. a = np.arange(10, dtype=np.longdouble).newbyteorder().byteswap()
  1698. i = nditer(a, ['buffered', 'external_loop'],
  1699. [['readwrite', 'nbo', 'aligned']],
  1700. casting='same_kind',
  1701. op_dtypes=[np.dtype('f4')],
  1702. buffersize=7)
  1703. with i:
  1704. for v in i:
  1705. v[...] *= 2
  1706. assert_equal(a, 2*np.arange(10, dtype=np.longdouble))
  1707. def test_iter_buffered_cast_structured_type():
  1708. # Tests buffering of structured types
  1709. # simple -> struct type (duplicates the value)
  1710. sdt = [('a', 'f4'), ('b', 'i8'), ('c', 'c8', (2, 3)), ('d', 'O')]
  1711. a = np.arange(3, dtype='f4') + 0.5
  1712. i = nditer(a, ['buffered', 'refs_ok'], ['readonly'],
  1713. casting='unsafe',
  1714. op_dtypes=sdt)
  1715. vals = [np.array(x) for x in i]
  1716. assert_equal(vals[0]['a'], 0.5)
  1717. assert_equal(vals[0]['b'], 0)
  1718. assert_equal(vals[0]['c'], [[(0.5)]*3]*2)
  1719. assert_equal(vals[0]['d'], 0.5)
  1720. assert_equal(vals[1]['a'], 1.5)
  1721. assert_equal(vals[1]['b'], 1)
  1722. assert_equal(vals[1]['c'], [[(1.5)]*3]*2)
  1723. assert_equal(vals[1]['d'], 1.5)
  1724. assert_equal(vals[0].dtype, np.dtype(sdt))
  1725. # object -> struct type
  1726. sdt = [('a', 'f4'), ('b', 'i8'), ('c', 'c8', (2, 3)), ('d', 'O')]
  1727. a = np.zeros((3,), dtype='O')
  1728. a[0] = (0.5, 0.5, [[0.5, 0.5, 0.5], [0.5, 0.5, 0.5]], 0.5)
  1729. a[1] = (1.5, 1.5, [[1.5, 1.5, 1.5], [1.5, 1.5, 1.5]], 1.5)
  1730. a[2] = (2.5, 2.5, [[2.5, 2.5, 2.5], [2.5, 2.5, 2.5]], 2.5)
  1731. if HAS_REFCOUNT:
  1732. rc = sys.getrefcount(a[0])
  1733. i = nditer(a, ['buffered', 'refs_ok'], ['readonly'],
  1734. casting='unsafe',
  1735. op_dtypes=sdt)
  1736. vals = [x.copy() for x in i]
  1737. assert_equal(vals[0]['a'], 0.5)
  1738. assert_equal(vals[0]['b'], 0)
  1739. assert_equal(vals[0]['c'], [[(0.5)]*3]*2)
  1740. assert_equal(vals[0]['d'], 0.5)
  1741. assert_equal(vals[1]['a'], 1.5)
  1742. assert_equal(vals[1]['b'], 1)
  1743. assert_equal(vals[1]['c'], [[(1.5)]*3]*2)
  1744. assert_equal(vals[1]['d'], 1.5)
  1745. assert_equal(vals[0].dtype, np.dtype(sdt))
  1746. vals, i, x = [None]*3
  1747. if HAS_REFCOUNT:
  1748. assert_equal(sys.getrefcount(a[0]), rc)
  1749. # single-field struct type -> simple
  1750. sdt = [('a', 'f4')]
  1751. a = np.array([(5.5,), (8,)], dtype=sdt)
  1752. i = nditer(a, ['buffered', 'refs_ok'], ['readonly'],
  1753. casting='unsafe',
  1754. op_dtypes='i4')
  1755. assert_equal([x_[()] for x_ in i], [5, 8])
  1756. # make sure multi-field struct type -> simple doesn't work
  1757. sdt = [('a', 'f4'), ('b', 'i8'), ('d', 'O')]
  1758. a = np.array([(5.5, 7, 'test'), (8, 10, 11)], dtype=sdt)
  1759. assert_raises(TypeError, lambda: (
  1760. nditer(a, ['buffered', 'refs_ok'], ['readonly'],
  1761. casting='unsafe',
  1762. op_dtypes='i4')))
  1763. # struct type -> struct type (field-wise copy)
  1764. sdt1 = [('a', 'f4'), ('b', 'i8'), ('d', 'O')]
  1765. sdt2 = [('d', 'u2'), ('a', 'O'), ('b', 'f8')]
  1766. a = np.array([(1, 2, 3), (4, 5, 6)], dtype=sdt1)
  1767. i = nditer(a, ['buffered', 'refs_ok'], ['readonly'],
  1768. casting='unsafe',
  1769. op_dtypes=sdt2)
  1770. assert_equal(i[0].dtype, np.dtype(sdt2))
  1771. assert_equal([np.array(x_) for x_ in i],
  1772. [np.array((1, 2, 3), dtype=sdt2),
  1773. np.array((4, 5, 6), dtype=sdt2)])
  1774. def test_iter_buffered_cast_structured_type_failure_with_cleanup():
  1775. # make sure struct type -> struct type with different
  1776. # number of fields fails
  1777. sdt1 = [('a', 'f4'), ('b', 'i8'), ('d', 'O')]
  1778. sdt2 = [('b', 'O'), ('a', 'f8')]
  1779. a = np.array([(1, 2, 3), (4, 5, 6)], dtype=sdt1)
  1780. for intent in ["readwrite", "readonly", "writeonly"]:
  1781. # This test was initially designed to test an error at a different
  1782. # place, but will now raise earlier to to the cast not being possible:
  1783. # `assert np.can_cast(a.dtype, sdt2, casting="unsafe")` fails.
  1784. # Without a faulty DType, there is probably no reliable
  1785. # way to get the initial tested behaviour.
  1786. simple_arr = np.array([1, 2], dtype="i,i") # requires clean up
  1787. with pytest.raises(TypeError):
  1788. nditer((simple_arr, a), ['buffered', 'refs_ok'], [intent, intent],
  1789. casting='unsafe', op_dtypes=["f,f", sdt2])
  1790. def test_buffered_cast_error_paths():
  1791. with pytest.raises(ValueError):
  1792. # The input is cast into an `S3` buffer
  1793. np.nditer((np.array("a", dtype="S1"),), op_dtypes=["i"],
  1794. casting="unsafe", flags=["buffered"])
  1795. # The `M8[ns]` is cast into the `S3` output
  1796. it = np.nditer((np.array(1, dtype="i"),), op_dtypes=["S1"],
  1797. op_flags=["writeonly"], casting="unsafe", flags=["buffered"])
  1798. with pytest.raises(ValueError):
  1799. with it:
  1800. buf = next(it)
  1801. buf[...] = "a" # cannot be converted to int.
  1802. @pytest.mark.skipif(IS_WASM, reason="Cannot start subprocess")
  1803. @pytest.mark.skipif(not HAS_REFCOUNT, reason="PyPy seems to not hit this.")
  1804. def test_buffered_cast_error_paths_unraisable():
  1805. # The following gives an unraisable error. Pytest sometimes captures that
  1806. # (depending python and/or pytest version). So with Python>=3.8 this can
  1807. # probably be cleaned out in the future to check for
  1808. # pytest.PytestUnraisableExceptionWarning:
  1809. code = textwrap.dedent("""
  1810. import numpy as np
  1811. it = np.nditer((np.array(1, dtype="i"),), op_dtypes=["S1"],
  1812. op_flags=["writeonly"], casting="unsafe", flags=["buffered"])
  1813. buf = next(it)
  1814. buf[...] = "a"
  1815. del buf, it # Flushing only happens during deallocate right now.
  1816. """)
  1817. res = subprocess.check_output([sys.executable, "-c", code],
  1818. stderr=subprocess.STDOUT, text=True)
  1819. assert "ValueError" in res
  1820. def test_iter_buffered_cast_subarray():
  1821. # Tests buffering of subarrays
  1822. # one element -> many (copies it to all)
  1823. sdt1 = [('a', 'f4')]
  1824. sdt2 = [('a', 'f8', (3, 2, 2))]
  1825. a = np.zeros((6,), dtype=sdt1)
  1826. a['a'] = np.arange(6)
  1827. i = nditer(a, ['buffered', 'refs_ok'], ['readonly'],
  1828. casting='unsafe',
  1829. op_dtypes=sdt2)
  1830. assert_equal(i[0].dtype, np.dtype(sdt2))
  1831. for x, count in zip(i, list(range(6))):
  1832. assert_(np.all(x['a'] == count))
  1833. # one element -> many -> back (copies it to all)
  1834. sdt1 = [('a', 'O', (1, 1))]
  1835. sdt2 = [('a', 'O', (3, 2, 2))]
  1836. a = np.zeros((6,), dtype=sdt1)
  1837. a['a'][:, 0, 0] = np.arange(6)
  1838. i = nditer(a, ['buffered', 'refs_ok'], ['readwrite'],
  1839. casting='unsafe',
  1840. op_dtypes=sdt2)
  1841. with i:
  1842. assert_equal(i[0].dtype, np.dtype(sdt2))
  1843. count = 0
  1844. for x in i:
  1845. assert_(np.all(x['a'] == count))
  1846. x['a'][0] += 2
  1847. count += 1
  1848. assert_equal(a['a'], np.arange(6).reshape(6, 1, 1)+2)
  1849. # many -> one element -> back (copies just element 0)
  1850. sdt1 = [('a', 'O', (3, 2, 2))]
  1851. sdt2 = [('a', 'O', (1,))]
  1852. a = np.zeros((6,), dtype=sdt1)
  1853. a['a'][:, 0, 0, 0] = np.arange(6)
  1854. i = nditer(a, ['buffered', 'refs_ok'], ['readwrite'],
  1855. casting='unsafe',
  1856. op_dtypes=sdt2)
  1857. with i:
  1858. assert_equal(i[0].dtype, np.dtype(sdt2))
  1859. count = 0
  1860. for x in i:
  1861. assert_equal(x['a'], count)
  1862. x['a'] += 2
  1863. count += 1
  1864. assert_equal(a['a'], np.arange(6).reshape(6, 1, 1, 1)*np.ones((1, 3, 2, 2))+2)
  1865. # many -> one element -> back (copies just element 0)
  1866. sdt1 = [('a', 'f8', (3, 2, 2))]
  1867. sdt2 = [('a', 'O', (1,))]
  1868. a = np.zeros((6,), dtype=sdt1)
  1869. a['a'][:, 0, 0, 0] = np.arange(6)
  1870. i = nditer(a, ['buffered', 'refs_ok'], ['readonly'],
  1871. casting='unsafe',
  1872. op_dtypes=sdt2)
  1873. assert_equal(i[0].dtype, np.dtype(sdt2))
  1874. count = 0
  1875. for x in i:
  1876. assert_equal(x['a'], count)
  1877. count += 1
  1878. # many -> one element (copies just element 0)
  1879. sdt1 = [('a', 'O', (3, 2, 2))]
  1880. sdt2 = [('a', 'f4', (1,))]
  1881. a = np.zeros((6,), dtype=sdt1)
  1882. a['a'][:, 0, 0, 0] = np.arange(6)
  1883. i = nditer(a, ['buffered', 'refs_ok'], ['readonly'],
  1884. casting='unsafe',
  1885. op_dtypes=sdt2)
  1886. assert_equal(i[0].dtype, np.dtype(sdt2))
  1887. count = 0
  1888. for x in i:
  1889. assert_equal(x['a'], count)
  1890. count += 1
  1891. # many -> matching shape (straightforward copy)
  1892. sdt1 = [('a', 'O', (3, 2, 2))]
  1893. sdt2 = [('a', 'f4', (3, 2, 2))]
  1894. a = np.zeros((6,), dtype=sdt1)
  1895. a['a'] = np.arange(6*3*2*2).reshape(6, 3, 2, 2)
  1896. i = nditer(a, ['buffered', 'refs_ok'], ['readonly'],
  1897. casting='unsafe',
  1898. op_dtypes=sdt2)
  1899. assert_equal(i[0].dtype, np.dtype(sdt2))
  1900. count = 0
  1901. for x in i:
  1902. assert_equal(x['a'], a[count]['a'])
  1903. count += 1
  1904. # vector -> smaller vector (truncates)
  1905. sdt1 = [('a', 'f8', (6,))]
  1906. sdt2 = [('a', 'f4', (2,))]
  1907. a = np.zeros((6,), dtype=sdt1)
  1908. a['a'] = np.arange(6*6).reshape(6, 6)
  1909. i = nditer(a, ['buffered', 'refs_ok'], ['readonly'],
  1910. casting='unsafe',
  1911. op_dtypes=sdt2)
  1912. assert_equal(i[0].dtype, np.dtype(sdt2))
  1913. count = 0
  1914. for x in i:
  1915. assert_equal(x['a'], a[count]['a'][:2])
  1916. count += 1
  1917. # vector -> bigger vector (pads with zeros)
  1918. sdt1 = [('a', 'f8', (2,))]
  1919. sdt2 = [('a', 'f4', (6,))]
  1920. a = np.zeros((6,), dtype=sdt1)
  1921. a['a'] = np.arange(6*2).reshape(6, 2)
  1922. i = nditer(a, ['buffered', 'refs_ok'], ['readonly'],
  1923. casting='unsafe',
  1924. op_dtypes=sdt2)
  1925. assert_equal(i[0].dtype, np.dtype(sdt2))
  1926. count = 0
  1927. for x in i:
  1928. assert_equal(x['a'][:2], a[count]['a'])
  1929. assert_equal(x['a'][2:], [0, 0, 0, 0])
  1930. count += 1
  1931. # vector -> matrix (broadcasts)
  1932. sdt1 = [('a', 'f8', (2,))]
  1933. sdt2 = [('a', 'f4', (2, 2))]
  1934. a = np.zeros((6,), dtype=sdt1)
  1935. a['a'] = np.arange(6*2).reshape(6, 2)
  1936. i = nditer(a, ['buffered', 'refs_ok'], ['readonly'],
  1937. casting='unsafe',
  1938. op_dtypes=sdt2)
  1939. assert_equal(i[0].dtype, np.dtype(sdt2))
  1940. count = 0
  1941. for x in i:
  1942. assert_equal(x['a'][0], a[count]['a'])
  1943. assert_equal(x['a'][1], a[count]['a'])
  1944. count += 1
  1945. # vector -> matrix (broadcasts and zero-pads)
  1946. sdt1 = [('a', 'f8', (2, 1))]
  1947. sdt2 = [('a', 'f4', (3, 2))]
  1948. a = np.zeros((6,), dtype=sdt1)
  1949. a['a'] = np.arange(6*2).reshape(6, 2, 1)
  1950. i = nditer(a, ['buffered', 'refs_ok'], ['readonly'],
  1951. casting='unsafe',
  1952. op_dtypes=sdt2)
  1953. assert_equal(i[0].dtype, np.dtype(sdt2))
  1954. count = 0
  1955. for x in i:
  1956. assert_equal(x['a'][:2, 0], a[count]['a'][:, 0])
  1957. assert_equal(x['a'][:2, 1], a[count]['a'][:, 0])
  1958. assert_equal(x['a'][2,:], [0, 0])
  1959. count += 1
  1960. # matrix -> matrix (truncates and zero-pads)
  1961. sdt1 = [('a', 'f8', (2, 3))]
  1962. sdt2 = [('a', 'f4', (3, 2))]
  1963. a = np.zeros((6,), dtype=sdt1)
  1964. a['a'] = np.arange(6*2*3).reshape(6, 2, 3)
  1965. i = nditer(a, ['buffered', 'refs_ok'], ['readonly'],
  1966. casting='unsafe',
  1967. op_dtypes=sdt2)
  1968. assert_equal(i[0].dtype, np.dtype(sdt2))
  1969. count = 0
  1970. for x in i:
  1971. assert_equal(x['a'][:2, 0], a[count]['a'][:, 0])
  1972. assert_equal(x['a'][:2, 1], a[count]['a'][:, 1])
  1973. assert_equal(x['a'][2,:], [0, 0])
  1974. count += 1
  1975. def test_iter_buffering_badwriteback():
  1976. # Writing back from a buffer cannot combine elements
  1977. # a needs write buffering, but had a broadcast dimension
  1978. a = np.arange(6).reshape(2, 3, 1)
  1979. b = np.arange(12).reshape(2, 3, 2)
  1980. assert_raises(ValueError, nditer, [a, b],
  1981. ['buffered', 'external_loop'],
  1982. [['readwrite'], ['writeonly']],
  1983. order='C')
  1984. # But if a is readonly, it's fine
  1985. nditer([a, b], ['buffered', 'external_loop'],
  1986. [['readonly'], ['writeonly']],
  1987. order='C')
  1988. # If a has just one element, it's fine too (constant 0 stride, a reduction)
  1989. a = np.arange(1).reshape(1, 1, 1)
  1990. nditer([a, b], ['buffered', 'external_loop', 'reduce_ok'],
  1991. [['readwrite'], ['writeonly']],
  1992. order='C')
  1993. # check that it fails on other dimensions too
  1994. a = np.arange(6).reshape(1, 3, 2)
  1995. assert_raises(ValueError, nditer, [a, b],
  1996. ['buffered', 'external_loop'],
  1997. [['readwrite'], ['writeonly']],
  1998. order='C')
  1999. a = np.arange(4).reshape(2, 1, 2)
  2000. assert_raises(ValueError, nditer, [a, b],
  2001. ['buffered', 'external_loop'],
  2002. [['readwrite'], ['writeonly']],
  2003. order='C')
  2004. def test_iter_buffering_string():
  2005. # Safe casting disallows shrinking strings
  2006. a = np.array(['abc', 'a', 'abcd'], dtype=np.bytes_)
  2007. assert_equal(a.dtype, np.dtype('S4'))
  2008. assert_raises(TypeError, nditer, a, ['buffered'], ['readonly'],
  2009. op_dtypes='S2')
  2010. i = nditer(a, ['buffered'], ['readonly'], op_dtypes='S6')
  2011. assert_equal(i[0], b'abc')
  2012. assert_equal(i[0].dtype, np.dtype('S6'))
  2013. a = np.array(['abc', 'a', 'abcd'], dtype=np.unicode_)
  2014. assert_equal(a.dtype, np.dtype('U4'))
  2015. assert_raises(TypeError, nditer, a, ['buffered'], ['readonly'],
  2016. op_dtypes='U2')
  2017. i = nditer(a, ['buffered'], ['readonly'], op_dtypes='U6')
  2018. assert_equal(i[0], 'abc')
  2019. assert_equal(i[0].dtype, np.dtype('U6'))
  2020. def test_iter_buffering_growinner():
  2021. # Test that the inner loop grows when no buffering is needed
  2022. a = np.arange(30)
  2023. i = nditer(a, ['buffered', 'growinner', 'external_loop'],
  2024. buffersize=5)
  2025. # Should end up with just one inner loop here
  2026. assert_equal(i[0].size, a.size)
  2027. @pytest.mark.slow
  2028. def test_iter_buffered_reduce_reuse():
  2029. # large enough array for all views, including negative strides.
  2030. a = np.arange(2*3**5)[3**5:3**5+1]
  2031. flags = ['buffered', 'delay_bufalloc', 'multi_index', 'reduce_ok', 'refs_ok']
  2032. op_flags = [('readonly',), ('readwrite', 'allocate')]
  2033. op_axes_list = [[(0, 1, 2), (0, 1, -1)], [(0, 1, 2), (0, -1, -1)]]
  2034. # wrong dtype to force buffering
  2035. op_dtypes = [float, a.dtype]
  2036. def get_params():
  2037. for xs in range(-3**2, 3**2 + 1):
  2038. for ys in range(xs, 3**2 + 1):
  2039. for op_axes in op_axes_list:
  2040. # last stride is reduced and because of that not
  2041. # important for this test, as it is the inner stride.
  2042. strides = (xs * a.itemsize, ys * a.itemsize, a.itemsize)
  2043. arr = np.lib.stride_tricks.as_strided(a, (3, 3, 3), strides)
  2044. for skip in [0, 1]:
  2045. yield arr, op_axes, skip
  2046. for arr, op_axes, skip in get_params():
  2047. nditer2 = np.nditer([arr.copy(), None],
  2048. op_axes=op_axes, flags=flags, op_flags=op_flags,
  2049. op_dtypes=op_dtypes)
  2050. with nditer2:
  2051. nditer2.operands[-1][...] = 0
  2052. nditer2.reset()
  2053. nditer2.iterindex = skip
  2054. for (a2_in, b2_in) in nditer2:
  2055. b2_in += a2_in.astype(np.int_)
  2056. comp_res = nditer2.operands[-1]
  2057. for bufsize in range(0, 3**3):
  2058. nditer1 = np.nditer([arr, None],
  2059. op_axes=op_axes, flags=flags, op_flags=op_flags,
  2060. buffersize=bufsize, op_dtypes=op_dtypes)
  2061. with nditer1:
  2062. nditer1.operands[-1][...] = 0
  2063. nditer1.reset()
  2064. nditer1.iterindex = skip
  2065. for (a1_in, b1_in) in nditer1:
  2066. b1_in += a1_in.astype(np.int_)
  2067. res = nditer1.operands[-1]
  2068. assert_array_equal(res, comp_res)
  2069. def test_iter_no_broadcast():
  2070. # Test that the no_broadcast flag works
  2071. a = np.arange(24).reshape(2, 3, 4)
  2072. b = np.arange(6).reshape(2, 3, 1)
  2073. c = np.arange(12).reshape(3, 4)
  2074. nditer([a, b, c], [],
  2075. [['readonly', 'no_broadcast'],
  2076. ['readonly'], ['readonly']])
  2077. assert_raises(ValueError, nditer, [a, b, c], [],
  2078. [['readonly'], ['readonly', 'no_broadcast'], ['readonly']])
  2079. assert_raises(ValueError, nditer, [a, b, c], [],
  2080. [['readonly'], ['readonly'], ['readonly', 'no_broadcast']])
  2081. class TestIterNested:
  2082. def test_basic(self):
  2083. # Test nested iteration basic usage
  2084. a = arange(12).reshape(2, 3, 2)
  2085. i, j = np.nested_iters(a, [[0], [1, 2]])
  2086. vals = [list(j) for _ in i]
  2087. assert_equal(vals, [[0, 1, 2, 3, 4, 5], [6, 7, 8, 9, 10, 11]])
  2088. i, j = np.nested_iters(a, [[0, 1], [2]])
  2089. vals = [list(j) for _ in i]
  2090. assert_equal(vals, [[0, 1], [2, 3], [4, 5], [6, 7], [8, 9], [10, 11]])
  2091. i, j = np.nested_iters(a, [[0, 2], [1]])
  2092. vals = [list(j) for _ in i]
  2093. assert_equal(vals, [[0, 2, 4], [1, 3, 5], [6, 8, 10], [7, 9, 11]])
  2094. def test_reorder(self):
  2095. # Test nested iteration basic usage
  2096. a = arange(12).reshape(2, 3, 2)
  2097. # In 'K' order (default), it gets reordered
  2098. i, j = np.nested_iters(a, [[0], [2, 1]])
  2099. vals = [list(j) for _ in i]
  2100. assert_equal(vals, [[0, 1, 2, 3, 4, 5], [6, 7, 8, 9, 10, 11]])
  2101. i, j = np.nested_iters(a, [[1, 0], [2]])
  2102. vals = [list(j) for _ in i]
  2103. assert_equal(vals, [[0, 1], [2, 3], [4, 5], [6, 7], [8, 9], [10, 11]])
  2104. i, j = np.nested_iters(a, [[2, 0], [1]])
  2105. vals = [list(j) for _ in i]
  2106. assert_equal(vals, [[0, 2, 4], [1, 3, 5], [6, 8, 10], [7, 9, 11]])
  2107. # In 'C' order, it doesn't
  2108. i, j = np.nested_iters(a, [[0], [2, 1]], order='C')
  2109. vals = [list(j) for _ in i]
  2110. assert_equal(vals, [[0, 2, 4, 1, 3, 5], [6, 8, 10, 7, 9, 11]])
  2111. i, j = np.nested_iters(a, [[1, 0], [2]], order='C')
  2112. vals = [list(j) for _ in i]
  2113. assert_equal(vals, [[0, 1], [6, 7], [2, 3], [8, 9], [4, 5], [10, 11]])
  2114. i, j = np.nested_iters(a, [[2, 0], [1]], order='C')
  2115. vals = [list(j) for _ in i]
  2116. assert_equal(vals, [[0, 2, 4], [6, 8, 10], [1, 3, 5], [7, 9, 11]])
  2117. def test_flip_axes(self):
  2118. # Test nested iteration with negative axes
  2119. a = arange(12).reshape(2, 3, 2)[::-1, ::-1, ::-1]
  2120. # In 'K' order (default), the axes all get flipped
  2121. i, j = np.nested_iters(a, [[0], [1, 2]])
  2122. vals = [list(j) for _ in i]
  2123. assert_equal(vals, [[0, 1, 2, 3, 4, 5], [6, 7, 8, 9, 10, 11]])
  2124. i, j = np.nested_iters(a, [[0, 1], [2]])
  2125. vals = [list(j) for _ in i]
  2126. assert_equal(vals, [[0, 1], [2, 3], [4, 5], [6, 7], [8, 9], [10, 11]])
  2127. i, j = np.nested_iters(a, [[0, 2], [1]])
  2128. vals = [list(j) for _ in i]
  2129. assert_equal(vals, [[0, 2, 4], [1, 3, 5], [6, 8, 10], [7, 9, 11]])
  2130. # In 'C' order, flipping axes is disabled
  2131. i, j = np.nested_iters(a, [[0], [1, 2]], order='C')
  2132. vals = [list(j) for _ in i]
  2133. assert_equal(vals, [[11, 10, 9, 8, 7, 6], [5, 4, 3, 2, 1, 0]])
  2134. i, j = np.nested_iters(a, [[0, 1], [2]], order='C')
  2135. vals = [list(j) for _ in i]
  2136. assert_equal(vals, [[11, 10], [9, 8], [7, 6], [5, 4], [3, 2], [1, 0]])
  2137. i, j = np.nested_iters(a, [[0, 2], [1]], order='C')
  2138. vals = [list(j) for _ in i]
  2139. assert_equal(vals, [[11, 9, 7], [10, 8, 6], [5, 3, 1], [4, 2, 0]])
  2140. def test_broadcast(self):
  2141. # Test nested iteration with broadcasting
  2142. a = arange(2).reshape(2, 1)
  2143. b = arange(3).reshape(1, 3)
  2144. i, j = np.nested_iters([a, b], [[0], [1]])
  2145. vals = [list(j) for _ in i]
  2146. assert_equal(vals, [[[0, 0], [0, 1], [0, 2]], [[1, 0], [1, 1], [1, 2]]])
  2147. i, j = np.nested_iters([a, b], [[1], [0]])
  2148. vals = [list(j) for _ in i]
  2149. assert_equal(vals, [[[0, 0], [1, 0]], [[0, 1], [1, 1]], [[0, 2], [1, 2]]])
  2150. def test_dtype_copy(self):
  2151. # Test nested iteration with a copy to change dtype
  2152. # copy
  2153. a = arange(6, dtype='i4').reshape(2, 3)
  2154. i, j = np.nested_iters(a, [[0], [1]],
  2155. op_flags=['readonly', 'copy'],
  2156. op_dtypes='f8')
  2157. assert_equal(j[0].dtype, np.dtype('f8'))
  2158. vals = [list(j) for _ in i]
  2159. assert_equal(vals, [[0, 1, 2], [3, 4, 5]])
  2160. vals = None
  2161. # writebackifcopy - using context manager
  2162. a = arange(6, dtype='f4').reshape(2, 3)
  2163. i, j = np.nested_iters(a, [[0], [1]],
  2164. op_flags=['readwrite', 'updateifcopy'],
  2165. casting='same_kind',
  2166. op_dtypes='f8')
  2167. with i, j:
  2168. assert_equal(j[0].dtype, np.dtype('f8'))
  2169. for x in i:
  2170. for y in j:
  2171. y[...] += 1
  2172. assert_equal(a, [[0, 1, 2], [3, 4, 5]])
  2173. assert_equal(a, [[1, 2, 3], [4, 5, 6]])
  2174. # writebackifcopy - using close()
  2175. a = arange(6, dtype='f4').reshape(2, 3)
  2176. i, j = np.nested_iters(a, [[0], [1]],
  2177. op_flags=['readwrite', 'updateifcopy'],
  2178. casting='same_kind',
  2179. op_dtypes='f8')
  2180. assert_equal(j[0].dtype, np.dtype('f8'))
  2181. for x in i:
  2182. for y in j:
  2183. y[...] += 1
  2184. assert_equal(a, [[0, 1, 2], [3, 4, 5]])
  2185. i.close()
  2186. j.close()
  2187. assert_equal(a, [[1, 2, 3], [4, 5, 6]])
  2188. def test_dtype_buffered(self):
  2189. # Test nested iteration with buffering to change dtype
  2190. a = arange(6, dtype='f4').reshape(2, 3)
  2191. i, j = np.nested_iters(a, [[0], [1]],
  2192. flags=['buffered'],
  2193. op_flags=['readwrite'],
  2194. casting='same_kind',
  2195. op_dtypes='f8')
  2196. assert_equal(j[0].dtype, np.dtype('f8'))
  2197. for x in i:
  2198. for y in j:
  2199. y[...] += 1
  2200. assert_equal(a, [[1, 2, 3], [4, 5, 6]])
  2201. def test_0d(self):
  2202. a = np.arange(12).reshape(2, 3, 2)
  2203. i, j = np.nested_iters(a, [[], [1, 0, 2]])
  2204. vals = [list(j) for _ in i]
  2205. assert_equal(vals, [[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]])
  2206. i, j = np.nested_iters(a, [[1, 0, 2], []])
  2207. vals = [list(j) for _ in i]
  2208. assert_equal(vals, [[0], [1], [2], [3], [4], [5], [6], [7], [8], [9], [10], [11]])
  2209. i, j, k = np.nested_iters(a, [[2, 0], [], [1]])
  2210. vals = []
  2211. for x in i:
  2212. for y in j:
  2213. vals.append([z for z in k])
  2214. assert_equal(vals, [[0, 2, 4], [1, 3, 5], [6, 8, 10], [7, 9, 11]])
  2215. def test_iter_nested_iters_dtype_buffered(self):
  2216. # Test nested iteration with buffering to change dtype
  2217. a = arange(6, dtype='f4').reshape(2, 3)
  2218. i, j = np.nested_iters(a, [[0], [1]],
  2219. flags=['buffered'],
  2220. op_flags=['readwrite'],
  2221. casting='same_kind',
  2222. op_dtypes='f8')
  2223. with i, j:
  2224. assert_equal(j[0].dtype, np.dtype('f8'))
  2225. for x in i:
  2226. for y in j:
  2227. y[...] += 1
  2228. assert_equal(a, [[1, 2, 3], [4, 5, 6]])
  2229. def test_iter_reduction_error():
  2230. a = np.arange(6)
  2231. assert_raises(ValueError, nditer, [a, None], [],
  2232. [['readonly'], ['readwrite', 'allocate']],
  2233. op_axes=[[0], [-1]])
  2234. a = np.arange(6).reshape(2, 3)
  2235. assert_raises(ValueError, nditer, [a, None], ['external_loop'],
  2236. [['readonly'], ['readwrite', 'allocate']],
  2237. op_axes=[[0, 1], [-1, -1]])
  2238. def test_iter_reduction():
  2239. # Test doing reductions with the iterator
  2240. a = np.arange(6)
  2241. i = nditer([a, None], ['reduce_ok'],
  2242. [['readonly'], ['readwrite', 'allocate']],
  2243. op_axes=[[0], [-1]])
  2244. # Need to initialize the output operand to the addition unit
  2245. with i:
  2246. i.operands[1][...] = 0
  2247. # Do the reduction
  2248. for x, y in i:
  2249. y[...] += x
  2250. # Since no axes were specified, should have allocated a scalar
  2251. assert_equal(i.operands[1].ndim, 0)
  2252. assert_equal(i.operands[1], np.sum(a))
  2253. a = np.arange(6).reshape(2, 3)
  2254. i = nditer([a, None], ['reduce_ok', 'external_loop'],
  2255. [['readonly'], ['readwrite', 'allocate']],
  2256. op_axes=[[0, 1], [-1, -1]])
  2257. # Need to initialize the output operand to the addition unit
  2258. with i:
  2259. i.operands[1][...] = 0
  2260. # Reduction shape/strides for the output
  2261. assert_equal(i[1].shape, (6,))
  2262. assert_equal(i[1].strides, (0,))
  2263. # Do the reduction
  2264. for x, y in i:
  2265. # Use a for loop instead of ``y[...] += x``
  2266. # (equivalent to ``y[...] = y[...].copy() + x``),
  2267. # because y has zero strides we use for the reduction
  2268. for j in range(len(y)):
  2269. y[j] += x[j]
  2270. # Since no axes were specified, should have allocated a scalar
  2271. assert_equal(i.operands[1].ndim, 0)
  2272. assert_equal(i.operands[1], np.sum(a))
  2273. # This is a tricky reduction case for the buffering double loop
  2274. # to handle
  2275. a = np.ones((2, 3, 5))
  2276. it1 = nditer([a, None], ['reduce_ok', 'external_loop'],
  2277. [['readonly'], ['readwrite', 'allocate']],
  2278. op_axes=[None, [0, -1, 1]])
  2279. it2 = nditer([a, None], ['reduce_ok', 'external_loop',
  2280. 'buffered', 'delay_bufalloc'],
  2281. [['readonly'], ['readwrite', 'allocate']],
  2282. op_axes=[None, [0, -1, 1]], buffersize=10)
  2283. with it1, it2:
  2284. it1.operands[1].fill(0)
  2285. it2.operands[1].fill(0)
  2286. it2.reset()
  2287. for x in it1:
  2288. x[1][...] += x[0]
  2289. for x in it2:
  2290. x[1][...] += x[0]
  2291. assert_equal(it1.operands[1], it2.operands[1])
  2292. assert_equal(it2.operands[1].sum(), a.size)
  2293. def test_iter_buffering_reduction():
  2294. # Test doing buffered reductions with the iterator
  2295. a = np.arange(6)
  2296. b = np.array(0., dtype='f8').byteswap().newbyteorder()
  2297. i = nditer([a, b], ['reduce_ok', 'buffered'],
  2298. [['readonly'], ['readwrite', 'nbo']],
  2299. op_axes=[[0], [-1]])
  2300. with i:
  2301. assert_equal(i[1].dtype, np.dtype('f8'))
  2302. assert_(i[1].dtype != b.dtype)
  2303. # Do the reduction
  2304. for x, y in i:
  2305. y[...] += x
  2306. # Since no axes were specified, should have allocated a scalar
  2307. assert_equal(b, np.sum(a))
  2308. a = np.arange(6).reshape(2, 3)
  2309. b = np.array([0, 0], dtype='f8').byteswap().newbyteorder()
  2310. i = nditer([a, b], ['reduce_ok', 'external_loop', 'buffered'],
  2311. [['readonly'], ['readwrite', 'nbo']],
  2312. op_axes=[[0, 1], [0, -1]])
  2313. # Reduction shape/strides for the output
  2314. with i:
  2315. assert_equal(i[1].shape, (3,))
  2316. assert_equal(i[1].strides, (0,))
  2317. # Do the reduction
  2318. for x, y in i:
  2319. # Use a for loop instead of ``y[...] += x``
  2320. # (equivalent to ``y[...] = y[...].copy() + x``),
  2321. # because y has zero strides we use for the reduction
  2322. for j in range(len(y)):
  2323. y[j] += x[j]
  2324. assert_equal(b, np.sum(a, axis=1))
  2325. # Iterator inner double loop was wrong on this one
  2326. p = np.arange(2) + 1
  2327. it = np.nditer([p, None],
  2328. ['delay_bufalloc', 'reduce_ok', 'buffered', 'external_loop'],
  2329. [['readonly'], ['readwrite', 'allocate']],
  2330. op_axes=[[-1, 0], [-1, -1]],
  2331. itershape=(2, 2))
  2332. with it:
  2333. it.operands[1].fill(0)
  2334. it.reset()
  2335. assert_equal(it[0], [1, 2, 1, 2])
  2336. # Iterator inner loop should take argument contiguity into account
  2337. x = np.ones((7, 13, 8), np.int8)[4:6,1:11:6,1:5].transpose(1, 2, 0)
  2338. x[...] = np.arange(x.size).reshape(x.shape)
  2339. y_base = np.arange(4*4, dtype=np.int8).reshape(4, 4)
  2340. y_base_copy = y_base.copy()
  2341. y = y_base[::2,:,None]
  2342. it = np.nditer([y, x],
  2343. ['buffered', 'external_loop', 'reduce_ok'],
  2344. [['readwrite'], ['readonly']])
  2345. with it:
  2346. for a, b in it:
  2347. a.fill(2)
  2348. assert_equal(y_base[1::2], y_base_copy[1::2])
  2349. assert_equal(y_base[::2], 2)
  2350. def test_iter_buffering_reduction_reuse_reduce_loops():
  2351. # There was a bug triggering reuse of the reduce loop inappropriately,
  2352. # which caused processing to happen in unnecessarily small chunks
  2353. # and overran the buffer.
  2354. a = np.zeros((2, 7))
  2355. b = np.zeros((1, 7))
  2356. it = np.nditer([a, b], flags=['reduce_ok', 'external_loop', 'buffered'],
  2357. op_flags=[['readonly'], ['readwrite']],
  2358. buffersize=5)
  2359. with it:
  2360. bufsizes = [x.shape[0] for x, y in it]
  2361. assert_equal(bufsizes, [5, 2, 5, 2])
  2362. assert_equal(sum(bufsizes), a.size)
  2363. def test_iter_writemasked_badinput():
  2364. a = np.zeros((2, 3))
  2365. b = np.zeros((3,))
  2366. m = np.array([[True, True, False], [False, True, False]])
  2367. m2 = np.array([True, True, False])
  2368. m3 = np.array([0, 1, 1], dtype='u1')
  2369. mbad1 = np.array([0, 1, 1], dtype='i1')
  2370. mbad2 = np.array([0, 1, 1], dtype='f4')
  2371. # Need an 'arraymask' if any operand is 'writemasked'
  2372. assert_raises(ValueError, nditer, [a, m], [],
  2373. [['readwrite', 'writemasked'], ['readonly']])
  2374. # A 'writemasked' operand must not be readonly
  2375. assert_raises(ValueError, nditer, [a, m], [],
  2376. [['readonly', 'writemasked'], ['readonly', 'arraymask']])
  2377. # 'writemasked' and 'arraymask' may not be used together
  2378. assert_raises(ValueError, nditer, [a, m], [],
  2379. [['readonly'], ['readwrite', 'arraymask', 'writemasked']])
  2380. # 'arraymask' may only be specified once
  2381. assert_raises(ValueError, nditer, [a, m, m2], [],
  2382. [['readwrite', 'writemasked'],
  2383. ['readonly', 'arraymask'],
  2384. ['readonly', 'arraymask']])
  2385. # An 'arraymask' with nothing 'writemasked' also doesn't make sense
  2386. assert_raises(ValueError, nditer, [a, m], [],
  2387. [['readwrite'], ['readonly', 'arraymask']])
  2388. # A writemasked reduction requires a similarly smaller mask
  2389. assert_raises(ValueError, nditer, [a, b, m], ['reduce_ok'],
  2390. [['readonly'],
  2391. ['readwrite', 'writemasked'],
  2392. ['readonly', 'arraymask']])
  2393. # But this should work with a smaller/equal mask to the reduction operand
  2394. np.nditer([a, b, m2], ['reduce_ok'],
  2395. [['readonly'],
  2396. ['readwrite', 'writemasked'],
  2397. ['readonly', 'arraymask']])
  2398. # The arraymask itself cannot be a reduction
  2399. assert_raises(ValueError, nditer, [a, b, m2], ['reduce_ok'],
  2400. [['readonly'],
  2401. ['readwrite', 'writemasked'],
  2402. ['readwrite', 'arraymask']])
  2403. # A uint8 mask is ok too
  2404. np.nditer([a, m3], ['buffered'],
  2405. [['readwrite', 'writemasked'],
  2406. ['readonly', 'arraymask']],
  2407. op_dtypes=['f4', None],
  2408. casting='same_kind')
  2409. # An int8 mask isn't ok
  2410. assert_raises(TypeError, np.nditer, [a, mbad1], ['buffered'],
  2411. [['readwrite', 'writemasked'],
  2412. ['readonly', 'arraymask']],
  2413. op_dtypes=['f4', None],
  2414. casting='same_kind')
  2415. # A float32 mask isn't ok
  2416. assert_raises(TypeError, np.nditer, [a, mbad2], ['buffered'],
  2417. [['readwrite', 'writemasked'],
  2418. ['readonly', 'arraymask']],
  2419. op_dtypes=['f4', None],
  2420. casting='same_kind')
  2421. def _is_buffered(iterator):
  2422. try:
  2423. iterator.itviews
  2424. except ValueError:
  2425. return True
  2426. return False
  2427. @pytest.mark.parametrize("a",
  2428. [np.zeros((3,), dtype='f8'),
  2429. np.zeros((9876, 3*5), dtype='f8')[::2, :],
  2430. np.zeros((4, 312, 124, 3), dtype='f8')[::2, :, ::2, :],
  2431. # Also test with the last dimension strided (so it does not fit if
  2432. # there is repeated access)
  2433. np.zeros((9,), dtype='f8')[::3],
  2434. np.zeros((9876, 3*10), dtype='f8')[::2, ::5],
  2435. np.zeros((4, 312, 124, 3), dtype='f8')[::2, :, ::2, ::-1]])
  2436. def test_iter_writemasked(a):
  2437. # Note, the slicing above is to ensure that nditer cannot combine multiple
  2438. # axes into one. The repetition is just to make things a bit more
  2439. # interesting.
  2440. shape = a.shape
  2441. reps = shape[-1] // 3
  2442. msk = np.empty(shape, dtype=bool)
  2443. msk[...] = [True, True, False] * reps
  2444. # When buffering is unused, 'writemasked' effectively does nothing.
  2445. # It's up to the user of the iterator to obey the requested semantics.
  2446. it = np.nditer([a, msk], [],
  2447. [['readwrite', 'writemasked'],
  2448. ['readonly', 'arraymask']])
  2449. with it:
  2450. for x, m in it:
  2451. x[...] = 1
  2452. # Because we violated the semantics, all the values became 1
  2453. assert_equal(a, np.broadcast_to([1, 1, 1] * reps, shape))
  2454. # Even if buffering is enabled, we still may be accessing the array
  2455. # directly.
  2456. it = np.nditer([a, msk], ['buffered'],
  2457. [['readwrite', 'writemasked'],
  2458. ['readonly', 'arraymask']])
  2459. # @seberg: I honestly don't currently understand why a "buffered" iterator
  2460. # would end up not using a buffer for the small array here at least when
  2461. # "writemasked" is used, that seems confusing... Check by testing for
  2462. # actual memory overlap!
  2463. is_buffered = True
  2464. with it:
  2465. for x, m in it:
  2466. x[...] = 2.5
  2467. if np.may_share_memory(x, a):
  2468. is_buffered = False
  2469. if not is_buffered:
  2470. # Because we violated the semantics, all the values became 2.5
  2471. assert_equal(a, np.broadcast_to([2.5, 2.5, 2.5] * reps, shape))
  2472. else:
  2473. # For large sizes, the iterator may be buffered:
  2474. assert_equal(a, np.broadcast_to([2.5, 2.5, 1] * reps, shape))
  2475. a[...] = 2.5
  2476. # If buffering will definitely happening, for instance because of
  2477. # a cast, only the items selected by the mask will be copied back from
  2478. # the buffer.
  2479. it = np.nditer([a, msk], ['buffered'],
  2480. [['readwrite', 'writemasked'],
  2481. ['readonly', 'arraymask']],
  2482. op_dtypes=['i8', None],
  2483. casting='unsafe')
  2484. with it:
  2485. for x, m in it:
  2486. x[...] = 3
  2487. # Even though we violated the semantics, only the selected values
  2488. # were copied back
  2489. assert_equal(a, np.broadcast_to([3, 3, 2.5] * reps, shape))
  2490. @pytest.mark.parametrize(["mask", "mask_axes"], [
  2491. # Allocated operand (only broadcasts with -1)
  2492. (None, [-1, 0]),
  2493. # Reduction along the first dimension (with and without op_axes)
  2494. (np.zeros((1, 4), dtype="bool"), [0, 1]),
  2495. (np.zeros((1, 4), dtype="bool"), None),
  2496. # Test 0-D and -1 op_axes
  2497. (np.zeros(4, dtype="bool"), [-1, 0]),
  2498. (np.zeros((), dtype="bool"), [-1, -1]),
  2499. (np.zeros((), dtype="bool"), None)])
  2500. def test_iter_writemasked_broadcast_error(mask, mask_axes):
  2501. # This assumes that a readwrite mask makes sense. This is likely not the
  2502. # case and should simply be deprecated.
  2503. arr = np.zeros((3, 4))
  2504. itflags = ["reduce_ok"]
  2505. mask_flags = ["arraymask", "readwrite", "allocate"]
  2506. a_flags = ["writeonly", "writemasked"]
  2507. if mask_axes is None:
  2508. op_axes = None
  2509. else:
  2510. op_axes = [mask_axes, [0, 1]]
  2511. with assert_raises(ValueError):
  2512. np.nditer((mask, arr), flags=itflags, op_flags=[mask_flags, a_flags],
  2513. op_axes=op_axes)
  2514. def test_iter_writemasked_decref():
  2515. # force casting (to make it interesting) by using a structured dtype.
  2516. arr = np.arange(10000).astype(">i,O")
  2517. original = arr.copy()
  2518. mask = np.random.randint(0, 2, size=10000).astype(bool)
  2519. it = np.nditer([arr, mask], ['buffered', "refs_ok"],
  2520. [['readwrite', 'writemasked'],
  2521. ['readonly', 'arraymask']],
  2522. op_dtypes=["<i,O", "?"])
  2523. singleton = object()
  2524. if HAS_REFCOUNT:
  2525. count = sys.getrefcount(singleton)
  2526. for buf, mask_buf in it:
  2527. buf[...] = (3, singleton)
  2528. del buf, mask_buf, it # delete everything to ensure correct cleanup
  2529. if HAS_REFCOUNT:
  2530. # The buffer would have included additional items, they must be
  2531. # cleared correctly:
  2532. assert sys.getrefcount(singleton) - count == np.count_nonzero(mask)
  2533. assert_array_equal(arr[~mask], original[~mask])
  2534. assert (arr[mask] == np.array((3, singleton), arr.dtype)).all()
  2535. del arr
  2536. if HAS_REFCOUNT:
  2537. assert sys.getrefcount(singleton) == count
  2538. def test_iter_non_writable_attribute_deletion():
  2539. it = np.nditer(np.ones(2))
  2540. attr = ["value", "shape", "operands", "itviews", "has_delayed_bufalloc",
  2541. "iterationneedsapi", "has_multi_index", "has_index", "dtypes",
  2542. "ndim", "nop", "itersize", "finished"]
  2543. for s in attr:
  2544. assert_raises(AttributeError, delattr, it, s)
  2545. def test_iter_writable_attribute_deletion():
  2546. it = np.nditer(np.ones(2))
  2547. attr = [ "multi_index", "index", "iterrange", "iterindex"]
  2548. for s in attr:
  2549. assert_raises(AttributeError, delattr, it, s)
  2550. def test_iter_element_deletion():
  2551. it = np.nditer(np.ones(3))
  2552. try:
  2553. del it[1]
  2554. del it[1:2]
  2555. except TypeError:
  2556. pass
  2557. except Exception:
  2558. raise AssertionError
  2559. def test_iter_allocated_array_dtypes():
  2560. # If the dtype of an allocated output has a shape, the shape gets
  2561. # tacked onto the end of the result.
  2562. it = np.nditer(([1, 3, 20], None), op_dtypes=[None, ('i4', (2,))])
  2563. for a, b in it:
  2564. b[0] = a - 1
  2565. b[1] = a + 1
  2566. assert_equal(it.operands[1], [[0, 2], [2, 4], [19, 21]])
  2567. # Check the same (less sensitive) thing when `op_axes` with -1 is given.
  2568. it = np.nditer(([[1, 3, 20]], None), op_dtypes=[None, ('i4', (2,))],
  2569. flags=["reduce_ok"], op_axes=[None, (-1, 0)])
  2570. for a, b in it:
  2571. b[0] = a - 1
  2572. b[1] = a + 1
  2573. assert_equal(it.operands[1], [[0, 2], [2, 4], [19, 21]])
  2574. # Make sure this works for scalars too
  2575. it = np.nditer((10, 2, None), op_dtypes=[None, None, ('i4', (2, 2))])
  2576. for a, b, c in it:
  2577. c[0, 0] = a - b
  2578. c[0, 1] = a + b
  2579. c[1, 0] = a * b
  2580. c[1, 1] = a / b
  2581. assert_equal(it.operands[2], [[8, 12], [20, 5]])
  2582. def test_0d_iter():
  2583. # Basic test for iteration of 0-d arrays:
  2584. i = nditer([2, 3], ['multi_index'], [['readonly']]*2)
  2585. assert_equal(i.ndim, 0)
  2586. assert_equal(next(i), (2, 3))
  2587. assert_equal(i.multi_index, ())
  2588. assert_equal(i.iterindex, 0)
  2589. assert_raises(StopIteration, next, i)
  2590. # test reset:
  2591. i.reset()
  2592. assert_equal(next(i), (2, 3))
  2593. assert_raises(StopIteration, next, i)
  2594. # test forcing to 0-d
  2595. i = nditer(np.arange(5), ['multi_index'], [['readonly']], op_axes=[()])
  2596. assert_equal(i.ndim, 0)
  2597. assert_equal(len(i), 1)
  2598. i = nditer(np.arange(5), ['multi_index'], [['readonly']],
  2599. op_axes=[()], itershape=())
  2600. assert_equal(i.ndim, 0)
  2601. assert_equal(len(i), 1)
  2602. # passing an itershape alone is not enough, the op_axes are also needed
  2603. with assert_raises(ValueError):
  2604. nditer(np.arange(5), ['multi_index'], [['readonly']], itershape=())
  2605. # Test a more complex buffered casting case (same as another test above)
  2606. sdt = [('a', 'f4'), ('b', 'i8'), ('c', 'c8', (2, 3)), ('d', 'O')]
  2607. a = np.array(0.5, dtype='f4')
  2608. i = nditer(a, ['buffered', 'refs_ok'], ['readonly'],
  2609. casting='unsafe', op_dtypes=sdt)
  2610. vals = next(i)
  2611. assert_equal(vals['a'], 0.5)
  2612. assert_equal(vals['b'], 0)
  2613. assert_equal(vals['c'], [[(0.5)]*3]*2)
  2614. assert_equal(vals['d'], 0.5)
  2615. def test_object_iter_cleanup():
  2616. # see gh-18450
  2617. # object arrays can raise a python exception in ufunc inner loops using
  2618. # nditer, which should cause iteration to stop & cleanup. There were bugs
  2619. # in the nditer cleanup when decref'ing object arrays.
  2620. # This test would trigger valgrind "uninitialized read" before the bugfix.
  2621. assert_raises(TypeError, lambda: np.zeros((17000, 2), dtype='f4') * None)
  2622. # this more explicit code also triggers the invalid access
  2623. arr = np.arange(np.BUFSIZE * 10).reshape(10, -1).astype(str)
  2624. oarr = arr.astype(object)
  2625. oarr[:, -1] = None
  2626. assert_raises(TypeError, lambda: np.add(oarr[:, ::-1], arr[:, ::-1]))
  2627. # followup: this tests for a bug introduced in the first pass of gh-18450,
  2628. # caused by an incorrect fallthrough of the TypeError
  2629. class T:
  2630. def __bool__(self):
  2631. raise TypeError("Ambiguous")
  2632. assert_raises(TypeError, np.logical_or.reduce,
  2633. np.array([T(), T()], dtype='O'))
  2634. def test_object_iter_cleanup_reduce():
  2635. # Similar as above, but a complex reduction case that was previously
  2636. # missed (see gh-18810).
  2637. # The following array is special in that it cannot be flattened:
  2638. arr = np.array([[None, 1], [-1, -1], [None, 2], [-1, -1]])[::2]
  2639. with pytest.raises(TypeError):
  2640. np.sum(arr)
  2641. @pytest.mark.parametrize("arr", [
  2642. np.ones((8000, 4, 2), dtype=object)[:, ::2, :],
  2643. np.ones((8000, 4, 2), dtype=object, order="F")[:, ::2, :],
  2644. np.ones((8000, 4, 2), dtype=object)[:, ::2, :].copy("F")])
  2645. def test_object_iter_cleanup_large_reduce(arr):
  2646. # More complicated calls are possible for large arrays:
  2647. out = np.ones(8000, dtype=np.intp)
  2648. # force casting with `dtype=object`
  2649. res = np.sum(arr, axis=(1, 2), dtype=object, out=out)
  2650. assert_array_equal(res, np.full(8000, 4, dtype=object))
  2651. def test_iter_too_large():
  2652. # The total size of the iterator must not exceed the maximum intp due
  2653. # to broadcasting. Dividing by 1024 will keep it small enough to
  2654. # give a legal array.
  2655. size = np.iinfo(np.intp).max // 1024
  2656. arr = np.lib.stride_tricks.as_strided(np.zeros(1), (size,), (0,))
  2657. assert_raises(ValueError, nditer, (arr, arr[:, None]))
  2658. # test the same for multiindex. That may get more interesting when
  2659. # removing 0 dimensional axis is allowed (since an iterator can grow then)
  2660. assert_raises(ValueError, nditer,
  2661. (arr, arr[:, None]), flags=['multi_index'])
  2662. def test_iter_too_large_with_multiindex():
  2663. # When a multi index is being tracked, the error is delayed this
  2664. # checks the delayed error messages and getting below that by
  2665. # removing an axis.
  2666. base_size = 2**10
  2667. num = 1
  2668. while base_size**num < np.iinfo(np.intp).max:
  2669. num += 1
  2670. shape_template = [1, 1] * num
  2671. arrays = []
  2672. for i in range(num):
  2673. shape = shape_template[:]
  2674. shape[i * 2] = 2**10
  2675. arrays.append(np.empty(shape))
  2676. arrays = tuple(arrays)
  2677. # arrays are now too large to be broadcast. The different modes test
  2678. # different nditer functionality with or without GIL.
  2679. for mode in range(6):
  2680. with assert_raises(ValueError):
  2681. _multiarray_tests.test_nditer_too_large(arrays, -1, mode)
  2682. # but if we do nothing with the nditer, it can be constructed:
  2683. _multiarray_tests.test_nditer_too_large(arrays, -1, 7)
  2684. # When an axis is removed, things should work again (half the time):
  2685. for i in range(num):
  2686. for mode in range(6):
  2687. # an axis with size 1024 is removed:
  2688. _multiarray_tests.test_nditer_too_large(arrays, i*2, mode)
  2689. # an axis with size 1 is removed:
  2690. with assert_raises(ValueError):
  2691. _multiarray_tests.test_nditer_too_large(arrays, i*2 + 1, mode)
  2692. def test_writebacks():
  2693. a = np.arange(6, dtype='f4')
  2694. au = a.byteswap().newbyteorder()
  2695. assert_(a.dtype.byteorder != au.dtype.byteorder)
  2696. it = nditer(au, [], [['readwrite', 'updateifcopy']],
  2697. casting='equiv', op_dtypes=[np.dtype('f4')])
  2698. with it:
  2699. it.operands[0][:] = 100
  2700. assert_equal(au, 100)
  2701. # do it again, this time raise an error,
  2702. it = nditer(au, [], [['readwrite', 'updateifcopy']],
  2703. casting='equiv', op_dtypes=[np.dtype('f4')])
  2704. try:
  2705. with it:
  2706. assert_equal(au.flags.writeable, False)
  2707. it.operands[0][:] = 0
  2708. raise ValueError('exit context manager on exception')
  2709. except:
  2710. pass
  2711. assert_equal(au, 0)
  2712. assert_equal(au.flags.writeable, True)
  2713. # cannot reuse i outside context manager
  2714. assert_raises(ValueError, getattr, it, 'operands')
  2715. it = nditer(au, [], [['readwrite', 'updateifcopy']],
  2716. casting='equiv', op_dtypes=[np.dtype('f4')])
  2717. with it:
  2718. x = it.operands[0]
  2719. x[:] = 6
  2720. assert_(x.flags.writebackifcopy)
  2721. assert_equal(au, 6)
  2722. assert_(not x.flags.writebackifcopy)
  2723. x[:] = 123 # x.data still valid
  2724. assert_equal(au, 6) # but not connected to au
  2725. it = nditer(au, [],
  2726. [['readwrite', 'updateifcopy']],
  2727. casting='equiv', op_dtypes=[np.dtype('f4')])
  2728. # reentering works
  2729. with it:
  2730. with it:
  2731. for x in it:
  2732. x[...] = 123
  2733. it = nditer(au, [],
  2734. [['readwrite', 'updateifcopy']],
  2735. casting='equiv', op_dtypes=[np.dtype('f4')])
  2736. # make sure exiting the inner context manager closes the iterator
  2737. with it:
  2738. with it:
  2739. for x in it:
  2740. x[...] = 123
  2741. assert_raises(ValueError, getattr, it, 'operands')
  2742. # do not crash if original data array is decrefed
  2743. it = nditer(au, [],
  2744. [['readwrite', 'updateifcopy']],
  2745. casting='equiv', op_dtypes=[np.dtype('f4')])
  2746. del au
  2747. with it:
  2748. for x in it:
  2749. x[...] = 123
  2750. # make sure we cannot reenter the closed iterator
  2751. enter = it.__enter__
  2752. assert_raises(RuntimeError, enter)
  2753. def test_close_equivalent():
  2754. ''' using a context amanger and using nditer.close are equivalent
  2755. '''
  2756. def add_close(x, y, out=None):
  2757. addop = np.add
  2758. it = np.nditer([x, y, out], [],
  2759. [['readonly'], ['readonly'], ['writeonly','allocate']])
  2760. for (a, b, c) in it:
  2761. addop(a, b, out=c)
  2762. ret = it.operands[2]
  2763. it.close()
  2764. return ret
  2765. def add_context(x, y, out=None):
  2766. addop = np.add
  2767. it = np.nditer([x, y, out], [],
  2768. [['readonly'], ['readonly'], ['writeonly','allocate']])
  2769. with it:
  2770. for (a, b, c) in it:
  2771. addop(a, b, out=c)
  2772. return it.operands[2]
  2773. z = add_close(range(5), range(5))
  2774. assert_equal(z, range(0, 10, 2))
  2775. z = add_context(range(5), range(5))
  2776. assert_equal(z, range(0, 10, 2))
  2777. def test_close_raises():
  2778. it = np.nditer(np.arange(3))
  2779. assert_equal (next(it), 0)
  2780. it.close()
  2781. assert_raises(StopIteration, next, it)
  2782. assert_raises(ValueError, getattr, it, 'operands')
  2783. def test_close_parameters():
  2784. it = np.nditer(np.arange(3))
  2785. assert_raises(TypeError, it.close, 1)
  2786. @pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts")
  2787. def test_warn_noclose():
  2788. a = np.arange(6, dtype='f4')
  2789. au = a.byteswap().newbyteorder()
  2790. with suppress_warnings() as sup:
  2791. sup.record(RuntimeWarning)
  2792. it = np.nditer(au, [], [['readwrite', 'updateifcopy']],
  2793. casting='equiv', op_dtypes=[np.dtype('f4')])
  2794. del it
  2795. assert len(sup.log) == 1
  2796. @pytest.mark.skipif(sys.version_info[:2] == (3, 9) and sys.platform == "win32",
  2797. reason="Errors with Python 3.9 on Windows")
  2798. @pytest.mark.parametrize(["in_dtype", "buf_dtype"],
  2799. [("i", "O"), ("O", "i"), # most simple cases
  2800. ("i,O", "O,O"), # structured partially only copying O
  2801. ("O,i", "i,O"), # structured casting to and from O
  2802. ])
  2803. @pytest.mark.parametrize("steps", [1, 2, 3])
  2804. def test_partial_iteration_cleanup(in_dtype, buf_dtype, steps):
  2805. """
  2806. Checks for reference counting leaks during cleanup. Using explicit
  2807. reference counts lead to occasional false positives (at least in parallel
  2808. test setups). This test now should still test leaks correctly when
  2809. run e.g. with pytest-valgrind or pytest-leaks
  2810. """
  2811. value = 2**30 + 1 # just a random value that Python won't intern
  2812. arr = np.full(int(np.BUFSIZE * 2.5), value).astype(in_dtype)
  2813. it = np.nditer(arr, op_dtypes=[np.dtype(buf_dtype)],
  2814. flags=["buffered", "external_loop", "refs_ok"], casting="unsafe")
  2815. for step in range(steps):
  2816. # The iteration finishes in 3 steps, the first two are partial
  2817. next(it)
  2818. del it # not necessary, but we test the cleanup
  2819. # Repeat the test with `iternext`
  2820. it = np.nditer(arr, op_dtypes=[np.dtype(buf_dtype)],
  2821. flags=["buffered", "external_loop", "refs_ok"], casting="unsafe")
  2822. for step in range(steps):
  2823. it.iternext()
  2824. del it # not necessary, but we test the cleanup
  2825. @pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts")
  2826. @pytest.mark.parametrize(["in_dtype", "buf_dtype"],
  2827. [("O", "i"), # most simple cases
  2828. ("O,i", "i,O"), # structured casting to and from O
  2829. ])
  2830. def test_partial_iteration_error(in_dtype, buf_dtype):
  2831. value = 123 # relies on python cache (leak-check will still find it)
  2832. arr = np.full(int(np.BUFSIZE * 2.5), value).astype(in_dtype)
  2833. if in_dtype == "O":
  2834. arr[int(np.BUFSIZE * 1.5)] = None
  2835. else:
  2836. arr[int(np.BUFSIZE * 1.5)]["f0"] = None
  2837. count = sys.getrefcount(value)
  2838. it = np.nditer(arr, op_dtypes=[np.dtype(buf_dtype)],
  2839. flags=["buffered", "external_loop", "refs_ok"], casting="unsafe")
  2840. with pytest.raises(TypeError):
  2841. # pytest.raises seems to have issues with the error originating
  2842. # in the for loop, so manually unravel:
  2843. next(it)
  2844. next(it) # raises TypeError
  2845. # Repeat the test with `iternext` after resetting, the buffers should
  2846. # already be cleared from any references, so resetting is sufficient.
  2847. it.reset()
  2848. with pytest.raises(TypeError):
  2849. it.iternext()
  2850. it.iternext()
  2851. assert count == sys.getrefcount(value)
  2852. def test_debug_print(capfd):
  2853. """
  2854. Matches the expected output of a debug print with the actual output.
  2855. Note that the iterator dump should not be considered stable API,
  2856. this test is mainly to ensure the print does not crash.
  2857. Currently uses a subprocess to avoid dealing with the C level `printf`s.
  2858. """
  2859. # the expected output with all addresses and sizes stripped (they vary
  2860. # and/or are platform dependent).
  2861. expected = """
  2862. ------ BEGIN ITERATOR DUMP ------
  2863. | Iterator Address:
  2864. | ItFlags: BUFFER REDUCE REUSE_REDUCE_LOOPS
  2865. | NDim: 2
  2866. | NOp: 2
  2867. | IterSize: 50
  2868. | IterStart: 0
  2869. | IterEnd: 50
  2870. | IterIndex: 0
  2871. | Iterator SizeOf:
  2872. | BufferData SizeOf:
  2873. | AxisData SizeOf:
  2874. |
  2875. | Perm: 0 1
  2876. | DTypes:
  2877. | DTypes: dtype('float64') dtype('int32')
  2878. | InitDataPtrs:
  2879. | BaseOffsets: 0 0
  2880. | Operands:
  2881. | Operand DTypes: dtype('int64') dtype('float64')
  2882. | OpItFlags:
  2883. | Flags[0]: READ CAST ALIGNED
  2884. | Flags[1]: READ WRITE CAST ALIGNED REDUCE
  2885. |
  2886. | BufferData:
  2887. | BufferSize: 50
  2888. | Size: 5
  2889. | BufIterEnd: 5
  2890. | REDUCE Pos: 0
  2891. | REDUCE OuterSize: 10
  2892. | REDUCE OuterDim: 1
  2893. | Strides: 8 4
  2894. | Ptrs:
  2895. | REDUCE Outer Strides: 40 0
  2896. | REDUCE Outer Ptrs:
  2897. | ReadTransferFn:
  2898. | ReadTransferData:
  2899. | WriteTransferFn:
  2900. | WriteTransferData:
  2901. | Buffers:
  2902. |
  2903. | AxisData[0]:
  2904. | Shape: 5
  2905. | Index: 0
  2906. | Strides: 16 8
  2907. | Ptrs:
  2908. | AxisData[1]:
  2909. | Shape: 10
  2910. | Index: 0
  2911. | Strides: 80 0
  2912. | Ptrs:
  2913. ------- END ITERATOR DUMP -------
  2914. """.strip().splitlines()
  2915. arr1 = np.arange(100, dtype=np.int64).reshape(10, 10)[:, ::2]
  2916. arr2 = np.arange(5.)
  2917. it = np.nditer((arr1, arr2), op_dtypes=["d", "i4"], casting="unsafe",
  2918. flags=["reduce_ok", "buffered"],
  2919. op_flags=[["readonly"], ["readwrite"]])
  2920. it.debug_print()
  2921. res = capfd.readouterr().out
  2922. res = res.strip().splitlines()
  2923. assert len(res) == len(expected)
  2924. for res_line, expected_line in zip(res, expected):
  2925. # The actual output may have additional pointers listed that are
  2926. # stripped from the example output:
  2927. assert res_line.startswith(expected_line.strip())