tests: add benchmarking test for cat picking

scikit-hep · Sep 15, 2021 · b07caa0 · b07caa0
1 parent 4c7e151
commit b07caa0
Showing 1 changed file with 54 additions and 0 deletions.
diff --git a/tests/test_benchmark_category_axis.py b/tests/test_benchmark_category_axis.py
@@ -34,3 +34,57 @@ def run(h, data):
         h.fill(data)
 
     benchmark(run, h, tuple(values) if dtype is tuple else values.astype(dtype))
+
+
+@pytest.mark.benchmark(group="Pick")
+def test_pick_only(benchmark):
+
+    h = bh.Histogram(
+        bh.axis.StrCategory([str(i) for i in range(100)]),
+        bh.axis.StrCategory([str(i) for i in range(100)]),
+        bh.axis.StrCategory([str(i) for i in range(100)]),
+        bh.axis.Regular(100, 0, 500),
+    )
+
+    h[...] = 1.0
+
+    def run(h):
+        return h[bh.loc("42"), bh.loc("42"), bh.loc("42"), :].view()
+
+    benchmark(get_slice, h)
+
+
+@pytest.mark.benchmark(group="Pick")
+def test_pick_only(benchmark):
+
+    h = bh.Histogram(
+        bh.axis.StrCategory([str(i) for i in range(32)]),
+        bh.axis.StrCategory([str(i) for i in range(32)]),
+        bh.axis.StrCategory([str(i) for i in range(32)]),
+        bh.axis.Regular(32, 0, 320),
+    )
+
+    h[...] = 1.0
+
+    def run(h):
+        return h[bh.loc("13"), bh.loc("13"), bh.loc("13"), :].view()
+
+    benchmark(run, h)
+
+
+@pytest.mark.benchmark(group="Pick")
+def test_pick_and_slice(benchmark):
+
+    h = bh.Histogram(
+        bh.axis.StrCategory([str(i) for i in range(32)]),
+        bh.axis.StrCategory([str(i) for i in range(32)]),
+        bh.axis.StrCategory([str(i) for i in range(32)]),
+        bh.axis.Regular(32, 0, 320),
+    )
+
+    h[...] = 1.0
+
+    def run(h):
+        return h[3:29, bh.loc("13"), bh.loc("13"), :].view()
+
+    benchmark(run, h)