codekingpro/portable-devtools
114k
1import inspect
2import subprocess
3import sys
4import textwrap
5import warnings
6
7import pytest
8
9import numpy as np
10import numpy._core._multiarray_tests as _multiarray_tests
11import numpy._core.umath as ncu
12from numpy import all, arange, array, nditer
13from numpy.testing import (
14 HAS_REFCOUNT,
15 IS_64BIT,
16 IS_PYPY,
17 IS_WASM,
18 assert_,
19 assert_array_equal,
20 assert_equal,
21 assert_raises,
22)
23from numpy.testing._private.utils import requires_memory
24
25
26def iter_multi_index(i):
27 ret = []
28 while not i.finished:
29 ret.append(i.multi_index)
30 i.iternext()
31 return ret
32
33def iter_indices(i):
34 ret = []
35 while not i.finished:
36 ret.append(i.index)
37 i.iternext()
38 return ret
39
40def iter_iterindices(i):
41 ret = []
42 while not i.finished:
43 ret.append(i.iterindex)
44 i.iternext()
45 return ret
46
47@pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts")
48def test_iter_refcount():
49 # Make sure the iterator doesn't leak
50
51 # Basic
52 a = arange(6)
53 dt = np.dtype('f4').newbyteorder()
54 rc_a = sys.getrefcount(a)
55 rc_dt = sys.getrefcount(dt)
56 with nditer(a, [],
57 [['readwrite', 'updateifcopy']],
58 casting='unsafe',
59 op_dtypes=[dt]) as it:
60 assert_(not it.iterationneedsapi)
61 assert_(sys.getrefcount(a) > rc_a)
62 assert_(sys.getrefcount(dt) > rc_dt)
63 # del 'it'
64 it = None
65 assert_equal(sys.getrefcount(a), rc_a)
66 assert_equal(sys.getrefcount(dt), rc_dt)
67
68 # With a copy
69 a = arange(6, dtype='f4')
70 dt = np.dtype('f4')
71 rc_a = sys.getrefcount(a)
72 rc_dt = sys.getrefcount(dt)
73 it = nditer(a, [],
74 [['readwrite']],
75 op_dtypes=[dt])
76 rc2_a = sys.getrefcount(a)
77 rc2_dt = sys.getrefcount(dt)
78 it2 = it.copy()
79 assert_(sys.getrefcount(a) > rc2_a)
80 if sys.version_info < (3, 13):
81 # np.dtype('f4') is immortal after Python 3.13
82 assert_(sys.getrefcount(dt) > rc2_dt)
83 it = None
84 assert_equal(sys.getrefcount(a), rc2_a)
85 assert_equal(sys.getrefcount(dt), rc2_dt)
86 it2 = None
87 assert_equal(sys.getrefcount(a), rc_a)
88 assert_equal(sys.getrefcount(dt), rc_dt)
89
90def test_iter_best_order():
91 # The iterator should always find the iteration order
92 # with increasing memory addresses
93
94 # Test the ordering for 1-D to 5-D shapes
95 for shape in [(5,), (3, 4), (2, 3, 4), (2, 3, 4, 3), (2, 3, 2, 2, 3)]:
96 a = arange(np.prod(shape))
97 # Test each combination of positive and negative strides
98 for dirs in range(2**len(shape)):
99 dirs_index = [slice(None)] * len(shape)
100 for bit in range(len(shape)):
101 if ((2**bit) & dirs):
102 dirs_index[bit] = slice(None, None, -1)
103 dirs_index = tuple(dirs_index)
104
105 aview = a.reshape(shape)[dirs_index]
106 # C-order
107 i = nditer(aview, [], [['readonly']])
108 assert_equal(list(i), a)
109 # Fortran-order
110 i = nditer(aview.T, [], [['readonly']])
111 assert_equal(list(i), a)
112 # Other order
113 if len(shape) > 2:
114 i = nditer(aview.swapaxes(0, 1), [], [['readonly']])
115 assert_equal(list(i), a)
116
117def test_iter_c_order():
118 # Test forcing C order
119
120 # Test the ordering for 1-D to 5-D shapes
121 for shape in [(5,), (3, 4), (2, 3, 4), (2, 3, 4, 3), (2, 3, 2, 2, 3)]:
122 a = arange(np.prod(shape))
123 # Test each combination of positive and negative strides
124 for dirs in range(2**len(shape)):
125 dirs_index = [slice(None)] * len(shape)
126 for bit in range(len(shape)):
127 if ((2**bit) & dirs):
128 dirs_index[bit] = slice(None, None, -1)
129 dirs_index = tuple(dirs_index)
130
131 aview = a.reshape(shape)[dirs_index]
132 # C-order
133 i = nditer(aview, order='C')
134 assert_equal(list(i), aview.ravel(order='C'))
135 # Fortran-order
136 i = nditer(aview.T, order='C')
137 assert_equal(list(i), aview.T.ravel(order='C'))
138 # Other order
139 if len(shape) > 2:
140 i = nditer(aview.swapaxes(0, 1), order='C')
141 assert_equal(list(i),
142 aview.swapaxes(0, 1).ravel(order='C'))
143
144def test_iter_f_order():
145 # Test forcing F order
146
147 # Test the ordering for 1-D to 5-D shapes
148 for shape in [(5,), (3, 4), (2, 3, 4), (2, 3, 4, 3), (2, 3, 2, 2, 3)]:
149 a = arange(np.prod(shape))
150 # Test each combination of positive and negative strides
151 for dirs in range(2**len(shape)):
152 dirs_index = [slice(None)] * len(shape)
153 for bit in range(len(shape)):
154 if ((2**bit) & dirs):
155 dirs_index[bit] = slice(None, None, -1)
156 dirs_index = tuple(dirs_index)
157
158 aview = a.reshape(shape)[dirs_index]
159 # C-order
160 i = nditer(aview, order='F')
161 assert_equal(list(i), aview.ravel(order='F'))
162 # Fortran-order
163 i = nditer(aview.T, order='F')
164 assert_equal(list(i), aview.T.ravel(order='F'))
165 # Other order
166 if len(shape) > 2:
167 i = nditer(aview.swapaxes(0, 1), order='F')
168 assert_equal(list(i),
169 aview.swapaxes(0, 1).ravel(order='F'))
170
171def test_iter_c_or_f_order():
172 # Test forcing any contiguous (C or F) order
173
174 # Test the ordering for 1-D to 5-D shapes
175 for shape in [(5,), (3, 4), (2, 3, 4), (2, 3, 4, 3), (2, 3, 2, 2, 3)]:
176 a = arange(np.prod(shape))
177 # Test each combination of positive and negative strides
178 for dirs in range(2**len(shape)):
179 dirs_index = [slice(None)] * len(shape)
180 for bit in range(len(shape)):
181 if ((2**bit) & dirs):
182 dirs_index[bit] = slice(None, None, -1)
183 dirs_index = tuple(dirs_index)
184
185 aview = a.reshape(shape)[dirs_index]
186 # C-order
187 i = nditer(aview, order='A')
188 assert_equal(list(i), aview.ravel(order='A'))
189 # Fortran-order
190 i = nditer(aview.T, order='A')
191 assert_equal(list(i), aview.T.ravel(order='A'))
192 # Other order
193 if len(shape) > 2:
194 i = nditer(aview.swapaxes(0, 1), order='A')
195 assert_equal(list(i),
196 aview.swapaxes(0, 1).ravel(order='A'))
197
198def test_nditer_multi_index_set():
199 # Test the multi_index set
200 a = np.arange(6).reshape(2, 3)
201 it = np.nditer(a, flags=['multi_index'])
202
203 # Removes the iteration on two first elements of a[0]
204 it.multi_index = (0, 2,)
205
206 assert_equal(list(it), [2, 3, 4, 5])
207
208@pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts")
209def test_nditer_multi_index_set_refcount():
210 # Test if the reference count on index variable is decreased
211
212 index = 0
213 i = np.nditer(np.array([111, 222, 333, 444]), flags=['multi_index'])
214
215 start_count = sys.getrefcount(index)
216 i.multi_index = (index,)
217 end_count = sys.getrefcount(index)
218
219 assert_equal(start_count, end_count)
220
221def test_iter_best_order_multi_index_1d():
222 # The multi-indices should be correct with any reordering
223
224 a = arange(4)
225 # 1D order
226 i = nditer(a, ['multi_index'], [['readonly']])
227 assert_equal(iter_multi_index(i), [(0,), (1,), (2,), (3,)])
228 # 1D reversed order
229 i = nditer(a[::-1], ['multi_index'], [['readonly']])
230 assert_equal(iter_multi_index(i), [(3,), (2,), (1,), (0,)])
231
232def test_iter_best_order_multi_index_2d():
233 # The multi-indices should be correct with any reordering
234
235 a = arange(6)
236 # 2D C-order
237 i = nditer(a.reshape(2, 3), ['multi_index'], [['readonly']])
238 assert_equal(iter_multi_index(i), [(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2)])
239 # 2D Fortran-order
240 i = nditer(a.reshape(2, 3).copy(order='F'), ['multi_index'], [['readonly']])
241 assert_equal(iter_multi_index(i), [(0, 0), (1, 0), (0, 1), (1, 1), (0, 2), (1, 2)])
242 # 2D reversed C-order
243 i = nditer(a.reshape(2, 3)[::-1], ['multi_index'], [['readonly']])
244 assert_equal(iter_multi_index(i), [(1, 0), (1, 1), (1, 2), (0, 0), (0, 1), (0, 2)])
245 i = nditer(a.reshape(2, 3)[:, ::-1], ['multi_index'], [['readonly']])
246 assert_equal(iter_multi_index(i), [(0, 2), (0, 1), (0, 0), (1, 2), (1, 1), (1, 0)])
247 i = nditer(a.reshape(2, 3)[::-1, ::-1], ['multi_index'], [['readonly']])
248 assert_equal(iter_multi_index(i), [(1, 2), (1, 1), (1, 0), (0, 2), (0, 1), (0, 0)])
249 # 2D reversed Fortran-order
250 i = nditer(a.reshape(2, 3).copy(order='F')[::-1], ['multi_index'], [['readonly']])
251 assert_equal(iter_multi_index(i), [(1, 0), (0, 0), (1, 1), (0, 1), (1, 2), (0, 2)])
252 i = nditer(a.reshape(2, 3).copy(order='F')[:, ::-1],
253 ['multi_index'], [['readonly']])
254 assert_equal(iter_multi_index(i), [(0, 2), (1, 2), (0, 1), (1, 1), (0, 0), (1, 0)])
255 i = nditer(a.reshape(2, 3).copy(order='F')[::-1, ::-1],
256 ['multi_index'], [['readonly']])
257 assert_equal(iter_multi_index(i), [(1, 2), (0, 2), (1, 1), (0, 1), (1, 0), (0, 0)])
258
259def test_iter_best_order_multi_index_3d():
260 # The multi-indices should be correct with any reordering
261
262 a = arange(12)
263 # 3D C-order
264 i = nditer(a.reshape(2, 3, 2), ['multi_index'], [['readonly']])
265 assert_equal(iter_multi_index(i),
266 [(0, 0, 0), (0, 0, 1), (0, 1, 0), (0, 1, 1), (0, 2, 0), (0, 2, 1),
267 (1, 0, 0), (1, 0, 1), (1, 1, 0), (1, 1, 1), (1, 2, 0), (1, 2, 1)])
268 # 3D Fortran-order
269 i = nditer(a.reshape(2, 3, 2).copy(order='F'), ['multi_index'], [['readonly']])
270 assert_equal(iter_multi_index(i),
271 [(0, 0, 0), (1, 0, 0), (0, 1, 0), (1, 1, 0), (0, 2, 0), (1, 2, 0),
272 (0, 0, 1), (1, 0, 1), (0, 1, 1), (1, 1, 1), (0, 2, 1), (1, 2, 1)])
273 # 3D reversed C-order
274 i = nditer(a.reshape(2, 3, 2)[::-1], ['multi_index'], [['readonly']])
275 assert_equal(iter_multi_index(i),
276 [(1, 0, 0), (1, 0, 1), (1, 1, 0), (1, 1, 1), (1, 2, 0), (1, 2, 1),
277 (0, 0, 0), (0, 0, 1), (0, 1, 0), (0, 1, 1), (0, 2, 0), (0, 2, 1)])
278 i = nditer(a.reshape(2, 3, 2)[:, ::-1], ['multi_index'], [['readonly']])
279 assert_equal(iter_multi_index(i),
280 [(0, 2, 0), (0, 2, 1), (0, 1, 0), (0, 1, 1), (0, 0, 0), (0, 0, 1),
281 (1, 2, 0), (1, 2, 1), (1, 1, 0), (1, 1, 1), (1, 0, 0), (1, 0, 1)])
282 i = nditer(a.reshape(2, 3, 2)[:, :, ::-1], ['multi_index'], [['readonly']])
283 assert_equal(iter_multi_index(i),
284 [(0, 0, 1), (0, 0, 0), (0, 1, 1), (0, 1, 0), (0, 2, 1), (0, 2, 0),
285 (1, 0, 1), (1, 0, 0), (1, 1, 1), (1, 1, 0), (1, 2, 1), (1, 2, 0)])
286 # 3D reversed Fortran-order
287 i = nditer(a.reshape(2, 3, 2).copy(order='F')[::-1],
288 ['multi_index'], [['readonly']])
289 assert_equal(iter_multi_index(i),
290 [(1, 0, 0), (0, 0, 0), (1, 1, 0), (0, 1, 0), (1, 2, 0), (0, 2, 0),
291 (1, 0, 1), (0, 0, 1), (1, 1, 1), (0, 1, 1), (1, 2, 1), (0, 2, 1)])
292 i = nditer(a.reshape(2, 3, 2).copy(order='F')[:, ::-1],
293 ['multi_index'], [['readonly']])
294 assert_equal(iter_multi_index(i),
295 [(0, 2, 0), (1, 2, 0), (0, 1, 0), (1, 1, 0), (0, 0, 0), (1, 0, 0),
296 (0, 2, 1), (1, 2, 1), (0, 1, 1), (1, 1, 1), (0, 0, 1), (1, 0, 1)])
297 i = nditer(a.reshape(2, 3, 2).copy(order='F')[:, :, ::-1],
298 ['multi_index'], [['readonly']])
299 assert_equal(iter_multi_index(i),
300 [(0, 0, 1), (1, 0, 1), (0, 1, 1), (1, 1, 1), (0, 2, 1), (1, 2, 1),
301 (0, 0, 0), (1, 0, 0), (0, 1, 0), (1, 1, 0), (0, 2, 0), (1, 2, 0)])
302
303def test_iter_best_order_c_index_1d():
304 # The C index should be correct with any reordering
305
306 a = arange(4)
307 # 1D order
308 i = nditer(a, ['c_index'], [['readonly']])
309 assert_equal(iter_indices(i), [0, 1, 2, 3])
310 # 1D reversed order
311 i = nditer(a[::-1], ['c_index'], [['readonly']])
312 assert_equal(iter_indices(i), [3, 2, 1, 0])
313
314def test_iter_best_order_c_index_2d():
315 # The C index should be correct with any reordering
316
317 a = arange(6)
318 # 2D C-order
319 i = nditer(a.reshape(2, 3), ['c_index'], [['readonly']])
320 assert_equal(iter_indices(i), [0, 1, 2, 3, 4, 5])
321 # 2D Fortran-order
322 i = nditer(a.reshape(2, 3).copy(order='F'),
323 ['c_index'], [['readonly']])
324 assert_equal(iter_indices(i), [0, 3, 1, 4, 2, 5])
325 # 2D reversed C-order
326 i = nditer(a.reshape(2, 3)[::-1], ['c_index'], [['readonly']])
327 assert_equal(iter_indices(i), [3, 4, 5, 0, 1, 2])
328 i = nditer(a.reshape(2, 3)[:, ::-1], ['c_index'], [['readonly']])
329 assert_equal(iter_indices(i), [2, 1, 0, 5, 4, 3])
330 i = nditer(a.reshape(2, 3)[::-1, ::-1], ['c_index'], [['readonly']])
331 assert_equal(iter_indices(i), [5, 4, 3, 2, 1, 0])
332 # 2D reversed Fortran-order
333 i = nditer(a.reshape(2, 3).copy(order='F')[::-1],
334 ['c_index'], [['readonly']])
335 assert_equal(iter_indices(i), [3, 0, 4, 1, 5, 2])
336 i = nditer(a.reshape(2, 3).copy(order='F')[:, ::-1],
337 ['c_index'], [['readonly']])
338 assert_equal(iter_indices(i), [2, 5, 1, 4, 0, 3])
339 i = nditer(a.reshape(2, 3).copy(order='F')[::-1, ::-1],
340 ['c_index'], [['readonly']])
341 assert_equal(iter_indices(i), [5, 2, 4, 1, 3, 0])
342
343def test_iter_best_order_c_index_3d():
344 # The C index should be correct with any reordering
345
346 a = arange(12)
347 # 3D C-order
348 i = nditer(a.reshape(2, 3, 2), ['c_index'], [['readonly']])
349 assert_equal(iter_indices(i),
350 [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11])
351 # 3D Fortran-order
352 i = nditer(a.reshape(2, 3, 2).copy(order='F'),
353 ['c_index'], [['readonly']])
354 assert_equal(iter_indices(i),
355 [0, 6, 2, 8, 4, 10, 1, 7, 3, 9, 5, 11])
356 # 3D reversed C-order
357 i = nditer(a.reshape(2, 3, 2)[::-1], ['c_index'], [['readonly']])
358 assert_equal(iter_indices(i),
359 [6, 7, 8, 9, 10, 11, 0, 1, 2, 3, 4, 5])
360 i = nditer(a.reshape(2, 3, 2)[:, ::-1], ['c_index'], [['readonly']])
361 assert_equal(iter_indices(i),
362 [4, 5, 2, 3, 0, 1, 10, 11, 8, 9, 6, 7])
363 i = nditer(a.reshape(2, 3, 2)[:, :, ::-1], ['c_index'], [['readonly']])
364 assert_equal(iter_indices(i),
365 [1, 0, 3, 2, 5, 4, 7, 6, 9, 8, 11, 10])
366 # 3D reversed Fortran-order
367 i = nditer(a.reshape(2, 3, 2).copy(order='F')[::-1],
368 ['c_index'], [['readonly']])
369 assert_equal(iter_indices(i),
370 [6, 0, 8, 2, 10, 4, 7, 1, 9, 3, 11, 5])
371 i = nditer(a.reshape(2, 3, 2).copy(order='F')[:, ::-1],
372 ['c_index'], [['readonly']])
373 assert_equal(iter_indices(i),
374 [4, 10, 2, 8, 0, 6, 5, 11, 3, 9, 1, 7])
375 i = nditer(a.reshape(2, 3, 2).copy(order='F')[:, :, ::-1],
376 ['c_index'], [['readonly']])
377 assert_equal(iter_indices(i),
378 [1, 7, 3, 9, 5, 11, 0, 6, 2, 8, 4, 10])
379
380def test_iter_best_order_f_index_1d():
381 # The Fortran index should be correct with any reordering
382
383 a = arange(4)
384 # 1D order
385 i = nditer(a, ['f_index'], [['readonly']])
386 assert_equal(iter_indices(i), [0, 1, 2, 3])
387 # 1D reversed order
388 i = nditer(a[::-1], ['f_index'], [['readonly']])
389 assert_equal(iter_indices(i), [3, 2, 1, 0])
390
391def test_iter_best_order_f_index_2d():
392 # The Fortran index should be correct with any reordering
393
394 a = arange(6)
395 # 2D C-order
396 i = nditer(a.reshape(2, 3), ['f_index'], [['readonly']])
397 assert_equal(iter_indices(i), [0, 2, 4, 1, 3, 5])
398 # 2D Fortran-order
399 i = nditer(a.reshape(2, 3).copy(order='F'),
400 ['f_index'], [['readonly']])
401 assert_equal(iter_indices(i), [0, 1, 2, 3, 4, 5])
402 # 2D reversed C-order
403 i = nditer(a.reshape(2, 3)[::-1], ['f_index'], [['readonly']])
404 assert_equal(iter_indices(i), [1, 3, 5, 0, 2, 4])
405 i = nditer(a.reshape(2, 3)[:, ::-1], ['f_index'], [['readonly']])
406 assert_equal(iter_indices(i), [4, 2, 0, 5, 3, 1])
407 i = nditer(a.reshape(2, 3)[::-1, ::-1], ['f_index'], [['readonly']])
408 assert_equal(iter_indices(i), [5, 3, 1, 4, 2, 0])
409 # 2D reversed Fortran-order
410 i = nditer(a.reshape(2, 3).copy(order='F')[::-1],
411 ['f_index'], [['readonly']])
412 assert_equal(iter_indices(i), [1, 0, 3, 2, 5, 4])
413 i = nditer(a.reshape(2, 3).copy(order='F')[:, ::-1],
414 ['f_index'], [['readonly']])
415 assert_equal(iter_indices(i), [4, 5, 2, 3, 0, 1])
416 i = nditer(a.reshape(2, 3).copy(order='F')[::-1, ::-1],
417 ['f_index'], [['readonly']])
418 assert_equal(iter_indices(i), [5, 4, 3, 2, 1, 0])
419
420def test_iter_best_order_f_index_3d():
421 # The Fortran index should be correct with any reordering
422
423 a = arange(12)
424 # 3D C-order
425 i = nditer(a.reshape(2, 3, 2), ['f_index'], [['readonly']])
426 assert_equal(iter_indices(i),
427 [0, 6, 2, 8, 4, 10, 1, 7, 3, 9, 5, 11])
428 # 3D Fortran-order
429 i = nditer(a.reshape(2, 3, 2).copy(order='F'),
430 ['f_index'], [['readonly']])
431 assert_equal(iter_indices(i),
432 [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11])
433 # 3D reversed C-order
434 i = nditer(a.reshape(2, 3, 2)[::-1], ['f_index'], [['readonly']])
435 assert_equal(iter_indices(i),
436 [1, 7, 3, 9, 5, 11, 0, 6, 2, 8, 4, 10])
437 i = nditer(a.reshape(2, 3, 2)[:, ::-1], ['f_index'], [['readonly']])
438 assert_equal(iter_indices(i),
439 [4, 10, 2, 8, 0, 6, 5, 11, 3, 9, 1, 7])
440 i = nditer(a.reshape(2, 3, 2)[:, :, ::-1], ['f_index'], [['readonly']])
441 assert_equal(iter_indices(i),
442 [6, 0, 8, 2, 10, 4, 7, 1, 9, 3, 11, 5])
443 # 3D reversed Fortran-order
444 i = nditer(a.reshape(2, 3, 2).copy(order='F')[::-1],
445 ['f_index'], [['readonly']])
446 assert_equal(iter_indices(i),
447 [1, 0, 3, 2, 5, 4, 7, 6, 9, 8, 11, 10])
448 i = nditer(a.reshape(2, 3, 2).copy(order='F')[:, ::-1],
449 ['f_index'], [['readonly']])
450 assert_equal(iter_indices(i),
451 [4, 5, 2, 3, 0, 1, 10, 11, 8, 9, 6, 7])
452 i = nditer(a.reshape(2, 3, 2).copy(order='F')[:, :, ::-1],
453 ['f_index'], [['readonly']])
454 assert_equal(iter_indices(i),
455 [6, 7, 8, 9, 10, 11, 0, 1, 2, 3, 4, 5])
456
457def test_iter_no_inner_full_coalesce():
458 # Check no_inner iterators which coalesce into a single inner loop
459
460 for shape in [(5,), (3, 4), (2, 3, 4), (2, 3, 4, 3), (2, 3, 2, 2, 3)]:
461 size = np.prod(shape)
462 a = arange(size)
463 # Test each combination of forward and backwards indexing
464 for dirs in range(2**len(shape)):
465 dirs_index = [slice(None)] * len(shape)
466 for bit in range(len(shape)):
467 if ((2**bit) & dirs):
468 dirs_index[bit] = slice(None, None, -1)
469 dirs_index = tuple(dirs_index)
470
471 aview = a.reshape(shape)[dirs_index]
472 # C-order
473 i = nditer(aview, ['external_loop'], [['readonly']])
474 assert_equal(i.ndim, 1)
475 assert_equal(i[0].shape, (size,))
476 # Fortran-order
477 i = nditer(aview.T, ['external_loop'], [['readonly']])
478 assert_equal(i.ndim, 1)
479 assert_equal(i[0].shape, (size,))
480 # Other order
481 if len(shape) > 2:
482 i = nditer(aview.swapaxes(0, 1),
483 ['external_loop'], [['readonly']])
484 assert_equal(i.ndim, 1)
485 assert_equal(i[0].shape, (size,))
486
487def test_iter_no_inner_dim_coalescing():
488 # Check no_inner iterators whose dimensions may not coalesce completely
489
490 # Skipping the last element in a dimension prevents coalescing
491 # with the next-bigger dimension
492 a = arange(24).reshape(2, 3, 4)[:, :, :-1]
493 i = nditer(a, ['external_loop'], [['readonly']])
494 assert_equal(i.ndim, 2)
495 assert_equal(i[0].shape, (3,))
496 a = arange(24).reshape(2, 3, 4)[:, :-1, :]
497 i = nditer(a, ['external_loop'], [['readonly']])
498 assert_equal(i.ndim, 2)
499 assert_equal(i[0].shape, (8,))
500 a = arange(24).reshape(2, 3, 4)[:-1, :, :]
501 i = nditer(a, ['external_loop'], [['readonly']])
502 assert_equal(i.ndim, 1)
503 assert_equal(i[0].shape, (12,))
504
505 # Even with lots of 1-sized dimensions, should still coalesce
506 a = arange(24).reshape(1, 1, 2, 1, 1, 3, 1, 1, 4, 1, 1)
507 i = nditer(a, ['external_loop'], [['readonly']])
508 assert_equal(i.ndim, 1)
509 assert_equal(i[0].shape, (24,))
510
511def test_iter_dim_coalescing():
512 # Check that the correct number of dimensions are coalesced
513
514 # Tracking a multi-index disables coalescing
515 a = arange(24).reshape(2, 3, 4)
516 i = nditer(a, ['multi_index'], [['readonly']])
517 assert_equal(i.ndim, 3)
518
519 # A tracked index can allow coalescing if it's compatible with the array
520 a3d = arange(24).reshape(2, 3, 4)
521 i = nditer(a3d, ['c_index'], [['readonly']])
522 assert_equal(i.ndim, 1)
523 i = nditer(a3d.swapaxes(0, 1), ['c_index'], [['readonly']])
524 assert_equal(i.ndim, 3)
525 i = nditer(a3d.T, ['c_index'], [['readonly']])
526 assert_equal(i.ndim, 3)
527 i = nditer(a3d.T, ['f_index'], [['readonly']])
528 assert_equal(i.ndim, 1)
529 i = nditer(a3d.T.swapaxes(0, 1), ['f_index'], [['readonly']])
530 assert_equal(i.ndim, 3)
531
532 # When C or F order is forced, coalescing may still occur
533 a3d = arange(24).reshape(2, 3, 4)
534 i = nditer(a3d, order='C')
535 assert_equal(i.ndim, 1)
536 i = nditer(a3d.T, order='C')
537 assert_equal(i.ndim, 3)
538 i = nditer(a3d, order='F')
539 assert_equal(i.ndim, 3)
540 i = nditer(a3d.T, order='F')
541 assert_equal(i.ndim, 1)
542 i = nditer(a3d, order='A')
543 assert_equal(i.ndim, 1)
544 i = nditer(a3d.T, order='A')
545 assert_equal(i.ndim, 1)
546
547def test_iter_broadcasting():
548 # Standard NumPy broadcasting rules
549
550 # 1D with scalar
551 i = nditer([arange(6), np.int32(2)], ['multi_index'], [['readonly']] * 2)
552 assert_equal(i.itersize, 6)
553 assert_equal(i.shape, (6,))
554
555 # 2D with scalar
556 i = nditer([arange(6).reshape(2, 3), np.int32(2)],
557 ['multi_index'], [['readonly']] * 2)
558 assert_equal(i.itersize, 6)
559 assert_equal(i.shape, (2, 3))
560 # 2D with 1D
561 i = nditer([arange(6).reshape(2, 3), arange(3)],
562 ['multi_index'], [['readonly']] * 2)
563 assert_equal(i.itersize, 6)
564 assert_equal(i.shape, (2, 3))
565 i = nditer([arange(2).reshape(2, 1), arange(3)],
566 ['multi_index'], [['readonly']] * 2)
567 assert_equal(i.itersize, 6)
568 assert_equal(i.shape, (2, 3))
569 # 2D with 2D
570 i = nditer([arange(2).reshape(2, 1), arange(3).reshape(1, 3)],
571 ['multi_index'], [['readonly']] * 2)
572 assert_equal(i.itersize, 6)
573 assert_equal(i.shape, (2, 3))
574
575 # 3D with scalar
576 i = nditer([np.int32(2), arange(24).reshape(4, 2, 3)],
577 ['multi_index'], [['readonly']] * 2)
578 assert_equal(i.itersize, 24)
579 assert_equal(i.shape, (4, 2, 3))
580 # 3D with 1D
581 i = nditer([arange(3), arange(24).reshape(4, 2, 3)],
582 ['multi_index'], [['readonly']] * 2)
583 assert_equal(i.itersize, 24)
584 assert_equal(i.shape, (4, 2, 3))
585 i = nditer([arange(3), arange(8).reshape(4, 2, 1)],
586 ['multi_index'], [['readonly']] * 2)
587 assert_equal(i.itersize, 24)
588 assert_equal(i.shape, (4, 2, 3))
589 # 3D with 2D
590 i = nditer([arange(6).reshape(2, 3), arange(24).reshape(4, 2, 3)],
591 ['multi_index'], [['readonly']] * 2)
592 assert_equal(i.itersize, 24)
593 assert_equal(i.shape, (4, 2, 3))
594 i = nditer([arange(2).reshape(2, 1), arange(24).reshape(4, 2, 3)],
595 ['multi_index'], [['readonly']] * 2)
596 assert_equal(i.itersize, 24)
597 assert_equal(i.shape, (4, 2, 3))
598 i = nditer([arange(3).reshape(1, 3), arange(8).reshape(4, 2, 1)],
599 ['multi_index'], [['readonly']] * 2)
600 assert_equal(i.itersize, 24)
601 assert_equal(i.shape, (4, 2, 3))
602 # 3D with 3D
603 i = nditer([arange(2).reshape(1, 2, 1), arange(3).reshape(1, 1, 3),
604 arange(4).reshape(4, 1, 1)],
605 ['multi_index'], [['readonly']] * 3)
606 assert_equal(i.itersize, 24)
607 assert_equal(i.shape, (4, 2, 3))
608 i = nditer([arange(6).reshape(1, 2, 3), arange(4).reshape(4, 1, 1)],
609 ['multi_index'], [['readonly']] * 2)
610 assert_equal(i.itersize, 24)
611 assert_equal(i.shape, (4, 2, 3))
612 i = nditer([arange(24).reshape(4, 2, 3), arange(12).reshape(4, 1, 3)],
613 ['multi_index'], [['readonly']] * 2)
614 assert_equal(i.itersize, 24)
615 assert_equal(i.shape, (4, 2, 3))
616
617def test_iter_itershape():
618 # Check that allocated outputs work with a specified shape
619 a = np.arange(6, dtype='i2').reshape(2, 3)
620 i = nditer([a, None], [], [['readonly'], ['writeonly', 'allocate']],
621 op_axes=[[0, 1, None], None],
622 itershape=(-1, -1, 4))
623 assert_equal(i.operands[1].shape, (2, 3, 4))
624 assert_equal(i.operands[1].strides, (24, 8, 2))
625
626 i = nditer([a.T, None], [], [['readonly'], ['writeonly', 'allocate']],
627 op_axes=[[0, 1, None], None],
628 itershape=(-1, -1, 4))
629 assert_equal(i.operands[1].shape, (3, 2, 4))
630 assert_equal(i.operands[1].strides, (8, 24, 2))
631
632 i = nditer([a.T, None], [], [['readonly'], ['writeonly', 'allocate']],
633 order='F',
634 op_axes=[[0, 1, None], None],
635 itershape=(-1, -1, 4))
636 assert_equal(i.operands[1].shape, (3, 2, 4))
637 assert_equal(i.operands[1].strides, (2, 6, 12))
638
639 # If we specify 1 in the itershape, it shouldn't allow broadcasting
640 # of that dimension to a bigger value
641 assert_raises(ValueError, nditer, [a, None], [],
642 [['readonly'], ['writeonly', 'allocate']],
643 op_axes=[[0, 1, None], None],
644 itershape=(-1, 1, 4))
645 # Test bug that for no op_axes but itershape, they are NULLed correctly
646 i = np.nditer([np.ones(2), None, None], itershape=(2,))
647
648def test_iter_broadcasting_errors():
649 # Check that errors are thrown for bad broadcasting shapes
650
651 # 1D with 1D
652 assert_raises(ValueError, nditer, [arange(2), arange(3)],
653 [], [['readonly']] * 2)
654 # 2D with 1D
655 assert_raises(ValueError, nditer,
656 [arange(6).reshape(2, 3), arange(2)],
657 [], [['readonly']] * 2)
658 # 2D with 2D
659 assert_raises(ValueError, nditer,
660 [arange(6).reshape(2, 3), arange(9).reshape(3, 3)],
661 [], [['readonly']] * 2)
662 assert_raises(ValueError, nditer,
663 [arange(6).reshape(2, 3), arange(4).reshape(2, 2)],
664 [], [['readonly']] * 2)
665 # 3D with 3D
666 assert_raises(ValueError, nditer,
667 [arange(36).reshape(3, 3, 4), arange(24).reshape(2, 3, 4)],
668 [], [['readonly']] * 2)
669 assert_raises(ValueError, nditer,
670 [arange(8).reshape(2, 4, 1), arange(24).reshape(2, 3, 4)],
671 [], [['readonly']] * 2)
672
673 # Verify that the error message mentions the right shapes
674 try:
675 nditer([arange(2).reshape(1, 2, 1),
676 arange(3).reshape(1, 3),
677 arange(6).reshape(2, 3)],
678 [],
679 [['readonly'], ['readonly'], ['writeonly', 'no_broadcast']])
680 raise AssertionError('Should have raised a broadcast error')
681 except ValueError as e:
682 msg = str(e)
683 # The message should contain the shape of the 3rd operand
684 assert_(msg.find('(2,3)') >= 0,
685 f'Message "{msg}" doesn\'t contain operand shape (2,3)')
686 # The message should contain the broadcast shape
687 assert_(msg.find('(1,2,3)') >= 0,
688 f'Message "{msg}" doesn\'t contain broadcast shape (1,2,3)')
689
690 try:
691 nditer([arange(6).reshape(2, 3), arange(2)],
692 [],
693 [['readonly'], ['readonly']],
694 op_axes=[[0, 1], [0, np.newaxis]],
695 itershape=(4, 3))
696 raise AssertionError('Should have raised a broadcast error')
697 except ValueError as e:
698 msg = str(e)
699 # The message should contain "shape->remappedshape" for each operand
700 assert_(msg.find('(2,3)->(2,3)') >= 0,
701 f'Message "{msg}" doesn\'t contain operand shape (2,3)->(2,3)')
702 assert_(msg.find('(2,)->(2,newaxis)') >= 0,
703 ('Message "%s" doesn\'t contain remapped operand shape'
704 '(2,)->(2,newaxis)') % msg)
705 # The message should contain the itershape parameter
706 assert_(msg.find('(4,3)') >= 0,
707 f'Message "{msg}" doesn\'t contain itershape parameter (4,3)')
708
709 try:
710 nditer([np.zeros((2, 1, 1)), np.zeros((2,))],
711 [],
712 [['writeonly', 'no_broadcast'], ['readonly']])
713 raise AssertionError('Should have raised a broadcast error')
714 except ValueError as e:
715 msg = str(e)
716 # The message should contain the shape of the bad operand
717 assert_(msg.find('(2,1,1)') >= 0,
718 f'Message "{msg}" doesn\'t contain operand shape (2,1,1)')
719 # The message should contain the broadcast shape
720 assert_(msg.find('(2,1,2)') >= 0,
721 f'Message "{msg}" doesn\'t contain the broadcast shape (2,1,2)')
722
723def test_iter_flags_errors():
724 # Check that bad combinations of flags produce errors
725
726 a = arange(6)
727
728 # Not enough operands
729 assert_raises(ValueError, nditer, [], [], [])
730 # Bad global flag
731 assert_raises(ValueError, nditer, [a], ['bad flag'], [['readonly']])
732 # Bad op flag
733 assert_raises(ValueError, nditer, [a], [], [['readonly', 'bad flag']])
734 # Bad order parameter
735 assert_raises(ValueError, nditer, [a], [], [['readonly']], order='G')
736 # Bad casting parameter
737 assert_raises(ValueError, nditer, [a], [], [['readonly']], casting='noon')
738 # op_flags must match ops
739 assert_raises(ValueError, nditer, [a] * 3, [], [['readonly']] * 2)
740 # Cannot track both a C and an F index
741 assert_raises(ValueError, nditer, a,
742 ['c_index', 'f_index'], [['readonly']])
743 # Inner iteration and multi-indices/indices are incompatible
744 assert_raises(ValueError, nditer, a,
745 ['external_loop', 'multi_index'], [['readonly']])
746 assert_raises(ValueError, nditer, a,
747 ['external_loop', 'c_index'], [['readonly']])
748 assert_raises(ValueError, nditer, a,
749 ['external_loop', 'f_index'], [['readonly']])
750 # Must specify exactly one of readwrite/readonly/writeonly per operand
751 assert_raises(ValueError, nditer, a, [], [[]])
752 assert_raises(ValueError, nditer, a, [], [['readonly', 'writeonly']])
753 assert_raises(ValueError, nditer, a, [], [['readonly', 'readwrite']])
754 assert_raises(ValueError, nditer, a, [], [['writeonly', 'readwrite']])
755 assert_raises(ValueError, nditer, a,
756 [], [['readonly', 'writeonly', 'readwrite']])
757 # Python scalars are always readonly
758 assert_raises(TypeError, nditer, 1.5, [], [['writeonly']])
759 assert_raises(TypeError, nditer, 1.5, [], [['readwrite']])
760 # Array scalars are always readonly
761 assert_raises(TypeError, nditer, np.int32(1), [], [['writeonly']])
762 assert_raises(TypeError, nditer, np.int32(1), [], [['readwrite']])
763 # Check readonly array
764 a.flags.writeable = False
765 assert_raises(ValueError, nditer, a, [], [['writeonly']])
766 assert_raises(ValueError, nditer, a, [], [['readwrite']])
767 a.flags.writeable = True
768 # Multi-indices available only with the multi_index flag
769 i = nditer(arange(6), [], [['readonly']])
770 assert_raises(ValueError, lambda i: i.multi_index, i)
771 # Index available only with an index flag
772 assert_raises(ValueError, lambda i: i.index, i)
773 # GotoCoords and GotoIndex incompatible with buffering or no_inner
774
775 def assign_multi_index(i):
776 i.multi_index = (0,)
777
778 def assign_index(i):
779 i.index = 0
780
781 def assign_iterindex(i):
782 i.iterindex = 0
783
784 def assign_iterrange(i):
785 i.iterrange = (0, 1)
786 i = nditer(arange(6), ['external_loop'])
787 assert_raises(ValueError, assign_multi_index, i)
788 assert_raises(ValueError, assign_index, i)
789 assert_raises(ValueError, assign_iterindex, i)
790 assert_raises(ValueError, assign_iterrange, i)
791 i = nditer(arange(6), ['buffered'])
792 assert_raises(ValueError, assign_multi_index, i)
793 assert_raises(ValueError, assign_index, i)
794 assert_raises(ValueError, assign_iterrange, i)
795 # Can't iterate if size is zero
796 assert_raises(ValueError, nditer, np.array([]))
797
798def test_iter_slice():
799 a, b, c = np.arange(3), np.arange(3), np.arange(3.)
800 i = nditer([a, b, c], [], ['readwrite'])
801 with i:
802 i[0:2] = (3, 3)
803 assert_equal(a, [3, 1, 2])
804 assert_equal(b, [3, 1, 2])
805 assert_equal(c, [0, 1, 2])
806 i[1] = 12
807 assert_equal(i[0:2], [3, 12])
808
809def test_iter_assign_mapping():
810 a = np.arange(24, dtype='f8').reshape(2, 3, 4).T
811 it = np.nditer(a, [], [['readwrite', 'updateifcopy']],
812 casting='same_kind', op_dtypes=[np.dtype('f4')])
813 with it:
814 it.operands[0][...] = 3
815 it.operands[0][...] = 14
816 assert_equal(a, 14)
817 it = np.nditer(a, [], [['readwrite', 'updateifcopy']],
818 casting='same_kind', op_dtypes=[np.dtype('f4')])
819 with it:
820 x = it.operands[0][-1:1]
821 x[...] = 14
822 it.operands[0][...] = -1234
823 assert_equal(a, -1234)
824 # check for no warnings on dealloc
825 x = None
826 it = None
827
828def test_iter_nbo_align_contig():
829 # Check that byte order, alignment, and contig changes work
830
831 # Byte order change by requesting a specific dtype
832 a = np.arange(6, dtype='f4')
833 au = a.byteswap()
834 au = au.view(au.dtype.newbyteorder())
835 assert_(a.dtype.byteorder != au.dtype.byteorder)
836 i = nditer(au, [], [['readwrite', 'updateifcopy']],
837 casting='equiv',
838 op_dtypes=[np.dtype('f4')])
839 with i:
840 # context manager triggers WRITEBACKIFCOPY on i at exit
841 assert_equal(i.dtypes[0].byteorder, a.dtype.byteorder)
842 assert_equal(i.operands[0].dtype.byteorder, a.dtype.byteorder)
843 assert_equal(i.operands[0], a)
844 i.operands[0][:] = 2
845 assert_equal(au, [2] * 6)
846 del i # should not raise a warning
847 # Byte order change by requesting NBO
848 a = np.arange(6, dtype='f4')
849 au = a.byteswap()
850 au = au.view(au.dtype.newbyteorder())
851 assert_(a.dtype.byteorder != au.dtype.byteorder)
852 with nditer(au, [], [['readwrite', 'updateifcopy', 'nbo']],
853 casting='equiv') as i:
854 # context manager triggers UPDATEIFCOPY on i at exit
855 assert_equal(i.dtypes[0].byteorder, a.dtype.byteorder)
856 assert_equal(i.operands[0].dtype.byteorder, a.dtype.byteorder)
857 assert_equal(i.operands[0], a)
858 i.operands[0][:] = 12345
859 i.operands[0][:] = 2
860 assert_equal(au, [2] * 6)
861
862 # Unaligned input
863 a = np.zeros((6 * 4 + 1,), dtype='i1')[1:]
864 a = a.view('f4')
865 a[:] = np.arange(6, dtype='f4')
866 assert_(not a.flags.aligned)
867 # Without 'aligned', shouldn't copy
868 i = nditer(a, [], [['readonly']])
869 assert_(not i.operands[0].flags.aligned)
870 assert_equal(i.operands[0], a)
871 # With 'aligned', should make a copy
872 with nditer(a, [], [['readwrite', 'updateifcopy', 'aligned']]) as i:
873 assert_(i.operands[0].flags.aligned)
874 # context manager triggers UPDATEIFCOPY on i at exit
875 assert_equal(i.operands[0], a)
876 i.operands[0][:] = 3
877 assert_equal(a, [3] * 6)
878
879 # Discontiguous input
880 a = arange(12)
881 # If it is contiguous, shouldn't copy
882 i = nditer(a[:6], [], [['readonly']])
883 assert_(i.operands[0].flags.contiguous)
884 assert_equal(i.operands[0], a[:6])
885 # If it isn't contiguous, should buffer
886 i = nditer(a[::2], ['buffered', 'external_loop'],
887 [['readonly', 'contig']],
888 buffersize=10)
889 assert_(i[0].flags.contiguous)
890 assert_equal(i[0], a[::2])
891
892def test_iter_array_cast():
893 # Check that arrays are cast as requested
894
895 # No cast 'f4' -> 'f4'
896 a = np.arange(6, dtype='f4').reshape(2, 3)
897 i = nditer(a, [], [['readwrite']], op_dtypes=[np.dtype('f4')])
898 with i:
899 assert_equal(i.operands[0], a)
900 assert_equal(i.operands[0].dtype, np.dtype('f4'))
901
902 # Byte-order cast '<f4' -> '>f4'
903 a = np.arange(6, dtype='<f4').reshape(2, 3)
904 with nditer(a, [], [['readwrite', 'updateifcopy']],
905 casting='equiv',
906 op_dtypes=[np.dtype('>f4')]) as i:
907 assert_equal(i.operands[0], a)
908 assert_equal(i.operands[0].dtype, np.dtype('>f4'))
909
910 # Safe case 'f4' -> 'f8'
911 a = np.arange(24, dtype='f4').reshape(2, 3, 4).swapaxes(1, 2)
912 i = nditer(a, [], [['readonly', 'copy']],
913 casting='safe',
914 op_dtypes=[np.dtype('f8')])
915 assert_equal(i.operands[0], a)
916 assert_equal(i.operands[0].dtype, np.dtype('f8'))
917 # The memory layout of the temporary should match a (a is (48,4,16))
918 # except negative strides get flipped to positive strides.
919 assert_equal(i.operands[0].strides, (96, 8, 32))
920 a = a[::-1, :, ::-1]
921 i = nditer(a, [], [['readonly', 'copy']],
922 casting='safe',
923 op_dtypes=[np.dtype('f8')])
924 assert_equal(i.operands[0], a)
925 assert_equal(i.operands[0].dtype, np.dtype('f8'))
926 assert_equal(i.operands[0].strides, (96, 8, 32))
927
928 # Same-kind cast 'f8' -> 'f4' -> 'f8'
929 a = np.arange(24, dtype='f8').reshape(2, 3, 4).T
930 with nditer(a, [],
931 [['readwrite', 'updateifcopy']],
932 casting='same_kind',
933 op_dtypes=[np.dtype('f4')]) as i:
934 assert_equal(i.operands[0], a)
935 assert_equal(i.operands[0].dtype, np.dtype('f4'))
936 assert_equal(i.operands[0].strides, (4, 16, 48))
937 # Check that WRITEBACKIFCOPY is activated at exit
938 i.operands[0][2, 1, 1] = -12.5
939 assert_(a[2, 1, 1] != -12.5)
940 assert_equal(a[2, 1, 1], -12.5)
941
942 a = np.arange(6, dtype='i4')[::-2]
943 with nditer(a, [],
944 [['writeonly', 'updateifcopy']],
945 casting='unsafe',
946 op_dtypes=[np.dtype('f4')]) as i:
947 assert_equal(i.operands[0].dtype, np.dtype('f4'))
948 # Even though the stride was negative in 'a', it
949 # becomes positive in the temporary
950 assert_equal(i.operands[0].strides, (4,))
951 i.operands[0][:] = [1, 2, 3]
952 assert_equal(a, [1, 2, 3])
953
954def test_iter_array_cast_errors():
955 # Check that invalid casts are caught
956
957 # Need to enable copying for casts to occur
958 assert_raises(TypeError, nditer, arange(2, dtype='f4'), [],
959 [['readonly']], op_dtypes=[np.dtype('f8')])
960 # Also need to allow casting for casts to occur
961 assert_raises(TypeError, nditer, arange(2, dtype='f4'), [],
962 [['readonly', 'copy']], casting='no',
963 op_dtypes=[np.dtype('f8')])
964 assert_raises(TypeError, nditer, arange(2, dtype='f4'), [],
965 [['readonly', 'copy']], casting='equiv',
966 op_dtypes=[np.dtype('f8')])
967 assert_raises(TypeError, nditer, arange(2, dtype='f8'), [],
968 [['writeonly', 'updateifcopy']],
969 casting='no',
970 op_dtypes=[np.dtype('f4')])
971 assert_raises(TypeError, nditer, arange(2, dtype='f8'), [],
972 [['writeonly', 'updateifcopy']],
973 casting='equiv',
974 op_dtypes=[np.dtype('f4')])
975 # '<f4' -> '>f4' should not work with casting='no'
976 assert_raises(TypeError, nditer, arange(2, dtype='<f4'), [],
977 [['readonly', 'copy']], casting='no',
978 op_dtypes=[np.dtype('>f4')])
979 # 'f4' -> 'f8' is a safe cast, but 'f8' -> 'f4' isn't
980 assert_raises(TypeError, nditer, arange(2, dtype='f4'), [],
981 [['readwrite', 'updateifcopy']],
982 casting='safe',
983 op_dtypes=[np.dtype('f8')])
984 assert_raises(TypeError, nditer, arange(2, dtype='f8'), [],
985 [['readwrite', 'updateifcopy']],
986 casting='safe',
987 op_dtypes=[np.dtype('f4')])
988 # 'f4' -> 'i4' is neither a safe nor a same-kind cast
989 assert_raises(TypeError, nditer, arange(2, dtype='f4'), [],
990 [['readonly', 'copy']],
991 casting='same_kind',
992 op_dtypes=[np.dtype('i4')])
993 assert_raises(TypeError, nditer, arange(2, dtype='i4'), [],
994 [['writeonly', 'updateifcopy']],
995 casting='same_kind',
996 op_dtypes=[np.dtype('f4')])
997
998def test_iter_scalar_cast():
999 # Check that scalars are cast as requested
1000
1001 # No cast 'f4' -> 'f4'
1002 i = nditer(np.float32(2.5), [], [['readonly']],
1003 op_dtypes=[np.dtype('f4')])
1004 assert_equal(i.dtypes[0], np.dtype('f4'))
1005 assert_equal(i.value.dtype, np.dtype('f4'))
1006 assert_equal(i.value, 2.5)
1007 # Safe cast 'f4' -> 'f8'
1008 i = nditer(np.float32(2.5), [],
1009 [['readonly', 'copy']],
1010 casting='safe',
1011 op_dtypes=[np.dtype('f8')])
1012 assert_equal(i.dtypes[0], np.dtype('f8'))
1013 assert_equal(i.value.dtype, np.dtype('f8'))
1014 assert_equal(i.value, 2.5)
1015 # Same-kind cast 'f8' -> 'f4'
1016 i = nditer(np.float64(2.5), [],
1017 [['readonly', 'copy']],
1018 casting='same_kind',
1019 op_dtypes=[np.dtype('f4')])
1020 assert_equal(i.dtypes[0], np.dtype('f4'))
1021 assert_equal(i.value.dtype, np.dtype('f4'))
1022 assert_equal(i.value, 2.5)
1023 # Unsafe cast 'f8' -> 'i4'
1024 i = nditer(np.float64(3.0), [],
1025 [['readonly', 'copy']],
1026 casting='unsafe',
1027 op_dtypes=[np.dtype('i4')])
1028 assert_equal(i.dtypes[0], np.dtype('i4'))
1029 assert_equal(i.value.dtype, np.dtype('i4'))
1030 assert_equal(i.value, 3)
1031 # Readonly scalars may be cast even without setting COPY or BUFFERED
1032 i = nditer(3, [], [['readonly']], op_dtypes=[np.dtype('f8')])
1033 assert_equal(i[0].dtype, np.dtype('f8'))
1034 assert_equal(i[0], 3.)
1035
1036def test_iter_scalar_cast_errors():
1037 # Check that invalid casts are caught
1038
1039 # Need to allow copying/buffering for write casts of scalars to occur
1040 assert_raises(TypeError, nditer, np.float32(2), [],
1041 [['readwrite']], op_dtypes=[np.dtype('f8')])
1042 assert_raises(TypeError, nditer, 2.5, [],
1043 [['readwrite']], op_dtypes=[np.dtype('f4')])
1044 # 'f8' -> 'f4' isn't a safe cast if the value would overflow
1045 assert_raises(TypeError, nditer, np.float64(1e60), [],
1046 [['readonly']],
1047 casting='safe',
1048 op_dtypes=[np.dtype('f4')])
1049 # 'f4' -> 'i4' is neither a safe nor a same-kind cast
1050 assert_raises(TypeError, nditer, np.float32(2), [],
1051 [['readonly']],
1052 casting='same_kind',
1053 op_dtypes=[np.dtype('i4')])
1054
1055def test_iter_object_arrays_basic():
1056 # Check that object arrays work
1057
1058 obj = {'a': 3, 'b': 'd'}
1059 a = np.array([[1, 2, 3], None, obj, None], dtype='O')
1060 if HAS_REFCOUNT:
1061 rc = sys.getrefcount(obj)
1062
1063 # Need to allow references for object arrays
1064 assert_raises(TypeError, nditer, a)
1065 if HAS_REFCOUNT:
1066 assert_equal(sys.getrefcount(obj), rc)
1067
1068 i = nditer(a, ['refs_ok'], ['readonly'])
1069 vals = [x_[()] for x_ in i]
1070 assert_equal(np.array(vals, dtype='O'), a)
1071 vals, i, x = [None] * 3
1072 if HAS_REFCOUNT:
1073 assert_equal(sys.getrefcount(obj), rc)
1074
1075 i = nditer(a.reshape(2, 2).T, ['refs_ok', 'buffered'],
1076 ['readonly'], order='C')
1077 assert_(i.iterationneedsapi)
1078 vals = [x_[()] for x_ in i]
1079 assert_equal(np.array(vals, dtype='O'), a.reshape(2, 2).ravel(order='F'))
1080 vals, i, x = [None] * 3
1081 if HAS_REFCOUNT:
1082 assert_equal(sys.getrefcount(obj), rc)
1083
1084 i = nditer(a.reshape(2, 2).T, ['refs_ok', 'buffered'],
1085 ['readwrite'], order='C')
1086 with i:
1087 for x in i:
1088 x[...] = None
1089 vals, i, x = [None] * 3
1090 if HAS_REFCOUNT:
1091 assert_(sys.getrefcount(obj) == rc - 1)
1092 assert_equal(a, np.array([None] * 4, dtype='O'))
1093
1094def test_iter_object_arrays_conversions():
1095 # Conversions to/from objects
1096 a = np.arange(6, dtype='O')
1097 i = nditer(a, ['refs_ok', 'buffered'], ['readwrite'],
1098 casting='unsafe', op_dtypes='i4')
1099 with i:
1100 for x in i:
1101 x[...] += 1
1102 assert_equal(a, np.arange(6) + 1)
1103
1104 a = np.arange(6, dtype='i4')
1105 i = nditer(a, ['refs_ok', 'buffered'], ['readwrite'],
1106 casting='unsafe', op_dtypes='O')
1107 with i:
1108 for x in i:
1109 x[...] += 1
1110 assert_equal(a, np.arange(6) + 1)
1111
1112 # Non-contiguous object array
1113 a = np.zeros((6,), dtype=[('p', 'i1'), ('a', 'O')])
1114 a = a['a']
1115 a[:] = np.arange(6)
1116 i = nditer(a, ['refs_ok', 'buffered'], ['readwrite'],
1117 casting='unsafe', op_dtypes='i4')
1118 with i:
1119 for x in i:
1120 x[...] += 1
1121 assert_equal(a, np.arange(6) + 1)
1122
1123 # Non-contiguous value array
1124 a = np.zeros((6,), dtype=[('p', 'i1'), ('a', 'i4')])
1125 a = a['a']
1126 a[:] = np.arange(6) + 98172488
1127 i = nditer(a, ['refs_ok', 'buffered'], ['readwrite'],
1128 casting='unsafe', op_dtypes='O')
1129 with i:
1130 ob = i[0][()]
1131 if HAS_REFCOUNT:
1132 rc = sys.getrefcount(ob)
1133 for x in i:
1134 x[...] += 1
1135 if HAS_REFCOUNT:
1136 newrc = sys.getrefcount(ob)
1137 assert_(newrc == rc - 1)
1138 assert_equal(a, np.arange(6) + 98172489)
1139
1140def test_iter_common_dtype():
1141 # Check that the iterator finds a common data type correctly
1142 # (some checks are somewhat duplicate after adopting NEP 50)
1143
1144 i = nditer([array([3], dtype='f4'), array([0], dtype='f8')],
1145 ['common_dtype'],
1146 [['readonly', 'copy']] * 2,
1147 casting='safe')
1148 assert_equal(i.dtypes[0], np.dtype('f8'))
1149 assert_equal(i.dtypes[1], np.dtype('f8'))
1150 i = nditer([array([3], dtype='i4'), array([0], dtype='f4')],
1151 ['common_dtype'],
1152 [['readonly', 'copy']] * 2,
1153 casting='safe')
1154 assert_equal(i.dtypes[0], np.dtype('f8'))
1155 assert_equal(i.dtypes[1], np.dtype('f8'))
1156 i = nditer([array([3], dtype='f4'), array(0, dtype='f8')],
1157 ['common_dtype'],
1158 [['readonly', 'copy']] * 2,
1159 casting='same_kind')
1160 assert_equal(i.dtypes[0], np.dtype('f8'))
1161 assert_equal(i.dtypes[1], np.dtype('f8'))
1162 i = nditer([array([3], dtype='u4'), array(0, dtype='i4')],
1163 ['common_dtype'],
1164 [['readonly', 'copy']] * 2,
1165 casting='safe')
1166 assert_equal(i.dtypes[0], np.dtype('i8'))
1167 assert_equal(i.dtypes[1], np.dtype('i8'))
1168 i = nditer([array([3], dtype='u4'), array(-12, dtype='i4')],
1169 ['common_dtype'],
1170 [['readonly', 'copy']] * 2,
1171 casting='safe')
1172 assert_equal(i.dtypes[0], np.dtype('i8'))
1173 assert_equal(i.dtypes[1], np.dtype('i8'))
1174 i = nditer([array([3], dtype='u4'), array(-12, dtype='i4'),
1175 array([2j], dtype='c8'), array([9], dtype='f8')],
1176 ['common_dtype'],
1177 [['readonly', 'copy']] * 4,
1178 casting='safe')
1179 assert_equal(i.dtypes[0], np.dtype('c16'))
1180 assert_equal(i.dtypes[1], np.dtype('c16'))
1181 assert_equal(i.dtypes[2], np.dtype('c16'))
1182 assert_equal(i.dtypes[3], np.dtype('c16'))
1183 assert_equal(i.value, (3, -12, 2j, 9))
1184
1185 # When allocating outputs, other outputs aren't factored in
1186 i = nditer([array([3], dtype='i4'), None, array([2j], dtype='c16')], [],
1187 [['readonly', 'copy'],
1188 ['writeonly', 'allocate'],
1189 ['writeonly']],
1190 casting='safe')
1191 assert_equal(i.dtypes[0], np.dtype('i4'))
1192 assert_equal(i.dtypes[1], np.dtype('i4'))
1193 assert_equal(i.dtypes[2], np.dtype('c16'))
1194 # But, if common data types are requested, they are
1195 i = nditer([array([3], dtype='i4'), None, array([2j], dtype='c16')],
1196 ['common_dtype'],
1197 [['readonly', 'copy'],
1198 ['writeonly', 'allocate'],
1199 ['writeonly']],
1200 casting='safe')