Skip to content

Commit

Permalink
feat: added bar type + correct legend for step type
Browse files Browse the repository at this point in the history
  • Loading branch information
0ctagon committed Oct 20, 2024
1 parent d9746dd commit 2301c7d
Show file tree
Hide file tree
Showing 9 changed files with 200 additions and 51 deletions.
3 changes: 3 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
@@ -1,5 +1,8 @@
# Project specific
src/mplhep/_version.py
*.root
result_images/
test/

# Byte-compiled / optimized / DLL files
__pycache__/
Expand Down
175 changes: 138 additions & 37 deletions src/mplhep/plot.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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]
Expand All @@ -423,53 +430,146 @@ 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)
)
if isinstance(yerr, bool) and yerr and plottables[i].variances is not None:
do_errors = True
else:
do_errors = yerr is not False and (yerr is not None or w2 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]
else:
if _kwargs.get("color") is not None:
_kwargs["edgecolor"] = _kwargs["color"]
else:
_kwargs["edgecolor"] = ax._get_lines.get_next_color() # type: ignore[attr-defined]
_kwargs["color"] = _kwargs["edgecolor"]

if histtype == "step":
_kwargs["fill"] = True
_kwargs["facecolor"] = "None"

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,
)
)
_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,
)
return_artists.append(
StairsArtists(
_s,
_e if do_errors else None,
_e_leg if do_errors else None,

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)):
Expand Down Expand Up @@ -531,9 +631,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:
Expand Down
Binary file added tests/baseline/test_histplot_bar.png
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Binary file modified tests/baseline/test_histplot_kwargs.png
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Binary file modified tests/baseline/test_histplot_real.png
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Binary file modified tests/baseline/test_histplot_types.png
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Binary file modified tests/baseline/test_simple_xerr.png
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
66 changes: 55 additions & 11 deletions tests/test_basic.py
Original file line number Diff line number Diff line change
Expand Up @@ -78,7 +78,7 @@ def test_onebin_hist():
fig, axs = plt.subplots()
h = hist.Hist(hist.axis.Regular(1, 0, 1))
h.fill([-1, 0.5])
hep.histplot(h, ax=axs)
hep.histplot(h, yerr=True, ax=axs)
return fig


Expand Down Expand Up @@ -138,10 +138,10 @@ def test_histplot_flow():
fig, axs = plt.subplots(2, 2, sharey=True, figsize=(10, 10))
axs = axs.flatten()

hep.histplot(h, ax=axs[0], flow="hint")
hep.histplot(h, ax=axs[1], flow="show")
hep.histplot(h, ax=axs[2], flow="sum")
hep.histplot(h, ax=axs[3], flow=None)
hep.histplot(h, ax=axs[0], yerr=True, flow="hint")
hep.histplot(h, ax=axs[1], yerr=True, flow="show")
hep.histplot(h, ax=axs[2], yerr=True, flow="sum")
hep.histplot(h, ax=axs[3], yerr=True, flow=None)

axs[0].set_title("Default(hint)", fontsize=18)
axs[1].set_title("Show", fontsize=18)
Expand Down Expand Up @@ -213,10 +213,10 @@ def test_histplot_uproot_flow():
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")
hep.histplot(h, ax=axs[0], yerr=True, flow="show")
hep.histplot(h2, ax=axs[1], yerr=True, flow="show")
hep.histplot(h3, ax=axs[2], yerr=True, flow="show")
hep.histplot(h4, ax=axs[3], yerr=True, flow="show")

axs[0].set_title("Two-side overflow", fontsize=18)
axs[1].set_title("Left-side overflow", fontsize=18)
Expand Down Expand Up @@ -615,16 +615,60 @@ 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", 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)


Expand Down
7 changes: 4 additions & 3 deletions tests/test_mock.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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,
)
Expand All @@ -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()

Expand Down

0 comments on commit 2301c7d

Please sign in to comment.