fix: array copy and logic for obtaining variances and values (#488)

* fix: array copy and logic

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: new logic with values and flow

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* tests: expand test to incorporate hists _without_ variance

* tests: fix parametrized histplot flow variances test

* ci: fix mypy

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: andrzejnovak <[email protected]>

Author: 3 people

Diff for: src/mplhep/plot.py

@@ -31,7 +31,6 @@
 ErrorBarArtists = namedtuple("ErrorBarArtists", "errorbar")
 ColormeshArtists = namedtuple("ColormeshArtists", "pcolormesh cbar text")
 
-
 Hist1DArtists = Union[StairsArtists, ErrorBarArtists]
 Hist2DArtists = ColormeshArtists
 
@@ -188,36 +187,40 @@ def histplot(
     plottables = []
     flow_bins = final_bins
     for i, h in enumerate(hists):
-        value, variance = h.values().copy(), h.variances()
-        if variance is not None:
-            variance = variance.copy()
-        underflow, overflow = 0, 0
-        underflowv, overflowv = 0, 0
-        # One sided flow bins - hist
-        if hasattr(h, "axes") and hasattr(h.axes[0], "traits"):
-            if hasattr(h.axes[0].traits, "overflow") and h.axes[0].traits.overflow:
-                overflow = h.values(flow=True)[-1].copy()
-                overflowv = h.variances(flow=True)[-1].copy()
-            if hasattr(h.axes[0].traits, "underflow") and h.axes[0].traits.underflow:
-                underflow = h.values(flow=True)[0].copy()
-                underflowv = h.variances(flow=True)[0].copy()
+        value, variance = np.copy(h.values()), h.variances()
+        if has_variances := variance is not None:
+            variance = np.copy(variance)
+        underflow, overflow = 0.0, 0.0
+        underflowv, overflowv = 0.0, 0.0
+        # One sided flow bins - hist (uproot hist does not have the over- or underflow traits)
+        if (
+            hasattr(h, "axes")
+            and (traits := getattr(h.axes[0], "traits", None)) is not None
+            and hasattr(traits, "underflow")
+            and hasattr(traits, "overflow")
+        ):
+            if traits.overflow:
+                overflow = np.copy(h.values(flow=True))[-1]
+                if has_variances:
+                    overflowv = np.copy(h.variances(flow=True))[-1]
+            if traits.underflow:
+                underflow = np.copy(h.values(flow=True))[0]
+                if has_variances:
+                    underflowv = np.copy(h.variances(flow=True))[0]
         # Both flow bins exist - uproot
-        if hasattr(h, "values") and "flow" in inspect.getfullargspec(h.values).args:
+        elif hasattr(h, "values") and "flow" in inspect.getfullargspec(h.values).args:
             if len(h.values()) + 2 == len(
                 h.values(flow=True)
             ):  # easy case, both over/under
                 underflow, overflow = (
-                    h.values(flow=True)[0].copy(),
-                    h.values(flow=True)[-1].copy(),
-                )
-                underflowv, overflowv = (
-                    h.variances(flow=True)[0].copy(),
-                    h.variances(flow=True)[-1].copy(),
+                    np.copy(h.values(flow=True))[0],
+                    np.copy(h.values(flow=True))[-1],
                 )
-        # Bins cannot be parsed
-        else:
-            underflow, overflow = 0, 0
-            underflowv, overflowv = 0, 0
+                if has_variances:
+                    underflowv, overflowv = (
+                        np.copy(h.variances(flow=True))[0],
+                        np.copy(h.variances(flow=True))[-1],
+                    )
 
         # Set plottables
         if flow == "none":
@@ -228,23 +231,27 @@ def histplot(
             _flow_bin_size = np.max(
                 [0.05 * (final_bins[-1] - final_bins[0]), np.mean(np.diff(final_bins))]
             )
-            flow_bins = final_bins.copy()
+            flow_bins = np.copy(final_bins)
             if underflow > 0:
                 flow_bins = np.r_[flow_bins[0] - _flow_bin_size, flow_bins]
                 value = np.r_[underflow, value]
-                variance = np.r_[underflowv, variance]
+                if has_variances:
+                    variance = np.r_[underflowv, variance]
             if overflow > 0:
                 flow_bins = np.r_[flow_bins, flow_bins[-1] + _flow_bin_size]
                 value = np.r_[value, overflow]
-                variance = np.r_[variance, overflowv]
+                if has_variances:
+                    variance = np.r_[variance, overflowv]
             plottables.append(Plottable(value, edges=flow_bins, variances=variance))
         elif flow == "sum":
             if underflow > 0:
                 value[0] += underflow
-                variance[0] += underflowv
+                if has_variances:
+                    variance[0] += underflowv
             if overflow > 0:
                 value[-1] += overflow
-                variance[-1] += overflowv
+                if has_variances:
+                    variance[-1] += overflowv
             plottables.append(Plottable(value, edges=final_bins, variances=variance))
         else:
             plottables.append(Plottable(value, edges=final_bins, variances=variance))
@@ -675,7 +682,7 @@ def hist2dplot(
 
     # TODO: use Histogram everywhere
 
-    H = h.values().copy()
+    H = np.copy(h.values())
     xbins, xtick_labels = get_plottable_protocol_bins(h.axes[0])
     ybins, ytick_labels = get_plottable_protocol_bins(h.axes[1])
     # Show under/overflow bins
@@ -715,7 +722,7 @@ def hist2dplot(
             H = padded
             xbins, ybins = pxbins, pybins
     elif flow == "sum":
-        H = h.values().copy()
+        H = np.copy(h.values())
         # Sum borders
         H[0], H[-1] = (
             H[0] + h.values(flow=True)[0, 1:-1],

Diff for: tests/baseline/test_histplot_hist_flow_no_variances.png

@@ -142,26 +142,29 @@ def test_histplot_flow():
     return fig
 
 
+@pytest.mark.parametrize("variances", [True, False], ids=["variances", "no_variances"])
 @pytest.mark.mpl_image_compare(style="default")
-def test_histplot_hist_flow():
+def test_histplot_hist_flow(variances):
     np.random.seed(0)
     entries = np.random.normal(10, 3, 400)
-    h = hist.new.Reg(20, 5, 15, name="x", flow=True).Weight()
-    h2 = hist.new.Reg(20, 5, 15, name="x", underflow=True, overflow=False).Weight()
-    h3 = hist.new.Reg(20, 5, 15, name="x", underflow=False, overflow=True).Weight()
-    h4 = hist.new.Reg(20, 5, 15, name="x", flow=False).Weight()
+    hist_constr = [
+        hist.new.Reg(20, 5, 15, name="x", flow=True),
+        hist.new.Reg(20, 5, 15, name="x", underflow=True, overflow=False),
+        hist.new.Reg(20, 5, 15, name="x", underflow=False, overflow=True),
+        hist.new.Reg(20, 5, 15, name="x", flow=False),
+    ]
+    if variances:
+        hists = [h.Weight() for h in hist_constr]
+    else:
+        hists = [h.Double() for h in hist_constr]
+    for h in hists:
+        h.fill(entries, weight=np.ones_like(entries))
 
-    h.fill(entries)
-    h2.fill(entries)
-    h3.fill(entries)
-    h4.fill(entries)
     fig, axs = plt.subplots(2, 2, sharey=True, figsize=(10, 10))
     axs = axs.flatten()
 
-    hep.histplot(h, ax=axs[0], flow="show")
-    hep.histplot(h2, ax=axs[1], flow="show")
-    hep.histplot(h3, ax=axs[2], flow="show")
-    hep.histplot(h4, ax=axs[3], flow="show")
+    for i, h in enumerate(hists):
+        hep.histplot(h, ax=axs[i], flow="show", yerr=variances)
 
     axs[0].set_title("Two-side overflow", fontsize=18)
     axs[1].set_title("Left-side overflow", fontsize=18)
@@ -187,17 +190,17 @@ def test_histplot_uproot_flow():
     h4.fill(entries[(entries > 5) & (entries < 15)])
     import uproot
 
-    f = uproot.recreate("flow_th1.root")
-    f["h"] = h
-    f["h2"] = h2
-    f["h3"] = h3
-    f["h4"] = h4
-
-    f = uproot.open("flow_th1.root")
-    h = f["h"]
-    h2 = f["h2"]
-    h3 = f["h3"]
-    h4 = f["h4"]
+    with uproot.recreate("flow_th1.root") as f:
+        f["h"] = h
+        f["h2"] = h2
+        f["h3"] = h3
+        f["h4"] = h4
+
+    with uproot.open("flow_th1.root") as f:
+        h = f["h"]
+        h2 = f["h2"]
+        h3 = f["h3"]
+        h4 = f["h4"]
 
     fig, axs = plt.subplots(2, 2, sharey=True, figsize=(10, 10))
     axs = axs.flatten()