fix: support setting with histograms (#1036)

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

Author: henryiii

pyproject.toml

@@ -145,7 +145,6 @@ minversion = "6.0"
 junit_family = "xunit2"
 addopts = [
   "-ra",
-  "--showlocals",
   "--strict-markers",
   "--strict-config",
   "--import-mode=importlib",

src/boost_histogram/axis/__init__.py

@@ -189,6 +189,16 @@ def __iter__(
     ) -> Iterator[float] | Iterator[str] | Iterator[tuple[float, float]]:
         return self._ax.__iter__()  # type: ignore[no-any-return]
 
+    def _process_callable(self, value: AxCallOrInt | None, *, default: int) -> int:
+        """
+        This processes a callable in start or stop. None gets replaced by default.
+        """
+        if value is None:
+            return default
+        if callable(value):
+            return value(self)
+        return value
+
     def _process_loc(
         self, start: AxCallOrInt | None, stop: AxCallOrInt | None
     ) -> tuple[int, int]:
@@ -201,18 +211,15 @@ def _process_loc(
         is turned off if underflow is not None.
         """
 
-        def _process_internal(item: AxCallOrInt | None, default: int) -> int:
-            return default if item is None else item(self) if callable(item) else item
-
         underflow = -1 if self._ax.traits_underflow else 0
         overflow = 1 if self._ax.traits_overflow else 0
 
         # Non-ordered axes only use flow if integrating from None to None
         if not self._ax.traits_ordered and not (start is None and stop is None):
             overflow = 0
 
-        begin = _process_internal(start, underflow)
-        end = _process_internal(stop, len(self) + overflow)
+        begin = self._process_callable(start, default=underflow)
+        end = self._process_callable(stop, default=len(self) + overflow)
 
         return begin, end
 

src/boost_histogram/histogram.py

@@ -1128,37 +1128,33 @@ def __setitem__(self, index: IndexingExpr, value: ArrayLike | Accumulator) -> No
         If an array is given that does not match, if it does match the
         with-overflow size, it fills that.
 
-        PLANNED (not yet supported):
-
             h[a:] = h2
 
         If another histogram is given, that must either match with or without
         overflow, where the overflow bins must be overflow bins (that is,
         you cannot set a histogram's flow bins from another histogram that
-        is 2 larger). Bin edges must be a close match, as well. If you don't
-        want this level of type safety, just use ``h[...] = h2.view()``.
+        is 2 larger). If you don't want this level of type safety, just use
+        ``h[...] = h2.view()``.
         """
         indexes = self._compute_commonindex(index)
 
-        if isinstance(value, Histogram):
-            raise TypeError("Not supported yet")
-
-        value = np.asarray(value)
+        in_array = np.asarray(value)
         view: Any = self.view(flow=True)
 
         value_shape: tuple[int, ...]
+
         # Support raw arrays for accumulators, the final dimension is the constructor values
         if (
-            value.ndim > 0
+            in_array.ndim > 0
             and len(view.dtype) > 0
-            and len(value.dtype) == 0
-            and len(view.dtype) == value.shape[-1]
+            and len(in_array.dtype) == 0
+            and len(view.dtype) == in_array.shape[-1]
         ):
-            value_shape = value.shape[:-1]
-            value_ndim = value.ndim - 1
+            value_shape = in_array.shape[:-1]
+            value_ndim = in_array.ndim - 1
         else:
-            value_shape = value.shape
-            value_ndim = value.ndim
+            value_shape = in_array.shape
+            value_ndim = in_array.ndim
 
         # NumPy does not broadcast partial slices, but we would need
         # to allow it (because we do allow broadcasting up dimensions)
@@ -1174,43 +1170,67 @@ def __setitem__(self, index: IndexingExpr, value: ArrayLike | Accumulator) -> No
             has_overflow = self.axes[n].traits.overflow
 
             if isinstance(request, slice):
+                # This ensures that callable start/stop are handled
+                start, stop = self.axes[n]._process_loc(request.start, request.stop)
+
                 # Only consider underflow/overflow if the endpoints are not given
-                use_underflow = has_underflow and request.start is None
-                use_overflow = has_overflow and request.stop is None
+                use_underflow = has_underflow and start < 0
+                use_overflow = has_overflow and stop > len(self.axes[n])
+
+                # If the input is a histogram, we need to exactly match underflow/overflow
+                if isinstance(value, Histogram):
+                    in_underflow = value.axes[n].traits.underflow
+                    in_overflow = value.axes[n].traits.overflow
+
+                    if use_underflow != in_underflow or use_overflow != in_overflow:
+                        msg = (
+                            f"Cannot set histogram with underflow={in_underflow} and overflow={in_overflow} "
+                            f"to a histogram slice with underflow={use_underflow} and overflow={use_overflow}"
+                        )
+                        raise ValueError(msg)
+
+                # Convert to non-flow coordinates
+                start_real = start + 1 if has_underflow else start
+                stop_real = stop + 1 if has_underflow else stop
 
-                # Make the limits explicit since we may need to shift them
-                start = 0 if request.start is None else request.start
-                stop = len(self.axes[n]) if request.stop is None else request.stop
-                request_len = stop - start
+                # This is the total requested length without flow bins
+                request_len = min(stop, len(self.axes[n])) - max(start, 0)
 
                 # If set to a scalar, then treat it like broadcasting without flow bins
                 # Normal requests here too
-                if value_ndim == 0 or request_len == value_shape[value_n]:
-                    start += has_underflow
-                    stop += has_underflow
+                # Also single element broadcasting
+                if (
+                    value_ndim == 0
+                    or request_len == value_shape[value_n]
+                    or value_shape[value_n] == 1
+                ):
+                    start_real += 1 if start < 0 else 0
+                    stop_real -= 1 if stop > len(self.axes[n]) else 0
 
                 # Expanded setting
                 elif request_len + use_underflow + use_overflow == value_shape[value_n]:
-                    start += has_underflow and not use_underflow
-                    stop += has_underflow + (has_overflow and use_overflow)
-
-                # Single element broadcasting
-                elif value_shape[value_n] == 1:
-                    start += has_underflow
-                    stop += has_underflow
+                    pass
 
                 else:
                     msg = f"Mismatched shapes in dimension {n}"
                     msg += f", {value_shape[n]} != {request_len}"
                     if use_underflow or use_overflow:
                         msg += f" or {request_len + use_underflow + use_overflow}"
                     raise ValueError(msg)
-                indexes[n] = slice(start, stop, request.step)
+                logger.debug(
+                    "__setitem__: axis %i, start: %i (actual %i), stop: %i (actual %i)",
+                    n,
+                    start,
+                    start_real,
+                    stop,
+                    stop_real,
+                )
+                indexes[n] = slice(start_real, stop_real, request.step)
                 value_n += 1
             else:
                 indexes[n] = request + has_underflow
 
-        view[tuple(indexes)] = value
+        view[tuple(indexes)] = in_array
 
     def project(self, *args: int) -> Self | float | Accumulator:
         """

tests/CMakeLists.txt

@@ -21,7 +21,7 @@ foreach(TEST_FILE IN LISTS BOOST_HIST_PY_TESTS)
   get_filename_component(TEST_NAME "${TEST_FILE}" NAME_WE)
   add_test(
     NAME ${TEST_NAME}
-    COMMAND ${Python_EXECUTABLE} -m pytest "${TEST_FILE}" --rootdir=.
+    COMMAND ${Python_EXECUTABLE} -m pytest "${TEST_FILE}" --rootdir=. --showlocals
     WORKING_DIRECTORY "${PROJECT_BINARY_DIR}")
   set_tests_properties(${TEST_NAME} PROPERTIES SKIP_RETURN_CODE 5)
 endforeach()

tests/test_histogram_indexing.py

@@ -295,6 +295,18 @@ def test_set_range_with_scalar():
     assert h[5] == 0
 
 
+def test_set_range_with_scalar_callable():
+    h = bh.Histogram(bh.axis.Integer(0, 10))
+    h[2:len] = 42
+
+    assert h[1] == 0
+    assert h[2] == 42
+    assert h[3] == 42
+    assert h[4] == 42
+    assert h[5] == 42
+    assert h[bh.overflow] == 0
+
+
 def test_set_all_with_scalar():
     h = bh.Histogram(bh.axis.Integer(0, 10))
     h[:] = 42
@@ -483,3 +495,62 @@ def test_large_index():
     )
     assert h.axes[0].value(6) == 99_999_001
     assert h.axes[0].index(99_999_001) == 6
+
+
+def test_scaling_slice():
+    h = bh.Histogram(bh.axis.Regular(3, 0, 3), bh.axis.StrCategory(["a", "b"]))
+    h.fill([1, 1, 2], "a")
+    h.fill([0], "b")
+
+    h[:, bh.loc("a")] *= 2
+
+    assert h[1, 0] == approx(4)
+    assert h[2, 0] == approx(2)
+    assert h[0, 1] == approx(1)
+
+
+def test_scaling_slice_weight():
+    h = bh.Histogram(
+        bh.axis.Regular(3, 0, 3),
+        bh.axis.StrCategory(["a", "b"]),
+        storage=bh.storage.Weight(),
+    )
+    h.fill([1, 1, 2], "a")
+    h.fill([0], "b")
+
+    h[:, bh.loc("a")] *= 2
+
+    assert h[1, 0].value == approx(4)
+    assert h[2, 0].value == approx(2)
+    assert h[0, 1].value == approx(1)
+
+
+def test_setting_histogram_mismatch():
+    h = bh.Histogram(bh.axis.Regular(10, 0, 10), bh.axis.Integer(0, 20))
+
+    h[:, 0] = bh.Histogram(bh.axis.Regular(10, 0, 10))
+    h[0:, 0] = bh.Histogram(bh.axis.Regular(10, 0, 10, underflow=False))
+    h[:len, 0] = bh.Histogram(bh.axis.Regular(10, 0, 10, overflow=False))
+    h[0:len, 0] = bh.Histogram(
+        bh.axis.Regular(10, 0, 10, underflow=False, overflow=False)
+    )
+    with pytest.raises(ValueError, match="Cannot set histogram with underflow"):
+        h[0:, 0] = bh.Histogram(bh.axis.Regular(10, 0, 10))
+    with pytest.raises(ValueError, match="Cannot set histogram with underflow"):
+        h[:len, 0] = bh.Histogram(bh.axis.Regular(10, 0, 10))
+    with pytest.raises(ValueError, match="Cannot set histogram with underflow"):
+        h[:, 0] = bh.Histogram(bh.axis.Regular(10, 0, 10, underflow=False))
+    with pytest.raises(ValueError, match="Cannot set histogram with underflow"):
+        h[:, 0] = bh.Histogram(bh.axis.Regular(10, 0, 10, overflow=False))
+    with pytest.raises(ValueError, match="Cannot set histogram with underflow"):
+        h[:, 0] = bh.Histogram(
+            bh.axis.Regular(10, 0, 10, underflow=False, overflow=False)
+        )
+    with pytest.raises(ValueError, match="Cannot set histogram with underflow"):
+        h[0:, 0] = bh.Histogram(
+            bh.axis.Regular(10, 0, 10, underflow=False, overflow=False)
+        )
+    with pytest.raises(ValueError, match="Cannot set histogram with underflow"):
+        h[:len, 0] = bh.Histogram(
+            bh.axis.Regular(10, 0, 10, underflow=False, overflow=False)
+        )