fix: faster picking when slices are made

Author: henryiii

docs/CHANGELOG.md

@@ -11,7 +11,7 @@
 
 #### Bug fixes
 * Support custom setters on AxesTuple subclasses. [#627][]
-* Faster picking if slices are not also used [#645][]
+* Faster picking if slices are not also used [#645][] or if they are [#648][] (1000x or more in some cases)
 * Throw an error when an AxesTuple setter is the wrong length (inspired by zip strict in Python 3.10) [#627][]
 * Fix error thrown on comparison with axis and non-axis object [#631][]
 * Static typing no longer thinks `storage=` is required [#604][]
@@ -29,6 +29,7 @@
 [#636]: https://github.com/scikit-hep/boost-histogram/pull/636
 [#645]: https://github.com/scikit-hep/boost-histogram/pull/645
 [#647]: https://github.com/scikit-hep/boost-histogram/pull/647
+[#648]: https://github.com/scikit-hep/boost-histogram/pull/648
 
 ## Version 1.1
 

pyproject.toml

@@ -11,7 +11,7 @@ write_to = "src/boost_histogram/version.py"
 
 [tool.pytest.ini_options]
 junit_family = "xunit2"
-addopts = "--benchmark-disable -Wd --strict-markers"
+addopts = "--benchmark-disable -Wd --strict-markers -ra"
 xfail_strict = true
 testpaths = ["tests"]
 required_plugins = ["pytest-benchmark"]

src/boost_histogram/_core/algorithm.pyi

@@ -2,6 +2,7 @@ import enum
 import typing
 
 class reduce_command:
+    iaxis: int
     def __repr__(self) -> str: ...
 
 class slice_mode(enum.Enum):

src/boost_histogram/_internal/hist.py

@@ -841,17 +841,33 @@ def __getitem__(  # noqa: C901
                 assert isinstance(stop, int)
                 slices.append(_core.algorithm.slice_and_rebin(i, start, stop, merge))
 
-        if slices:
-            logger.debug("Reduce with %s", slices)
-            reduced = self._hist.reduce(*slices)
-        elif pick_set or pick_each or integrations:
-            # Can avoid a copy in these cases, will be copied anyway
-            logger.debug("Reduce is empty, but picking or slicing, so no copy needed")
+        # Will be updated below
+        if slices or pick_set or pick_each or integrations:
             reduced = self._hist
         else:
-            logger.debug("Reduce is empty, just making a copy")
+            logger.debug("Reduce actions are all empty, just making a copy")
             reduced = copy.copy(self._hist)
 
+        if pick_each:
+            tuple_slice = tuple(
+                pick_each.get(i, slice(None)) for i in range(reduced.rank())
+            )
+            logger.debug("Slices for pick each: %s", tuple_slice)
+            axes = [
+                reduced.axis(i) for i in range(reduced.rank()) if i not in pick_each
+            ]
+            logger.debug("Axes: %s", axes)
+            new_reduced = reduced.__class__(axes)
+            new_reduced.view(flow=True)[...] = reduced.view(flow=True)[tuple_slice]
+            reduced = new_reduced
+            integrations = {i - sum(j <= i for j in pick_each) for i in integrations}
+            for slice_ in slices:
+                slice_.iaxis -= sum(j <= slice_.iaxis for j in pick_each)
+
+        if slices:
+            logger.debug("Reduce with %s", slices)
+            reduced = reduced.reduce(*slices)
+
         if pick_set:
             warnings.warn(
                 "List indexing selection is experimental. Removed bins are not placed in overflow."
@@ -880,20 +896,6 @@ def __getitem__(  # noqa: C901
             new_reduced.view(flow=True)[...] = reduced_view
             reduced = new_reduced
 
-        if pick_each:
-            tuple_slice = tuple(
-                pick_each.get(i, slice(None)) for i in range(reduced.rank())
-            )
-            logger.debug("Slices for pick each: %s", tuple_slice)
-            axes = [
-                reduced.axis(i) for i in range(reduced.rank()) if i not in pick_each
-            ]
-            logger.debug("Axes: %s", axes)
-            new_reduced = reduced.__class__(axes)
-            new_reduced.view(flow=True)[...] = reduced.view(flow=True)[tuple_slice]
-            reduced = new_reduced
-            integrations = {i - sum(j <= i for j in pick_each) for i in integrations}
-
         if integrations:
             projections = [i for i in range(reduced.rank()) if i not in integrations]
             reduced = reduced.project(*projections)

src/register_algorithm.cpp

@@ -10,6 +10,7 @@
 void register_algorithms(py::module& algorithm) {
     py::class_<bh::algorithm::reduce_command>(algorithm, "reduce_command")
         .def(py::init<bh::algorithm::reduce_command>())
+        .def_readwrite("iaxis", &bh::algorithm::reduce_command::iaxis)
         .def("__repr__", [](const bh::algorithm::reduce_command& self) {
             using range_t = bh::algorithm::reduce_command::range_t;