fix : flow plot bug on axis range for histplot and hist2dplot (#432)

* fix : flow plot bug on axis range for histplot and hist2dplot * feat : add test for one bin hist with flow bins #431
scikit-hep · Jul 10, 2023 · 2aa455e · 2aa455e
1 parent bd90705
commit 2aa455e
Show file tree

Hide file tree

Showing 8 changed files with 35 additions and 13 deletions.
diff --git a/src/mplhep/plot.py b/src/mplhep/plot.py
@@ -171,6 +171,8 @@ def histplot(
     # Process input
     hists = list(process_histogram_parts(H, bins))
     final_bins, xtick_labels = get_plottable_protocol_bins(hists[0].axes[0])
+    _bin_widths = np.diff(final_bins)
+    _bin_centers = final_bins[1:] - _bin_widths / float(2)
     assert final_bins.ndim == 1, "bins need to be 1 dimensional"
     _x_axes_label = ax.get_xlabel()
     x_axes_label = (
@@ -237,8 +239,8 @@ def histplot(
                         final_bins,
                         0,
                         [
-                            final_bins[0] - (final_bins[-1] - final_bins[0]) * 0.08,
-                            final_bins[0] - (final_bins[-1] - final_bins[0]) * 0.03,
+                            final_bins[0] - _bin_widths[0] * len(_bin_widths) * 0.08,
+                            final_bins[0] - _bin_widths[0] * len(_bin_widths) * 0.03,
                         ],
                     )
                 value, variance = np.insert(value, 0, np.nan), np.insert(
@@ -252,8 +254,8 @@ def histplot(
                     flow_bins = np.append(
                         flow_bins,
                         [
-                            final_bins[-1] + (final_bins[-1] - final_bins[0]) * 0.03,
-                            final_bins[-1] + (final_bins[-1] - final_bins[0]) * 0.08,
+                            final_bins[-1] + _bin_widths[-1] * len(_bin_widths) * 0.03,
+                            final_bins[-1] + _bin_widths[-1] * len(_bin_widths) * 0.08,
                         ],
                     )
                 value, variance = np.append(value, np.nan), np.append(variance, np.nan)
@@ -314,9 +316,6 @@ def iterable_not_string(arg):
             for i in range(len(_chunked_kwargs)):
                 _chunked_kwargs[i][kwarg] = kwargs[kwarg]
 
-    _bin_widths = np.diff(final_bins)
-    _bin_centers = final_bins[1:] - _bin_widths / float(2)
-
     ############################
     # # yerr calculation
     _yerr: np.ndarray | None
@@ -529,7 +528,7 @@ def iterable_not_string(arg):
             if flow == "hint":
                 ax.plot(
                     [
-                        final_bins[0] - (final_bins[-3] - final_bins[2]) * 0.03,
+                        final_bins[0] - _bin_widths[0] * len(_bin_widths) * 0.03,
                         final_bins[0],
                     ],
                     [0, 0],
@@ -550,7 +549,7 @@ def iterable_not_string(arg):
                 ax.plot(
                     [
                         final_bins[-1],
-                        final_bins[-1] + (final_bins[-3] - final_bins[2]) * 0.03,
+                        final_bins[-1] + _bin_widths[-1] * len(_bin_widths) * 0.03,
                     ],
                     [0, 0],
                     **kwargs,
@@ -798,7 +797,10 @@ def hist2dplot(
         if any(h.values(flow=True)[0] > 0):
             if flow == "hint":
                 ax.plot(
-                    [xbins[0] - (xbins[-3] - xbins[2]) * 0.03, xbins[0]],
+                    [
+                        xbins[0] - np.diff(xbins)[0] * len(np.diff(xbins)) * 0.03,
+                        xbins[0],
+                    ],
                     [0, 0],
                     transform=trans,
                     **kwargs,
@@ -809,7 +811,10 @@ def hist2dplot(
         if any(h.values(flow=True)[:, 0] > 0):
             if flow == "hint":
                 ax.plot(
-                    [xbins[-1] + (xbins[-3] - xbins[2]) * 0.03, xbins[-1]],
+                    [
+                        xbins[-1] + np.diff(xbins)[-1] * len(np.diff(xbins)) * 0.03,
+                        xbins[-1],
+                    ],
                     [0, 0],
                     transform=trans,
                     **kwargs,
@@ -820,7 +825,10 @@ def hist2dplot(
         if any(h.values(flow=True)[-1] > 0):
             if flow == "hint":
                 ax.plot(
-                    [xbins[0], xbins[0] - (xbins[-3] - xbins[2]) * 0.03],
+                    [
+                        xbins[0],
+                        xbins[0] - np.diff(xbins)[0] * len(np.diff(xbins)) * 0.03,
+                    ],
                     [1, 1],
                     transform=trans,
                     **kwargs,
@@ -832,7 +840,10 @@ def hist2dplot(
         if any(h.values(flow=True)[:, -1] > 0):
             if flow == "hint":
                 ax.plot(
-                    [xbins[-1] + (xbins[-3] - xbins[2]) * 0.03, xbins[-1]],
+                    [
+                        xbins[-1] + np.diff(xbins)[-1] * len(np.diff(xbins)) * 0.03,
+                        xbins[-1],
+                    ],
                     [1, 1],
                     transform=trans,
                     **kwargs,

diff --git a/tests/baseline/test_hist2dplot_flow.png b/tests/baseline/test_hist2dplot_flow.png
diff --git a/tests/baseline/test_histplot_flow.png b/tests/baseline/test_histplot_flow.png
diff --git a/tests/baseline/test_inputs_bh.png b/tests/baseline/test_inputs_bh.png
diff --git a/tests/baseline/test_inputs_bh_cat.png b/tests/baseline/test_inputs_bh_cat.png
diff --git a/tests/baseline/test_inputs_uproot.png b/tests/baseline/test_inputs_uproot.png
diff --git a/tests/baseline/test_onebin_hist.png b/tests/baseline/test_onebin_hist.png
diff --git a/tests/test_basic.py b/tests/test_basic.py
@@ -63,6 +63,17 @@ def test_log():
     return fig
 
 
+@pytest.mark.mpl_image_compare(style="default", remove_text=True)
+def test_onebin_hist():
+    import hist
+
+    fig, axs = plt.subplots()
+    h = hist.Hist(hist.axis.Regular(1, 0, 1))
+    h.fill([-1, 0.5])
+    hep.histplot(h, ax=axs)
+    return fig
+
+
 @pytest.mark.mpl_image_compare(style="default", remove_text=True)
 def test_histplot():
     np.random.seed(0)