fix: axis passthrough (#980)

Signed-off-by: Henry Schreiner <[email protected]>

docs/changelog.md

@@ -2,9 +2,31 @@
 
 ## Version 1.5
 
+### Version 1.5.2
+
+Fix for axis metadata not passing though non-uniform rebinning correctly. Flow
+bins are now preserved when doing a non-uniform rebinning. Also adds the
+ability to rebin by edges or an existing axis.
+
+#### Features
+
+- Support `edges=` and `axis=` in `bh.rebin` [#977][]
+- Allow axis to be passed through `bh.rebin` [#980][]
+
+#### Bug fixes
+
+- Axis metadata was broken when rebinning [#978][]
+- Flow bins were lost when using variable rebinning [#977][]
+
+[#977]: https://github.com/scikit-hep/boost-histogram/pull/977
+[#978]: https://github.com/scikit-hep/boost-histogram/pull/978
+[#980]: https://github.com/scikit-hep/boost-histogram/pull/980
+
 ### Version 1.5.1
 
-Make non-uniform rebinning work for Weight() and friends [#972][]
+#### Bug fixes
+
+- Make non-uniform rebinning work for Weight() and friends [#972][]
 
 [#972]: https://github.com/scikit-hep/boost-histogram/pull/972
 

pyproject.toml

@@ -66,6 +66,7 @@ dev = [
     "numpy",
     "setuptools_scm",
     "typer",
+    "uhi",
 ]
 docs = [
     "ipython",
@@ -112,14 +113,12 @@ sdist.exclude = [
   "extern/*/README.md",
 ]
 
-
 [[tool.scikit-build.generate]]
 path = "boost_histogram/version.py"
 template = '''
 __version__ = version = "$version"
 '''
 
-
 [tool.setuptools_scm]
 
 

src/boost_histogram/histogram.py

@@ -32,7 +32,7 @@
 from ._utils import cast, register
 from .axis import AxesTuple, Axis, Variable
 from .storage import Double, Storage
-from .typing import Accumulator, ArrayLike, CppHistogram
+from .typing import Accumulator, ArrayLike, CppHistogram, RebinProtocol
 from .view import MeanView, WeightedMeanView, WeightedSumView, _to_view
 
 if TYPE_CHECKING:
@@ -108,7 +108,7 @@ def __dir__() -> list[str]:
 
 CppAxis = NewType("CppAxis", object)
 
-SimpleIndexing = Union[SupportsIndex, slice]
+SimpleIndexing = Union[SupportsIndex, slice, RebinProtocol]
 InnerIndexing = Union[SimpleIndexing, Callable[[Axis], int]]
 FullInnerIndexing = Union[InnerIndexing, List[InnerIndexing]]
 IndexingWithMapping = Union[FullInnerIndexing, Mapping[int, FullInnerIndexing]]
@@ -908,11 +908,18 @@ def __getitem__(self: H, index: IndexingExpr) -> H | float | Accumulator:
             start, stop = self.axes[i]._process_loc(ind.start, ind.stop)
 
             groups = []
+            new_axis = None
             if ind != slice(None):
                 merge = 1
                 if ind.step is not None:
                     if getattr(ind.step, "factor", None) is not None:
                         merge = ind.step.factor
+                    elif (
+                        hasattr(ind.step, "axis_mapping")
+                        and (tmp_both := ind.step.axis_mapping(self.axes[i]))
+                        is not None
+                    ):
+                        groups, new_axis = tmp_both
                     elif (
                         hasattr(ind.step, "group_mapping")
                         and (tmp_groups := ind.step.group_mapping(self.axes[i]))
@@ -958,23 +965,31 @@ def __getitem__(self: H, index: IndexingExpr) -> H | float | Accumulator:
                         new_axes_indices += [axes[i].edges[j + group]]
                         j += group
 
-                    variable_axis = Variable(
-                        new_axes_indices, __dict__=axes[i].metadata
-                    )
-                    axes[i] = variable_axis._ax
+                    if new_axis is None:
+                        new_axis = Variable(
+                            new_axes_indices,
+                            __dict__=axes[i].metadata,
+                            underflow=axes[i].traits_underflow,
+                            overflow=axes[i].traits_overflow,
+                        )
+                    old_axis = axes[i]
+                    axes[i] = new_axis._ax
 
                     logger.debug("Axes: %s", axes)
 
                     new_reduced = reduced.__class__(axes)
                     new_view = new_reduced.view(flow=True)
                     j = 0
                     new_j_base = 0
-                    if self.axes[i].traits.underflow:
+
+                    if old_axis.traits_underflow and axes[i].traits_underflow:
                         groups.insert(0, 1)
-                    else:
+                    elif axes[i].traits_underflow:
                         new_j_base = 1
-                    if self.axes[i].traits.overflow:
+
+                    if old_axis.traits_overflow and axes[i].traits_overflow:
                         groups.append(1)
+
                     for new_j, group in enumerate(groups):
                         for _ in range(group):
                             pos = [slice(None)] * (i)

src/boost_histogram/tag.py

@@ -158,6 +158,15 @@ def __repr__(self) -> str:
                 break
         return return_str
 
+    # Note: this preserves the input type of `self.axis`, so is safe within a
+    # single UHI library, but not cross-library. Returns None for the axis if
+    # an axis is not provided, the caller should make an axis if that's the
+    # case.
+    def axis_mapping(
+        self, axis: PlottableAxis
+    ) -> tuple[Sequence[int], PlottableAxis | None]:
+        return (self.group_mapping(axis), self.axis)
+
     def group_mapping(self, axis: PlottableAxis) -> Sequence[int]:
         if self.groups is not None:
             if sum(self.groups) != len(axis):

src/boost_histogram/typing.py

@@ -4,9 +4,11 @@
 
 if TYPE_CHECKING:
     from builtins import ellipsis
+    from collections.abc import Sequence
 
     from numpy import ufunc as Ufunc
     from numpy.typing import ArrayLike
+    from uhi.typing.plottable import PlottableAxis
 
     from boost_histogram._core.accumulators import Mean, WeightedMean, WeightedSum
     from boost_histogram._core.hist import _BaseHistogram as CppHistogram
@@ -23,6 +25,7 @@
     "ArrayLike",
     "AxisLike",
     "CppHistogram",
+    "RebinProtocol",
     "StdIndex",
     "StrIndex",
     "Ufunc",
@@ -39,3 +42,9 @@ def __len__(self) -> int: ...
 StrIndex = Union[
     int, slice, str, "ellipsis", Tuple[Union[slice, int, str, "ellipsis"], ...]
 ]
+
+
+class RebinProtocol(Protocol):
+    def axis_mapping(
+        self, axis: PlottableAxis
+    ) -> tuple[Sequence[int], PlottableAxis | None]: ...

tests/test_histogram.py

@@ -653,10 +653,12 @@ def test_rebin_1d(metadata):
     hs = h[bh.rebin(edges=[1.0, 1.2, 1.6, 2.2, 5.0])]
     assert_array_equal(hs.view(), [1, 0, 0, 3])
     assert_array_equal(hs.axes.edges[0], [1.0, 1.2, 1.6, 2.2, 5.0])
+    assert h.axes[0].metadata is hs.axes[0].metadata
 
-    hs = h[bh.rebin(axis=bh.axis.Variable([1.0, 1.2, 1.6, 2.2, 5.0]))]
+    hs = h[bh.rebin(axis=bh.axis.Variable([1.0, 1.2, 1.6, 2.2, 5.0], metadata="hi"))]
     assert_array_equal(hs.view(), [1, 0, 0, 3])
     assert_array_equal(hs.axes.edges[0], [1.0, 1.2, 1.6, 2.2, 5.0])
+    assert hs.axes[0].metadata == "hi"
 
 
 def test_rebin_1d_flow():
@@ -667,15 +669,23 @@ def test_rebin_1d_flow():
     assert_array_equal(hs.view(flow=True), [1, 2, 2, 1])
     assert_array_equal(hs.axes.edges[0], [0.0, 3.0, 5.0])
 
+    # Flow bins are kept from the original
     h = bh.Histogram(bh.axis.Regular(5, 0, 5, underflow=False, overflow=False))
     h.fill([-1, 1.1, 2.2, 3.3, 4.4, 5.5])
     hs = h[bh.rebin(edges=[0, 3, 5.0])]
-    assert_array_equal(hs.view(flow=True), [0, 2, 2, 0])
+    assert_array_equal(hs.view(flow=True), [2, 2])
 
     h = bh.Histogram(bh.axis.Regular(5, 0, 5, underflow=True, overflow=False))
     h.fill([-1, 1.1, 2.2, 3.3, 4.4, 5.5])
     hs = h[bh.rebin(edges=[0, 3, 5.0])]
-    assert_array_equal(hs.view(flow=True), [1, 2, 2, 0])
+    assert_array_equal(hs.view(flow=True), [1, 2, 2])
+
+    h = bh.Histogram(bh.axis.Regular(5, 0, 5, underflow=True, overflow=True))
+    h.fill([-1, 1.1, 2.2, 3.3, 4.4, 5.5])
+    hs = h[
+        bh.rebin(axis=bh.axis.Variable([0, 3, 5.0], underflow=False, overflow=False))
+    ]
+    assert_array_equal(hs.view(flow=True), [2, 2])
 
 
 def test_shrink_rebin_1d():