Skip to content

Commit

Permalink
feat: update hist2d flow style
Browse files Browse the repository at this point in the history
  • Loading branch information
@andrzejnovak
andrzejnovak committed Apr 1, 2024
1 parent c94dc7c commit ba1aab8
Show file tree
Hide file tree
Showing 6 changed files with 149 additions and 116 deletions.
237 changes: 121 additions & 116 deletions src/mplhep/plot.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
isLight,
process_histogram_parts,
align_marker,
to_padded2d,
)

if TYPE_CHECKING:
Expand Down Expand Up @@ -501,7 +502,7 @@ def iterable_not_string(arg):

# Flow extra styling
if flow == "hint":
s = (
_marker_size = (
30
* ax.get_window_extent()
.transformed(ax.figure.dpi_scale_trans.inverted())
Expand All @@ -511,7 +512,7 @@ def iterable_not_string(arg):
ax.scatter(
final_bins[0],
0,
s,
_marker_size,
marker=align_marker("<", halign="right"),
edgecolor="black",
zorder=5,
Expand All @@ -523,7 +524,7 @@ def iterable_not_string(arg):
ax.scatter(
final_bins[-1],
0,
s,
_marker_size,
marker=align_marker(">", halign="left"),
edgecolor="black",
zorder=5,
Expand Down Expand Up @@ -674,7 +675,7 @@ def hist2dplot(

# TODO: use Histogram everywhere

H = h.values()
H = h.values().copy()
xbins, xtick_labels = get_plottable_protocol_bins(h.axes[0])
ybins, ytick_labels = get_plottable_protocol_bins(h.axes[1])
# Show under/overflow bins
Expand All @@ -697,49 +698,31 @@ def hist2dplot(
):
flow = None
print(f"Warning: you don't have flow bins stored in {h}")
elif flow == "show":
H = h.values(flow=True)
if any(h.values(flow=True)[0] > 0):
xbins = np.array(
[
xbins[0] - (xbins[-1] - xbins[0]) * 0.08,
xbins[0] - (xbins[-1] - xbins[0]) * 0.03,
*xbins,
]
)
if any(h.values(flow=True)[-1] > 0):
xbins = np.array(
[
*xbins,
xbins[-1] + (xbins[-1] - xbins[0]) * 0.03,
xbins[-1] + (xbins[-1] - xbins[0]) * 0.08,
]
)
if any(h.values(flow=True)[:, 0] > 0):
ybins = np.array(
[
ybins[0] - (ybins[-1] - ybins[0]) * 0.08,
ybins[0] - (ybins[-1] - ybins[0]) * 0.03,
*ybins,
]
)
if any(h.values(flow=True)[:, -1] > 0):
ybins = np.array(
[
*ybins,
ybins[-1] + (ybins[-1] - ybins[0]) * 0.03,
ybins[-1] + (ybins[-1] - ybins[0]) * 0.08,
]
)

if any(h.values(flow=True)[0] > 0.0):
H = np.insert(H, (1), np.nan, axis=-1)
if any(h.values(flow=True)[-1] > 0.0):
H = np.insert(H, (-1), np.nan, axis=-1)
if any(h.values(flow=True)[:, 0] > 0):
H = np.insert(H, (1), np.full(np.shape(H)[1], np.nan), axis=0)
if any(h.values(flow=True)[:, -1] > 0):
H = np.insert(H, (-1), np.full(np.shape(H)[1], np.nan), axis=0)
elif flow in ["hint", "show"]:
xwidth, ywidth = (xbins[-1] - xbins[0]) * 0.05, (ybins[-1] - ybins[0]) * 0.05
pxbins = np.r_[xbins[0] - xwidth, xbins, xbins[-1] + xwidth]
pybins = np.r_[ybins[0] - ywidth, ybins, ybins[-1] + ywidth]
padded = to_padded2d(h)
hint_xlo, hint_xhi, hint_ylo, hint_yhi = True, True, True, True
if np.all(padded[0, :] == 0):
padded = padded[1:, :]
pxbins = pxbins[1:]
hint_xlo = False
if np.all(padded[-1, :] == 0):
padded = padded[:-1, :]
pxbins = pxbins[:-1]
hint_xhi = False
if np.all(padded[:, 0] == 0):
padded = padded[:, 1:]
pybins = pybins[1:]
hint_ylo = False
if np.all(padded[:, -1] == 0):
padded = padded[:, :-1]
pybins = pybins[:-1]
hint_yhi = False
if flow == "show":
H = padded
xbins, ybins = pxbins, pybins
elif flow == "sum":
H = h.values().copy()
# Sum borders
Expand Down Expand Up @@ -812,76 +795,98 @@ def hist2dplot(
cb_obj = None

plt.sca(ax)
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)
ax_h = ax.bbox.height
kwargs = dict(
marker=[(-0.5, -d), (0.5, d)],
markersize=ax_h * 0.05,
linestyle="none",
color="k",
mec="k",
mew=1,
clip_on=False,
if flow == "show":
if hint_xlo:
ax.plot(
[xbins[1]] * 2,
[0, 1],
ls="--",
color="lightgrey",
clip_on=False,
transform=ax.get_xaxis_transform(),
)
if hint_xhi:
ax.plot(
[xbins[-2]] * 2,
[0, 1],
ls="--",
color="lightgrey",
clip_on=False,
transform=ax.get_xaxis_transform(),
)
if hint_ylo:
ax.plot(
[0, 1],
[ybins[1]] * 2,
ls="--",
color="lightgrey",
clip_on=False,
transform=ax.get_yaxis_transform(),
)
if hint_yhi:
ax.plot(
[0, 1],
[ybins[-2]] * 2,
ls="--",
color="lightgrey",
clip_on=False,
transform=ax.get_yaxis_transform(),
)
elif flow == "hint":
_marker_size = (
30
* ax.get_window_extent()
.transformed(ax.figure.dpi_scale_trans.inverted())
.width
)
if any(h.values(flow=True)[0] > 0):
if flow == "hint":
ax.plot(
[
xbins[0] - np.diff(xbins)[0] * len(np.diff(xbins)) * 0.03,
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)
if any(h.values(flow=True)[:, 0] > 0):
if flow == "hint":
ax.plot(
[
xbins[-1] + np.diff(xbins)[-1] * len(np.diff(xbins)) * 0.03,
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)
if any(h.values(flow=True)[-1] > 0):
if flow == "hint":
ax.plot(
[
xbins[0],
xbins[0] - np.diff(xbins)[0] * len(np.diff(xbins)) * 0.03,
],
[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)

if any(h.values(flow=True)[:, -1] > 0):
if flow == "hint":
ax.plot(
[
xbins[-1] + np.diff(xbins)[-1] * len(np.diff(xbins)) * 0.03,
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)
if hint_xlo:
ax.scatter(
0,
0,
_marker_size,
marker=align_marker("<", halign="right", valign="bottom"),
edgecolor="black",
zorder=5,
clip_on=False,
facecolor="white",
transform=ax.transAxes,
)
if hint_xhi:
ax.scatter(
1,
0,
_marker_size,
marker=align_marker(">", halign="left"),
edgecolor="black",
zorder=5,
clip_on=False,
facecolor="white",
transform=ax.transAxes,
)
if hint_ylo:
ax.scatter(
0,
0,
_marker_size,
marker=align_marker("v", valign="top", halign="left"),
edgecolor="black",
zorder=5,
clip_on=False,
facecolor="white",
transform=ax.transAxes,
)
if hint_yhi:
ax.scatter(
0,
1,
_marker_size,
marker=align_marker("^", valign="bottom"),
edgecolor="black",
zorder=5,
clip_on=False,
facecolor="white",
transform=ax.transAxes,
)

_labels: np.ndarray | None = None
if isinstance(labels, bool):
Expand Down
28 changes: 28 additions & 0 deletions src/mplhep/utils.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -375,3 +375,31 @@ def align_marker(
m_arr[:, 1] += valign / 2

return Path(m_arr, bm.get_path().codes)


def to_padded2d(h, variances=False):
if np.array_equal(
np.array(h.values().shape) + 2, np.array(h.values(flow=True).shape)
):
padded = h.values(flow=True)
padded_varis = h.variances(flow=True)
else:
vals_flow = h.values(flow=True)
variances_flow = h.variances(flow=True)
xpadlo, xpadhi = 1 - h.axes[0].traits.underflow, 1 - h.axes[0].traits.overflow
ypadlo, ypadhi = 1 - h.axes[1].traits.underflow, 1 - h.axes[1].traits.overflow
xpadhi_m, mypadhi_m = [-pad if pad != 0 else None for pad in [xpadhi, ypadhi]]

padded = np.zeros(
(
vals_flow.shape[0] + xpadlo + xpadhi,
(vals_flow.shape[1] + ypadlo + ypadhi),
)
)
padded_varis = padded.copy()
padded[xpadlo:xpadhi_m, ypadlo:mypadhi_m] = vals_flow
padded_varis[xpadlo:xpadhi_m, ypadlo:mypadhi_m] = variances_flow
if variances:
return padded, padded_varis
else:
return padded
Binary file modified tests/baseline/test_hist2dplot_flow.png
View file
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_bh.png
View file
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_bh_cat.png
View file
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_uproot.png
View file
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.

0 comments on commit ba1aab8

Please sign in to comment.