Add logic for updating bin contents

scikit-hep · Feb 28, 2024 · b963e2b · b963e2b
1 parent d734fa4
commit b963e2b
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 13 deletions.
diff --git a/src/boost_histogram/_internal/hist.py b/src/boost_histogram/_internal/hist.py
@@ -892,25 +892,30 @@ def __getitem__(self: H, index: IndexingExpr) -> H | float | Accumulator:
                     reduced_view = reduced.view(flow=True)
                     new_axes_indices = [axes[i].edges[0]]
 
-                    j: int = 0
+                    j = 0
                     for group in groups:
-                        new_axes_indices += [axes[i].edges[j + 1 : j + group + 1][-1]]
+                        new_axes_indices += [axes[i].edges[j + group]]
                         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)
 
                     new_reduced = reduced.__class__(axes)
-                    new_reduced.fill(
-                        np.add.reduceat(reduced.to_numpy()[0], new_axes_indices, axis=i)
-                    )
+                    new_view = new_reduced.view(flow=True)
+
+                    j = 1
+                    for new_j, group in enumerate(groups):
+                        for _ in range(group):
+                            pos = [slice] * (i)
+                            new_view[(*pos, new_j + 1, ...)] += reduced_view[  # type: ignore[arg-type]
+                                (*pos, j, ...)  # type: ignore[arg-type]
+                            ]
+                            j += 1
+
                     reduced = new_reduced
 
         # Will be updated below

diff --git a/tests/test_histogram.py b/tests/test_histogram.py
@@ -632,13 +632,18 @@ def test_shrink_1d():
 
 def test_rebin_1d():
     h = bh.Histogram(bh.axis.Regular(20, 1, 5))
-    h.fill(1.1)
+    h.fill([1.1, 2.2, 3.3, 4.4])
 
-    hs = h[{0: slice(None, None, bh.rebin(4))}]
-    assert_array_equal(hs.view(), [1, 0, 0, 0, 0])
+    hs = h[{0: slice(None, None, bh.tag.Rebinner(4))}]
+    assert_array_equal(hs.view(), [1, 1, 1, 0, 1])
+    print("Here")
 
-    hs = h[{0: bh.rebin(4)}]
-    assert_array_equal(hs.view(), [1, 0, 0, 0, 0])
+    hs = h[{0: bh.tag.Rebinner(4)}]
+    assert_array_equal(hs.view(), [1, 1, 1, 0, 1])
+
+    hs = h[{0: bh.tag.Rebinner(groups=[1, 2, 3, 4])}]
+    assert_array_equal(hs.view(), [1, 0, 0, 1])
+    assert_array_equal(hs.axes.edges[0], [1.0, 1.2, 1.6, 2.2, 3.0])
 
 
 def test_shrink_rebin_1d():