def plot_on(
self,
ax: plt.axis,
draw_legend: bool = True,
legend_inside: bool = True,
legend_kwargs: dict = {},
yaxis_scale=1.3,
hide_labels: bool = False,
) -> plt.axis:
for component in self.components:
if component.style == 'point':
ax.errorbar(
component.data.x_values,
component.data.y_values,
yerr=component.data.y_errors,
xerr=component.data.x_errors,
label=component.label,
color=component.color,
ls=component.ls,
marker=component.marker,
)
elif component.style == 'box':
ax.bar(component.data.x_values,
2 * component.data.y_errors,
width=2 * component.data.x_errors,
bottom=component.data.y_values -
component.data.y_errors,
label=component.label,
color=component.color,
alpha=0.5)
else:
raise NotImplementedError(
"Options: point|box. If you require a new kind of plot, report a feature request"
)
if not hide_labels:
ax.set_xlabel(self.variable.x_label, plot_style.xlabel_pos)
ax.set_ylabel(self.variable.y_label, plot_style.ylabel_pos)
if draw_legend:
if legend_inside:
ax.legend(frameon=False, **legend_kwargs)
ylims = ax.get_ylim()
ax.set_ylim(ylims[0], yaxis_scale * ylims[1])
else:
ax.legend(frameon=False,
bbox_to_anchor=(1, 1),
**legend_kwargs)
return ax
def get_auto_ylims(ax: plt.axis,
hMC,
*,
hdata=None,
log_y=False,
yaxis_scale='auto'
) -> Tuple[Union[int, float], Union[int, float]]:
"""Return the minimum and maximum values for the y axis. If yaxis_scale is 'auto', get a value such that the
histogram stays in bounds of the plot, and the legend does not overlap with the histogram.
"""
ylims = ax.get_ylim()
ymin = ylims[0]
if yaxis_scale == 'auto':
if not isinstance(hMC[0], float):
hMC = hMC[-1]
labelFraction = 1 / 4
legendFraction = 1 / 3
maxBelowLabel = max(hMC[:len(hMC) * 2 // 3])
maxBelowLegend = max(hMC[len(hMC) * 2 // 3:])
if hdata is not None:
maxBelowLabel = max(maxBelowLabel,
max(hdata[:len(hdata) * 2 // 3]))
maxBelowLegend = max(maxBelowLegend,
max(hdata[len(hdata) * 2 // 3:]))
if log_y:
ymax = max(maxBelowLabel**(1 / (1 - labelFraction)),
maxBelowLegend**(1 / (1 - legendFraction)))
else:
ymax = max(maxBelowLabel / (1 - labelFraction),
maxBelowLegend / (1 - legendFraction))
else:
if log_y:
ymax = ylims[1]**yaxis_scale
else:
ymax = ylims[1] * yaxis_scale
return ymin, ymax
def set_inset_spectrum(axis: plt.axis, data: np.ndarray, current_index: int,
peak_collection: PeakCollection) -> plt.axis:
axis.plot(data) # Show intensity spectrum
axis = plot_peaks(axis=axis, collection=peak_collection) # Show peaks
y_bot, y_top = axis.get_ylim()
text_height = y_bot + 0.6 * (y_top - y_bot) # Data coordinates
_margin = 50
x_lim = [
max(0, current_index - _margin),
min(len(data) - 1, current_index + _margin)
]
# Plot clusters
# for cluster in peak_collection.get_clusters:
# bound_left, bound_right = cluster.get_value_slice
# if bound_right > x_lim[0] or bound_left < x_lim[1]:
# axis.axvspan(bound_left, bound_right, alpha=0.5, color='green')
# if x_lim[0] < cluster.get_avg_x < x_lim[1]:
# axis.text(x=cluster.get_avg_x, y=text_height, s=r'$\tilde{m}$'+f'={cluster.get_transverse_mode_id}')
axis.axvline(x=current_index, color='r')
axis.set_xlim(x_lim)
return axis
def plot_on(self,
ax: plt.axis,
ylabel="Events",
draw_legend=True,
legend_inside=True,
hide_labels: bool = False):
bin_edges, bin_mids, bin_width = self._get_bin_edges()
self._bin_edges = bin_edges
self._bin_mids = bin_mids
self._bin_width = bin_width
ax.hist(
x=[comp.data for comp in self._mc_components['stacked']],
bins=bin_edges,
weights=[comp.weights for comp in self._mc_components['stacked']],
stacked=True,
edgecolor="black",
lw=0.3,
color=[comp.color for comp in self._mc_components['stacked']],
label=[comp.label for comp in self._mc_components['stacked']],
histtype='stepfilled')
if not hide_labels:
ax.set_xlabel(self._variable.x_label, plot_style.xlabel_pos)
y_label = self._get_y_label(False, bin_width, ylabel)
ax.set_ylabel(y_label, plot_style.ylabel_pos)
if draw_legend:
if legend_inside:
ax.legend(frameon=False)
ylims = ax.get_ylim()
ax.set_ylim(ylims[0], 1.4 * ylims[1])
else:
ax.legend(frameon=False, bbox_to_anchor=(1, 1))
return ax
def plot_on(
self,
ax1: plt.axis,
ax2,
style="stacked",
ylabel="Events",
sum_color=plot_style.KITColors.kit_purple,
draw_legend: bool = True,
legend_inside: bool = True,
):
bin_edges, bin_mids, bin_width = self._get_bin_edges()
self._bin_edges = bin_edges
self._bin_mids = bin_mids
self._bin_width = bin_width
sum_w = np.sum(np.array([
binned_statistic(comp.data,
comp.weights,
statistic="sum",
bins=bin_edges)[0]
for comp in self._mc_components["MC"]
]),
axis=0)
sum_w2 = np.sum(np.array([
binned_statistic(comp.data,
comp.weights**2,
statistic="sum",
bins=bin_edges)[0]
for comp in self._mc_components["MC"]
]),
axis=0)
hdata, _ = np.histogram(self._data_component.data, bins=bin_edges)
if style.lower() == "stacked":
ax1.hist(
x=[comp.data for comp in self._mc_components['MC']],
bins=bin_edges,
weights=[comp.weights for comp in self._mc_components['MC']],
stacked=True,
edgecolor="black",
lw=0.3,
color=[comp.color for comp in self._mc_components['MC']],
label=[comp.label for comp in self._mc_components['MC']],
histtype='stepfilled')
ax1.bar(x=bin_mids,
height=2 * np.sqrt(sum_w2),
width=self.bin_width,
bottom=sum_w - np.sqrt(sum_w2),
color="black",
hatch="///////",
fill=False,
lw=0,
label="MC stat. unc.")
if style.lower() == "summed":
ax1.bar(x=bin_mids,
height=2 * np.sqrt(sum_w2),
width=self.bin_width,
bottom=sum_w - np.sqrt(sum_w2),
color=sum_color,
lw=0,
label="MC")
ax1.errorbar(x=bin_mids,
y=hdata,
yerr=np.sqrt(hdata),
ls="",
marker=".",
color="black",
label=self._data_component.label)
y_label = self._get_y_label(False, bin_width, evts_or_cand=ylabel)
# ax1.legend(loc=0, bbox_to_anchor=(1,1))
ax1.set_ylabel(y_label, plot_style.ylabel_pos)
if draw_legend:
if legend_inside:
ax1.legend(frameon=False)
ylims = ax1.get_ylim()
ax1.set_ylim(ylims[0], 1.4 * ylims[1])
else:
ax1.legend(frameon=False, bbox_to_anchor=(1, 1))
ax2.set_ylabel(r"$\frac{\mathrm{Data - MC}}{\mathrm{Data}}$")
ax2.set_xlabel(self._variable.x_label, plot_style.xlabel_pos)
ax2.set_ylim((-1, 1))
try:
uhdata = unp.uarray(hdata, np.sqrt(hdata))
uhmc = unp.uarray(sum_w, np.sqrt(sum_w2))
ratio = (uhdata - uhmc) / uhdata
ax2.axhline(y=0, color=plot_style.KITColors.dark_grey, alpha=0.8)
ax2.errorbar(bin_mids,
unp.nominal_values(ratio),
yerr=unp.std_devs(ratio),
ls="",
marker=".",
color=plot_style.KITColors.kit_black)
except ZeroDivisionError:
ax2.axhline(y=0, color=plot_style.KITColors.dark_grey, alpha=0.8)
plt.subplots_adjust(hspace=0.08)
def plot_on(self,
ax: plt.axis,
draw_legend: bool = True,
legend_inside: bool = True,
yaxis_scale=1.3,
normed: bool = False,
ylabel="Events",
hide_labels: bool = False) -> plt.axis:
"""
Plots the component on a given matplotlib.pyplot.axis
:param ax: matplotlib.pyplot.axis where the histograms will be drawn
on.
:param draw_legend: Draw legend on axis if True.
:param normed: If true the histograms are normalized.
:return: matplotlib.pyplot.axis with histogram drawn on it
"""
bin_edges, bin_mids, bin_width = self._get_bin_edges()
self._bin_edges = bin_edges
self._bin_mids = bin_mids
self._bin_width = bin_width
for component in self._mc_components['single']:
if component.histtype == 'stepfilled':
alpha = 0.6
edge_color = 'black'
else:
edge_color = None
alpha = 1.0
ax.hist(x=component.data,
bins=bin_edges,
density=normed,
weights=component.weights,
histtype=component.histtype,
label=component.label,
edgecolor=edge_color
if edge_color is not None else component.color,
alpha=alpha,
lw=1.5,
ls=component.ls,
color=component.color)
if not hide_labels:
ax.set_xlabel(self._variable.x_label, plot_style.xlabel_pos)
y_label = self._get_y_label(normed=normed,
bin_width=bin_width,
evts_or_cand=ylabel)
ax.set_ylabel(y_label, plot_style.ylabel_pos)
if draw_legend:
if legend_inside:
ax.legend(frameon=False)
ylims = ax.get_ylim()
ax.set_ylim(ylims[0], yaxis_scale * ylims[1])
else:
ax.legend(frameon=False, bbox_to_anchor=(1, 1))
return ax
