12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730173117321733173417351736173717381739174017411742174317441745174617471748174917501751175217531754175517561757175817591760176117621763176417651766176717681769177017711772177317741775177617771778177917801781178217831784178517861787178817891790179117921793179417951796179717981799180018011802180318041805180618071808180918101811181218131814181518161817181818191820182118221823182418251826182718281829183018311832183318341835183618371838183918401841184218431844184518461847184818491850185118521853185418551856185718581859186018611862186318641865186618671868186918701871187218731874187518761877187818791880188118821883188418851886188718881889189018911892189318941895189618971898189919001901190219031904190519061907190819091910191119121913191419151916191719181919192019211922192319241925192619271928192919301931193219331934193519361937193819391940194119421943194419451946194719481949195019511952195319541955195619571958195919601961196219631964196519661967196819691970197119721973197419751976197719781979198019811982198319841985198619871988198919901991199219931994199519961997199819992000200120022003200420052006200720082009201020112012201320142015201620172018201920202021202220232024202520262027202820292030203120322033203420352036203720382039204020412042204320442045204620472048204920502051205220532054205520562057205820592060206120622063206420652066206720682069207020712072207320742075207620772078207920802081208220832084208520862087208820892090209120922093209420952096209720982099210021012102210321042105210621072108210921102111211221132114211521162117211821192120212121222123212421252126212721282129213021312132213321342135213621372138213921402141214221432144214521462147214821492150215121522153215421552156215721582159216021612162216321642165216621672168216921702171217221732174217521762177217821792180218121822183218421852186218721882189219021912192219321942195219621972198219922002201220222032204220522062207220822092210221122122213221422152216221722182219222022212222222322242225222622272228222922302231223222332234223522362237223822392240224122422243224422452246224722482249225022512252225322542255225622572258225922602261226222632264226522662267226822692270227122722273227422752276227722782279228022812282228322842285228622872288228922902291229222932294229522962297229822992300230123022303230423052306230723082309231023112312231323142315231623172318231923202321232223232324232523262327232823292330233123322333233423352336233723382339234023412342234323442345234623472348234923502351235223532354235523562357235823592360236123622363236423652366236723682369237023712372237323742375237623772378237923802381238223832384238523862387238823892390239123922393239423952396239723982399240024012402240324042405240624072408240924102411241224132414241524162417241824192420242124222423242424252426242724282429243024312432243324342435243624372438243924402441244224432444244524462447244824492450245124522453245424552456245724582459246024612462246324642465246624672468246924702471247224732474247524762477247824792480248124822483248424852486248724882489249024912492249324942495249624972498249925002501250225032504250525062507250825092510251125122513251425152516251725182519252025212522252325242525252625272528252925302531253225332534253525362537253825392540254125422543254425452546254725482549255025512552255325542555255625572558255925602561256225632564256525662567256825692570257125722573257425752576257725782579258025812582258325842585258625872588258925902591259225932594259525962597259825992600260126022603260426052606260726082609261026112612261326142615261626172618261926202621262226232624262526262627262826292630263126322633263426352636263726382639264026412642264326442645264626472648264926502651265226532654265526562657265826592660266126622663266426652666266726682669267026712672267326742675267626772678267926802681268226832684268526862687268826892690269126922693269426952696269726982699270027012702270327042705270627072708270927102711271227132714271527162717271827192720272127222723272427252726272727282729273027312732273327342735273627372738273927402741274227432744274527462747274827492750275127522753275427552756275727582759276027612762276327642765276627672768276927702771277227732774277527762777277827792780278127822783278427852786278727882789279027912792279327942795279627972798279928002801280228032804280528062807280828092810281128122813281428152816281728182819282028212822282328242825282628272828282928302831283228332834283528362837283828392840284128422843284428452846284728482849285028512852285328542855285628572858285928602861286228632864286528662867286828692870287128722873287428752876287728782879288028812882288328842885288628872888288928902891289228932894289528962897289828992900290129022903290429052906290729082909291029112912291329142915291629172918291929202921292229232924292529262927292829292930293129322933293429352936293729382939294029412942294329442945294629472948294929502951295229532954295529562957295829592960296129622963296429652966296729682969297029712972297329742975297629772978297929802981298229832984298529862987 |
- import sys
- import pytest
- import numpy as np
- import numpy.core._multiarray_tests as _multiarray_tests
- from numpy import array, arange, nditer, all
- from numpy.testing import (
- assert_, assert_equal, assert_array_equal, assert_raises,
- HAS_REFCOUNT, suppress_warnings
- )
- def iter_multi_index(i):
- ret = []
- while not i.finished:
- ret.append(i.multi_index)
- i.iternext()
- return ret
- def iter_indices(i):
- ret = []
- while not i.finished:
- ret.append(i.index)
- i.iternext()
- return ret
- def iter_iterindices(i):
- ret = []
- while not i.finished:
- ret.append(i.iterindex)
- i.iternext()
- return ret
- @pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts")
- def test_iter_refcount():
- # Make sure the iterator doesn't leak
- # Basic
- a = arange(6)
- dt = np.dtype('f4').newbyteorder()
- rc_a = sys.getrefcount(a)
- rc_dt = sys.getrefcount(dt)
- with nditer(a, [],
- [['readwrite', 'updateifcopy']],
- casting='unsafe',
- op_dtypes=[dt]) as it:
- assert_(not it.iterationneedsapi)
- assert_(sys.getrefcount(a) > rc_a)
- assert_(sys.getrefcount(dt) > rc_dt)
- # del 'it'
- it = None
- assert_equal(sys.getrefcount(a), rc_a)
- assert_equal(sys.getrefcount(dt), rc_dt)
- # With a copy
- a = arange(6, dtype='f4')
- dt = np.dtype('f4')
- rc_a = sys.getrefcount(a)
- rc_dt = sys.getrefcount(dt)
- it = nditer(a, [],
- [['readwrite']],
- op_dtypes=[dt])
- rc2_a = sys.getrefcount(a)
- rc2_dt = sys.getrefcount(dt)
- it2 = it.copy()
- assert_(sys.getrefcount(a) > rc2_a)
- assert_(sys.getrefcount(dt) > rc2_dt)
- it = None
- assert_equal(sys.getrefcount(a), rc2_a)
- assert_equal(sys.getrefcount(dt), rc2_dt)
- it2 = None
- assert_equal(sys.getrefcount(a), rc_a)
- assert_equal(sys.getrefcount(dt), rc_dt)
- del it2 # avoid pyflakes unused variable warning
- def test_iter_best_order():
- # The iterator should always find the iteration order
- # with increasing memory addresses
- # Test the ordering for 1-D to 5-D shapes
- for shape in [(5,), (3, 4), (2, 3, 4), (2, 3, 4, 3), (2, 3, 2, 2, 3)]:
- a = arange(np.prod(shape))
- # Test each combination of positive and negative strides
- for dirs in range(2**len(shape)):
- dirs_index = [slice(None)]*len(shape)
- for bit in range(len(shape)):
- if ((2**bit) & dirs):
- dirs_index[bit] = slice(None, None, -1)
- dirs_index = tuple(dirs_index)
- aview = a.reshape(shape)[dirs_index]
- # C-order
- i = nditer(aview, [], [['readonly']])
- assert_equal([x for x in i], a)
- # Fortran-order
- i = nditer(aview.T, [], [['readonly']])
- assert_equal([x for x in i], a)
- # Other order
- if len(shape) > 2:
- i = nditer(aview.swapaxes(0, 1), [], [['readonly']])
- assert_equal([x for x in i], a)
- def test_iter_c_order():
- # Test forcing C order
- # Test the ordering for 1-D to 5-D shapes
- for shape in [(5,), (3, 4), (2, 3, 4), (2, 3, 4, 3), (2, 3, 2, 2, 3)]:
- a = arange(np.prod(shape))
- # Test each combination of positive and negative strides
- for dirs in range(2**len(shape)):
- dirs_index = [slice(None)]*len(shape)
- for bit in range(len(shape)):
- if ((2**bit) & dirs):
- dirs_index[bit] = slice(None, None, -1)
- dirs_index = tuple(dirs_index)
- aview = a.reshape(shape)[dirs_index]
- # C-order
- i = nditer(aview, order='C')
- assert_equal([x for x in i], aview.ravel(order='C'))
- # Fortran-order
- i = nditer(aview.T, order='C')
- assert_equal([x for x in i], aview.T.ravel(order='C'))
- # Other order
- if len(shape) > 2:
- i = nditer(aview.swapaxes(0, 1), order='C')
- assert_equal([x for x in i],
- aview.swapaxes(0, 1).ravel(order='C'))
- def test_iter_f_order():
- # Test forcing F order
- # Test the ordering for 1-D to 5-D shapes
- for shape in [(5,), (3, 4), (2, 3, 4), (2, 3, 4, 3), (2, 3, 2, 2, 3)]:
- a = arange(np.prod(shape))
- # Test each combination of positive and negative strides
- for dirs in range(2**len(shape)):
- dirs_index = [slice(None)]*len(shape)
- for bit in range(len(shape)):
- if ((2**bit) & dirs):
- dirs_index[bit] = slice(None, None, -1)
- dirs_index = tuple(dirs_index)
- aview = a.reshape(shape)[dirs_index]
- # C-order
- i = nditer(aview, order='F')
- assert_equal([x for x in i], aview.ravel(order='F'))
- # Fortran-order
- i = nditer(aview.T, order='F')
- assert_equal([x for x in i], aview.T.ravel(order='F'))
- # Other order
- if len(shape) > 2:
- i = nditer(aview.swapaxes(0, 1), order='F')
- assert_equal([x for x in i],
- aview.swapaxes(0, 1).ravel(order='F'))
- def test_iter_c_or_f_order():
- # Test forcing any contiguous (C or F) order
- # Test the ordering for 1-D to 5-D shapes
- for shape in [(5,), (3, 4), (2, 3, 4), (2, 3, 4, 3), (2, 3, 2, 2, 3)]:
- a = arange(np.prod(shape))
- # Test each combination of positive and negative strides
- for dirs in range(2**len(shape)):
- dirs_index = [slice(None)]*len(shape)
- for bit in range(len(shape)):
- if ((2**bit) & dirs):
- dirs_index[bit] = slice(None, None, -1)
- dirs_index = tuple(dirs_index)
- aview = a.reshape(shape)[dirs_index]
- # C-order
- i = nditer(aview, order='A')
- assert_equal([x for x in i], aview.ravel(order='A'))
- # Fortran-order
- i = nditer(aview.T, order='A')
- assert_equal([x for x in i], aview.T.ravel(order='A'))
- # Other order
- if len(shape) > 2:
- i = nditer(aview.swapaxes(0, 1), order='A')
- assert_equal([x for x in i],
- aview.swapaxes(0, 1).ravel(order='A'))
- def test_iter_best_order_multi_index_1d():
- # The multi-indices should be correct with any reordering
- a = arange(4)
- # 1D order
- i = nditer(a, ['multi_index'], [['readonly']])
- assert_equal(iter_multi_index(i), [(0,), (1,), (2,), (3,)])
- # 1D reversed order
- i = nditer(a[::-1], ['multi_index'], [['readonly']])
- assert_equal(iter_multi_index(i), [(3,), (2,), (1,), (0,)])
- def test_iter_best_order_multi_index_2d():
- # The multi-indices should be correct with any reordering
- a = arange(6)
- # 2D C-order
- i = nditer(a.reshape(2, 3), ['multi_index'], [['readonly']])
- assert_equal(iter_multi_index(i), [(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2)])
- # 2D Fortran-order
- i = nditer(a.reshape(2, 3).copy(order='F'), ['multi_index'], [['readonly']])
- assert_equal(iter_multi_index(i), [(0, 0), (1, 0), (0, 1), (1, 1), (0, 2), (1, 2)])
- # 2D reversed C-order
- i = nditer(a.reshape(2, 3)[::-1], ['multi_index'], [['readonly']])
- assert_equal(iter_multi_index(i), [(1, 0), (1, 1), (1, 2), (0, 0), (0, 1), (0, 2)])
- i = nditer(a.reshape(2, 3)[:, ::-1], ['multi_index'], [['readonly']])
- assert_equal(iter_multi_index(i), [(0, 2), (0, 1), (0, 0), (1, 2), (1, 1), (1, 0)])
- i = nditer(a.reshape(2, 3)[::-1, ::-1], ['multi_index'], [['readonly']])
- assert_equal(iter_multi_index(i), [(1, 2), (1, 1), (1, 0), (0, 2), (0, 1), (0, 0)])
- # 2D reversed Fortran-order
- i = nditer(a.reshape(2, 3).copy(order='F')[::-1], ['multi_index'], [['readonly']])
- assert_equal(iter_multi_index(i), [(1, 0), (0, 0), (1, 1), (0, 1), (1, 2), (0, 2)])
- i = nditer(a.reshape(2, 3).copy(order='F')[:, ::-1],
- ['multi_index'], [['readonly']])
- assert_equal(iter_multi_index(i), [(0, 2), (1, 2), (0, 1), (1, 1), (0, 0), (1, 0)])
- i = nditer(a.reshape(2, 3).copy(order='F')[::-1, ::-1],
- ['multi_index'], [['readonly']])
- assert_equal(iter_multi_index(i), [(1, 2), (0, 2), (1, 1), (0, 1), (1, 0), (0, 0)])
- def test_iter_best_order_multi_index_3d():
- # The multi-indices should be correct with any reordering
- a = arange(12)
- # 3D C-order
- i = nditer(a.reshape(2, 3, 2), ['multi_index'], [['readonly']])
- assert_equal(iter_multi_index(i),
- [(0, 0, 0), (0, 0, 1), (0, 1, 0), (0, 1, 1), (0, 2, 0), (0, 2, 1),
- (1, 0, 0), (1, 0, 1), (1, 1, 0), (1, 1, 1), (1, 2, 0), (1, 2, 1)])
- # 3D Fortran-order
- i = nditer(a.reshape(2, 3, 2).copy(order='F'), ['multi_index'], [['readonly']])
- assert_equal(iter_multi_index(i),
- [(0, 0, 0), (1, 0, 0), (0, 1, 0), (1, 1, 0), (0, 2, 0), (1, 2, 0),
- (0, 0, 1), (1, 0, 1), (0, 1, 1), (1, 1, 1), (0, 2, 1), (1, 2, 1)])
- # 3D reversed C-order
- i = nditer(a.reshape(2, 3, 2)[::-1], ['multi_index'], [['readonly']])
- assert_equal(iter_multi_index(i),
- [(1, 0, 0), (1, 0, 1), (1, 1, 0), (1, 1, 1), (1, 2, 0), (1, 2, 1),
- (0, 0, 0), (0, 0, 1), (0, 1, 0), (0, 1, 1), (0, 2, 0), (0, 2, 1)])
- i = nditer(a.reshape(2, 3, 2)[:, ::-1], ['multi_index'], [['readonly']])
- assert_equal(iter_multi_index(i),
- [(0, 2, 0), (0, 2, 1), (0, 1, 0), (0, 1, 1), (0, 0, 0), (0, 0, 1),
- (1, 2, 0), (1, 2, 1), (1, 1, 0), (1, 1, 1), (1, 0, 0), (1, 0, 1)])
- i = nditer(a.reshape(2, 3, 2)[:,:, ::-1], ['multi_index'], [['readonly']])
- assert_equal(iter_multi_index(i),
- [(0, 0, 1), (0, 0, 0), (0, 1, 1), (0, 1, 0), (0, 2, 1), (0, 2, 0),
- (1, 0, 1), (1, 0, 0), (1, 1, 1), (1, 1, 0), (1, 2, 1), (1, 2, 0)])
- # 3D reversed Fortran-order
- i = nditer(a.reshape(2, 3, 2).copy(order='F')[::-1],
- ['multi_index'], [['readonly']])
- assert_equal(iter_multi_index(i),
- [(1, 0, 0), (0, 0, 0), (1, 1, 0), (0, 1, 0), (1, 2, 0), (0, 2, 0),
- (1, 0, 1), (0, 0, 1), (1, 1, 1), (0, 1, 1), (1, 2, 1), (0, 2, 1)])
- i = nditer(a.reshape(2, 3, 2).copy(order='F')[:, ::-1],
- ['multi_index'], [['readonly']])
- assert_equal(iter_multi_index(i),
- [(0, 2, 0), (1, 2, 0), (0, 1, 0), (1, 1, 0), (0, 0, 0), (1, 0, 0),
- (0, 2, 1), (1, 2, 1), (0, 1, 1), (1, 1, 1), (0, 0, 1), (1, 0, 1)])
- i = nditer(a.reshape(2, 3, 2).copy(order='F')[:,:, ::-1],
- ['multi_index'], [['readonly']])
- assert_equal(iter_multi_index(i),
- [(0, 0, 1), (1, 0, 1), (0, 1, 1), (1, 1, 1), (0, 2, 1), (1, 2, 1),
- (0, 0, 0), (1, 0, 0), (0, 1, 0), (1, 1, 0), (0, 2, 0), (1, 2, 0)])
- def test_iter_best_order_c_index_1d():
- # The C index should be correct with any reordering
- a = arange(4)
- # 1D order
- i = nditer(a, ['c_index'], [['readonly']])
- assert_equal(iter_indices(i), [0, 1, 2, 3])
- # 1D reversed order
- i = nditer(a[::-1], ['c_index'], [['readonly']])
- assert_equal(iter_indices(i), [3, 2, 1, 0])
- def test_iter_best_order_c_index_2d():
- # The C index should be correct with any reordering
- a = arange(6)
- # 2D C-order
- i = nditer(a.reshape(2, 3), ['c_index'], [['readonly']])
- assert_equal(iter_indices(i), [0, 1, 2, 3, 4, 5])
- # 2D Fortran-order
- i = nditer(a.reshape(2, 3).copy(order='F'),
- ['c_index'], [['readonly']])
- assert_equal(iter_indices(i), [0, 3, 1, 4, 2, 5])
- # 2D reversed C-order
- i = nditer(a.reshape(2, 3)[::-1], ['c_index'], [['readonly']])
- assert_equal(iter_indices(i), [3, 4, 5, 0, 1, 2])
- i = nditer(a.reshape(2, 3)[:, ::-1], ['c_index'], [['readonly']])
- assert_equal(iter_indices(i), [2, 1, 0, 5, 4, 3])
- i = nditer(a.reshape(2, 3)[::-1, ::-1], ['c_index'], [['readonly']])
- assert_equal(iter_indices(i), [5, 4, 3, 2, 1, 0])
- # 2D reversed Fortran-order
- i = nditer(a.reshape(2, 3).copy(order='F')[::-1],
- ['c_index'], [['readonly']])
- assert_equal(iter_indices(i), [3, 0, 4, 1, 5, 2])
- i = nditer(a.reshape(2, 3).copy(order='F')[:, ::-1],
- ['c_index'], [['readonly']])
- assert_equal(iter_indices(i), [2, 5, 1, 4, 0, 3])
- i = nditer(a.reshape(2, 3).copy(order='F')[::-1, ::-1],
- ['c_index'], [['readonly']])
- assert_equal(iter_indices(i), [5, 2, 4, 1, 3, 0])
- def test_iter_best_order_c_index_3d():
- # The C index should be correct with any reordering
- a = arange(12)
- # 3D C-order
- i = nditer(a.reshape(2, 3, 2), ['c_index'], [['readonly']])
- assert_equal(iter_indices(i),
- [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11])
- # 3D Fortran-order
- i = nditer(a.reshape(2, 3, 2).copy(order='F'),
- ['c_index'], [['readonly']])
- assert_equal(iter_indices(i),
- [0, 6, 2, 8, 4, 10, 1, 7, 3, 9, 5, 11])
- # 3D reversed C-order
- i = nditer(a.reshape(2, 3, 2)[::-1], ['c_index'], [['readonly']])
- assert_equal(iter_indices(i),
- [6, 7, 8, 9, 10, 11, 0, 1, 2, 3, 4, 5])
- i = nditer(a.reshape(2, 3, 2)[:, ::-1], ['c_index'], [['readonly']])
- assert_equal(iter_indices(i),
- [4, 5, 2, 3, 0, 1, 10, 11, 8, 9, 6, 7])
- i = nditer(a.reshape(2, 3, 2)[:,:, ::-1], ['c_index'], [['readonly']])
- assert_equal(iter_indices(i),
- [1, 0, 3, 2, 5, 4, 7, 6, 9, 8, 11, 10])
- # 3D reversed Fortran-order
- i = nditer(a.reshape(2, 3, 2).copy(order='F')[::-1],
- ['c_index'], [['readonly']])
- assert_equal(iter_indices(i),
- [6, 0, 8, 2, 10, 4, 7, 1, 9, 3, 11, 5])
- i = nditer(a.reshape(2, 3, 2).copy(order='F')[:, ::-1],
- ['c_index'], [['readonly']])
- assert_equal(iter_indices(i),
- [4, 10, 2, 8, 0, 6, 5, 11, 3, 9, 1, 7])
- i = nditer(a.reshape(2, 3, 2).copy(order='F')[:,:, ::-1],
- ['c_index'], [['readonly']])
- assert_equal(iter_indices(i),
- [1, 7, 3, 9, 5, 11, 0, 6, 2, 8, 4, 10])
- def test_iter_best_order_f_index_1d():
- # The Fortran index should be correct with any reordering
- a = arange(4)
- # 1D order
- i = nditer(a, ['f_index'], [['readonly']])
- assert_equal(iter_indices(i), [0, 1, 2, 3])
- # 1D reversed order
- i = nditer(a[::-1], ['f_index'], [['readonly']])
- assert_equal(iter_indices(i), [3, 2, 1, 0])
- def test_iter_best_order_f_index_2d():
- # The Fortran index should be correct with any reordering
- a = arange(6)
- # 2D C-order
- i = nditer(a.reshape(2, 3), ['f_index'], [['readonly']])
- assert_equal(iter_indices(i), [0, 2, 4, 1, 3, 5])
- # 2D Fortran-order
- i = nditer(a.reshape(2, 3).copy(order='F'),
- ['f_index'], [['readonly']])
- assert_equal(iter_indices(i), [0, 1, 2, 3, 4, 5])
- # 2D reversed C-order
- i = nditer(a.reshape(2, 3)[::-1], ['f_index'], [['readonly']])
- assert_equal(iter_indices(i), [1, 3, 5, 0, 2, 4])
- i = nditer(a.reshape(2, 3)[:, ::-1], ['f_index'], [['readonly']])
- assert_equal(iter_indices(i), [4, 2, 0, 5, 3, 1])
- i = nditer(a.reshape(2, 3)[::-1, ::-1], ['f_index'], [['readonly']])
- assert_equal(iter_indices(i), [5, 3, 1, 4, 2, 0])
- # 2D reversed Fortran-order
- i = nditer(a.reshape(2, 3).copy(order='F')[::-1],
- ['f_index'], [['readonly']])
- assert_equal(iter_indices(i), [1, 0, 3, 2, 5, 4])
- i = nditer(a.reshape(2, 3).copy(order='F')[:, ::-1],
- ['f_index'], [['readonly']])
- assert_equal(iter_indices(i), [4, 5, 2, 3, 0, 1])
- i = nditer(a.reshape(2, 3).copy(order='F')[::-1, ::-1],
- ['f_index'], [['readonly']])
- assert_equal(iter_indices(i), [5, 4, 3, 2, 1, 0])
- def test_iter_best_order_f_index_3d():
- # The Fortran index should be correct with any reordering
- a = arange(12)
- # 3D C-order
- i = nditer(a.reshape(2, 3, 2), ['f_index'], [['readonly']])
- assert_equal(iter_indices(i),
- [0, 6, 2, 8, 4, 10, 1, 7, 3, 9, 5, 11])
- # 3D Fortran-order
- i = nditer(a.reshape(2, 3, 2).copy(order='F'),
- ['f_index'], [['readonly']])
- assert_equal(iter_indices(i),
- [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11])
- # 3D reversed C-order
- i = nditer(a.reshape(2, 3, 2)[::-1], ['f_index'], [['readonly']])
- assert_equal(iter_indices(i),
- [1, 7, 3, 9, 5, 11, 0, 6, 2, 8, 4, 10])
- i = nditer(a.reshape(2, 3, 2)[:, ::-1], ['f_index'], [['readonly']])
- assert_equal(iter_indices(i),
- [4, 10, 2, 8, 0, 6, 5, 11, 3, 9, 1, 7])
- i = nditer(a.reshape(2, 3, 2)[:,:, ::-1], ['f_index'], [['readonly']])
- assert_equal(iter_indices(i),
- [6, 0, 8, 2, 10, 4, 7, 1, 9, 3, 11, 5])
- # 3D reversed Fortran-order
- i = nditer(a.reshape(2, 3, 2).copy(order='F')[::-1],
- ['f_index'], [['readonly']])
- assert_equal(iter_indices(i),
- [1, 0, 3, 2, 5, 4, 7, 6, 9, 8, 11, 10])
- i = nditer(a.reshape(2, 3, 2).copy(order='F')[:, ::-1],
- ['f_index'], [['readonly']])
- assert_equal(iter_indices(i),
- [4, 5, 2, 3, 0, 1, 10, 11, 8, 9, 6, 7])
- i = nditer(a.reshape(2, 3, 2).copy(order='F')[:,:, ::-1],
- ['f_index'], [['readonly']])
- assert_equal(iter_indices(i),
- [6, 7, 8, 9, 10, 11, 0, 1, 2, 3, 4, 5])
- def test_iter_no_inner_full_coalesce():
- # Check no_inner iterators which coalesce into a single inner loop
- for shape in [(5,), (3, 4), (2, 3, 4), (2, 3, 4, 3), (2, 3, 2, 2, 3)]:
- size = np.prod(shape)
- a = arange(size)
- # Test each combination of forward and backwards indexing
- for dirs in range(2**len(shape)):
- dirs_index = [slice(None)]*len(shape)
- for bit in range(len(shape)):
- if ((2**bit) & dirs):
- dirs_index[bit] = slice(None, None, -1)
- dirs_index = tuple(dirs_index)
- aview = a.reshape(shape)[dirs_index]
- # C-order
- i = nditer(aview, ['external_loop'], [['readonly']])
- assert_equal(i.ndim, 1)
- assert_equal(i[0].shape, (size,))
- # Fortran-order
- i = nditer(aview.T, ['external_loop'], [['readonly']])
- assert_equal(i.ndim, 1)
- assert_equal(i[0].shape, (size,))
- # Other order
- if len(shape) > 2:
- i = nditer(aview.swapaxes(0, 1),
- ['external_loop'], [['readonly']])
- assert_equal(i.ndim, 1)
- assert_equal(i[0].shape, (size,))
- def test_iter_no_inner_dim_coalescing():
- # Check no_inner iterators whose dimensions may not coalesce completely
- # Skipping the last element in a dimension prevents coalescing
- # with the next-bigger dimension
- a = arange(24).reshape(2, 3, 4)[:,:, :-1]
- i = nditer(a, ['external_loop'], [['readonly']])
- assert_equal(i.ndim, 2)
- assert_equal(i[0].shape, (3,))
- a = arange(24).reshape(2, 3, 4)[:, :-1,:]
- i = nditer(a, ['external_loop'], [['readonly']])
- assert_equal(i.ndim, 2)
- assert_equal(i[0].shape, (8,))
- a = arange(24).reshape(2, 3, 4)[:-1,:,:]
- i = nditer(a, ['external_loop'], [['readonly']])
- assert_equal(i.ndim, 1)
- assert_equal(i[0].shape, (12,))
- # Even with lots of 1-sized dimensions, should still coalesce
- a = arange(24).reshape(1, 1, 2, 1, 1, 3, 1, 1, 4, 1, 1)
- i = nditer(a, ['external_loop'], [['readonly']])
- assert_equal(i.ndim, 1)
- assert_equal(i[0].shape, (24,))
- def test_iter_dim_coalescing():
- # Check that the correct number of dimensions are coalesced
- # Tracking a multi-index disables coalescing
- a = arange(24).reshape(2, 3, 4)
- i = nditer(a, ['multi_index'], [['readonly']])
- assert_equal(i.ndim, 3)
- # A tracked index can allow coalescing if it's compatible with the array
- a3d = arange(24).reshape(2, 3, 4)
- i = nditer(a3d, ['c_index'], [['readonly']])
- assert_equal(i.ndim, 1)
- i = nditer(a3d.swapaxes(0, 1), ['c_index'], [['readonly']])
- assert_equal(i.ndim, 3)
- i = nditer(a3d.T, ['c_index'], [['readonly']])
- assert_equal(i.ndim, 3)
- i = nditer(a3d.T, ['f_index'], [['readonly']])
- assert_equal(i.ndim, 1)
- i = nditer(a3d.T.swapaxes(0, 1), ['f_index'], [['readonly']])
- assert_equal(i.ndim, 3)
- # When C or F order is forced, coalescing may still occur
- a3d = arange(24).reshape(2, 3, 4)
- i = nditer(a3d, order='C')
- assert_equal(i.ndim, 1)
- i = nditer(a3d.T, order='C')
- assert_equal(i.ndim, 3)
- i = nditer(a3d, order='F')
- assert_equal(i.ndim, 3)
- i = nditer(a3d.T, order='F')
- assert_equal(i.ndim, 1)
- i = nditer(a3d, order='A')
- assert_equal(i.ndim, 1)
- i = nditer(a3d.T, order='A')
- assert_equal(i.ndim, 1)
- def test_iter_broadcasting():
- # Standard NumPy broadcasting rules
- # 1D with scalar
- i = nditer([arange(6), np.int32(2)], ['multi_index'], [['readonly']]*2)
- assert_equal(i.itersize, 6)
- assert_equal(i.shape, (6,))
- # 2D with scalar
- i = nditer([arange(6).reshape(2, 3), np.int32(2)],
- ['multi_index'], [['readonly']]*2)
- assert_equal(i.itersize, 6)
- assert_equal(i.shape, (2, 3))
- # 2D with 1D
- i = nditer([arange(6).reshape(2, 3), arange(3)],
- ['multi_index'], [['readonly']]*2)
- assert_equal(i.itersize, 6)
- assert_equal(i.shape, (2, 3))
- i = nditer([arange(2).reshape(2, 1), arange(3)],
- ['multi_index'], [['readonly']]*2)
- assert_equal(i.itersize, 6)
- assert_equal(i.shape, (2, 3))
- # 2D with 2D
- i = nditer([arange(2).reshape(2, 1), arange(3).reshape(1, 3)],
- ['multi_index'], [['readonly']]*2)
- assert_equal(i.itersize, 6)
- assert_equal(i.shape, (2, 3))
- # 3D with scalar
- i = nditer([np.int32(2), arange(24).reshape(4, 2, 3)],
- ['multi_index'], [['readonly']]*2)
- assert_equal(i.itersize, 24)
- assert_equal(i.shape, (4, 2, 3))
- # 3D with 1D
- i = nditer([arange(3), arange(24).reshape(4, 2, 3)],
- ['multi_index'], [['readonly']]*2)
- assert_equal(i.itersize, 24)
- assert_equal(i.shape, (4, 2, 3))
- i = nditer([arange(3), arange(8).reshape(4, 2, 1)],
- ['multi_index'], [['readonly']]*2)
- assert_equal(i.itersize, 24)
- assert_equal(i.shape, (4, 2, 3))
- # 3D with 2D
- i = nditer([arange(6).reshape(2, 3), arange(24).reshape(4, 2, 3)],
- ['multi_index'], [['readonly']]*2)
- assert_equal(i.itersize, 24)
- assert_equal(i.shape, (4, 2, 3))
- i = nditer([arange(2).reshape(2, 1), arange(24).reshape(4, 2, 3)],
- ['multi_index'], [['readonly']]*2)
- assert_equal(i.itersize, 24)
- assert_equal(i.shape, (4, 2, 3))
- i = nditer([arange(3).reshape(1, 3), arange(8).reshape(4, 2, 1)],
- ['multi_index'], [['readonly']]*2)
- assert_equal(i.itersize, 24)
- assert_equal(i.shape, (4, 2, 3))
- # 3D with 3D
- i = nditer([arange(2).reshape(1, 2, 1), arange(3).reshape(1, 1, 3),
- arange(4).reshape(4, 1, 1)],
- ['multi_index'], [['readonly']]*3)
- assert_equal(i.itersize, 24)
- assert_equal(i.shape, (4, 2, 3))
- i = nditer([arange(6).reshape(1, 2, 3), arange(4).reshape(4, 1, 1)],
- ['multi_index'], [['readonly']]*2)
- assert_equal(i.itersize, 24)
- assert_equal(i.shape, (4, 2, 3))
- i = nditer([arange(24).reshape(4, 2, 3), arange(12).reshape(4, 1, 3)],
- ['multi_index'], [['readonly']]*2)
- assert_equal(i.itersize, 24)
- assert_equal(i.shape, (4, 2, 3))
- def test_iter_itershape():
- # Check that allocated outputs work with a specified shape
- a = np.arange(6, dtype='i2').reshape(2, 3)
- i = nditer([a, None], [], [['readonly'], ['writeonly', 'allocate']],
- op_axes=[[0, 1, None], None],
- itershape=(-1, -1, 4))
- assert_equal(i.operands[1].shape, (2, 3, 4))
- assert_equal(i.operands[1].strides, (24, 8, 2))
- i = nditer([a.T, None], [], [['readonly'], ['writeonly', 'allocate']],
- op_axes=[[0, 1, None], None],
- itershape=(-1, -1, 4))
- assert_equal(i.operands[1].shape, (3, 2, 4))
- assert_equal(i.operands[1].strides, (8, 24, 2))
- i = nditer([a.T, None], [], [['readonly'], ['writeonly', 'allocate']],
- order='F',
- op_axes=[[0, 1, None], None],
- itershape=(-1, -1, 4))
- assert_equal(i.operands[1].shape, (3, 2, 4))
- assert_equal(i.operands[1].strides, (2, 6, 12))
- # If we specify 1 in the itershape, it shouldn't allow broadcasting
- # of that dimension to a bigger value
- assert_raises(ValueError, nditer, [a, None], [],
- [['readonly'], ['writeonly', 'allocate']],
- op_axes=[[0, 1, None], None],
- itershape=(-1, 1, 4))
- # Test bug that for no op_axes but itershape, they are NULLed correctly
- i = np.nditer([np.ones(2), None, None], itershape=(2,))
- def test_iter_broadcasting_errors():
- # Check that errors are thrown for bad broadcasting shapes
- # 1D with 1D
- assert_raises(ValueError, nditer, [arange(2), arange(3)],
- [], [['readonly']]*2)
- # 2D with 1D
- assert_raises(ValueError, nditer,
- [arange(6).reshape(2, 3), arange(2)],
- [], [['readonly']]*2)
- # 2D with 2D
- assert_raises(ValueError, nditer,
- [arange(6).reshape(2, 3), arange(9).reshape(3, 3)],
- [], [['readonly']]*2)
- assert_raises(ValueError, nditer,
- [arange(6).reshape(2, 3), arange(4).reshape(2, 2)],
- [], [['readonly']]*2)
- # 3D with 3D
- assert_raises(ValueError, nditer,
- [arange(36).reshape(3, 3, 4), arange(24).reshape(2, 3, 4)],
- [], [['readonly']]*2)
- assert_raises(ValueError, nditer,
- [arange(8).reshape(2, 4, 1), arange(24).reshape(2, 3, 4)],
- [], [['readonly']]*2)
- # Verify that the error message mentions the right shapes
- try:
- nditer([arange(2).reshape(1, 2, 1),
- arange(3).reshape(1, 3),
- arange(6).reshape(2, 3)],
- [],
- [['readonly'], ['readonly'], ['writeonly', 'no_broadcast']])
- raise AssertionError('Should have raised a broadcast error')
- except ValueError as e:
- msg = str(e)
- # The message should contain the shape of the 3rd operand
- assert_(msg.find('(2,3)') >= 0,
- 'Message "%s" doesn\'t contain operand shape (2,3)' % msg)
- # The message should contain the broadcast shape
- assert_(msg.find('(1,2,3)') >= 0,
- 'Message "%s" doesn\'t contain broadcast shape (1,2,3)' % msg)
- try:
- nditer([arange(6).reshape(2, 3), arange(2)],
- [],
- [['readonly'], ['readonly']],
- op_axes=[[0, 1], [0, np.newaxis]],
- itershape=(4, 3))
- raise AssertionError('Should have raised a broadcast error')
- except ValueError as e:
- msg = str(e)
- # The message should contain "shape->remappedshape" for each operand
- assert_(msg.find('(2,3)->(2,3)') >= 0,
- 'Message "%s" doesn\'t contain operand shape (2,3)->(2,3)' % msg)
- assert_(msg.find('(2,)->(2,newaxis)') >= 0,
- ('Message "%s" doesn\'t contain remapped operand shape' +
- '(2,)->(2,newaxis)') % msg)
- # The message should contain the itershape parameter
- assert_(msg.find('(4,3)') >= 0,
- 'Message "%s" doesn\'t contain itershape parameter (4,3)' % msg)
- try:
- nditer([np.zeros((2, 1, 1)), np.zeros((2,))],
- [],
- [['writeonly', 'no_broadcast'], ['readonly']])
- raise AssertionError('Should have raised a broadcast error')
- except ValueError as e:
- msg = str(e)
- # The message should contain the shape of the bad operand
- assert_(msg.find('(2,1,1)') >= 0,
- 'Message "%s" doesn\'t contain operand shape (2,1,1)' % msg)
- # The message should contain the broadcast shape
- assert_(msg.find('(2,1,2)') >= 0,
- 'Message "%s" doesn\'t contain the broadcast shape (2,1,2)' % msg)
- def test_iter_flags_errors():
- # Check that bad combinations of flags produce errors
- a = arange(6)
- # Not enough operands
- assert_raises(ValueError, nditer, [], [], [])
- # Too many operands
- assert_raises(ValueError, nditer, [a]*100, [], [['readonly']]*100)
- # Bad global flag
- assert_raises(ValueError, nditer, [a], ['bad flag'], [['readonly']])
- # Bad op flag
- assert_raises(ValueError, nditer, [a], [], [['readonly', 'bad flag']])
- # Bad order parameter
- assert_raises(ValueError, nditer, [a], [], [['readonly']], order='G')
- # Bad casting parameter
- assert_raises(ValueError, nditer, [a], [], [['readonly']], casting='noon')
- # op_flags must match ops
- assert_raises(ValueError, nditer, [a]*3, [], [['readonly']]*2)
- # Cannot track both a C and an F index
- assert_raises(ValueError, nditer, a,
- ['c_index', 'f_index'], [['readonly']])
- # Inner iteration and multi-indices/indices are incompatible
- assert_raises(ValueError, nditer, a,
- ['external_loop', 'multi_index'], [['readonly']])
- assert_raises(ValueError, nditer, a,
- ['external_loop', 'c_index'], [['readonly']])
- assert_raises(ValueError, nditer, a,
- ['external_loop', 'f_index'], [['readonly']])
- # Must specify exactly one of readwrite/readonly/writeonly per operand
- assert_raises(ValueError, nditer, a, [], [[]])
- assert_raises(ValueError, nditer, a, [], [['readonly', 'writeonly']])
- assert_raises(ValueError, nditer, a, [], [['readonly', 'readwrite']])
- assert_raises(ValueError, nditer, a, [], [['writeonly', 'readwrite']])
- assert_raises(ValueError, nditer, a,
- [], [['readonly', 'writeonly', 'readwrite']])
- # Python scalars are always readonly
- assert_raises(TypeError, nditer, 1.5, [], [['writeonly']])
- assert_raises(TypeError, nditer, 1.5, [], [['readwrite']])
- # Array scalars are always readonly
- assert_raises(TypeError, nditer, np.int32(1), [], [['writeonly']])
- assert_raises(TypeError, nditer, np.int32(1), [], [['readwrite']])
- # Check readonly array
- a.flags.writeable = False
- assert_raises(ValueError, nditer, a, [], [['writeonly']])
- assert_raises(ValueError, nditer, a, [], [['readwrite']])
- a.flags.writeable = True
- # Multi-indices available only with the multi_index flag
- i = nditer(arange(6), [], [['readonly']])
- assert_raises(ValueError, lambda i:i.multi_index, i)
- # Index available only with an index flag
- assert_raises(ValueError, lambda i:i.index, i)
- # GotoCoords and GotoIndex incompatible with buffering or no_inner
- def assign_multi_index(i):
- i.multi_index = (0,)
- def assign_index(i):
- i.index = 0
- def assign_iterindex(i):
- i.iterindex = 0
- def assign_iterrange(i):
- i.iterrange = (0, 1)
- i = nditer(arange(6), ['external_loop'])
- assert_raises(ValueError, assign_multi_index, i)
- assert_raises(ValueError, assign_index, i)
- assert_raises(ValueError, assign_iterindex, i)
- assert_raises(ValueError, assign_iterrange, i)
- i = nditer(arange(6), ['buffered'])
- assert_raises(ValueError, assign_multi_index, i)
- assert_raises(ValueError, assign_index, i)
- assert_raises(ValueError, assign_iterrange, i)
- # Can't iterate if size is zero
- assert_raises(ValueError, nditer, np.array([]))
- def test_iter_slice():
- a, b, c = np.arange(3), np.arange(3), np.arange(3.)
- i = nditer([a, b, c], [], ['readwrite'])
- with i:
- i[0:2] = (3, 3)
- assert_equal(a, [3, 1, 2])
- assert_equal(b, [3, 1, 2])
- assert_equal(c, [0, 1, 2])
- i[1] = 12
- assert_equal(i[0:2], [3, 12])
- def test_iter_assign_mapping():
- a = np.arange(24, dtype='f8').reshape(2, 3, 4).T
- it = np.nditer(a, [], [['readwrite', 'updateifcopy']],
- casting='same_kind', op_dtypes=[np.dtype('f4')])
- with it:
- it.operands[0][...] = 3
- it.operands[0][...] = 14
- assert_equal(a, 14)
- it = np.nditer(a, [], [['readwrite', 'updateifcopy']],
- casting='same_kind', op_dtypes=[np.dtype('f4')])
- with it:
- x = it.operands[0][-1:1]
- x[...] = 14
- it.operands[0][...] = -1234
- assert_equal(a, -1234)
- # check for no warnings on dealloc
- x = None
- it = None
- def test_iter_nbo_align_contig():
- # Check that byte order, alignment, and contig changes work
- # Byte order change by requesting a specific dtype
- a = np.arange(6, dtype='f4')
- au = a.byteswap().newbyteorder()
- assert_(a.dtype.byteorder != au.dtype.byteorder)
- i = nditer(au, [], [['readwrite', 'updateifcopy']],
- casting='equiv',
- op_dtypes=[np.dtype('f4')])
- with i:
- # context manager triggers UPDATEIFCOPY on i at exit
- assert_equal(i.dtypes[0].byteorder, a.dtype.byteorder)
- assert_equal(i.operands[0].dtype.byteorder, a.dtype.byteorder)
- assert_equal(i.operands[0], a)
- i.operands[0][:] = 2
- assert_equal(au, [2]*6)
- del i # should not raise a warning
- # Byte order change by requesting NBO
- a = np.arange(6, dtype='f4')
- au = a.byteswap().newbyteorder()
- assert_(a.dtype.byteorder != au.dtype.byteorder)
- with nditer(au, [], [['readwrite', 'updateifcopy', 'nbo']],
- casting='equiv') as i:
- # context manager triggers UPDATEIFCOPY on i at exit
- assert_equal(i.dtypes[0].byteorder, a.dtype.byteorder)
- assert_equal(i.operands[0].dtype.byteorder, a.dtype.byteorder)
- assert_equal(i.operands[0], a)
- i.operands[0][:] = 12345
- i.operands[0][:] = 2
- assert_equal(au, [2]*6)
- # Unaligned input
- a = np.zeros((6*4+1,), dtype='i1')[1:]
- a.dtype = 'f4'
- a[:] = np.arange(6, dtype='f4')
- assert_(not a.flags.aligned)
- # Without 'aligned', shouldn't copy
- i = nditer(a, [], [['readonly']])
- assert_(not i.operands[0].flags.aligned)
- assert_equal(i.operands[0], a)
- # With 'aligned', should make a copy
- with nditer(a, [], [['readwrite', 'updateifcopy', 'aligned']]) as i:
- assert_(i.operands[0].flags.aligned)
- # context manager triggers UPDATEIFCOPY on i at exit
- assert_equal(i.operands[0], a)
- i.operands[0][:] = 3
- assert_equal(a, [3]*6)
- # Discontiguous input
- a = arange(12)
- # If it is contiguous, shouldn't copy
- i = nditer(a[:6], [], [['readonly']])
- assert_(i.operands[0].flags.contiguous)
- assert_equal(i.operands[0], a[:6])
- # If it isn't contiguous, should buffer
- i = nditer(a[::2], ['buffered', 'external_loop'],
- [['readonly', 'contig']],
- buffersize=10)
- assert_(i[0].flags.contiguous)
- assert_equal(i[0], a[::2])
- def test_iter_array_cast():
- # Check that arrays are cast as requested
- # No cast 'f4' -> 'f4'
- a = np.arange(6, dtype='f4').reshape(2, 3)
- i = nditer(a, [], [['readwrite']], op_dtypes=[np.dtype('f4')])
- with i:
- assert_equal(i.operands[0], a)
- assert_equal(i.operands[0].dtype, np.dtype('f4'))
- # Byte-order cast '<f4' -> '>f4'
- a = np.arange(6, dtype='<f4').reshape(2, 3)
- with nditer(a, [], [['readwrite', 'updateifcopy']],
- casting='equiv',
- op_dtypes=[np.dtype('>f4')]) as i:
- assert_equal(i.operands[0], a)
- assert_equal(i.operands[0].dtype, np.dtype('>f4'))
- # Safe case 'f4' -> 'f8'
- a = np.arange(24, dtype='f4').reshape(2, 3, 4).swapaxes(1, 2)
- i = nditer(a, [], [['readonly', 'copy']],
- casting='safe',
- op_dtypes=[np.dtype('f8')])
- assert_equal(i.operands[0], a)
- assert_equal(i.operands[0].dtype, np.dtype('f8'))
- # The memory layout of the temporary should match a (a is (48,4,16))
- # except negative strides get flipped to positive strides.
- assert_equal(i.operands[0].strides, (96, 8, 32))
- a = a[::-1,:, ::-1]
- i = nditer(a, [], [['readonly', 'copy']],
- casting='safe',
- op_dtypes=[np.dtype('f8')])
- assert_equal(i.operands[0], a)
- assert_equal(i.operands[0].dtype, np.dtype('f8'))
- assert_equal(i.operands[0].strides, (96, 8, 32))
- # Same-kind cast 'f8' -> 'f4' -> 'f8'
- a = np.arange(24, dtype='f8').reshape(2, 3, 4).T
- with nditer(a, [],
- [['readwrite', 'updateifcopy']],
- casting='same_kind',
- op_dtypes=[np.dtype('f4')]) as i:
- assert_equal(i.operands[0], a)
- assert_equal(i.operands[0].dtype, np.dtype('f4'))
- assert_equal(i.operands[0].strides, (4, 16, 48))
- # Check that WRITEBACKIFCOPY is activated at exit
- i.operands[0][2, 1, 1] = -12.5
- assert_(a[2, 1, 1] != -12.5)
- assert_equal(a[2, 1, 1], -12.5)
- a = np.arange(6, dtype='i4')[::-2]
- with nditer(a, [],
- [['writeonly', 'updateifcopy']],
- casting='unsafe',
- op_dtypes=[np.dtype('f4')]) as i:
- assert_equal(i.operands[0].dtype, np.dtype('f4'))
- # Even though the stride was negative in 'a', it
- # becomes positive in the temporary
- assert_equal(i.operands[0].strides, (4,))
- i.operands[0][:] = [1, 2, 3]
- assert_equal(a, [1, 2, 3])
- def test_iter_array_cast_errors():
- # Check that invalid casts are caught
- # Need to enable copying for casts to occur
- assert_raises(TypeError, nditer, arange(2, dtype='f4'), [],
- [['readonly']], op_dtypes=[np.dtype('f8')])
- # Also need to allow casting for casts to occur
- assert_raises(TypeError, nditer, arange(2, dtype='f4'), [],
- [['readonly', 'copy']], casting='no',
- op_dtypes=[np.dtype('f8')])
- assert_raises(TypeError, nditer, arange(2, dtype='f4'), [],
- [['readonly', 'copy']], casting='equiv',
- op_dtypes=[np.dtype('f8')])
- assert_raises(TypeError, nditer, arange(2, dtype='f8'), [],
- [['writeonly', 'updateifcopy']],
- casting='no',
- op_dtypes=[np.dtype('f4')])
- assert_raises(TypeError, nditer, arange(2, dtype='f8'), [],
- [['writeonly', 'updateifcopy']],
- casting='equiv',
- op_dtypes=[np.dtype('f4')])
- # '<f4' -> '>f4' should not work with casting='no'
- assert_raises(TypeError, nditer, arange(2, dtype='<f4'), [],
- [['readonly', 'copy']], casting='no',
- op_dtypes=[np.dtype('>f4')])
- # 'f4' -> 'f8' is a safe cast, but 'f8' -> 'f4' isn't
- assert_raises(TypeError, nditer, arange(2, dtype='f4'), [],
- [['readwrite', 'updateifcopy']],
- casting='safe',
- op_dtypes=[np.dtype('f8')])
- assert_raises(TypeError, nditer, arange(2, dtype='f8'), [],
- [['readwrite', 'updateifcopy']],
- casting='safe',
- op_dtypes=[np.dtype('f4')])
- # 'f4' -> 'i4' is neither a safe nor a same-kind cast
- assert_raises(TypeError, nditer, arange(2, dtype='f4'), [],
- [['readonly', 'copy']],
- casting='same_kind',
- op_dtypes=[np.dtype('i4')])
- assert_raises(TypeError, nditer, arange(2, dtype='i4'), [],
- [['writeonly', 'updateifcopy']],
- casting='same_kind',
- op_dtypes=[np.dtype('f4')])
- def test_iter_scalar_cast():
- # Check that scalars are cast as requested
- # No cast 'f4' -> 'f4'
- i = nditer(np.float32(2.5), [], [['readonly']],
- op_dtypes=[np.dtype('f4')])
- assert_equal(i.dtypes[0], np.dtype('f4'))
- assert_equal(i.value.dtype, np.dtype('f4'))
- assert_equal(i.value, 2.5)
- # Safe cast 'f4' -> 'f8'
- i = nditer(np.float32(2.5), [],
- [['readonly', 'copy']],
- casting='safe',
- op_dtypes=[np.dtype('f8')])
- assert_equal(i.dtypes[0], np.dtype('f8'))
- assert_equal(i.value.dtype, np.dtype('f8'))
- assert_equal(i.value, 2.5)
- # Same-kind cast 'f8' -> 'f4'
- i = nditer(np.float64(2.5), [],
- [['readonly', 'copy']],
- casting='same_kind',
- op_dtypes=[np.dtype('f4')])
- assert_equal(i.dtypes[0], np.dtype('f4'))
- assert_equal(i.value.dtype, np.dtype('f4'))
- assert_equal(i.value, 2.5)
- # Unsafe cast 'f8' -> 'i4'
- i = nditer(np.float64(3.0), [],
- [['readonly', 'copy']],
- casting='unsafe',
- op_dtypes=[np.dtype('i4')])
- assert_equal(i.dtypes[0], np.dtype('i4'))
- assert_equal(i.value.dtype, np.dtype('i4'))
- assert_equal(i.value, 3)
- # Readonly scalars may be cast even without setting COPY or BUFFERED
- i = nditer(3, [], [['readonly']], op_dtypes=[np.dtype('f8')])
- assert_equal(i[0].dtype, np.dtype('f8'))
- assert_equal(i[0], 3.)
- def test_iter_scalar_cast_errors():
- # Check that invalid casts are caught
- # Need to allow copying/buffering for write casts of scalars to occur
- assert_raises(TypeError, nditer, np.float32(2), [],
- [['readwrite']], op_dtypes=[np.dtype('f8')])
- assert_raises(TypeError, nditer, 2.5, [],
- [['readwrite']], op_dtypes=[np.dtype('f4')])
- # 'f8' -> 'f4' isn't a safe cast if the value would overflow
- assert_raises(TypeError, nditer, np.float64(1e60), [],
- [['readonly']],
- casting='safe',
- op_dtypes=[np.dtype('f4')])
- # 'f4' -> 'i4' is neither a safe nor a same-kind cast
- assert_raises(TypeError, nditer, np.float32(2), [],
- [['readonly']],
- casting='same_kind',
- op_dtypes=[np.dtype('i4')])
- def test_iter_object_arrays_basic():
- # Check that object arrays work
- obj = {'a':3,'b':'d'}
- a = np.array([[1, 2, 3], None, obj, None], dtype='O')
- if HAS_REFCOUNT:
- rc = sys.getrefcount(obj)
- # Need to allow references for object arrays
- assert_raises(TypeError, nditer, a)
- if HAS_REFCOUNT:
- assert_equal(sys.getrefcount(obj), rc)
- i = nditer(a, ['refs_ok'], ['readonly'])
- vals = [x_[()] for x_ in i]
- assert_equal(np.array(vals, dtype='O'), a)
- vals, i, x = [None]*3
- if HAS_REFCOUNT:
- assert_equal(sys.getrefcount(obj), rc)
- i = nditer(a.reshape(2, 2).T, ['refs_ok', 'buffered'],
- ['readonly'], order='C')
- assert_(i.iterationneedsapi)
- vals = [x_[()] for x_ in i]
- assert_equal(np.array(vals, dtype='O'), a.reshape(2, 2).ravel(order='F'))
- vals, i, x = [None]*3
- if HAS_REFCOUNT:
- assert_equal(sys.getrefcount(obj), rc)
- i = nditer(a.reshape(2, 2).T, ['refs_ok', 'buffered'],
- ['readwrite'], order='C')
- with i:
- for x in i:
- x[...] = None
- vals, i, x = [None]*3
- if HAS_REFCOUNT:
- assert_(sys.getrefcount(obj) == rc-1)
- assert_equal(a, np.array([None]*4, dtype='O'))
- def test_iter_object_arrays_conversions():
- # Conversions to/from objects
- a = np.arange(6, dtype='O')
- i = nditer(a, ['refs_ok', 'buffered'], ['readwrite'],
- casting='unsafe', op_dtypes='i4')
- with i:
- for x in i:
- x[...] += 1
- assert_equal(a, np.arange(6)+1)
- a = np.arange(6, dtype='i4')
- i = nditer(a, ['refs_ok', 'buffered'], ['readwrite'],
- casting='unsafe', op_dtypes='O')
- with i:
- for x in i:
- x[...] += 1
- assert_equal(a, np.arange(6)+1)
- # Non-contiguous object array
- a = np.zeros((6,), dtype=[('p', 'i1'), ('a', 'O')])
- a = a['a']
- a[:] = np.arange(6)
- i = nditer(a, ['refs_ok', 'buffered'], ['readwrite'],
- casting='unsafe', op_dtypes='i4')
- with i:
- for x in i:
- x[...] += 1
- assert_equal(a, np.arange(6)+1)
- #Non-contiguous value array
- a = np.zeros((6,), dtype=[('p', 'i1'), ('a', 'i4')])
- a = a['a']
- a[:] = np.arange(6) + 98172488
- i = nditer(a, ['refs_ok', 'buffered'], ['readwrite'],
- casting='unsafe', op_dtypes='O')
- with i:
- ob = i[0][()]
- if HAS_REFCOUNT:
- rc = sys.getrefcount(ob)
- for x in i:
- x[...] += 1
- if HAS_REFCOUNT:
- assert_(sys.getrefcount(ob) == rc-1)
- assert_equal(a, np.arange(6)+98172489)
- def test_iter_common_dtype():
- # Check that the iterator finds a common data type correctly
- i = nditer([array([3], dtype='f4'), array([0], dtype='f8')],
- ['common_dtype'],
- [['readonly', 'copy']]*2,
- casting='safe')
- assert_equal(i.dtypes[0], np.dtype('f8'))
- assert_equal(i.dtypes[1], np.dtype('f8'))
- i = nditer([array([3], dtype='i4'), array([0], dtype='f4')],
- ['common_dtype'],
- [['readonly', 'copy']]*2,
- casting='safe')
- assert_equal(i.dtypes[0], np.dtype('f8'))
- assert_equal(i.dtypes[1], np.dtype('f8'))
- i = nditer([array([3], dtype='f4'), array(0, dtype='f8')],
- ['common_dtype'],
- [['readonly', 'copy']]*2,
- casting='same_kind')
- assert_equal(i.dtypes[0], np.dtype('f4'))
- assert_equal(i.dtypes[1], np.dtype('f4'))
- i = nditer([array([3], dtype='u4'), array(0, dtype='i4')],
- ['common_dtype'],
- [['readonly', 'copy']]*2,
- casting='safe')
- assert_equal(i.dtypes[0], np.dtype('u4'))
- assert_equal(i.dtypes[1], np.dtype('u4'))
- i = nditer([array([3], dtype='u4'), array(-12, dtype='i4')],
- ['common_dtype'],
- [['readonly', 'copy']]*2,
- casting='safe')
- assert_equal(i.dtypes[0], np.dtype('i8'))
- assert_equal(i.dtypes[1], np.dtype('i8'))
- i = nditer([array([3], dtype='u4'), array(-12, dtype='i4'),
- array([2j], dtype='c8'), array([9], dtype='f8')],
- ['common_dtype'],
- [['readonly', 'copy']]*4,
- casting='safe')
- assert_equal(i.dtypes[0], np.dtype('c16'))
- assert_equal(i.dtypes[1], np.dtype('c16'))
- assert_equal(i.dtypes[2], np.dtype('c16'))
- assert_equal(i.dtypes[3], np.dtype('c16'))
- assert_equal(i.value, (3, -12, 2j, 9))
- # When allocating outputs, other outputs aren't factored in
- i = nditer([array([3], dtype='i4'), None, array([2j], dtype='c16')], [],
- [['readonly', 'copy'],
- ['writeonly', 'allocate'],
- ['writeonly']],
- casting='safe')
- assert_equal(i.dtypes[0], np.dtype('i4'))
- assert_equal(i.dtypes[1], np.dtype('i4'))
- assert_equal(i.dtypes[2], np.dtype('c16'))
- # But, if common data types are requested, they are
- i = nditer([array([3], dtype='i4'), None, array([2j], dtype='c16')],
- ['common_dtype'],
- [['readonly', 'copy'],
- ['writeonly', 'allocate'],
- ['writeonly']],
- casting='safe')
- assert_equal(i.dtypes[0], np.dtype('c16'))
- assert_equal(i.dtypes[1], np.dtype('c16'))
- assert_equal(i.dtypes[2], np.dtype('c16'))
- def test_iter_copy_if_overlap():
- # Ensure the iterator makes copies on read/write overlap, if requested
- # Copy not needed, 1 op
- for flag in ['readonly', 'writeonly', 'readwrite']:
- a = arange(10)
- i = nditer([a], ['copy_if_overlap'], [[flag]])
- with i:
- assert_(i.operands[0] is a)
- # Copy needed, 2 ops, read-write overlap
- x = arange(10)
- a = x[1:]
- b = x[:-1]
- with nditer([a, b], ['copy_if_overlap'], [['readonly'], ['readwrite']]) as i:
- assert_(not np.shares_memory(*i.operands))
- # Copy not needed with elementwise, 2 ops, exactly same arrays
- x = arange(10)
- a = x
- b = x
- i = nditer([a, b], ['copy_if_overlap'], [['readonly', 'overlap_assume_elementwise'],
- ['readwrite', 'overlap_assume_elementwise']])
- with i:
- assert_(i.operands[0] is a and i.operands[1] is b)
- with nditer([a, b], ['copy_if_overlap'], [['readonly'], ['readwrite']]) as i:
- assert_(i.operands[0] is a and not np.shares_memory(i.operands[1], b))
- # Copy not needed, 2 ops, no overlap
- x = arange(10)
- a = x[::2]
- b = x[1::2]
- i = nditer([a, b], ['copy_if_overlap'], [['readonly'], ['writeonly']])
- assert_(i.operands[0] is a and i.operands[1] is b)
- # Copy needed, 2 ops, read-write overlap
- x = arange(4, dtype=np.int8)
- a = x[3:]
- b = x.view(np.int32)[:1]
- with nditer([a, b], ['copy_if_overlap'], [['readonly'], ['writeonly']]) as i:
- assert_(not np.shares_memory(*i.operands))
- # Copy needed, 3 ops, read-write overlap
- for flag in ['writeonly', 'readwrite']:
- x = np.ones([10, 10])
- a = x
- b = x.T
- c = x
- with nditer([a, b, c], ['copy_if_overlap'],
- [['readonly'], ['readonly'], [flag]]) as i:
- a2, b2, c2 = i.operands
- assert_(not np.shares_memory(a2, c2))
- assert_(not np.shares_memory(b2, c2))
- # Copy not needed, 3 ops, read-only overlap
- x = np.ones([10, 10])
- a = x
- b = x.T
- c = x
- i = nditer([a, b, c], ['copy_if_overlap'],
- [['readonly'], ['readonly'], ['readonly']])
- a2, b2, c2 = i.operands
- assert_(a is a2)
- assert_(b is b2)
- assert_(c is c2)
- # Copy not needed, 3 ops, read-only overlap
- x = np.ones([10, 10])
- a = x
- b = np.ones([10, 10])
- c = x.T
- i = nditer([a, b, c], ['copy_if_overlap'],
- [['readonly'], ['writeonly'], ['readonly']])
- a2, b2, c2 = i.operands
- assert_(a is a2)
- assert_(b is b2)
- assert_(c is c2)
- # Copy not needed, 3 ops, write-only overlap
- x = np.arange(7)
- a = x[:3]
- b = x[3:6]
- c = x[4:7]
- i = nditer([a, b, c], ['copy_if_overlap'],
- [['readonly'], ['writeonly'], ['writeonly']])
- a2, b2, c2 = i.operands
- assert_(a is a2)
- assert_(b is b2)
- assert_(c is c2)
- def test_iter_op_axes():
- # Check that custom axes work
- # Reverse the axes
- a = arange(6).reshape(2, 3)
- i = nditer([a, a.T], [], [['readonly']]*2, op_axes=[[0, 1], [1, 0]])
- assert_(all([x == y for (x, y) in i]))
- a = arange(24).reshape(2, 3, 4)
- i = nditer([a.T, a], [], [['readonly']]*2, op_axes=[[2, 1, 0], None])
- assert_(all([x == y for (x, y) in i]))
- # Broadcast 1D to any dimension
- a = arange(1, 31).reshape(2, 3, 5)
- b = arange(1, 3)
- i = nditer([a, b], [], [['readonly']]*2, op_axes=[None, [0, -1, -1]])
- assert_equal([x*y for (x, y) in i], (a*b.reshape(2, 1, 1)).ravel())
- b = arange(1, 4)
- i = nditer([a, b], [], [['readonly']]*2, op_axes=[None, [-1, 0, -1]])
- assert_equal([x*y for (x, y) in i], (a*b.reshape(1, 3, 1)).ravel())
- b = arange(1, 6)
- i = nditer([a, b], [], [['readonly']]*2,
- op_axes=[None, [np.newaxis, np.newaxis, 0]])
- assert_equal([x*y for (x, y) in i], (a*b.reshape(1, 1, 5)).ravel())
- # Inner product-style broadcasting
- a = arange(24).reshape(2, 3, 4)
- b = arange(40).reshape(5, 2, 4)
- i = nditer([a, b], ['multi_index'], [['readonly']]*2,
- op_axes=[[0, 1, -1, -1], [-1, -1, 0, 1]])
- assert_equal(i.shape, (2, 3, 5, 2))
- # Matrix product-style broadcasting
- a = arange(12).reshape(3, 4)
- b = arange(20).reshape(4, 5)
- i = nditer([a, b], ['multi_index'], [['readonly']]*2,
- op_axes=[[0, -1], [-1, 1]])
- assert_equal(i.shape, (3, 5))
- def test_iter_op_axes_errors():
- # Check that custom axes throws errors for bad inputs
- # Wrong number of items in op_axes
- a = arange(6).reshape(2, 3)
- assert_raises(ValueError, nditer, [a, a], [], [['readonly']]*2,
- op_axes=[[0], [1], [0]])
- # Out of bounds items in op_axes
- assert_raises(ValueError, nditer, [a, a], [], [['readonly']]*2,
- op_axes=[[2, 1], [0, 1]])
- assert_raises(ValueError, nditer, [a, a], [], [['readonly']]*2,
- op_axes=[[0, 1], [2, -1]])
- # Duplicate items in op_axes
- assert_raises(ValueError, nditer, [a, a], [], [['readonly']]*2,
- op_axes=[[0, 0], [0, 1]])
- assert_raises(ValueError, nditer, [a, a], [], [['readonly']]*2,
- op_axes=[[0, 1], [1, 1]])
- # Different sized arrays in op_axes
- assert_raises(ValueError, nditer, [a, a], [], [['readonly']]*2,
- op_axes=[[0, 1], [0, 1, 0]])
- # Non-broadcastable dimensions in the result
- assert_raises(ValueError, nditer, [a, a], [], [['readonly']]*2,
- op_axes=[[0, 1], [1, 0]])
- def test_iter_copy():
- # Check that copying the iterator works correctly
- a = arange(24).reshape(2, 3, 4)
- # Simple iterator
- i = nditer(a)
- j = i.copy()
- assert_equal([x[()] for x in i], [x[()] for x in j])
- i.iterindex = 3
- j = i.copy()
- assert_equal([x[()] for x in i], [x[()] for x in j])
- # Buffered iterator
- i = nditer(a, ['buffered', 'ranged'], order='F', buffersize=3)
- j = i.copy()
- assert_equal([x[()] for x in i], [x[()] for x in j])
- i.iterindex = 3
- j = i.copy()
- assert_equal([x[()] for x in i], [x[()] for x in j])
- i.iterrange = (3, 9)
- j = i.copy()
- assert_equal([x[()] for x in i], [x[()] for x in j])
- i.iterrange = (2, 18)
- next(i)
- next(i)
- j = i.copy()
- assert_equal([x[()] for x in i], [x[()] for x in j])
- # Casting iterator
- with nditer(a, ['buffered'], order='F', casting='unsafe',
- op_dtypes='f8', buffersize=5) as i:
- j = i.copy()
- assert_equal([x[()] for x in j], a.ravel(order='F'))
- a = arange(24, dtype='<i4').reshape(2, 3, 4)
- with nditer(a, ['buffered'], order='F', casting='unsafe',
- op_dtypes='>f8', buffersize=5) as i:
- j = i.copy()
- assert_equal([x[()] for x in j], a.ravel(order='F'))
- def test_iter_allocate_output_simple():
- # Check that the iterator will properly allocate outputs
- # Simple case
- a = arange(6)
- i = nditer([a, None], [], [['readonly'], ['writeonly', 'allocate']],
- op_dtypes=[None, np.dtype('f4')])
- assert_equal(i.operands[1].shape, a.shape)
- assert_equal(i.operands[1].dtype, np.dtype('f4'))
- def test_iter_allocate_output_buffered_readwrite():
- # Allocated output with buffering + delay_bufalloc
- a = arange(6)
- i = nditer([a, None], ['buffered', 'delay_bufalloc'],
- [['readonly'], ['allocate', 'readwrite']])
- with i:
- i.operands[1][:] = 1
- i.reset()
- for x in i:
- x[1][...] += x[0][...]
- assert_equal(i.operands[1], a+1)
- def test_iter_allocate_output_itorder():
- # The allocated output should match the iteration order
- # C-order input, best iteration order
- a = arange(6, dtype='i4').reshape(2, 3)
- i = nditer([a, None], [], [['readonly'], ['writeonly', 'allocate']],
- op_dtypes=[None, np.dtype('f4')])
- assert_equal(i.operands[1].shape, a.shape)
- assert_equal(i.operands[1].strides, a.strides)
- assert_equal(i.operands[1].dtype, np.dtype('f4'))
- # F-order input, best iteration order
- a = arange(24, dtype='i4').reshape(2, 3, 4).T
- i = nditer([a, None], [], [['readonly'], ['writeonly', 'allocate']],
- op_dtypes=[None, np.dtype('f4')])
- assert_equal(i.operands[1].shape, a.shape)
- assert_equal(i.operands[1].strides, a.strides)
- assert_equal(i.operands[1].dtype, np.dtype('f4'))
- # Non-contiguous input, C iteration order
- a = arange(24, dtype='i4').reshape(2, 3, 4).swapaxes(0, 1)
- i = nditer([a, None], [],
- [['readonly'], ['writeonly', 'allocate']],
- order='C',
- op_dtypes=[None, np.dtype('f4')])
- assert_equal(i.operands[1].shape, a.shape)
- assert_equal(i.operands[1].strides, (32, 16, 4))
- assert_equal(i.operands[1].dtype, np.dtype('f4'))
- def test_iter_allocate_output_opaxes():
- # Specifying op_axes should work
- a = arange(24, dtype='i4').reshape(2, 3, 4)
- i = nditer([None, a], [], [['writeonly', 'allocate'], ['readonly']],
- op_dtypes=[np.dtype('u4'), None],
- op_axes=[[1, 2, 0], None])
- assert_equal(i.operands[0].shape, (4, 2, 3))
- assert_equal(i.operands[0].strides, (4, 48, 16))
- assert_equal(i.operands[0].dtype, np.dtype('u4'))
- def test_iter_allocate_output_types_promotion():
- # Check type promotion of automatic outputs
- i = nditer([array([3], dtype='f4'), array([0], dtype='f8'), None], [],
- [['readonly']]*2+[['writeonly', 'allocate']])
- assert_equal(i.dtypes[2], np.dtype('f8'))
- i = nditer([array([3], dtype='i4'), array([0], dtype='f4'), None], [],
- [['readonly']]*2+[['writeonly', 'allocate']])
- assert_equal(i.dtypes[2], np.dtype('f8'))
- i = nditer([array([3], dtype='f4'), array(0, dtype='f8'), None], [],
- [['readonly']]*2+[['writeonly', 'allocate']])
- assert_equal(i.dtypes[2], np.dtype('f4'))
- i = nditer([array([3], dtype='u4'), array(0, dtype='i4'), None], [],
- [['readonly']]*2+[['writeonly', 'allocate']])
- assert_equal(i.dtypes[2], np.dtype('u4'))
- i = nditer([array([3], dtype='u4'), array(-12, dtype='i4'), None], [],
- [['readonly']]*2+[['writeonly', 'allocate']])
- assert_equal(i.dtypes[2], np.dtype('i8'))
- def test_iter_allocate_output_types_byte_order():
- # Verify the rules for byte order changes
- # When there's just one input, the output type exactly matches
- a = array([3], dtype='u4').newbyteorder()
- i = nditer([a, None], [],
- [['readonly'], ['writeonly', 'allocate']])
- assert_equal(i.dtypes[0], i.dtypes[1])
- # With two or more inputs, the output type is in native byte order
- i = nditer([a, a, None], [],
- [['readonly'], ['readonly'], ['writeonly', 'allocate']])
- assert_(i.dtypes[0] != i.dtypes[2])
- assert_equal(i.dtypes[0].newbyteorder('='), i.dtypes[2])
- def test_iter_allocate_output_types_scalar():
- # If the inputs are all scalars, the output should be a scalar
- i = nditer([None, 1, 2.3, np.float32(12), np.complex128(3)], [],
- [['writeonly', 'allocate']] + [['readonly']]*4)
- assert_equal(i.operands[0].dtype, np.dtype('complex128'))
- assert_equal(i.operands[0].ndim, 0)
- def test_iter_allocate_output_subtype():
- # Make sure that the subtype with priority wins
- class MyNDArray(np.ndarray):
- __array_priority__ = 15
- # subclass vs ndarray
- a = np.array([[1, 2], [3, 4]]).view(MyNDArray)
- b = np.arange(4).reshape(2, 2).T
- i = nditer([a, b, None], [],
- [['readonly'], ['readonly'], ['writeonly', 'allocate']])
- assert_equal(type(a), type(i.operands[2]))
- assert_(type(b) is not type(i.operands[2]))
- assert_equal(i.operands[2].shape, (2, 2))
- # If subtypes are disabled, we should get back an ndarray.
- i = nditer([a, b, None], [],
- [['readonly'], ['readonly'],
- ['writeonly', 'allocate', 'no_subtype']])
- assert_equal(type(b), type(i.operands[2]))
- assert_(type(a) is not type(i.operands[2]))
- assert_equal(i.operands[2].shape, (2, 2))
- def test_iter_allocate_output_errors():
- # Check that the iterator will throw errors for bad output allocations
- # Need an input if no output data type is specified
- a = arange(6)
- assert_raises(TypeError, nditer, [a, None], [],
- [['writeonly'], ['writeonly', 'allocate']])
- # Allocated output should be flagged for writing
- assert_raises(ValueError, nditer, [a, None], [],
- [['readonly'], ['allocate', 'readonly']])
- # Allocated output can't have buffering without delayed bufalloc
- assert_raises(ValueError, nditer, [a, None], ['buffered'],
- ['allocate', 'readwrite'])
- # Must specify at least one input
- assert_raises(ValueError, nditer, [None, None], [],
- [['writeonly', 'allocate'],
- ['writeonly', 'allocate']],
- op_dtypes=[np.dtype('f4'), np.dtype('f4')])
- # If using op_axes, must specify all the axes
- a = arange(24, dtype='i4').reshape(2, 3, 4)
- assert_raises(ValueError, nditer, [a, None], [],
- [['readonly'], ['writeonly', 'allocate']],
- op_dtypes=[None, np.dtype('f4')],
- op_axes=[None, [0, np.newaxis, 1]])
- # If using op_axes, the axes must be within bounds
- assert_raises(ValueError, nditer, [a, None], [],
- [['readonly'], ['writeonly', 'allocate']],
- op_dtypes=[None, np.dtype('f4')],
- op_axes=[None, [0, 3, 1]])
- # If using op_axes, there can't be duplicates
- assert_raises(ValueError, nditer, [a, None], [],
- [['readonly'], ['writeonly', 'allocate']],
- op_dtypes=[None, np.dtype('f4')],
- op_axes=[None, [0, 2, 1, 0]])
- # Not all axes may be specified if a reduction. If there is a hole
- # in op_axes, this is an error.
- a = arange(24, dtype='i4').reshape(2, 3, 4)
- assert_raises(ValueError, nditer, [a, None], ["reduce_ok"],
- [['readonly'], ['readwrite', 'allocate']],
- op_dtypes=[None, np.dtype('f4')],
- op_axes=[None, [0, np.newaxis, 2]])
- def test_iter_remove_axis():
- a = arange(24).reshape(2, 3, 4)
- i = nditer(a, ['multi_index'])
- i.remove_axis(1)
- assert_equal([x for x in i], a[:, 0,:].ravel())
- a = a[::-1,:,:]
- i = nditer(a, ['multi_index'])
- i.remove_axis(0)
- assert_equal([x for x in i], a[0,:,:].ravel())
- def test_iter_remove_multi_index_inner_loop():
- # Check that removing multi-index support works
- a = arange(24).reshape(2, 3, 4)
- i = nditer(a, ['multi_index'])
- assert_equal(i.ndim, 3)
- assert_equal(i.shape, (2, 3, 4))
- assert_equal(i.itviews[0].shape, (2, 3, 4))
- # Removing the multi-index tracking causes all dimensions to coalesce
- before = [x for x in i]
- i.remove_multi_index()
- after = [x for x in i]
- assert_equal(before, after)
- assert_equal(i.ndim, 1)
- assert_raises(ValueError, lambda i:i.shape, i)
- assert_equal(i.itviews[0].shape, (24,))
- # Removing the inner loop means there's just one iteration
- i.reset()
- assert_equal(i.itersize, 24)
- assert_equal(i[0].shape, tuple())
- i.enable_external_loop()
- assert_equal(i.itersize, 24)
- assert_equal(i[0].shape, (24,))
- assert_equal(i.value, arange(24))
- def test_iter_iterindex():
- # Make sure iterindex works
- buffersize = 5
- a = arange(24).reshape(4, 3, 2)
- for flags in ([], ['buffered']):
- i = nditer(a, flags, buffersize=buffersize)
- assert_equal(iter_iterindices(i), list(range(24)))
- i.iterindex = 2
- assert_equal(iter_iterindices(i), list(range(2, 24)))
- i = nditer(a, flags, order='F', buffersize=buffersize)
- assert_equal(iter_iterindices(i), list(range(24)))
- i.iterindex = 5
- assert_equal(iter_iterindices(i), list(range(5, 24)))
- i = nditer(a[::-1], flags, order='F', buffersize=buffersize)
- assert_equal(iter_iterindices(i), list(range(24)))
- i.iterindex = 9
- assert_equal(iter_iterindices(i), list(range(9, 24)))
- i = nditer(a[::-1, ::-1], flags, order='C', buffersize=buffersize)
- assert_equal(iter_iterindices(i), list(range(24)))
- i.iterindex = 13
- assert_equal(iter_iterindices(i), list(range(13, 24)))
- i = nditer(a[::1, ::-1], flags, buffersize=buffersize)
- assert_equal(iter_iterindices(i), list(range(24)))
- i.iterindex = 23
- assert_equal(iter_iterindices(i), list(range(23, 24)))
- i.reset()
- i.iterindex = 2
- assert_equal(iter_iterindices(i), list(range(2, 24)))
- def test_iter_iterrange():
- # Make sure getting and resetting the iterrange works
- buffersize = 5
- a = arange(24, dtype='i4').reshape(4, 3, 2)
- a_fort = a.ravel(order='F')
- i = nditer(a, ['ranged'], ['readonly'], order='F',
- buffersize=buffersize)
- assert_equal(i.iterrange, (0, 24))
- assert_equal([x[()] for x in i], a_fort)
- for r in [(0, 24), (1, 2), (3, 24), (5, 5), (0, 20), (23, 24)]:
- i.iterrange = r
- assert_equal(i.iterrange, r)
- assert_equal([x[()] for x in i], a_fort[r[0]:r[1]])
- i = nditer(a, ['ranged', 'buffered'], ['readonly'], order='F',
- op_dtypes='f8', buffersize=buffersize)
- assert_equal(i.iterrange, (0, 24))
- assert_equal([x[()] for x in i], a_fort)
- for r in [(0, 24), (1, 2), (3, 24), (5, 5), (0, 20), (23, 24)]:
- i.iterrange = r
- assert_equal(i.iterrange, r)
- assert_equal([x[()] for x in i], a_fort[r[0]:r[1]])
- def get_array(i):
- val = np.array([], dtype='f8')
- for x in i:
- val = np.concatenate((val, x))
- return val
- i = nditer(a, ['ranged', 'buffered', 'external_loop'],
- ['readonly'], order='F',
- op_dtypes='f8', buffersize=buffersize)
- assert_equal(i.iterrange, (0, 24))
- assert_equal(get_array(i), a_fort)
- for r in [(0, 24), (1, 2), (3, 24), (5, 5), (0, 20), (23, 24)]:
- i.iterrange = r
- assert_equal(i.iterrange, r)
- assert_equal(get_array(i), a_fort[r[0]:r[1]])
- def test_iter_buffering():
- # Test buffering with several buffer sizes and types
- arrays = []
- # F-order swapped array
- arrays.append(np.arange(24,
- dtype='c16').reshape(2, 3, 4).T.newbyteorder().byteswap())
- # Contiguous 1-dimensional array
- arrays.append(np.arange(10, dtype='f4'))
- # Unaligned array
- a = np.zeros((4*16+1,), dtype='i1')[1:]
- a.dtype = 'i4'
- a[:] = np.arange(16, dtype='i4')
- arrays.append(a)
- # 4-D F-order array
- arrays.append(np.arange(120, dtype='i4').reshape(5, 3, 2, 4).T)
- for a in arrays:
- for buffersize in (1, 2, 3, 5, 8, 11, 16, 1024):
- vals = []
- i = nditer(a, ['buffered', 'external_loop'],
- [['readonly', 'nbo', 'aligned']],
- order='C',
- casting='equiv',
- buffersize=buffersize)
- while not i.finished:
- assert_(i[0].size <= buffersize)
- vals.append(i[0].copy())
- i.iternext()
- assert_equal(np.concatenate(vals), a.ravel(order='C'))
- def test_iter_write_buffering():
- # Test that buffering of writes is working
- # F-order swapped array
- a = np.arange(24).reshape(2, 3, 4).T.newbyteorder().byteswap()
- i = nditer(a, ['buffered'],
- [['readwrite', 'nbo', 'aligned']],
- casting='equiv',
- order='C',
- buffersize=16)
- x = 0
- with i:
- while not i.finished:
- i[0] = x
- x += 1
- i.iternext()
- assert_equal(a.ravel(order='C'), np.arange(24))
- def test_iter_buffering_delayed_alloc():
- # Test that delaying buffer allocation works
- a = np.arange(6)
- b = np.arange(1, dtype='f4')
- i = nditer([a, b], ['buffered', 'delay_bufalloc', 'multi_index', 'reduce_ok'],
- ['readwrite'],
- casting='unsafe',
- op_dtypes='f4')
- assert_(i.has_delayed_bufalloc)
- assert_raises(ValueError, lambda i:i.multi_index, i)
- assert_raises(ValueError, lambda i:i[0], i)
- assert_raises(ValueError, lambda i:i[0:2], i)
- def assign_iter(i):
- i[0] = 0
- assert_raises(ValueError, assign_iter, i)
- i.reset()
- assert_(not i.has_delayed_bufalloc)
- assert_equal(i.multi_index, (0,))
- with i:
- assert_equal(i[0], 0)
- i[1] = 1
- assert_equal(i[0:2], [0, 1])
- assert_equal([[x[0][()], x[1][()]] for x in i], list(zip(range(6), [1]*6)))
- def test_iter_buffered_cast_simple():
- # Test that buffering can handle a simple cast
- a = np.arange(10, dtype='f4')
- i = nditer(a, ['buffered', 'external_loop'],
- [['readwrite', 'nbo', 'aligned']],
- casting='same_kind',
- op_dtypes=[np.dtype('f8')],
- buffersize=3)
- with i:
- for v in i:
- v[...] *= 2
- assert_equal(a, 2*np.arange(10, dtype='f4'))
- def test_iter_buffered_cast_byteswapped():
- # Test that buffering can handle a cast which requires swap->cast->swap
- a = np.arange(10, dtype='f4').newbyteorder().byteswap()
- i = nditer(a, ['buffered', 'external_loop'],
- [['readwrite', 'nbo', 'aligned']],
- casting='same_kind',
- op_dtypes=[np.dtype('f8').newbyteorder()],
- buffersize=3)
- with i:
- for v in i:
- v[...] *= 2
- assert_equal(a, 2*np.arange(10, dtype='f4'))
- with suppress_warnings() as sup:
- sup.filter(np.ComplexWarning)
- a = np.arange(10, dtype='f8').newbyteorder().byteswap()
- i = nditer(a, ['buffered', 'external_loop'],
- [['readwrite', 'nbo', 'aligned']],
- casting='unsafe',
- op_dtypes=[np.dtype('c8').newbyteorder()],
- buffersize=3)
- with i:
- for v in i:
- v[...] *= 2
- assert_equal(a, 2*np.arange(10, dtype='f8'))
- def test_iter_buffered_cast_byteswapped_complex():
- # Test that buffering can handle a cast which requires swap->cast->copy
- a = np.arange(10, dtype='c8').newbyteorder().byteswap()
- a += 2j
- i = nditer(a, ['buffered', 'external_loop'],
- [['readwrite', 'nbo', 'aligned']],
- casting='same_kind',
- op_dtypes=[np.dtype('c16')],
- buffersize=3)
- with i:
- for v in i:
- v[...] *= 2
- assert_equal(a, 2*np.arange(10, dtype='c8') + 4j)
- a = np.arange(10, dtype='c8')
- a += 2j
- i = nditer(a, ['buffered', 'external_loop'],
- [['readwrite', 'nbo', 'aligned']],
- casting='same_kind',
- op_dtypes=[np.dtype('c16').newbyteorder()],
- buffersize=3)
- with i:
- for v in i:
- v[...] *= 2
- assert_equal(a, 2*np.arange(10, dtype='c8') + 4j)
- a = np.arange(10, dtype=np.clongdouble).newbyteorder().byteswap()
- a += 2j
- i = nditer(a, ['buffered', 'external_loop'],
- [['readwrite', 'nbo', 'aligned']],
- casting='same_kind',
- op_dtypes=[np.dtype('c16')],
- buffersize=3)
- with i:
- for v in i:
- v[...] *= 2
- assert_equal(a, 2*np.arange(10, dtype=np.clongdouble) + 4j)
- a = np.arange(10, dtype=np.longdouble).newbyteorder().byteswap()
- i = nditer(a, ['buffered', 'external_loop'],
- [['readwrite', 'nbo', 'aligned']],
- casting='same_kind',
- op_dtypes=[np.dtype('f4')],
- buffersize=7)
- with i:
- for v in i:
- v[...] *= 2
- assert_equal(a, 2*np.arange(10, dtype=np.longdouble))
- def test_iter_buffered_cast_structured_type():
- # Tests buffering of structured types
- # simple -> struct type (duplicates the value)
- sdt = [('a', 'f4'), ('b', 'i8'), ('c', 'c8', (2, 3)), ('d', 'O')]
- a = np.arange(3, dtype='f4') + 0.5
- i = nditer(a, ['buffered', 'refs_ok'], ['readonly'],
- casting='unsafe',
- op_dtypes=sdt)
- vals = [np.array(x) for x in i]
- assert_equal(vals[0]['a'], 0.5)
- assert_equal(vals[0]['b'], 0)
- assert_equal(vals[0]['c'], [[(0.5)]*3]*2)
- assert_equal(vals[0]['d'], 0.5)
- assert_equal(vals[1]['a'], 1.5)
- assert_equal(vals[1]['b'], 1)
- assert_equal(vals[1]['c'], [[(1.5)]*3]*2)
- assert_equal(vals[1]['d'], 1.5)
- assert_equal(vals[0].dtype, np.dtype(sdt))
- # object -> struct type
- sdt = [('a', 'f4'), ('b', 'i8'), ('c', 'c8', (2, 3)), ('d', 'O')]
- a = np.zeros((3,), dtype='O')
- a[0] = (0.5, 0.5, [[0.5, 0.5, 0.5], [0.5, 0.5, 0.5]], 0.5)
- a[1] = (1.5, 1.5, [[1.5, 1.5, 1.5], [1.5, 1.5, 1.5]], 1.5)
- a[2] = (2.5, 2.5, [[2.5, 2.5, 2.5], [2.5, 2.5, 2.5]], 2.5)
- if HAS_REFCOUNT:
- rc = sys.getrefcount(a[0])
- i = nditer(a, ['buffered', 'refs_ok'], ['readonly'],
- casting='unsafe',
- op_dtypes=sdt)
- vals = [x.copy() for x in i]
- assert_equal(vals[0]['a'], 0.5)
- assert_equal(vals[0]['b'], 0)
- assert_equal(vals[0]['c'], [[(0.5)]*3]*2)
- assert_equal(vals[0]['d'], 0.5)
- assert_equal(vals[1]['a'], 1.5)
- assert_equal(vals[1]['b'], 1)
- assert_equal(vals[1]['c'], [[(1.5)]*3]*2)
- assert_equal(vals[1]['d'], 1.5)
- assert_equal(vals[0].dtype, np.dtype(sdt))
- vals, i, x = [None]*3
- if HAS_REFCOUNT:
- assert_equal(sys.getrefcount(a[0]), rc)
- # single-field struct type -> simple
- sdt = [('a', 'f4')]
- a = np.array([(5.5,), (8,)], dtype=sdt)
- i = nditer(a, ['buffered', 'refs_ok'], ['readonly'],
- casting='unsafe',
- op_dtypes='i4')
- assert_equal([x_[()] for x_ in i], [5, 8])
- # make sure multi-field struct type -> simple doesn't work
- sdt = [('a', 'f4'), ('b', 'i8'), ('d', 'O')]
- a = np.array([(5.5, 7, 'test'), (8, 10, 11)], dtype=sdt)
- assert_raises(TypeError, lambda: (
- nditer(a, ['buffered', 'refs_ok'], ['readonly'],
- casting='unsafe',
- op_dtypes='i4')))
- # struct type -> struct type (field-wise copy)
- sdt1 = [('a', 'f4'), ('b', 'i8'), ('d', 'O')]
- sdt2 = [('d', 'u2'), ('a', 'O'), ('b', 'f8')]
- a = np.array([(1, 2, 3), (4, 5, 6)], dtype=sdt1)
- i = nditer(a, ['buffered', 'refs_ok'], ['readonly'],
- casting='unsafe',
- op_dtypes=sdt2)
- assert_equal(i[0].dtype, np.dtype(sdt2))
- assert_equal([np.array(x_) for x_ in i],
- [np.array((1, 2, 3), dtype=sdt2),
- np.array((4, 5, 6), dtype=sdt2)])
- # make sure struct type -> struct type with different
- # number of fields fails
- sdt1 = [('a', 'f4'), ('b', 'i8'), ('d', 'O')]
- sdt2 = [('b', 'O'), ('a', 'f8')]
- a = np.array([(1, 2, 3), (4, 5, 6)], dtype=sdt1)
- assert_raises(ValueError, lambda : (
- nditer(a, ['buffered', 'refs_ok'], ['readwrite'],
- casting='unsafe',
- op_dtypes=sdt2)))
- def test_iter_buffered_cast_subarray():
- # Tests buffering of subarrays
- # one element -> many (copies it to all)
- sdt1 = [('a', 'f4')]
- sdt2 = [('a', 'f8', (3, 2, 2))]
- a = np.zeros((6,), dtype=sdt1)
- a['a'] = np.arange(6)
- i = nditer(a, ['buffered', 'refs_ok'], ['readonly'],
- casting='unsafe',
- op_dtypes=sdt2)
- assert_equal(i[0].dtype, np.dtype(sdt2))
- for x, count in zip(i, list(range(6))):
- assert_(np.all(x['a'] == count))
- # one element -> many -> back (copies it to all)
- sdt1 = [('a', 'O', (1, 1))]
- sdt2 = [('a', 'O', (3, 2, 2))]
- a = np.zeros((6,), dtype=sdt1)
- a['a'][:, 0, 0] = np.arange(6)
- i = nditer(a, ['buffered', 'refs_ok'], ['readwrite'],
- casting='unsafe',
- op_dtypes=sdt2)
- with i:
- assert_equal(i[0].dtype, np.dtype(sdt2))
- count = 0
- for x in i:
- assert_(np.all(x['a'] == count))
- x['a'][0] += 2
- count += 1
- assert_equal(a['a'], np.arange(6).reshape(6, 1, 1)+2)
- # many -> one element -> back (copies just element 0)
- sdt1 = [('a', 'O', (3, 2, 2))]
- sdt2 = [('a', 'O', (1,))]
- a = np.zeros((6,), dtype=sdt1)
- a['a'][:, 0, 0, 0] = np.arange(6)
- i = nditer(a, ['buffered', 'refs_ok'], ['readwrite'],
- casting='unsafe',
- op_dtypes=sdt2)
- with i:
- assert_equal(i[0].dtype, np.dtype(sdt2))
- count = 0
- for x in i:
- assert_equal(x['a'], count)
- x['a'] += 2
- count += 1
- assert_equal(a['a'], np.arange(6).reshape(6, 1, 1, 1)*np.ones((1, 3, 2, 2))+2)
- # many -> one element -> back (copies just element 0)
- sdt1 = [('a', 'f8', (3, 2, 2))]
- sdt2 = [('a', 'O', (1,))]
- a = np.zeros((6,), dtype=sdt1)
- a['a'][:, 0, 0, 0] = np.arange(6)
- i = nditer(a, ['buffered', 'refs_ok'], ['readonly'],
- casting='unsafe',
- op_dtypes=sdt2)
- assert_equal(i[0].dtype, np.dtype(sdt2))
- count = 0
- for x in i:
- assert_equal(x['a'], count)
- count += 1
- # many -> one element (copies just element 0)
- sdt1 = [('a', 'O', (3, 2, 2))]
- sdt2 = [('a', 'f4', (1,))]
- a = np.zeros((6,), dtype=sdt1)
- a['a'][:, 0, 0, 0] = np.arange(6)
- i = nditer(a, ['buffered', 'refs_ok'], ['readonly'],
- casting='unsafe',
- op_dtypes=sdt2)
- assert_equal(i[0].dtype, np.dtype(sdt2))
- count = 0
- for x in i:
- assert_equal(x['a'], count)
- count += 1
- # many -> matching shape (straightforward copy)
- sdt1 = [('a', 'O', (3, 2, 2))]
- sdt2 = [('a', 'f4', (3, 2, 2))]
- a = np.zeros((6,), dtype=sdt1)
- a['a'] = np.arange(6*3*2*2).reshape(6, 3, 2, 2)
- i = nditer(a, ['buffered', 'refs_ok'], ['readonly'],
- casting='unsafe',
- op_dtypes=sdt2)
- assert_equal(i[0].dtype, np.dtype(sdt2))
- count = 0
- for x in i:
- assert_equal(x['a'], a[count]['a'])
- count += 1
- # vector -> smaller vector (truncates)
- sdt1 = [('a', 'f8', (6,))]
- sdt2 = [('a', 'f4', (2,))]
- a = np.zeros((6,), dtype=sdt1)
- a['a'] = np.arange(6*6).reshape(6, 6)
- i = nditer(a, ['buffered', 'refs_ok'], ['readonly'],
- casting='unsafe',
- op_dtypes=sdt2)
- assert_equal(i[0].dtype, np.dtype(sdt2))
- count = 0
- for x in i:
- assert_equal(x['a'], a[count]['a'][:2])
- count += 1
- # vector -> bigger vector (pads with zeros)
- sdt1 = [('a', 'f8', (2,))]
- sdt2 = [('a', 'f4', (6,))]
- a = np.zeros((6,), dtype=sdt1)
- a['a'] = np.arange(6*2).reshape(6, 2)
- i = nditer(a, ['buffered', 'refs_ok'], ['readonly'],
- casting='unsafe',
- op_dtypes=sdt2)
- assert_equal(i[0].dtype, np.dtype(sdt2))
- count = 0
- for x in i:
- assert_equal(x['a'][:2], a[count]['a'])
- assert_equal(x['a'][2:], [0, 0, 0, 0])
- count += 1
- # vector -> matrix (broadcasts)
- sdt1 = [('a', 'f8', (2,))]
- sdt2 = [('a', 'f4', (2, 2))]
- a = np.zeros((6,), dtype=sdt1)
- a['a'] = np.arange(6*2).reshape(6, 2)
- i = nditer(a, ['buffered', 'refs_ok'], ['readonly'],
- casting='unsafe',
- op_dtypes=sdt2)
- assert_equal(i[0].dtype, np.dtype(sdt2))
- count = 0
- for x in i:
- assert_equal(x['a'][0], a[count]['a'])
- assert_equal(x['a'][1], a[count]['a'])
- count += 1
- # vector -> matrix (broadcasts and zero-pads)
- sdt1 = [('a', 'f8', (2, 1))]
- sdt2 = [('a', 'f4', (3, 2))]
- a = np.zeros((6,), dtype=sdt1)
- a['a'] = np.arange(6*2).reshape(6, 2, 1)
- i = nditer(a, ['buffered', 'refs_ok'], ['readonly'],
- casting='unsafe',
- op_dtypes=sdt2)
- assert_equal(i[0].dtype, np.dtype(sdt2))
- count = 0
- for x in i:
- assert_equal(x['a'][:2, 0], a[count]['a'][:, 0])
- assert_equal(x['a'][:2, 1], a[count]['a'][:, 0])
- assert_equal(x['a'][2,:], [0, 0])
- count += 1
- # matrix -> matrix (truncates and zero-pads)
- sdt1 = [('a', 'f8', (2, 3))]
- sdt2 = [('a', 'f4', (3, 2))]
- a = np.zeros((6,), dtype=sdt1)
- a['a'] = np.arange(6*2*3).reshape(6, 2, 3)
- i = nditer(a, ['buffered', 'refs_ok'], ['readonly'],
- casting='unsafe',
- op_dtypes=sdt2)
- assert_equal(i[0].dtype, np.dtype(sdt2))
- count = 0
- for x in i:
- assert_equal(x['a'][:2, 0], a[count]['a'][:, 0])
- assert_equal(x['a'][:2, 1], a[count]['a'][:, 1])
- assert_equal(x['a'][2,:], [0, 0])
- count += 1
- def test_iter_buffering_badwriteback():
- # Writing back from a buffer cannot combine elements
- # a needs write buffering, but had a broadcast dimension
- a = np.arange(6).reshape(2, 3, 1)
- b = np.arange(12).reshape(2, 3, 2)
- assert_raises(ValueError, nditer, [a, b],
- ['buffered', 'external_loop'],
- [['readwrite'], ['writeonly']],
- order='C')
- # But if a is readonly, it's fine
- nditer([a, b], ['buffered', 'external_loop'],
- [['readonly'], ['writeonly']],
- order='C')
- # If a has just one element, it's fine too (constant 0 stride, a reduction)
- a = np.arange(1).reshape(1, 1, 1)
- nditer([a, b], ['buffered', 'external_loop', 'reduce_ok'],
- [['readwrite'], ['writeonly']],
- order='C')
- # check that it fails on other dimensions too
- a = np.arange(6).reshape(1, 3, 2)
- assert_raises(ValueError, nditer, [a, b],
- ['buffered', 'external_loop'],
- [['readwrite'], ['writeonly']],
- order='C')
- a = np.arange(4).reshape(2, 1, 2)
- assert_raises(ValueError, nditer, [a, b],
- ['buffered', 'external_loop'],
- [['readwrite'], ['writeonly']],
- order='C')
- def test_iter_buffering_string():
- # Safe casting disallows shrinking strings
- a = np.array(['abc', 'a', 'abcd'], dtype=np.bytes_)
- assert_equal(a.dtype, np.dtype('S4'))
- assert_raises(TypeError, nditer, a, ['buffered'], ['readonly'],
- op_dtypes='S2')
- i = nditer(a, ['buffered'], ['readonly'], op_dtypes='S6')
- assert_equal(i[0], b'abc')
- assert_equal(i[0].dtype, np.dtype('S6'))
- a = np.array(['abc', 'a', 'abcd'], dtype=np.unicode_)
- assert_equal(a.dtype, np.dtype('U4'))
- assert_raises(TypeError, nditer, a, ['buffered'], ['readonly'],
- op_dtypes='U2')
- i = nditer(a, ['buffered'], ['readonly'], op_dtypes='U6')
- assert_equal(i[0], u'abc')
- assert_equal(i[0].dtype, np.dtype('U6'))
- def test_iter_buffering_growinner():
- # Test that the inner loop grows when no buffering is needed
- a = np.arange(30)
- i = nditer(a, ['buffered', 'growinner', 'external_loop'],
- buffersize=5)
- # Should end up with just one inner loop here
- assert_equal(i[0].size, a.size)
- @pytest.mark.slow
- def test_iter_buffered_reduce_reuse():
- # large enough array for all views, including negative strides.
- a = np.arange(2*3**5)[3**5:3**5+1]
- flags = ['buffered', 'delay_bufalloc', 'multi_index', 'reduce_ok', 'refs_ok']
- op_flags = [('readonly',), ('readwrite', 'allocate')]
- op_axes_list = [[(0, 1, 2), (0, 1, -1)], [(0, 1, 2), (0, -1, -1)]]
- # wrong dtype to force buffering
- op_dtypes = [float, a.dtype]
- def get_params():
- for xs in range(-3**2, 3**2 + 1):
- for ys in range(xs, 3**2 + 1):
- for op_axes in op_axes_list:
- # last stride is reduced and because of that not
- # important for this test, as it is the inner stride.
- strides = (xs * a.itemsize, ys * a.itemsize, a.itemsize)
- arr = np.lib.stride_tricks.as_strided(a, (3, 3, 3), strides)
- for skip in [0, 1]:
- yield arr, op_axes, skip
- for arr, op_axes, skip in get_params():
- nditer2 = np.nditer([arr.copy(), None],
- op_axes=op_axes, flags=flags, op_flags=op_flags,
- op_dtypes=op_dtypes)
- with nditer2:
- nditer2.operands[-1][...] = 0
- nditer2.reset()
- nditer2.iterindex = skip
- for (a2_in, b2_in) in nditer2:
- b2_in += a2_in.astype(np.int_)
- comp_res = nditer2.operands[-1]
- for bufsize in range(0, 3**3):
- nditer1 = np.nditer([arr, None],
- op_axes=op_axes, flags=flags, op_flags=op_flags,
- buffersize=bufsize, op_dtypes=op_dtypes)
- with nditer1:
- nditer1.operands[-1][...] = 0
- nditer1.reset()
- nditer1.iterindex = skip
- for (a1_in, b1_in) in nditer1:
- b1_in += a1_in.astype(np.int_)
- res = nditer1.operands[-1]
- assert_array_equal(res, comp_res)
- def test_iter_no_broadcast():
- # Test that the no_broadcast flag works
- a = np.arange(24).reshape(2, 3, 4)
- b = np.arange(6).reshape(2, 3, 1)
- c = np.arange(12).reshape(3, 4)
- nditer([a, b, c], [],
- [['readonly', 'no_broadcast'],
- ['readonly'], ['readonly']])
- assert_raises(ValueError, nditer, [a, b, c], [],
- [['readonly'], ['readonly', 'no_broadcast'], ['readonly']])
- assert_raises(ValueError, nditer, [a, b, c], [],
- [['readonly'], ['readonly'], ['readonly', 'no_broadcast']])
- class TestIterNested:
- def test_basic(self):
- # Test nested iteration basic usage
- a = arange(12).reshape(2, 3, 2)
- i, j = np.nested_iters(a, [[0], [1, 2]])
- vals = [list(j) for _ in i]
- assert_equal(vals, [[0, 1, 2, 3, 4, 5], [6, 7, 8, 9, 10, 11]])
- i, j = np.nested_iters(a, [[0, 1], [2]])
- vals = [list(j) for _ in i]
- assert_equal(vals, [[0, 1], [2, 3], [4, 5], [6, 7], [8, 9], [10, 11]])
- i, j = np.nested_iters(a, [[0, 2], [1]])
- vals = [list(j) for _ in i]
- assert_equal(vals, [[0, 2, 4], [1, 3, 5], [6, 8, 10], [7, 9, 11]])
- def test_reorder(self):
- # Test nested iteration basic usage
- a = arange(12).reshape(2, 3, 2)
- # In 'K' order (default), it gets reordered
- i, j = np.nested_iters(a, [[0], [2, 1]])
- vals = [list(j) for _ in i]
- assert_equal(vals, [[0, 1, 2, 3, 4, 5], [6, 7, 8, 9, 10, 11]])
- i, j = np.nested_iters(a, [[1, 0], [2]])
- vals = [list(j) for _ in i]
- assert_equal(vals, [[0, 1], [2, 3], [4, 5], [6, 7], [8, 9], [10, 11]])
- i, j = np.nested_iters(a, [[2, 0], [1]])
- vals = [list(j) for _ in i]
- assert_equal(vals, [[0, 2, 4], [1, 3, 5], [6, 8, 10], [7, 9, 11]])
- # In 'C' order, it doesn't
- i, j = np.nested_iters(a, [[0], [2, 1]], order='C')
- vals = [list(j) for _ in i]
- assert_equal(vals, [[0, 2, 4, 1, 3, 5], [6, 8, 10, 7, 9, 11]])
- i, j = np.nested_iters(a, [[1, 0], [2]], order='C')
- vals = [list(j) for _ in i]
- assert_equal(vals, [[0, 1], [6, 7], [2, 3], [8, 9], [4, 5], [10, 11]])
- i, j = np.nested_iters(a, [[2, 0], [1]], order='C')
- vals = [list(j) for _ in i]
- assert_equal(vals, [[0, 2, 4], [6, 8, 10], [1, 3, 5], [7, 9, 11]])
- def test_flip_axes(self):
- # Test nested iteration with negative axes
- a = arange(12).reshape(2, 3, 2)[::-1, ::-1, ::-1]
- # In 'K' order (default), the axes all get flipped
- i, j = np.nested_iters(a, [[0], [1, 2]])
- vals = [list(j) for _ in i]
- assert_equal(vals, [[0, 1, 2, 3, 4, 5], [6, 7, 8, 9, 10, 11]])
- i, j = np.nested_iters(a, [[0, 1], [2]])
- vals = [list(j) for _ in i]
- assert_equal(vals, [[0, 1], [2, 3], [4, 5], [6, 7], [8, 9], [10, 11]])
- i, j = np.nested_iters(a, [[0, 2], [1]])
- vals = [list(j) for _ in i]
- assert_equal(vals, [[0, 2, 4], [1, 3, 5], [6, 8, 10], [7, 9, 11]])
- # In 'C' order, flipping axes is disabled
- i, j = np.nested_iters(a, [[0], [1, 2]], order='C')
- vals = [list(j) for _ in i]
- assert_equal(vals, [[11, 10, 9, 8, 7, 6], [5, 4, 3, 2, 1, 0]])
- i, j = np.nested_iters(a, [[0, 1], [2]], order='C')
- vals = [list(j) for _ in i]
- assert_equal(vals, [[11, 10], [9, 8], [7, 6], [5, 4], [3, 2], [1, 0]])
- i, j = np.nested_iters(a, [[0, 2], [1]], order='C')
- vals = [list(j) for _ in i]
- assert_equal(vals, [[11, 9, 7], [10, 8, 6], [5, 3, 1], [4, 2, 0]])
- def test_broadcast(self):
- # Test nested iteration with broadcasting
- a = arange(2).reshape(2, 1)
- b = arange(3).reshape(1, 3)
- i, j = np.nested_iters([a, b], [[0], [1]])
- vals = [list(j) for _ in i]
- assert_equal(vals, [[[0, 0], [0, 1], [0, 2]], [[1, 0], [1, 1], [1, 2]]])
- i, j = np.nested_iters([a, b], [[1], [0]])
- vals = [list(j) for _ in i]
- assert_equal(vals, [[[0, 0], [1, 0]], [[0, 1], [1, 1]], [[0, 2], [1, 2]]])
- def test_dtype_copy(self):
- # Test nested iteration with a copy to change dtype
- # copy
- a = arange(6, dtype='i4').reshape(2, 3)
- i, j = np.nested_iters(a, [[0], [1]],
- op_flags=['readonly', 'copy'],
- op_dtypes='f8')
- assert_equal(j[0].dtype, np.dtype('f8'))
- vals = [list(j) for _ in i]
- assert_equal(vals, [[0, 1, 2], [3, 4, 5]])
- vals = None
- # writebackifcopy - using context manager
- a = arange(6, dtype='f4').reshape(2, 3)
- i, j = np.nested_iters(a, [[0], [1]],
- op_flags=['readwrite', 'updateifcopy'],
- casting='same_kind',
- op_dtypes='f8')
- with i, j:
- assert_equal(j[0].dtype, np.dtype('f8'))
- for x in i:
- for y in j:
- y[...] += 1
- assert_equal(a, [[0, 1, 2], [3, 4, 5]])
- assert_equal(a, [[1, 2, 3], [4, 5, 6]])
- # writebackifcopy - using close()
- a = arange(6, dtype='f4').reshape(2, 3)
- i, j = np.nested_iters(a, [[0], [1]],
- op_flags=['readwrite', 'updateifcopy'],
- casting='same_kind',
- op_dtypes='f8')
- assert_equal(j[0].dtype, np.dtype('f8'))
- for x in i:
- for y in j:
- y[...] += 1
- assert_equal(a, [[0, 1, 2], [3, 4, 5]])
- i.close()
- j.close()
- assert_equal(a, [[1, 2, 3], [4, 5, 6]])
- def test_dtype_buffered(self):
- # Test nested iteration with buffering to change dtype
- a = arange(6, dtype='f4').reshape(2, 3)
- i, j = np.nested_iters(a, [[0], [1]],
- flags=['buffered'],
- op_flags=['readwrite'],
- casting='same_kind',
- op_dtypes='f8')
- assert_equal(j[0].dtype, np.dtype('f8'))
- for x in i:
- for y in j:
- y[...] += 1
- assert_equal(a, [[1, 2, 3], [4, 5, 6]])
- def test_0d(self):
- a = np.arange(12).reshape(2, 3, 2)
- i, j = np.nested_iters(a, [[], [1, 0, 2]])
- vals = [list(j) for _ in i]
- assert_equal(vals, [[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]])
- i, j = np.nested_iters(a, [[1, 0, 2], []])
- vals = [list(j) for _ in i]
- assert_equal(vals, [[0], [1], [2], [3], [4], [5], [6], [7], [8], [9], [10], [11]])
- i, j, k = np.nested_iters(a, [[2, 0], [], [1]])
- vals = []
- for x in i:
- for y in j:
- vals.append([z for z in k])
- assert_equal(vals, [[0, 2, 4], [1, 3, 5], [6, 8, 10], [7, 9, 11]])
- def test_iter_nested_iters_dtype_buffered(self):
- # Test nested iteration with buffering to change dtype
- a = arange(6, dtype='f4').reshape(2, 3)
- i, j = np.nested_iters(a, [[0], [1]],
- flags=['buffered'],
- op_flags=['readwrite'],
- casting='same_kind',
- op_dtypes='f8')
- with i, j:
- assert_equal(j[0].dtype, np.dtype('f8'))
- for x in i:
- for y in j:
- y[...] += 1
- assert_equal(a, [[1, 2, 3], [4, 5, 6]])
- def test_iter_reduction_error():
- a = np.arange(6)
- assert_raises(ValueError, nditer, [a, None], [],
- [['readonly'], ['readwrite', 'allocate']],
- op_axes=[[0], [-1]])
- a = np.arange(6).reshape(2, 3)
- assert_raises(ValueError, nditer, [a, None], ['external_loop'],
- [['readonly'], ['readwrite', 'allocate']],
- op_axes=[[0, 1], [-1, -1]])
- def test_iter_reduction():
- # Test doing reductions with the iterator
- a = np.arange(6)
- i = nditer([a, None], ['reduce_ok'],
- [['readonly'], ['readwrite', 'allocate']],
- op_axes=[[0], [-1]])
- # Need to initialize the output operand to the addition unit
- with i:
- i.operands[1][...] = 0
- # Do the reduction
- for x, y in i:
- y[...] += x
- # Since no axes were specified, should have allocated a scalar
- assert_equal(i.operands[1].ndim, 0)
- assert_equal(i.operands[1], np.sum(a))
- a = np.arange(6).reshape(2, 3)
- i = nditer([a, None], ['reduce_ok', 'external_loop'],
- [['readonly'], ['readwrite', 'allocate']],
- op_axes=[[0, 1], [-1, -1]])
- # Need to initialize the output operand to the addition unit
- with i:
- i.operands[1][...] = 0
- # Reduction shape/strides for the output
- assert_equal(i[1].shape, (6,))
- assert_equal(i[1].strides, (0,))
- # Do the reduction
- for x, y in i:
- # Use a for loop instead of ``y[...] += x``
- # (equivalent to ``y[...] = y[...].copy() + x``),
- # because y has zero strides we use for the reduction
- for j in range(len(y)):
- y[j] += x[j]
- # Since no axes were specified, should have allocated a scalar
- assert_equal(i.operands[1].ndim, 0)
- assert_equal(i.operands[1], np.sum(a))
- # This is a tricky reduction case for the buffering double loop
- # to handle
- a = np.ones((2, 3, 5))
- it1 = nditer([a, None], ['reduce_ok', 'external_loop'],
- [['readonly'], ['readwrite', 'allocate']],
- op_axes=[None, [0, -1, 1]])
- it2 = nditer([a, None], ['reduce_ok', 'external_loop',
- 'buffered', 'delay_bufalloc'],
- [['readonly'], ['readwrite', 'allocate']],
- op_axes=[None, [0, -1, 1]], buffersize=10)
- with it1, it2:
- it1.operands[1].fill(0)
- it2.operands[1].fill(0)
- it2.reset()
- for x in it1:
- x[1][...] += x[0]
- for x in it2:
- x[1][...] += x[0]
- assert_equal(it1.operands[1], it2.operands[1])
- assert_equal(it2.operands[1].sum(), a.size)
- def test_iter_buffering_reduction():
- # Test doing buffered reductions with the iterator
- a = np.arange(6)
- b = np.array(0., dtype='f8').byteswap().newbyteorder()
- i = nditer([a, b], ['reduce_ok', 'buffered'],
- [['readonly'], ['readwrite', 'nbo']],
- op_axes=[[0], [-1]])
- with i:
- assert_equal(i[1].dtype, np.dtype('f8'))
- assert_(i[1].dtype != b.dtype)
- # Do the reduction
- for x, y in i:
- y[...] += x
- # Since no axes were specified, should have allocated a scalar
- assert_equal(b, np.sum(a))
- a = np.arange(6).reshape(2, 3)
- b = np.array([0, 0], dtype='f8').byteswap().newbyteorder()
- i = nditer([a, b], ['reduce_ok', 'external_loop', 'buffered'],
- [['readonly'], ['readwrite', 'nbo']],
- op_axes=[[0, 1], [0, -1]])
- # Reduction shape/strides for the output
- with i:
- assert_equal(i[1].shape, (3,))
- assert_equal(i[1].strides, (0,))
- # Do the reduction
- for x, y in i:
- # Use a for loop instead of ``y[...] += x``
- # (equivalent to ``y[...] = y[...].copy() + x``),
- # because y has zero strides we use for the reduction
- for j in range(len(y)):
- y[j] += x[j]
- assert_equal(b, np.sum(a, axis=1))
- # Iterator inner double loop was wrong on this one
- p = np.arange(2) + 1
- it = np.nditer([p, None],
- ['delay_bufalloc', 'reduce_ok', 'buffered', 'external_loop'],
- [['readonly'], ['readwrite', 'allocate']],
- op_axes=[[-1, 0], [-1, -1]],
- itershape=(2, 2))
- with it:
- it.operands[1].fill(0)
- it.reset()
- assert_equal(it[0], [1, 2, 1, 2])
- # Iterator inner loop should take argument contiguity into account
- x = np.ones((7, 13, 8), np.int8)[4:6,1:11:6,1:5].transpose(1, 2, 0)
- x[...] = np.arange(x.size).reshape(x.shape)
- y_base = np.arange(4*4, dtype=np.int8).reshape(4, 4)
- y_base_copy = y_base.copy()
- y = y_base[::2,:,None]
- it = np.nditer([y, x],
- ['buffered', 'external_loop', 'reduce_ok'],
- [['readwrite'], ['readonly']])
- with it:
- for a, b in it:
- a.fill(2)
- assert_equal(y_base[1::2], y_base_copy[1::2])
- assert_equal(y_base[::2], 2)
- def test_iter_buffering_reduction_reuse_reduce_loops():
- # There was a bug triggering reuse of the reduce loop inappropriately,
- # which caused processing to happen in unnecessarily small chunks
- # and overran the buffer.
- a = np.zeros((2, 7))
- b = np.zeros((1, 7))
- it = np.nditer([a, b], flags=['reduce_ok', 'external_loop', 'buffered'],
- op_flags=[['readonly'], ['readwrite']],
- buffersize=5)
- with it:
- bufsizes = [x.shape[0] for x, y in it]
- assert_equal(bufsizes, [5, 2, 5, 2])
- assert_equal(sum(bufsizes), a.size)
- def test_iter_writemasked_badinput():
- a = np.zeros((2, 3))
- b = np.zeros((3,))
- m = np.array([[True, True, False], [False, True, False]])
- m2 = np.array([True, True, False])
- m3 = np.array([0, 1, 1], dtype='u1')
- mbad1 = np.array([0, 1, 1], dtype='i1')
- mbad2 = np.array([0, 1, 1], dtype='f4')
- # Need an 'arraymask' if any operand is 'writemasked'
- assert_raises(ValueError, nditer, [a, m], [],
- [['readwrite', 'writemasked'], ['readonly']])
- # A 'writemasked' operand must not be readonly
- assert_raises(ValueError, nditer, [a, m], [],
- [['readonly', 'writemasked'], ['readonly', 'arraymask']])
- # 'writemasked' and 'arraymask' may not be used together
- assert_raises(ValueError, nditer, [a, m], [],
- [['readonly'], ['readwrite', 'arraymask', 'writemasked']])
- # 'arraymask' may only be specified once
- assert_raises(ValueError, nditer, [a, m, m2], [],
- [['readwrite', 'writemasked'],
- ['readonly', 'arraymask'],
- ['readonly', 'arraymask']])
- # An 'arraymask' with nothing 'writemasked' also doesn't make sense
- assert_raises(ValueError, nditer, [a, m], [],
- [['readwrite'], ['readonly', 'arraymask']])
- # A writemasked reduction requires a similarly smaller mask
- assert_raises(ValueError, nditer, [a, b, m], ['reduce_ok'],
- [['readonly'],
- ['readwrite', 'writemasked'],
- ['readonly', 'arraymask']])
- # But this should work with a smaller/equal mask to the reduction operand
- np.nditer([a, b, m2], ['reduce_ok'],
- [['readonly'],
- ['readwrite', 'writemasked'],
- ['readonly', 'arraymask']])
- # The arraymask itself cannot be a reduction
- assert_raises(ValueError, nditer, [a, b, m2], ['reduce_ok'],
- [['readonly'],
- ['readwrite', 'writemasked'],
- ['readwrite', 'arraymask']])
- # A uint8 mask is ok too
- np.nditer([a, m3], ['buffered'],
- [['readwrite', 'writemasked'],
- ['readonly', 'arraymask']],
- op_dtypes=['f4', None],
- casting='same_kind')
- # An int8 mask isn't ok
- assert_raises(TypeError, np.nditer, [a, mbad1], ['buffered'],
- [['readwrite', 'writemasked'],
- ['readonly', 'arraymask']],
- op_dtypes=['f4', None],
- casting='same_kind')
- # A float32 mask isn't ok
- assert_raises(TypeError, np.nditer, [a, mbad2], ['buffered'],
- [['readwrite', 'writemasked'],
- ['readonly', 'arraymask']],
- op_dtypes=['f4', None],
- casting='same_kind')
- def test_iter_writemasked():
- a = np.zeros((3,), dtype='f8')
- msk = np.array([True, True, False])
- # When buffering is unused, 'writemasked' effectively does nothing.
- # It's up to the user of the iterator to obey the requested semantics.
- it = np.nditer([a, msk], [],
- [['readwrite', 'writemasked'],
- ['readonly', 'arraymask']])
- with it:
- for x, m in it:
- x[...] = 1
- # Because we violated the semantics, all the values became 1
- assert_equal(a, [1, 1, 1])
- # Even if buffering is enabled, we still may be accessing the array
- # directly.
- it = np.nditer([a, msk], ['buffered'],
- [['readwrite', 'writemasked'],
- ['readonly', 'arraymask']])
- with it:
- for x, m in it:
- x[...] = 2.5
- # Because we violated the semantics, all the values became 2.5
- assert_equal(a, [2.5, 2.5, 2.5])
- # If buffering will definitely happening, for instance because of
- # a cast, only the items selected by the mask will be copied back from
- # the buffer.
- it = np.nditer([a, msk], ['buffered'],
- [['readwrite', 'writemasked'],
- ['readonly', 'arraymask']],
- op_dtypes=['i8', None],
- casting='unsafe')
- with it:
- for x, m in it:
- x[...] = 3
- # Even though we violated the semantics, only the selected values
- # were copied back
- assert_equal(a, [3, 3, 2.5])
- def test_iter_non_writable_attribute_deletion():
- it = np.nditer(np.ones(2))
- attr = ["value", "shape", "operands", "itviews", "has_delayed_bufalloc",
- "iterationneedsapi", "has_multi_index", "has_index", "dtypes",
- "ndim", "nop", "itersize", "finished"]
- for s in attr:
- assert_raises(AttributeError, delattr, it, s)
- def test_iter_writable_attribute_deletion():
- it = np.nditer(np.ones(2))
- attr = [ "multi_index", "index", "iterrange", "iterindex"]
- for s in attr:
- assert_raises(AttributeError, delattr, it, s)
- def test_iter_element_deletion():
- it = np.nditer(np.ones(3))
- try:
- del it[1]
- del it[1:2]
- except TypeError:
- pass
- except Exception:
- raise AssertionError
- def test_iter_allocated_array_dtypes():
- # If the dtype of an allocated output has a shape, the shape gets
- # tacked onto the end of the result.
- it = np.nditer(([1, 3, 20], None), op_dtypes=[None, ('i4', (2,))])
- for a, b in it:
- b[0] = a - 1
- b[1] = a + 1
- assert_equal(it.operands[1], [[0, 2], [2, 4], [19, 21]])
- # Check the same (less sensitive) thing when `op_axes` with -1 is given.
- it = np.nditer(([[1, 3, 20]], None), op_dtypes=[None, ('i4', (2,))],
- flags=["reduce_ok"], op_axes=[None, (-1, 0)])
- for a, b in it:
- b[0] = a - 1
- b[1] = a + 1
- assert_equal(it.operands[1], [[0, 2], [2, 4], [19, 21]])
- # Make sure this works for scalars too
- it = np.nditer((10, 2, None), op_dtypes=[None, None, ('i4', (2, 2))])
- for a, b, c in it:
- c[0, 0] = a - b
- c[0, 1] = a + b
- c[1, 0] = a * b
- c[1, 1] = a / b
- assert_equal(it.operands[2], [[8, 12], [20, 5]])
- def test_0d_iter():
- # Basic test for iteration of 0-d arrays:
- i = nditer([2, 3], ['multi_index'], [['readonly']]*2)
- assert_equal(i.ndim, 0)
- assert_equal(next(i), (2, 3))
- assert_equal(i.multi_index, ())
- assert_equal(i.iterindex, 0)
- assert_raises(StopIteration, next, i)
- # test reset:
- i.reset()
- assert_equal(next(i), (2, 3))
- assert_raises(StopIteration, next, i)
- # test forcing to 0-d
- i = nditer(np.arange(5), ['multi_index'], [['readonly']], op_axes=[()])
- assert_equal(i.ndim, 0)
- assert_equal(len(i), 1)
- i = nditer(np.arange(5), ['multi_index'], [['readonly']],
- op_axes=[()], itershape=())
- assert_equal(i.ndim, 0)
- assert_equal(len(i), 1)
- # passing an itershape alone is not enough, the op_axes are also needed
- with assert_raises(ValueError):
- nditer(np.arange(5), ['multi_index'], [['readonly']], itershape=())
- # Test a more complex buffered casting case (same as another test above)
- sdt = [('a', 'f4'), ('b', 'i8'), ('c', 'c8', (2, 3)), ('d', 'O')]
- a = np.array(0.5, dtype='f4')
- i = nditer(a, ['buffered', 'refs_ok'], ['readonly'],
- casting='unsafe', op_dtypes=sdt)
- vals = next(i)
- assert_equal(vals['a'], 0.5)
- assert_equal(vals['b'], 0)
- assert_equal(vals['c'], [[(0.5)]*3]*2)
- assert_equal(vals['d'], 0.5)
- def test_object_iter_cleanup():
- # see gh-18450
- # object arrays can raise a python exception in ufunc inner loops using
- # nditer, which should cause iteration to stop & cleanup. There were bugs
- # in the nditer cleanup when decref'ing object arrays.
- # This test would trigger valgrind "uninitialized read" before the bugfix.
- assert_raises(TypeError, lambda: np.zeros((17000, 2), dtype='f4') * None)
- # this more explicit code also triggers the invalid access
- arr = np.arange(np.BUFSIZE * 10).reshape(10, -1).astype(str)
- oarr = arr.astype(object)
- oarr[:, -1] = None
- assert_raises(TypeError, lambda: np.add(oarr[:, ::-1], arr[:, ::-1]))
- # followup: this tests for a bug introduced in the first pass of gh-18450,
- # caused by an incorrect fallthrough of the TypeError
- class T:
- def __bool__(self):
- raise TypeError("Ambiguous")
- assert_raises(TypeError, np.logical_or.reduce,
- np.array([T(), T()], dtype='O'))
- def test_object_iter_cleanup_reduce():
- # Similar as above, but a complex reduction case that was previously
- # missed (see gh-18810).
- # The following array is special in that it cannot be flattened:
- arr = np.array([[None, 1], [-1, -1], [None, 2], [-1, -1]])[::2]
- with pytest.raises(TypeError):
- np.sum(arr)
- @pytest.mark.parametrize("arr", [
- np.ones((8000, 4, 2), dtype=object)[:, ::2, :],
- np.ones((8000, 4, 2), dtype=object, order="F")[:, ::2, :],
- np.ones((8000, 4, 2), dtype=object)[:, ::2, :].copy("F")])
- def test_object_iter_cleanup_large_reduce(arr):
- # More complicated calls are possible for large arrays:
- out = np.ones(8000, dtype=np.intp)
- # force casting with `dtype=object`
- res = np.sum(arr, axis=(1, 2), dtype=object, out=out)
- assert_array_equal(res, np.full(8000, 4, dtype=object))
- def test_iter_too_large():
- # The total size of the iterator must not exceed the maximum intp due
- # to broadcasting. Dividing by 1024 will keep it small enough to
- # give a legal array.
- size = np.iinfo(np.intp).max // 1024
- arr = np.lib.stride_tricks.as_strided(np.zeros(1), (size,), (0,))
- assert_raises(ValueError, nditer, (arr, arr[:, None]))
- # test the same for multiindex. That may get more interesting when
- # removing 0 dimensional axis is allowed (since an iterator can grow then)
- assert_raises(ValueError, nditer,
- (arr, arr[:, None]), flags=['multi_index'])
- def test_iter_too_large_with_multiindex():
- # When a multi index is being tracked, the error is delayed this
- # checks the delayed error messages and getting below that by
- # removing an axis.
- base_size = 2**10
- num = 1
- while base_size**num < np.iinfo(np.intp).max:
- num += 1
- shape_template = [1, 1] * num
- arrays = []
- for i in range(num):
- shape = shape_template[:]
- shape[i * 2] = 2**10
- arrays.append(np.empty(shape))
- arrays = tuple(arrays)
- # arrays are now too large to be broadcast. The different modes test
- # different nditer functionality with or without GIL.
- for mode in range(6):
- with assert_raises(ValueError):
- _multiarray_tests.test_nditer_too_large(arrays, -1, mode)
- # but if we do nothing with the nditer, it can be constructed:
- _multiarray_tests.test_nditer_too_large(arrays, -1, 7)
- # When an axis is removed, things should work again (half the time):
- for i in range(num):
- for mode in range(6):
- # an axis with size 1024 is removed:
- _multiarray_tests.test_nditer_too_large(arrays, i*2, mode)
- # an axis with size 1 is removed:
- with assert_raises(ValueError):
- _multiarray_tests.test_nditer_too_large(arrays, i*2 + 1, mode)
- def test_writebacks():
- a = np.arange(6, dtype='f4')
- au = a.byteswap().newbyteorder()
- assert_(a.dtype.byteorder != au.dtype.byteorder)
- it = nditer(au, [], [['readwrite', 'updateifcopy']],
- casting='equiv', op_dtypes=[np.dtype('f4')])
- with it:
- it.operands[0][:] = 100
- assert_equal(au, 100)
- # do it again, this time raise an error,
- it = nditer(au, [], [['readwrite', 'updateifcopy']],
- casting='equiv', op_dtypes=[np.dtype('f4')])
- try:
- with it:
- assert_equal(au.flags.writeable, False)
- it.operands[0][:] = 0
- raise ValueError('exit context manager on exception')
- except:
- pass
- assert_equal(au, 0)
- assert_equal(au.flags.writeable, True)
- # cannot reuse i outside context manager
- assert_raises(ValueError, getattr, it, 'operands')
- it = nditer(au, [], [['readwrite', 'updateifcopy']],
- casting='equiv', op_dtypes=[np.dtype('f4')])
- with it:
- x = it.operands[0]
- x[:] = 6
- assert_(x.flags.writebackifcopy)
- assert_equal(au, 6)
- assert_(not x.flags.writebackifcopy)
- x[:] = 123 # x.data still valid
- assert_equal(au, 6) # but not connected to au
- it = nditer(au, [],
- [['readwrite', 'updateifcopy']],
- casting='equiv', op_dtypes=[np.dtype('f4')])
- # reentering works
- with it:
- with it:
- for x in it:
- x[...] = 123
- it = nditer(au, [],
- [['readwrite', 'updateifcopy']],
- casting='equiv', op_dtypes=[np.dtype('f4')])
- # make sure exiting the inner context manager closes the iterator
- with it:
- with it:
- for x in it:
- x[...] = 123
- assert_raises(ValueError, getattr, it, 'operands')
- # do not crash if original data array is decrefed
- it = nditer(au, [],
- [['readwrite', 'updateifcopy']],
- casting='equiv', op_dtypes=[np.dtype('f4')])
- del au
- with it:
- for x in it:
- x[...] = 123
- # make sure we cannot reenter the closed iterator
- enter = it.__enter__
- assert_raises(RuntimeError, enter)
- def test_close_equivalent():
- ''' using a context amanger and using nditer.close are equivalent
- '''
- def add_close(x, y, out=None):
- addop = np.add
- it = np.nditer([x, y, out], [],
- [['readonly'], ['readonly'], ['writeonly','allocate']])
- for (a, b, c) in it:
- addop(a, b, out=c)
- ret = it.operands[2]
- it.close()
- return ret
- def add_context(x, y, out=None):
- addop = np.add
- it = np.nditer([x, y, out], [],
- [['readonly'], ['readonly'], ['writeonly','allocate']])
- with it:
- for (a, b, c) in it:
- addop(a, b, out=c)
- return it.operands[2]
- z = add_close(range(5), range(5))
- assert_equal(z, range(0, 10, 2))
- z = add_context(range(5), range(5))
- assert_equal(z, range(0, 10, 2))
- def test_close_raises():
- it = np.nditer(np.arange(3))
- assert_equal (next(it), 0)
- it.close()
- assert_raises(StopIteration, next, it)
- assert_raises(ValueError, getattr, it, 'operands')
- @pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts")
- def test_warn_noclose():
- a = np.arange(6, dtype='f4')
- au = a.byteswap().newbyteorder()
- with suppress_warnings() as sup:
- sup.record(RuntimeWarning)
- it = np.nditer(au, [], [['readwrite', 'updateifcopy']],
- casting='equiv', op_dtypes=[np.dtype('f4')])
- del it
- assert len(sup.log) == 1
- @pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts")
- @pytest.mark.parametrize(["in_dtype", "buf_dtype"],
- [("i", "O"), ("O", "i"), # most simple cases
- ("i,O", "O,O"), # structured partially only copying O
- ("O,i", "i,O"), # structured casting to and from O
- ])
- @pytest.mark.parametrize("steps", [1, 2, 3])
- def test_partial_iteration_cleanup(in_dtype, buf_dtype, steps):
- value = 123 # relies on python cache (leak-check will still find it)
- arr = np.full(int(np.BUFSIZE * 2.5), value).astype(in_dtype)
- count = sys.getrefcount(value)
- it = np.nditer(arr, op_dtypes=[np.dtype(buf_dtype)],
- flags=["buffered", "external_loop", "refs_ok"], casting="unsafe")
- for step in range(steps):
- # The iteration finishes in 3 steps, the first two are partial
- next(it)
- # Note that resetting does not free references
- del it
- assert count == sys.getrefcount(value)
- # Repeat the test with `iternext`
- it = np.nditer(arr, op_dtypes=[np.dtype(buf_dtype)],
- flags=["buffered", "external_loop", "refs_ok"], casting="unsafe")
- for step in range(steps):
- it.iternext()
- del it # should ensure cleanup
- assert count == sys.getrefcount(value)
- @pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts")
- @pytest.mark.parametrize(["in_dtype", "buf_dtype"],
- [("O", "i"), # most simple cases
- ("O,i", "i,O"), # structured casting to and from O
- ])
- def test_partial_iteration_error(in_dtype, buf_dtype):
- value = 123 # relies on python cache (leak-check will still find it)
- arr = np.full(int(np.BUFSIZE * 2.5), value).astype(in_dtype)
- if in_dtype == "O":
- arr[int(np.BUFSIZE * 1.5)] = None
- else:
- arr[int(np.BUFSIZE * 1.5)]["f0"] = None
- count = sys.getrefcount(value)
- it = np.nditer(arr, op_dtypes=[np.dtype(buf_dtype)],
- flags=["buffered", "external_loop", "refs_ok"], casting="unsafe")
- with pytest.raises(TypeError):
- # pytest.raises seems to have issues with the error originating
- # in the for loop, so manually unravel:
- next(it)
- next(it) # raises TypeError
- # Repeat the test with `iternext` after resetting, the buffers should
- # already be cleared from any references, so resetting is sufficient.
- it.reset()
- with pytest.raises(TypeError):
- it.iternext()
- it.iternext()
- assert count == sys.getrefcount(value)
|