fix: correctly pass through w2method (#558)

Author: andrzejnovak

src/mplhep/plot.py

@@ -178,6 +178,17 @@ def histplot(
         "hint",
         "none",
     }, "flow must be show, sum, hint, or none"
+    if hasattr(H, "values") or hasattr(H[0], "values"):  # Check for hist-like inputs
+        assert bins is None, (
+            "When plotting hist(-like) objects, specifying bins is not allowed."
+        )
+        assert w2 is None, (
+            "When plotting hist(-like) objects, specifying w2 is not allowed."
+        )
+    if w2 is not None:
+        assert np.array(w2).shape == np.array(H).shape, (
+            "w2 must have the same shape as H"
+        )
 
     # Convert 1/0 etc to real bools
     stack = bool(stack)

src/mplhep/utils.py

@@ -273,11 +273,15 @@ def get_plottables(
 
     if w2 is not None:
         for _w2, _plottable in zip(
-            w2.reshape(len(plottables), len(final_bins) - 1), plottables
+            np.array(w2).reshape(len(plottables), len(final_bins) - 1), plottables
         ):
             _plottable.variances = _w2
             _plottable.method = w2method
 
+    for _plottable in plottables:
+        if _plottable.variances is not None:
+            _plottable.method = w2method
+
     if w2 is not None and yerr is not None:
         msg = "Can only supply errors or w2"
         raise ValueError(msg)
@@ -417,7 +421,9 @@ def norm_stack_plottables(plottables, bins, stack=False, density=False, binwnorm
 
 
 class Plottable:
-    def __init__(self, values, *, edges=None, variances=None, yerr=None):
+    def __init__(
+        self, values, *, edges=None, variances=None, yerr=None, w2method="poisson"
+    ):
         self._values = np.array(values).astype(float)
         self.variances = None
         self._variances = None
@@ -434,7 +440,7 @@ def __init__(self, values, *, edges=None, variances=None, yerr=None):
         if self.edges is None:
             self.edges = np.arange(len(values) + 1)
         self.centers = self.edges[:-1] + np.diff(self.edges) / 2
-        self.method = "poisson"
+        self.method = w2method
 
         self.yerr = yerr
         assert self.variances is None or self.yerr is None
@@ -470,12 +476,11 @@ def errors(self, method=None):
                     method = "poisson"
                 else:
                     method = "sqrt"
-
         if self._errors_present:
             return
 
-        def sqrt_method(values, _):
-            return values - np.sqrt(values), values + np.sqrt(values)
+        def sqrt_method(values, variances):
+            return values - np.sqrt(variances), values + np.sqrt(variances)
 
         def calculate_relative(method_fcn, variances):
             return np.abs(method_fcn(self.values, variances) - self.values)

tests/baseline/test_histplot_w2_methods.png

@@ -633,6 +633,46 @@ def test_histplot_w2():
     return fig
 
 
+@pytest.mark.mpl_image_compare(style="default", remove_text=True)
+def test_histplot_w2_methods():
+    htype1 = [10, 20, 30, 40, 30, 20, 10, 1, 0, 1, 0]
+    htype1_w2 = [10, 20, 30, 40, 30, 20, 50, 1, 0, 1, 0]
+    np.random.seed(0)
+    htype2 = hist.new.Reg(11, 0, 11).Weight().fill(np.random.normal(3, 2, 100))
+
+    def fcn1(w, _):
+        return np.maximum(0, w - np.ones_like(w) * 3), w + np.ones_like(w) * 3
+
+    def fcn2(w, _):
+        return w - np.ones_like(w) * 0.2 * np.mean(w), w + np.ones_like(
+            w
+        ) * 0.2 * np.mean(w)
+
+    fig, axs = plt.subplots(2, 3, figsize=(12, 8))
+    for ax, method in zip(axs.flatten(), [None, "poisson", "sqrt", fcn1, fcn2]):
+        hep.histplot(htype1, w2=htype1_w2, w2method=method, ax=ax, label="With w2")
+        hep.histplot(htype1, w2method=method, ax=ax, label="No w2 passed")
+        htype2.plot(w2method=method, ax=ax, label="Hist")
+        ax.set_title(str(method))
+        ax.legend()
+    return fig
+
+
+@pytest.mark.mpl_image_compare(style="default", remove_text=True)
+def test_histplot_w2_poisson_handling():
+    np.random.seed(0)
+    evts = np.random.normal(2, 2, 100)
+    weights = np.random.uniform(0.99, 1.01, 100)
+    htype1 = hist.new.Reg(5, 0, 5).Weight().fill(evts)
+    htype2 = hist.new.Reg(5, 0, 5).Weight().fill(evts, weight=weights)
+
+    fig, ax = plt.subplots()
+    htype1.plot(ax=ax, histtype="errorbar", capsize=4, label="Raw counts")
+    htype2.plot(ax=ax, label="Weighted counts")
+    ax.legend()
+    return fig
+
+
 @pytest.mark.mpl_image_compare(style="default", remove_text=True)
 def test_histplot_types():
     hs, bins = [[2, 3, 4], [5, 4, 3]], [0, 1, 2, 3]