From e7d6c764c7d49129b0f0980aeca27466efdcd245 Mon Sep 17 00:00:00 2001
From: Saransh Chopra <saransh0701@gmail.com>
Date: Fri, 2 Feb 2024 17:25:40 +0100
Subject: [PATCH] Refactor and move the rebinning logic in the loop

---
 src/boost_histogram/_internal/hist.py | 57 +++++++++++++++------------
 1 file changed, 31 insertions(+), 26 deletions(-)

diff --git a/src/boost_histogram/_internal/hist.py b/src/boost_histogram/_internal/hist.py
index 8e6f1676..ce9f0fb9 100644
--- a/src/boost_histogram/_internal/hist.py
+++ b/src/boost_histogram/_internal/hist.py
@@ -853,6 +853,7 @@ def __getitem__(self: H, index: IndexingExpr) -> H | float | Accumulator:
 
             if ind != slice(None):
                 merge = 1
+                groups = []
                 if ind.step is not None:
                     if getattr(ind.step, "factor", None) is not None:
                         merge = ind.step.factor
@@ -870,7 +871,8 @@ def __getitem__(self: H, index: IndexingExpr) -> H | float | Accumulator:
                                     i, start, stop, _core.algorithm.slice_mode.crop
                                 )
                             )
-                        continue
+                        if len(groups) == 0:
+                            continue
                     else:
                         raise IndexError(
                             "The third argument to a slice must be rebin or projection"
@@ -878,10 +880,37 @@ def __getitem__(self: H, index: IndexingExpr) -> H | float | Accumulator:
 
                 assert isinstance(start, int)
                 assert isinstance(stop, int)
-                if not (ind.step is not None and ind.step.factor is None):
+                # rebinning with factor
+                if len(groups) == 0:
                     slices.append(
                         _core.algorithm.slice_and_rebin(i, start, stop, merge)
                     )
+                # rebinning with groups
+                elif len(groups) != 0:
+                    reduced = self._hist
+                    axes = [reduced.axis(x) for x in range(reduced.rank())]
+                    reduced_view = reduced.view(flow=True)
+                    new_axes_indices = [axes[i].edges[0]]
+                    j: int = 0
+                    for group in groups:
+                        new_axes_indices += [axes[i].edges[j + 1 : j + group + 1][-1]]
+                        j = group
+
+                    variable_axis = Variable(
+                        new_axes_indices, metadata=axes[i].metadata
+                    )
+                    axes[i] = variable_axis._ax
+                    reduced_view = np.take(
+                        reduced_view, range(len(reduced_view)), axis=i
+                    )
+
+                    logger.debug("Axes: %s", axes)
+
+                    # redistribute the bin contents here
+
+                    new_reduced = reduced.__class__(axes)
+                    new_reduced.view(flow=True)[...] = reduced_view
+                    reduced = new_reduced
 
         # Will be updated below
         if slices or pick_set or pick_each or integrations:
@@ -890,30 +919,6 @@ def __getitem__(self: H, index: IndexingExpr) -> H | float | Accumulator:
             logger.debug("Reduce actions are all empty, just making a copy")
             reduced = copy.copy(self._hist)
 
-        # bin re-grouping
-        if (
-            hasattr(ind, "step")
-            and ind.step is not None
-            and ind.step.groups is not None
-        ):
-            axes = [reduced.axis(i) for i in range(reduced.rank())]
-            reduced_view = reduced.view(flow=True)
-            new_axes_indices = [axes[i].edges[0]]
-            j: int = 0
-            for group in groups:
-                new_axes_indices += [axes[i].edges[j + 1 : j + group + 1][-1]]
-                j = group
-
-            variable_axis = Variable(new_axes_indices, metadata=axes[i].metadata)
-            axes[i] = variable_axis._axis
-            reduced_view = np.take(reduced_view, range(len(reduced_view)), axis=i)
-
-            logger.debug("Axes: %s", axes)
-
-            new_reduced = reduced.__class__(axes)
-            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())