Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: added bar type #529

Merged
merged 7 commits into from
Oct 29, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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
158 changes: 124 additions & 34 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,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)):
Expand Down Expand Up @@ -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:
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_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_inputs_basic.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.
53 changes: 51 additions & 2 deletions tests/test_basic.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)


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
Loading