Skip to content

Commit

Permalink
feat : implemented broken-axis signed
Browse files Browse the repository at this point in the history
- add additional space for "show" option, change the tick labels
- using slash for "hint", "show" options
  • Loading branch information
Ming-Yan committed Mar 16, 2023
1 parent 4208d43 commit 2c3b83f
Showing 1 changed file with 163 additions and 55 deletions.
218 changes: 163 additions & 55 deletions src/mplhep/plot.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
from matplotlib import patches
from matplotlib.offsetbox import AnchoredText
from matplotlib.transforms import Bbox
from mpl_toolkits.axes_grid1 import axes_size, make_axes_locatable
Expand Down Expand Up @@ -182,19 +181,18 @@ def histplot(
plottables = []
final_bins = np.array(
[
final_bins[0] - (final_bins[1] - final_bins[0]) * 2,
final_bins[0] - (final_bins[1] - final_bins[0]) * 3,
final_bins[0] - (final_bins[1] - final_bins[0]),
*final_bins,
final_bins[-1] + (final_bins[1] - final_bins[0]) * 2,
final_bins[-1] + (final_bins[1] - final_bins[0]),
final_bins[-1] + (final_bins[1] - final_bins[0]) * 3,
]
)
for h in hists:
plottables.append(
Plottable(
h.view(flow=True)["value"],
edges=final_bins,
variances=h.view(flow=True)["variance"],
)
)
value, variance = h.view(flow=True)["value"], h.view(flow=True)["variance"]
value, variance = np.insert(value, -1, 0), np.insert(variance, -1, 0)
value, variance = np.insert(value, 1, 0), np.insert(variance, 1, 0)
plottables.append(Plottable(value, edges=final_bins, variances=variance))
# "sum": Add under/overflow bin to first/last bin
elif flow == "sum":
plottables = []
Expand Down Expand Up @@ -445,15 +443,89 @@ def iterable_not_string(arg):

if x_axes_label:
ax.set_xlabel(x_axes_label)
if flow == "hint":
if flow == "hint" or flow == "show":
underflow, overflow = 0.0, 0.0
for h in hists:
underflow = underflow + h.view(flow=True)["value"][0]
overflow = overflow + h.view(flow=True)["value"][-1]
d = 0.9 # proportion of vertical to horizontal extent of the slanted line
trans = mpl.transforms.blended_transform_factory(ax.transData, ax.transAxes)
kwargs = dict(
marker=[(-0.5, -d), (0.5, d)],
markersize=15,
linestyle="none",
color="k",
mec="k",
mew=1,
clip_on=False,
transform=trans,
)
xticks = ax.get_xticks().tolist()
if underflow > 0.0:
ax.axvspan(final_bins[0], final_bins[1], facecolor="red", alpha=0.5)
if flow == "hint":
ax.plot(
[
final_bins[0] - (final_bins[1] - final_bins[0]) / 2.0,
final_bins[0],
],
[0, 0],
**kwargs,
)
ax.plot(
[
final_bins[0] - (final_bins[1] - final_bins[0]) / 2.0,
final_bins[0],
],
[1, 1],
**kwargs,
)
if flow == "show":
ax.plot(
[final_bins[1], final_bins[2]],
[0, 0],
**kwargs,
)
ax.plot(
[final_bins[1], final_bins[2]],
[1, 1],
**kwargs,
)
xticks[0] = ""
xticks[1] = f"<{final_bins[2]}"

ax.set_xticklabels(xticks)
if overflow > 0.0:
ax.axvspan(final_bins[-2], final_bins[-1], facecolor="red", alpha=0.5)
if flow == "hint":
ax.plot(
[
final_bins[-1],
final_bins[-1] + (final_bins[1] - final_bins[0]) / 2.0,
],
[0, 0],
**kwargs,
)
ax.plot(
[
final_bins[-1],
final_bins[-1] + (final_bins[1] - final_bins[0]) / 2.0,
],
[1, 1],
**kwargs,
)
if flow == "show":
ax.plot(
[final_bins[-3], final_bins[-2]],
[0, 0],
**kwargs,
)
ax.plot(
[final_bins[-3], final_bins[-2]],
[1, 1],
**kwargs,
)
xticks[-1] = ""
xticks[-2] = f">{final_bins[-3]}"
ax.set_xticklabels(xticks)

return return_artists

Expand Down Expand Up @@ -542,18 +614,25 @@ def hist2dplot(
H = hist.view(flow=True)["value"]
xbins = np.array(
[
xbins[0] - (xbins[1] - xbins[0]) * 2,
xbins[0] - (xbins[1] - xbins[0]) * 3,
xbins[0] - (xbins[1] - xbins[0]),
*xbins,
xbins[-1] + (xbins[1] - xbins[0]) * 2,
xbins[-1] + (xbins[1] - xbins[0]),
xbins[-1] + (xbins[1] - xbins[0]) * 3,
]
)
ybins = np.array(
[
ybins[0] - (ybins[1] - ybins[0]) * 2,
ybins[0] - (ybins[1] - ybins[0]) * 3,
ybins[0] - (ybins[1] - ybins[0]),
*ybins,
ybins[-1] + (ybins[1] - ybins[0]) * 2,
ybins[-1] + (ybins[1] - ybins[0]),
ybins[-1] + (ybins[1] - ybins[0]) * 3,
]
)
H = np.insert(H, (1, -1), 0, axis=-1)
H = np.insert(H, (1, -1), np.zeros(np.shape(H)[1]), axis=0)

if flow == "sum":
H[0, 0], H[-1, -1], H[0, -1], H[-1, 0] = (
hist.view(flow=True)["value"][0, 0] + H[0, 0],
Expand Down Expand Up @@ -615,48 +694,77 @@ def hist2dplot(
cb_obj = None

plt.sca(ax)
if flow == "hint":
if flow == "hint" or flow == "show":
d = 0.9 # proportion of vertical to horizontal extent of the slanted line
trans = mpl.transforms.blended_transform_factory(ax.transData, ax.transAxes)
kwargs = dict(
marker=[(-0.5, -d), (0.5, d)],
markersize=15,
linestyle="none",
color="k",
mec="k",
mew=1,
clip_on=False,
)
xticks = ax.get_xticks().tolist()
yticks = ax.get_yticks().tolist()
if hist.view(flow=True)["value"][0, 0] > 0.0:
ret1 = patches.Rectangle(
(xbins[0], ybins[0]),
width=xbins[1] - xbins[0],
height=ybins[1] - ybins[0],
fc="none",
ec="r",
lw=2,
)
ax.add_patch(ret1)

if flow == "hint":
ax.plot(
[xbins[0] - (xbins[1] - xbins[0]) / 2.0, xbins[0]],
[0, 0],
transform=trans,
**kwargs,
)
if flow == "show":
ax.plot([xbins[1], xbins[2]], [0, 0], transform=trans, **kwargs)
ax.plot([xbins[0], xbins[0]], [ybins[1], ybins[2]], **kwargs)
xticks[0] = ""
xticks[1] = f"<{xbins[1]}"
ax.set_xticklabels(xticks)
if hist.view(flow=True)["value"][-1, 0] > 0.0:
ret2 = patches.Rectangle(
(xbins[-2], ybins[0]),
width=xbins[1] - xbins[0],
height=ybins[1] - ybins[0],
fc="none",
ec="r",
lw=2,
)
ax.add_patch(ret2)
if flow == "hint":
ax.plot(
[xbins[-1] + (xbins[1] - xbins[0]) / 2.0, xbins[-1]],
[0, 0],
transform=trans,
**kwargs,
)
if flow == "show":
ax.plot([xbins[-3], xbins[-2]], [0, 0], transform=trans, **kwargs)
ax.plot([xbins[-1], xbins[-1]], [ybins[1], ybins[2]], **kwargs)
xticks[-1] = ""
xticks[-2] = f">{xbins[-2]}"
ax.set_xticklabels(xticks)
if hist.view(flow=True)["value"][0, -1] > 0.0:
ret3 = patches.Rectangle(
(xbins[0], ybins[-2]),
width=xbins[1] - xbins[0],
height=ybins[1] - ybins[0],
fc="none",
ec="r",
lw=2,
)
ax.add_patch(ret3)
if flow == "hint":
ax.plot(
[xbins[0], xbins[0] - (xbins[1] - xbins[0]) / 2.0],
[1, 1],
transform=trans,
**kwargs,
)
if flow == "show":
ax.plot([xbins[1], xbins[2]], [1, 1], transform=trans, **kwargs)
ax.plot([xbins[0], xbins[0]], [ybins[-3], ybins[-2]], **kwargs)
yticks[0] = ""
yticks[1] = f"<{ybins[1]}"
ax.set_yticklabels(yticks)
if hist.view(flow=True)["value"][-1, -1] > 0.0:
ret4 = patches.Rectangle(
(xbins[-2], ybins[-2]),
width=xbins[1] - xbins[0],
height=ybins[1] - ybins[0],
fc="none",
ec="r",
lw=2,
)
ax.add_patch(ret4)
if flow == "hint":
ax.plot(
[xbins[-1] + (xbins[1] - xbins[0]) / 2.0, xbins[-1]],
[1, 1],
transform=trans,
**kwargs,
)
if flow == "show":
ax.plot([xbins[-3], xbins[-2]], [1, 1], transform=trans, **kwargs)
ax.plot([xbins[-1], xbins[-1]], [ybins[-3], ybins[-2]], **kwargs)
yticks[-1] = ""
yticks[-2] = f">{ybins[-2]}"
ax.set_yticklabels(yticks)

_labels: np.ndarray | None = None
if isinstance(labels, bool):
_labels = H if labels else None
Expand Down

0 comments on commit 2c3b83f

Please sign in to comment.