Skip to content

Commit

Permalink
Refactor and move the rebinning logic in the loop
Browse files Browse the repository at this point in the history
  • Loading branch information
Saransh-cpp committed Feb 2, 2024
1 parent c73a486 commit a38965c
Showing 1 changed file with 31 additions and 26 deletions.
57 changes: 31 additions & 26 deletions src/boost_histogram/_internal/hist.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -870,18 +871,46 @@ 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"
)

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:
Expand All @@ -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())
Expand Down

0 comments on commit a38965c

Please sign in to comment.