1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
|
from operator import mul
import numpy as np
from numpy.random import randint
from numpy.lib import Arrayterator
def test():
np.random.seed(np.arange(10))
# Create a random array
ndims = randint(5)+1
shape = tuple(randint(10)+1 for dim in range(ndims))
els = reduce(mul, shape)
a = np.arange(els)
a.shape = shape
buf_size = randint(2*els)
b = Arrayterator(a, buf_size)
# Check that each block has at most ``buf_size`` elements
for block in b:
assert len(block.flat) <= (buf_size or els)
# Check that all elements are iterated correctly
assert list(b.flat) == list(a.flat)
# Slice arrayterator
start = [randint(dim) for dim in shape]
stop = [randint(dim)+1 for dim in shape]
step = [randint(dim)+1 for dim in shape]
slice_ = tuple(slice(*t) for t in zip(start, stop, step))
c = b[slice_]
d = a[slice_]
# Check that each block has at most ``buf_size`` elements
for block in c:
assert len(block.flat) <= (buf_size or els)
# Check that the arrayterator is sliced correctly
assert np.all(c.__array__() == d)
# Check that all elements are iterated correctly
assert list(c.flat) == list(d.flat)
if __name__ == '__main__':
from numpy.testing import run_module_suite
run_module_suite()
|
