BENCH: Add some basic ma benchmarks

eric-wieser · eric-wieser · commit 78a1837e4abd · 2017-04-12T01:11:14.000+01:00
diff --git a/benchmarks/benchmarks/bench_ma.py b/benchmarks/benchmarks/bench_ma.py
@@ -18,3 +18,96 @@ def time_masked_array_l100(self):
 
     def time_masked_array_l100_t100(self):
         np.ma.masked_array(self.l100, self.t100)
+
+
+class Indexing(Benchmark):
+    param_names = ['masked', 'ndim', 'size']
+    params = [[True, False],
+              [1, 2],
+              [10, 100, 1000]]
+    def setup(self, masked, ndim, size):
+        x = np.arange(size**ndim).reshape(ndim * (size,))
+
+        if masked:
+            self.m = np.ma.array(x, mask=x%2 == 0)
+        else:
+            self.m = np.ma.array(x)
+
+        self.idx_scalar = (size//2,) * ndim
+        self.idx_0d = (size//2,) * ndim + (Ellipsis,)
+        self.idx_1d = (size//2,) * (ndim - 1)
+
+    def time_scalar(self, masked, ndim, size):
+        self.m[self.idx_scalar]
+
+    def time_0d(self, masked, ndim, size):
+        self.m[self.idx_0d]
+
+    def time_1d(self, masked, ndim, size):
+        self.m[self.idx_1d]
+
+
+class UFunc(Benchmark):
+    param_names = ['a_masked', 'b_masked', 'size']
+    params = [[True, False],
+              [True, False],
+              [10, 100, 1000]]
+
+    def setup(self, a_masked, b_masked, size):
+        x = np.arange(size).astype(np.uint8)
+
+        self.a_scalar = np.ma.masked if a_masked else 5
+        self.b_scalar = np.ma.masked if b_masked else 3
+
+        self.a_1d = np.ma.array(x, mask=x%2 == 0 if a_masked else np.ma.nomask)
+        self.b_1d = np.ma.array(x, mask=x%3 == 0 if b_masked else np.ma.nomask)
+
+        self.a_2d = self.a_1d.reshape(1, -1)
+        self.b_2d = self.a_1d.reshape(-1, 1)
+
+    def time_scalar(self, a_masked, b_masked, size):
+        np.ma.add(self.a_scalar, self.b_scalar)
+
+    def time_scalar_1d(self, a_masked, b_masked, size):
+        np.ma.add(self.a_scalar, self.b_1d)
+
+    def time_1d(self, a_masked, b_masked, size):
+        np.ma.add(self.a_1d, self.b_1d)
+
+    def time_2d(self, a_masked, b_masked, size):
+        # broadcasting happens this time
+        np.ma.add(self.a_2d, self.b_2d)
+
+
+class Concatenate(Benchmark):
+    param_names = ['mode', 'n']
+    params = [
+        ['ndarray', 'unmasked',
+         'ndarray+masked', 'unmasked+masked',
+         'masked'],
+        [2, 100, 2000]
+    ]
+
+    def setup(self, mode, n):
+        normal = np.zeros((n, n), int)
+        unmasked = np.ma.zeros((n, n), int)
+        masked = np.ma.array(normal, mask=True)
+
+        mode_parts = mode.split('+')
+        base = mode_parts[0]
+        promote = 'masked' in mode_parts[1:]
+
+        if base == 'ndarray':
+            args = 10 * (normal,)
+        elif base == 'unmasked':
+            args = 10 * (unmasked,)
+        else:
+            args = 10 * (masked,)
+
+        if promote:
+            args = args[:-1] + (masked,)
+
+        self.args = args
+
+    def time_it(self, mode, n):
+        np.ma.concatenate(self.args)
