From 8ddf4633de92679557db92b98fa383afa69d0c42 Mon Sep 17 00:00:00 2001
From: Henry Schreiner <henryschreineriii@gmail.com>
Date: Tue, 18 May 2021 15:31:26 -0400
Subject: [PATCH] fix: slicing supports under/overflow

---
 docs/CHANGELOG.md                     |  5 +++-
 src/boost_histogram/_internal/hist.py | 40 ++++++++++++++++++---------
 tests/test_histogram_indexing.py      | 25 +++++++++++++++++
 3 files changed, 56 insertions(+), 14 deletions(-)

diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md
index 602c9fb1f..8a7aee6f6 100644
--- a/docs/CHANGELOG.md
+++ b/docs/CHANGELOG.md
@@ -2,7 +2,10 @@
 
 ## UPCOMING
 
-No changes yet.
+* Fix "picking" on a flow bin [#576][]
+
+[#576]: https://github.com/scikit-hep/boost-histogram/pull/576
+
 
 ## Version 1.0
 
diff --git a/src/boost_histogram/_internal/hist.py b/src/boost_histogram/_internal/hist.py
index 702b3c554..659c672c9 100644
--- a/src/boost_histogram/_internal/hist.py
+++ b/src/boost_histogram/_internal/hist.py
@@ -726,18 +726,22 @@ def __getitem__(  # noqa: C901
             except RuntimeError:
                 pass
 
-        integrations = set()
-        slices = []
+        integrations: Set[int] = set()
+        slices: List[_core.algorithm.reduce_command] = []
+        pick_each: Dict[int, int] = dict()
 
         # Compute needed slices and projections
         for i, ind in enumerate(indexes):
             if hasattr(ind, "__index__"):
-                ind = slice(ind.__index__(), ind.__index__() + 1, sum)  # type: ignore
-
+                pick_each[i] = ind.__index__() + (  # type: ignore
+                    1 if self.axes[i].traits.underflow else 0
+                )
+                continue
             elif not isinstance(ind, slice):
                 raise IndexError(
                     "Must be a slice, an integer, or follow the locator protocol."
                 )
+
             # If the dictionary brackets are forgotten, it's easy to put a slice
             # into a slice - adding a nicer error message in that case
             if any(isinstance(v, slice) for v in (ind.start, ind.stop, ind.step)):
@@ -778,15 +782,25 @@ def __getitem__(  # noqa: C901
         logger.debug("Reduce with %s", slices)
         reduced = self._hist.reduce(*slices)
 
-        if not integrations:
-            return self._new_hist(reduced)
-        projections = [i for i in range(self.ndim) if i not in integrations]
-
-        return (
-            self._new_hist(reduced.project(*projections))
-            if projections
-            else reduced.sum(flow=True)
-        )
+        if pick_each:
+            my_slice = tuple(
+                pick_each.get(i, slice(None)) for i in range(reduced.rank())
+            )
+            logger.debug("Slices: %s", my_slice)
+            axes = [
+                reduced.axis(i) for i in range(reduced.rank()) if i not in pick_each
+            ]
+            logger.debug("Axes: %s", axes)
+            new_reduced = reduced.__class__(axes)
+            new_reduced.view(flow=True)[...] = reduced.view(flow=True)[my_slice]
+            reduced = new_reduced
+            integrations = {i - sum(j <= i for j in pick_each) for i in integrations}
+
+        if integrations:
+            projections = [i for i in range(reduced.rank()) if i not in integrations]
+            reduced = reduced.project(*projections)
+
+        return self._new_hist(reduced) if reduced.rank() > 0 else reduced.sum(flow=True)
 
     def __setitem__(
         self, index: IndexingExpr, value: Union[ArrayLike, Accumulator]
diff --git a/tests/test_histogram_indexing.py b/tests/test_histogram_indexing.py
index 3429b9fb1..aeb33937d 100644
--- a/tests/test_histogram_indexing.py
+++ b/tests/test_histogram_indexing.py
@@ -1,6 +1,7 @@
 import numpy as np
 import pytest
 from numpy.testing import assert_array_equal
+from pytest import approx
 
 import boost_histogram as bh
 
@@ -351,6 +352,30 @@ def test_pick_int_category():
     assert_array_equal(h[:, :, bh.loc(7)].view(), 0)
 
 
+@pytest.mark.parametrize(
+    "ax",
+    [bh.axis.Regular(3, 0, 1), bh.axis.Variable([0, 0.3, 0.6, 1])],
+    ids=["regular", "variable"],
+)
+def test_pick_flowbin(ax):
+    w = 1e-2  # e.g. a cross section for a process
+    x = [-0.1, -0.1, 0.1, 0.1, 0.1]
+    y = [-0.1, 0.1, -0.1, -0.1, 0.1]
+
+    h = bh.Histogram(
+        ax,
+        ax,
+        storage=bh.storage.Weight(),
+    )
+    h.fill(x, y, weight=w)
+
+    uf_slice = h[bh.tag.underflow, ...]
+    assert uf_slice.values(flow=True) == approx(np.array([1, 1, 0, 0, 0]) * w)
+
+    uf_slice = h[..., bh.tag.underflow]
+    assert uf_slice.values(flow=True) == approx(np.array([1, 2, 0, 0, 0]) * w)
+
+
 def test_axes_tuple():
     h = bh.Histogram(bh.axis.Regular(10, 0, 1))
     assert isinstance(h.axes[:1], bh._internal.axestuple.AxesTuple)