feat: added bar/barstep histtypes

* feat: added bar type + correct legend for step type * fix: if available, yerr is shown by default for step histtype * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: back to simple line in legend for step histtype --------- Co-authored-by: Andrzej Novak <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
scikit-hep · Oct 29, 2024 · 702a55b · 702a55b
1 parent be11ed5
commit 702a55b
Show file tree

Hide file tree

Showing 8 changed files with 182 additions and 39 deletions.
diff --git a/.gitignore b/.gitignore
@@ -1,5 +1,8 @@
 # Project specific
 src/mplhep/_version.py
+*.root
+result_images/
+test/
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

diff --git a/src/mplhep/plot.py b/src/mplhep/plot.py
@@ -129,13 +129,14 @@ def histplot(
         binwnorm : float, optional
             If true, convert sum weights to bin-width-normalized, with unit equal to
                 supplied value (usually you want to specify 1.)
-        histtype: {'step', 'fill', 'band', 'errorbar'}, optional, default: "step"
+        histtype: {'step', 'fill', 'errorbar', 'bar', 'barstep', 'band'}, optional, default: "step"
             Type of histogram to plot:
-
             - "step": skyline/step/outline of a histogram using `plt.stairs <https://matplotlib.org/stable/api/_as_gen/matplotlib.axes.Axes.stairs.html#matplotlib-axes-axes-stairs>`_
             - "fill": filled histogram using `plt.stairs <https://matplotlib.org/stable/api/_as_gen/matplotlib.axes.Axes.stairs.html#matplotlib-axes-axes-stairs>`_
-            - "band": filled band spanning the yerr range of the histogram using `plt.stairs <https://matplotlib.org/stable/api/_as_gen/matplotlib.axes.Axes.stairs.html#matplotlib-axes-axes-stairs>`_
             - "errorbar": single marker histogram using `plt.errorbar <https://matplotlib.org/stable/api/_as_gen/matplotlib.axes.Axes.errorbar.html#matplotlib-axes-axes-errorbar>`_
+            - "bar": If multiple data are given the bars are arranged side by side using `plt.bar <https://matplotlib.org/stable/api/_as_gen/matplotlib.axes.Axes.bar.html#matplotlib-axes-axes-bar>`_ If only one histogram is provided, it will be treated as "fill" histtype
+            - "barstep": If multiple data are given the steps are arranged side by side using `plt.stairs <https://matplotlib.org/stable/api/_as_gen/matplotlib.axes.Axes.stairs.html#matplotlib-axes-axes-stairs>`_ . Supports yerr representation. If one histogram is provided, it will be treated as "step" histtype.
+            - "band": filled band spanning the yerr range of the histogram using `plt.stairs <https://matplotlib.org/stable/api/_as_gen/matplotlib.axes.Axes.stairs.html#matplotlib-axes-axes-stairs>`_
         xerr:  bool or float, optional
             Size of xerr if ``histtype == 'errorbar'``. If ``True``, bin-width will be used.
         label : str or list, optional
@@ -168,8 +169,8 @@ def histplot(
         raise ValueError(msg)
 
     # arg check
-    _allowed_histtype = ["fill", "step", "errorbar", "band"]
-    _err_message = f"Select 'histtype' from: {_allowed_histtype}"
+    _allowed_histtype = ["fill", "step", "errorbar", "band", "bar", "barstep"]
+    _err_message = f"Select 'histtype' from: {_allowed_histtype}, got '{histtype}'"
     assert histtype in _allowed_histtype, _err_message
     assert flow is None or flow in {
         "show",
@@ -409,6 +410,12 @@ def iterable_not_string(arg):
     ##########
     # Plotting
     return_artists: list[StairsArtists | ErrorBarArtists] = []
+
+    if histtype == "bar" and len(plottables) == 1:
+        histtype = "fill"
+    elif histtype == "barstep" and len(plottables) == 1:
+        histtype = "step"
+
     # customize color cycle assignment when stacking to match legend
     if stack:
         plottables = plottables[::-1]
@@ -423,53 +430,135 @@ def iterable_not_string(arg):
             for i in range(len(plottables)):
                 _chunked_kwargs[i].update({"color": _colors[i]})
 
-    if histtype == "step":
+    if "bar" in histtype:
+        if kwargs.get("bin_width") is None:
+            _full_bin_width = 0.8
+        else:
+            _full_bin_width = kwargs.pop("bin_width")
+        _shift = np.linspace(
+            -(_full_bin_width / 2), _full_bin_width / 2, len(plottables), endpoint=False
+        )
+        _shift += _full_bin_width / (2 * len(plottables))
+
+    if "step" in histtype:
         for i in range(len(plottables)):
             do_errors = yerr is not False and (
                 (yerr is not None or w2 is not None)
                 or (plottables[i].variances is not None)
             )
 
             _kwargs = _chunked_kwargs[i]
+
+            if _kwargs.get("bin_width"):
+                _kwargs.pop("bin_width")
+
             _label = _labels[i] if do_errors else None
             _step_label = _labels[i] if not do_errors else None
 
             _kwargs = soft_update_kwargs(_kwargs, {"linewidth": 1.5})
 
             _plot_info = plottables[i].to_stairs()
             _plot_info["baseline"] = None if not edges else 0
-            _s = ax.stairs(
-                **_plot_info,
-                label=_step_label,
-                **_kwargs,
-            )
 
-            if do_errors:
-                _kwargs = soft_update_kwargs(_kwargs, {"color": _s.get_edgecolor()})
-                _ls = _kwargs.pop("linestyle", "-")
-                _kwargs["linestyle"] = "none"
-                _plot_info = plottables[i].to_errorbar()
-                _e = ax.errorbar(
+            if _kwargs.get("color") is None:
+                _kwargs["color"] = ax._get_lines.get_next_color()  # type: ignore[attr-defined]
+
+            if histtype == "step":
+                _s = ax.stairs(
                     **_plot_info,
+                    label=_step_label,
                     **_kwargs,
                 )
-                _e_leg = ax.errorbar(
-                    [],
-                    [],
-                    yerr=1,
-                    xerr=None,
-                    color=_s.get_edgecolor(),
-                    label=_label,
-                    linestyle=_ls,
+                if do_errors:
+                    _kwargs = soft_update_kwargs(_kwargs, {"color": _s.get_edgecolor()})
+                    _ls = _kwargs.pop("linestyle", "-")
+                    _kwargs["linestyle"] = "none"
+                    _plot_info = plottables[i].to_errorbar()
+                    _e = ax.errorbar(
+                        **_plot_info,
+                        **_kwargs,
+                    )
+                    _e_leg = ax.errorbar(
+                        [],
+                        [],
+                        yerr=1,
+                        xerr=None,
+                        color=_s.get_edgecolor(),
+                        label=_label,
+                        linestyle=_ls,
+                    )
+                return_artists.append(
+                    StairsArtists(
+                        _s,
+                        _e if do_errors else None,
+                        _e_leg if do_errors else None,
+                    )
                 )
-            return_artists.append(
-                StairsArtists(
-                    _s,
-                    _e if do_errors else None,
-                    _e_leg if do_errors else None,
+                _artist = _s
+
+            # histtype = barstep
+            else:
+                if _kwargs.get("edgecolor") is None:
+                    edgecolor = _kwargs.get("color")
+                else:
+                    edgecolor = _kwargs.pop("edgecolor")
+
+                _b = ax.bar(
+                    plottables[i].centers + _shift[i],
+                    plottables[i].values,
+                    width=_full_bin_width / len(plottables),
+                    label=_step_label,
+                    align="center",
+                    edgecolor=edgecolor,
+                    fill=False,
+                    **_kwargs,
+                )
+
+                if do_errors:
+                    _ls = _kwargs.pop("linestyle", "-")
+                    # _kwargs["linestyle"] = "none"
+                    _plot_info = plottables[i].to_errorbar()
+                    _e = ax.errorbar(
+                        _plot_info["x"] + _shift[i],
+                        _plot_info["y"],
+                        yerr=_plot_info["yerr"],
+                        linestyle="none",
+                        **_kwargs,
+                    )
+                    _e_leg = ax.errorbar(
+                        [],
+                        [],
+                        yerr=1,
+                        xerr=None,
+                        color=_kwargs.get("color"),
+                        label=_label,
+                        linestyle=_ls,
+                    )
+                return_artists.append(
+                    StairsArtists(
+                        _b, _e if do_errors else None, _e_leg if do_errors else None
+                    )
                 )
+                _artist = _b  # type: ignore[assignment]
+
+    elif histtype == "bar":
+        for i in range(len(plottables)):
+            _kwargs = _chunked_kwargs[i]
+
+            if _kwargs.get("bin_width"):
+                _kwargs.pop("bin_width")
+
+            _b = ax.bar(
+                plottables[i].centers + _shift[i],
+                plottables[i].values,
+                width=_full_bin_width / len(plottables),
+                label=_labels[i],
+                align="center",
+                fill=True,
+                **_kwargs,
             )
-        _artist = _s
+            return_artists.append(StairsArtists(_b, None, None))
+        _artist = _b  # type: ignore[assignment]
 
     elif histtype == "fill":
         for i in range(len(plottables)):
@@ -531,9 +620,10 @@ def iterable_not_string(arg):
         _artist = _e[0]
 
     # Add sticky edges for autoscale
-    listy = _artist.sticky_edges.y
-    assert hasattr(listy, "append"), "cannot append to sticky edges"
-    listy.append(0)
+    if "bar" not in histtype:
+        listy = _artist.sticky_edges.y
+        assert hasattr(listy, "append"), "cannot append to sticky edges"
+        listy.append(0)
 
     if xtick_labels is None or flow == "show":
         if binticks:

diff --git a/tests/baseline/test_histplot_bar.png b/tests/baseline/test_histplot_bar.png
diff --git a/tests/baseline/test_histplot_types.png b/tests/baseline/test_histplot_types.png
diff --git a/tests/baseline/test_inputs_basic.png b/tests/baseline/test_inputs_basic.png
diff --git a/tests/baseline/test_simple_xerr.png b/tests/baseline/test_simple_xerr.png
diff --git a/tests/test_basic.py b/tests/test_basic.py
@@ -615,16 +615,65 @@ def test_histplot_w2():
 @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]
-    fig, axs = plt.subplots(3, 2, figsize=(8, 12))
+    fig, axs = plt.subplots(5, 2, figsize=(8, 16))
     axs = axs.flatten()
 
-    for i, htype in enumerate(["step", "fill", "errorbar"]):
+    for i, htype in enumerate(["step", "fill", "errorbar", "bar", "barstep"]):
         hep.histplot(hs[0], bins, yerr=True, histtype=htype, ax=axs[i * 2], alpha=0.7)
         hep.histplot(hs, bins, yerr=True, histtype=htype, ax=axs[i * 2 + 1], alpha=0.7)
 
     return fig
 
 
+@pytest.mark.mpl_image_compare(style="default", remove_text=True)
+def test_histplot_bar():
+    bins = list(range(6))
+    h1 = [1, 2, 3, 2, 1]
+    h2 = [2, 2, 2, 2, 2]
+    h3 = [2, 1, 2, 1, 2]
+    h4 = [3, 1, 2, 1, 3]
+
+    fig, axs = plt.subplots(2, 2, sharex=True, sharey=True, figsize=(10, 10))
+    axs = axs.flatten()
+
+    axs[0].set_title("Histype bar", fontsize=18)
+    hep.histplot(
+        [h1, h2, h3, h4],
+        bins,
+        histtype="bar",
+        label=["h1", "h2", "h3", "h4"],
+        ax=axs[0],
+    )
+    axs[0].legend()
+
+    axs[1].set_title("Histtype barstep", fontsize=18)
+    hep.histplot(
+        [h1, h2, h3],
+        bins,
+        histtype="barstep",
+        yerr=False,
+        label=["h1", "h2", "h3"],
+        ax=axs[1],
+    )
+    axs[1].legend()
+
+    axs[2].set_title("Histtype barstep", fontsize=18)
+    hep.histplot(
+        [h1, h2], bins, histtype="barstep", yerr=True, label=["h1", "h2"], ax=axs[2]
+    )
+    axs[2].legend()
+
+    axs[3].set_title("Histype bar", fontsize=18)
+    hep.histplot(
+        [h1, h2], bins, histtype="bar", label=["h1", "h2"], bin_width=0.2, ax=axs[3]
+    )
+    axs[3].legend()
+
+    fig.subplots_adjust(wspace=0.1)
+
+    return fig
+
+
 h = np.geomspace(1, 10, 10)
 
 

diff --git a/tests/test_mock.py b/tests/test_mock.py
@@ -46,14 +46,15 @@ def test_simple(mock_matplotlib):
     bins = [0, 1, 2, 3]
     hep.histplot(h, bins, yerr=True, label="X", ax=ax)
 
-    assert len(ax.mock_calls) == 12
+    assert len(ax.mock_calls) == 13
 
     ax.stairs.assert_called_once_with(
         values=approx([1.0, 3.0, 2.0]),
         edges=approx([0.0, 1.0, 2.0, 3.0]),
         baseline=0,
         label=None,
         linewidth=1.5,
+        color="next-color",
     )
 
     assert ax.errorbar.call_count == 2
@@ -74,7 +75,7 @@ def test_simple(mock_matplotlib):
             approx([0.82724622, 1.63270469, 1.29181456]),
             approx([2.29952656, 2.91818583, 2.63785962]),
         ],
-        color=ax.stairs().get_edgecolor(),
+        color="next-color",
         linestyle="none",
         linewidth=1.5,
     )
@@ -90,7 +91,7 @@ def test_histplot_real(mock_matplotlib):
     hep.histplot([a, b, c], bins=bins, ax=ax, yerr=True, label=["MC1", "MC2", "Data"])
     ax.legend()
     ax.set_title("Raw")
-    assert len(ax.mock_calls) == 24
+    assert len(ax.mock_calls) == 27
 
     ax.reset_mock()