tests: add benchmarking test for cat picking

scikit-hep · Sep 15, 2021 · 83520f1 · 83520f1
1 parent f3be0d3
commit 83520f1
Showing 1 changed file with 36 additions and 0 deletions.
diff --git a/tests/test_benchmark_category_axis.py b/tests/test_benchmark_category_axis.py
@@ -34,3 +34,39 @@ 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(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)