1from .common import Benchmark
2
3import numpy as np
4
5
6class Block(Benchmark):
7    params = [1, 10, 100]
8    param_names = ['size']
9
10    def setup(self, n):
11        self.a_2d = np.ones((2 * n, 2 * n))
12        self.b_1d = np.ones(2 * n)
13        self.b_2d = 2 * self.a_2d
14
15        self.a = np.ones(3 * n)
16        self.b = np.ones(3 * n)
17
18        self.one_2d = np.ones((1 * n, 3 * n))
19        self.two_2d = np.ones((1 * n, 3 * n))
20        self.three_2d = np.ones((1 * n, 6 * n))
21        self.four_1d = np.ones(6 * n)
22        self.five_0d = np.ones(1 * n)
23        self.six_1d = np.ones(5 * n)
24        # avoid np.zeros's lazy allocation that might cause
25        # page faults during benchmark
26        self.zero_2d = np.full((2 * n, 6 * n), 0)
27
28        self.one = np.ones(3 * n)
29        self.two = 2 * np.ones((3, 3 * n))
30        self.three = 3 * np.ones(3 * n)
31        self.four = 4 * np.ones(3 * n)
32        self.five = 5 * np.ones(1 * n)
33        self.six = 6 * np.ones(5 * n)
34        # avoid np.zeros's lazy allocation that might cause
35        # page faults during benchmark
36        self.zero = np.full((2 * n, 6 * n), 0)
37
38    def time_block_simple_row_wise(self, n):
39        np.block([self.a_2d, self.b_2d])
40
41    def time_block_simple_column_wise(self, n):
42        np.block([[self.a_2d], [self.b_2d]])
43
44    def time_block_complicated(self, n):
45        np.block([[self.one_2d, self.two_2d],
46                  [self.three_2d],
47                  [self.four_1d],
48                  [self.five_0d, self.six_1d],
49                  [self.zero_2d]])
50
51    def time_nested(self, n):
52        np.block([
53            [
54                np.block([
55                   [self.one],
56                   [self.three],
57                   [self.four]
58                ]),
59                self.two
60            ],
61            [self.five, self.six],
62            [self.zero]
63        ])
64
65    def time_no_lists(self, n):
66        np.block(1)
67        np.block(np.eye(3 * n))
68
69
70class Block2D(Benchmark):
71    params = [[(16, 16), (32, 32), (64, 64), (128, 128), (256, 256), (512, 512), (1024, 1024)],
72              ['uint8', 'uint16', 'uint32', 'uint64'],
73              [(2, 2), (4, 4)]]
74    param_names = ['shape', 'dtype', 'n_chunks']
75
76    def setup(self, shape, dtype, n_chunks):
77
78        self.block_list = [
79             [np.full(shape=[s//n_chunk for s, n_chunk in zip(shape, n_chunks)],
80                     fill_value=1, dtype=dtype) for _ in range(n_chunks[1])]
81            for _ in range(n_chunks[0])
82        ]
83
84    def time_block2d(self, shape, dtype, n_chunks):
85        np.block(self.block_list)
86
87
88class Block3D(Benchmark):
89    """This benchmark concatenates an array of size ``(5n)^3``"""
90    # Having copy as a `mode` of the block3D
91    # allows us to directly compare the benchmark of block
92    # to that of a direct memory copy into new buffers with
93    # the ASV framework.
94    # block and copy will be plotted on the same graph
95    # as opposed to being displayed as separate benchmarks
96    params = [[1, 10, 100],
97              ['block', 'copy']]
98    param_names = ['n', 'mode']
99
100    def setup(self, n, mode):
101        # Slow setup method: hence separated from the others above
102        self.a000 = np.ones((2 * n, 2 * n, 2 * n), int) * 1
103
104        self.a100 = np.ones((3 * n, 2 * n, 2 * n), int) * 2
105        self.a010 = np.ones((2 * n, 3 * n, 2 * n), int) * 3
106        self.a001 = np.ones((2 * n, 2 * n, 3 * n), int) * 4
107
108        self.a011 = np.ones((2 * n, 3 * n, 3 * n), int) * 5
109        self.a101 = np.ones((3 * n, 2 * n, 3 * n), int) * 6
110        self.a110 = np.ones((3 * n, 3 * n, 2 * n), int) * 7
111
112        self.a111 = np.ones((3 * n, 3 * n, 3 * n), int) * 8
113
114        self.block = [
115            [
116                [self.a000, self.a001],
117                [self.a010, self.a011],
118            ],
119            [
120                [self.a100, self.a101],
121                [self.a110, self.a111],
122            ]
123        ]
124        self.arr_list = [a
125                         for two_d in self.block
126                         for one_d in two_d
127                         for a in one_d]
128
129    def time_3d(self, n, mode):
130        if mode == 'block':
131            np.block(self.block)
132        else:  # mode == 'copy'
133            [arr.copy() for arr in self.arr_list]
134
135    # Retain old benchmark name for backward compat
136    time_3d.benchmark_name = "bench_shape_base.Block.time_3d"
137