codekingpro/portable-devtools

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