diff --git a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_uhi.py b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_uhi.py index 23d11a6f3e897..5faebd583396a 100644 --- a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_uhi.py +++ b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_uhi.py @@ -169,9 +169,6 @@ def _process_index_for_axis(self, index, axis, include_flow_bins=False, is_slice if index == -1: return _overflow(self, axis) - 1 - if index == _overflow(self, axis): - return index + (1 if include_flow_bins else 0) - # Shift the indices by 1 to align with the UHI convention, # where 0 corresponds to the first bin, unlike ROOT where 0 represents underflow and 1 is the first bin. nbins = _get_axis_len(self, axis) + (1 if is_slice_stop else 0) @@ -379,9 +376,16 @@ def _setitem(self, index, value): _slice_set(self, uhi_index, index, value) +def _iter(self): + array = _values_by_copy(self, include_flow_bins=True) + for val in array.flat: + yield val.item() + + def _add_indexing_features(klass: Any) -> None: klass.__getitem__ = _getitem klass.__setitem__ = _setitem + klass.__iter__ = _iter """ @@ -492,18 +496,20 @@ def _values_default(self) -> np.typing.NDArray[Any]: # noqa: F821 # Special case for TH1K: we need the array length to correspond to the number of bins # according to the UHI plotting protocol -def _values_by_copy(self) -> np.typing.NDArray[Any]: # noqa: F821 +def _values_by_copy(self, include_flow_bins=False) -> np.typing.NDArray[Any]: # noqa: F821 from itertools import product import numpy as np + offset = 0 if include_flow_bins else 1 dimensions = [ - range(1, _get_axis_len(self, axis, include_flow_bins=False) + 1) for axis in range(self.GetDimension()) + range(offset, _get_axis_len(self, axis, include_flow_bins=include_flow_bins) + offset) + for axis in range(self.GetDimension()) ] bin_combinations = product(*dimensions) return np.array([self.GetBinContent(*bin) for bin in bin_combinations]).reshape( - _shape(self, include_flow_bins=False) + _shape(self, include_flow_bins=include_flow_bins) ) diff --git a/bindings/pyroot/pythonizations/test/uhi_indexing.py b/bindings/pyroot/pythonizations/test/uhi_indexing.py index 8ecba4b7d9026..4fb93986c0413 100644 --- a/bindings/pyroot/pythonizations/test/uhi_indexing.py +++ b/bindings/pyroot/pythonizations/test/uhi_indexing.py @@ -328,6 +328,16 @@ def test_equality(self, hist_setup): hist_full_slice = hist_setup[...] assert hist_setup != hist_full_slice + def test_list_iter(self, hist_setup): + import numpy as np + + if _special_setting(hist_setup): + pytest.skip("Setting cannot be tested here") + + expected = np.full(_shape(hist_setup), 3, dtype=np.int64) + hist_setup[...] = expected + assert list(hist_setup) == expected.flatten().tolist() + if __name__ == "__main__": pytest.main(args=[__file__])