BENCH: Add basic benchmarks for take and putmask

1 files changed, 45 insertions, 0 deletions

diff --git a/benchmarks/benchmarks/bench_itemselection.py b/benchmarks/benchmarks/bench_itemselection.py
new file mode 100644
index 000000000..27fc49e30
--- /dev/null
+++ b/benchmarks/benchmarks/bench_itemselection.py
@@ -0,0 +1,45 @@
+from __future__ import absolute_import, division, print_function
+
+from .common import Benchmark, TYPES1
+
+import numpy as np
+
+
+class Take(Benchmark):
+    params = [
+        [(1000, 1), (1000, 2), (2, 1000, 1), (1000, 3)],
+        ["raise", "wrap", "clip"],
+        TYPES1]
+    param_names = ["shape", "mode", "dtype"]
+
+    def setup(self, shape, mode, dtype):
+        self.arr = np.ones(shape, dtype)
+        self.indices = np.arange(1000)
+
+    def time_contiguous(self, shape, mode, dtype):
+        self.arr.take(self.indices, axis=-2, mode=mode)
+
+
+class PutMask(Benchmark):
+    params = [
+        [True, False],
+        TYPES1]
+    param_names = ["values_is_scalar", "dtype"]
+
+    def setup(self, values_is_scalar, dtype):
+        if values_is_scalar:
+            self.vals = np.array(1., dtype=dtype)
+        else:
+            self.vals = np.ones(1000, dtype=dtype)
+
+        self.arr = np.ones(1000, dtype=dtype)
+
+        self.dense_mask = np.ones(1000, dtype="bool")
+        self.sparse_mask = np.zeros(1000, dtype="bool")
+
+    def time_dense(self, values_is_scalar, dtype):
+        np.putmask(self.arr, self.dense_mask, self.vals)
+
+    def time_sparse(self, values_is_scalar, dtype):
+        np.putmask(self.arr, self.sparse_mask, self.vals)
+