def plot_RPF_regions(input_file: str, hist_name: str, output_prefix: str = ".", printing_extensions: Optional[List[str]] = None) -> None:
""" Main entry point for stand-alone highlight plotting functionality.
If this is being used as a library, call ``plot_RPF_fit_regions(...)`` directly instead.
Args:
input_file (str): Path to the input file.
hist_name (str): Name of the histogram to be highlighted.
output_prefix (str): Directory where the output file should be stored. Default: "."
printing_extensions (list): Printing extensions to be used. Default: None, which corresponds
to printing to ``.pdf``.
Returns:
None.
"""
# Argument validation
if printing_extensions is None:
printing_extensions = [".pdf"]
# Basic setup
# Create logger
logging.basicConfig(level=logging.DEBUG)
# Quiet down the matplotlib logging
logging.getLogger("matplotlib").setLevel(logging.INFO)
# Retrieve hist
with histogram.RootOpen(filename = input_file, mode = "READ") as f:
hist = f.Get(hist_name)
hist.SetDirectory(0)
# Increase the size of the fonts, etc
with sns.plotting_context(context = "notebook", font_scale = 1.5):
# See the possible arguments to highlight_region_of_surface(...)
# For example, to turn on the color overlay, it would be:
#highlight_args = {"use_color_screen" : True}
highlight_args: Dict[str, bool] = {}
# Call plotting functions
fig, ax = plot_RPF_fit_regions(
histogram.get_array_from_hist2D(hist),
highlight_regions = define_highlight_regions(),
colormap = "ROOT_kBird", view_angle = (35, 225),
**highlight_args,
)
# Modify axis labels
# Set the distance from axis to label in pixels.
# Having to set this by hand isn't ideal, but tight_layout doesn't work as well for 3D plots.
ax.xaxis.labelpad = 12
# Visually, delta eta looks closer
ax.yaxis.labelpad = 15
# If desired, add additional labeling here.
# Probably want to use, for examxple, `ax.text2D(0.1, 0.9, "label text", transform = ax.transAxes)`
# This would place the text in the upper left corner (it's like NDC)
# Save and finish up
# The figure will be saved at ``output_prefix/output_path.printing_extension``
output_wrapper = analysis_objects.PlottingOutputWrapper(output_prefix = output_prefix, printing_extensions = printing_extensions)
plot_base.save_plot(output_wrapper, fig, output_name = "highlightRPFRegions")
plt.close(fig)
def setup(self) -> bool:
""" Setup the analysis and the PYTHIA 6 generator.
The main task is to setup the efficiency histograms.
Usually, we'd like the generator to set itself up, but ``TPythia{6,8}`` is a singleton,
so we need to call set up immediately before we run the particular analysis.
"""
# Setup the efficiency histograms
# Distribution to sample for determining which efficiency hist to use.
# Based on year 14 ZDC rates (plot from Dan).
# 0 = 0-33 kHz -> 1/9
# 1 = 33-66 kHz -> 5/9
# 2 = 66-100 kHz -> 3/9
self.efficiency_sampling = [0, 1, 1, 1, 1, 1, 2, 2, 2]
# Retrieve the efficiency histograms
hists = histogram.get_histograms_in_file(
filename="inputData/AuAu/200/y14_efficiency_dca1.root")
centrality_index_map = {
# 0-10% in most analyses.
# 0 = 0-5%, 1 = 5-10%
params.EventActivity.central:
list(range(0, 2)),
# 20-50% in Joel's STAR analysis.
# 4 = 20-25%, 5 = 25-30%, 6 = 30-35%, 7 = 35-40%, 8 = 40-45%, 9 = 45-50%
params.EventActivity.semi_central:
list(range(4, 10)),
}
for interation_rate_index in range(3):
centrality_hists = []
for centrality_index in centrality_index_map[self.event_activity]:
h_root = hists[
f"efficiency_lumi_{interation_rate_index}_cent_{centrality_index}"]
_, _, h_temp = histogram.get_array_from_hist2D(
h_root, set_zero_to_NaN=False)
centrality_hists.append(h_temp)
# Average the efficiency over the centrality bins.
final_efficiency_hist = sum(centrality_hists) / len(
centrality_hists)
self.efficiency_hists.append(final_efficiency_hist)
if interation_rate_index == 0:
# h_root was set from the last iteration, so we take advantage of it.
# Take advantage of the last iteration
self.efficiency_pt_bin_edges = histogram.get_bin_edges_from_axis(
h_root.GetXaxis())
self.efficiency_eta_bin_edges = histogram.get_bin_edges_from_axis(
h_root.GetYaxis())
# Complete setup of PYTHIA6 if necessary.
if self.generator.initialized is False:
self.generator.setup()
return True
def plot_2D_efficiency_data(efficiency_hist: Hist, centrality_range: params.SelectedRange, output_info: analysis_objects.PlottingOutputWrapper) -> None:
""" Plot the 2D efficiency data
Args:
efficiency_hist: Efficiecny histogram.
centrality_range: Centrality range.
output_info: Output info for saving figures.
Returns:
None.
"""
X, Y, efficiency_data = histogram.get_array_from_hist2D(hist = efficiency_hist)
logger.debug(f"efficiency data min: {np.nanmin(efficiency_data)}, max: {np.nanmax(efficiency_data)}")
fig, ax = plt.subplots(figsize = (8, 6))
im = ax.imshow(
efficiency_data.T, extent = [np.nanmin(X), np.nanmax(X), np.nanmin(Y), np.nanmax(Y)],
interpolation = "nearest", aspect = "auto", origin = "lower",
norm = matplotlib.colors.Normalize(
vmin = np.nanmin(efficiency_data), vmax = np.nanmax(efficiency_data)
#vmin = 0.5, vmax = 1,
),
cmap = "viridis",
)
# Add the colorbar
color_bar = fig.colorbar(im, ax = ax)
color_bar.set_label("Efficiency")
# Labels
ax.set_xlabel(fr"${labels.pt_display_label()}\:({labels.momentum_units_label_gev()})$")
ax.set_ylabel(r"$\eta$")
title = f"{period} tracking efficiency data"
if system != params.CollisionSystem.pp:
title += rf", ${centrality_range.min} \textendash {centrality_range.max}\%$"
ax.set_title(title, size = 16)
# Final adjustments
fig.tight_layout()
name = f"efficiency_{period}"
if system != params.CollisionSystem.pp:
name += f"_centrality_{centrality_range.min}_{centrality_range.max}"
plot_base.save_plot(output_info, fig, name)
# Cleanup
plt.close(fig)
def test_get_array_from_hist2D(self, logging_mixin, use_bin_edges, set_zero_to_NaN, test_root_hists):
""" Test getting numpy arrays from a 2D hist. """
hist = test_root_hists.hist2D
x, y, hist_array = histogram.get_array_from_hist2D(hist = hist, set_zero_to_NaN = set_zero_to_NaN, return_bin_edges = use_bin_edges)
# Determine expected values
if use_bin_edges:
epsilon = 1e-9
x_bin_edges = np.empty(hist.GetXaxis().GetNbins() + 1)
x_bin_edges[:-1] = [hist.GetXaxis().GetBinLowEdge(i) for i in range(1, hist.GetXaxis().GetNbins() + 1)]
x_bin_edges[-1] = hist.GetXaxis().GetBinUpEdge(hist.GetXaxis().GetNbins())
y_bin_edges = np.empty(hist.GetYaxis().GetNbins() + 1)
y_bin_edges[:-1] = [hist.GetYaxis().GetBinLowEdge(i) for i in range(1, hist.GetYaxis().GetNbins() + 1)]
y_bin_edges[-1] = hist.GetYaxis().GetBinUpEdge(hist.GetYaxis().GetNbins())
x_mesh = np.arange(np.amin(x_bin_edges), np.amax(x_bin_edges) + epsilon, hist.GetXaxis().GetBinWidth(1))
y_mesh = np.arange(np.amin(y_bin_edges), np.amax(y_bin_edges) + epsilon, hist.GetYaxis().GetBinWidth(1))
else:
x_mesh = np.array([hist.GetXaxis().GetBinCenter(i) for i in range(1, hist.GetXaxis().GetNbins() + 1)])
y_mesh = np.array([hist.GetYaxis().GetBinCenter(i) for i in range(1, hist.GetYaxis().GetNbins() + 1)])
x_range = x_mesh
y_range = y_mesh
expected_x, expected_y = np.meshgrid(x_range, y_range)
expected_hist_array = np.array([
hist.GetBinContent(x, y)
for x in range(1, hist.GetXaxis().GetNbins() + 1)
for y in range(1, hist.GetYaxis().GetNbins() + 1)], dtype=np.float32
).reshape(hist.GetYaxis().GetNbins(), hist.GetXaxis().GetNbins())
if set_zero_to_NaN:
expected_hist_array[expected_hist_array == 0] = np.nan
assert np.allclose(x, expected_x)
assert np.allclose(y, expected_y)
assert np.allclose(hist_array, expected_hist_array, equal_nan = True)
# Check particular values for good measure.
assert np.isclose(hist_array[1][0], 1.0)
if set_zero_to_NaN:
assert np.isnan(hist_array[0][1])
else:
assert np.isclose(hist_array[0][1], 0.0)
def _plot_2D_hists(self, fig: matplotlib.figure.Figure,
ax: matplotlib.axes.Axes) -> None:
""" Plot a 2D hist using matplotlib.
Note:
This is is restricted to plotting only one hist (because how would you plot multiple
2D hists on one canvas)?
Args:
fig: Figure from matplotlib
ax: Axis from matplotlib.
Raises:
ValueError: If more than one histogram is stored in the ``hists``, which doesn't make sense for
2D hist, which doesn't make sense for 2D hists.
"""
if len(self.hists) > 1:
raise ValueError(
self.hists,
"Too many hists are included for a 2D hist. Should only be one!"
)
# We can do a few different types of plots:
# - Surface plot
# - Image plot (like colz), created using:
# - imshow() by default
# - pcolormesh() as an option
# NOTE: seaborn heatmap is not really as flexible here, so we'll handle it manually instead.
# Since the style is already applied, there isn't really anything lost
# Retrieve the array corresponding to the data
# NOTE: The convention for the array is arr[x_index][y_index]. Apparently this matches up with
# numpy axis conventions (?). We will have to transpose the data when we go to plot it.
# This is the same convention as used in root_numpy.
X, Y, hist_array = histogram.get_array_from_hist2D(
hist=self._first_hist(),
set_zero_to_NaN=True,
return_bin_edges=not self.surface)
# Define and fill kwargs
kwargs = {}
# Set the z axis normalization via this particular function
# Will be called when creating the arguments
normalizationFunction = matplotlib.colors.Normalize
if self.logz:
normalizationFunction = matplotlib.colors.LogNorm
# Evaluate
kwargs["norm"] = normalizationFunction(vmin=np.nanmin(hist_array),
vmax=np.nanmax(hist_array))
logger.debug("min: {}, max: {}".format(np.nanmin(hist_array),
np.nanmax(hist_array)))
# Colormap is the default from sns.heatmap
kwargs["cmap"] = plot_base.prepare_colormap(sns.cm.rocket)
# Label is included so we could use a legend if we want
kwargs["label"] = self._first_hist().GetTitle()
logger.debug("kwargs: {}".format(kwargs))
if self.surface:
logger.debug("Plotting surface")
# Need to retrieve a special 3D axis for the surface plot
ax = plt.axes(projection="3d")
# For the surface plot, we want to specify (X,Y) as bin centers. Edges of surfaces
# will be at these points.
# NOTE: There are n-1 faces for n points, so not every value will be represented by a face.
# However, the location of the points at the bin centers is still correct!
# Create the plot
axFromPlot = ax.plot_surface(X, Y, hist_array.T, **kwargs)
else:
# Assume an image plot if we don't explicitly select surface
logger.debug("Plotting as an image")
# pcolormesh is somewhat like a hist in that (X,Y) should define bin edges. So we need (X,Y)
# to define lower bin edges, with the last element in the array corresponding to the lower
# edge of the next (n+1) bin (which is the same as the upper edge of the nth bin).
#
# imshow also takes advantage of these limits to determine the extent, but only in a limited
# way, as it just takes the min and max of each range.
#
# Plot with either imshow or pcolormesh
# Anecdotally, I think pcolormesh is a bit more flexible, but imshow seems to be much faster.
# According to https://stackoverflow.com/a/21169703, either one can be fine with the proper
# options. So the plan is to stick with imshow unless pcolormesh is needed for some reason
if self.use_pcolor_mesh:
logger.debug("Plotting with pcolormesh")
axFromPlot = plt.pcolormesh(X, Y, hist_array.T, **kwargs)
else:
logger.debug("Plotting with imshow ")
# This imshow is quite similar to rplt.imshow (it is located in
# ``plotting.root2matplotlib.imshow``), which can be seen here:
# https://github.com/rootpy/rootpy/blob/master/rootpy/plotting/root2matplotlib.py#L743
# However, we don't copy it directly because we want to set the 0s to nan so they won't
# be plotted.
extent = [np.amin(X), np.amax(X), np.amin(Y), np.amax(Y)]
#logger.debug(f"Extent: {extent}, binEdges[1]: {binEdges[1]}")
axFromPlot = plt.imshow(hist_array.T,
extent=extent,
interpolation="nearest",
aspect="auto",
origin="lower",
**kwargs)
# Draw the colorbar based on the drawn axis above.
label = self.z_label if self.z_label else self._first_hist().GetZaxis(
).GetTitle()
# NOTE: The mappable argument must be the separate axes, not the overall one
# NOTE: This can cause the warning:
# '''
# matplotlib/colors.py:1031: RuntimeWarning: invalid value encountered in less_equal
# mask |= resdat <= 0"
# '''
# The warning is due to the nan we introduced above. It can safely be ignored
# See: https://stackoverflow.com/a/34955622
# (Could suppress, but I don't feel it's necessary at the moment)
fig.colorbar(axFromPlot, label=label)
def _plot_response_matrix_with_matplotlib(
name: str, x_label: str, y_label: str, output_name: str, hist: Hist,
plot_errors_hist: bool,
output_info: analysis_objects.PlottingOutputWrapper) -> None:
""" Underlying function to actually plot a response matrix with matplotlib.
Args:
name: Name of the histogram.
x_label: X axis label.
y_label: Y axis label.
output_name: Output name of the histogram.
hist: The response matrix related 2D hist.
errors_hist: True if the hist is the response matrix errors hist.
output_info: Output information.
Returns:
None
"""
# Setup
fig, ax = plt.subplots(figsize=(8, 6))
# Convert the histogram
X, Y, hist_array = histogram.get_array_from_hist2D(
hist=hist,
set_zero_to_NaN=True,
return_bin_edges=True,
)
# Determine and fill args
kwargs = {}
# Create a log z axis heat map.
kwargs["norm"] = matplotlib.colors.LogNorm(vmin=np.nanmin(hist_array),
vmax=np.nanmax(hist_array))
logger.debug(f"min: {np.nanmin(hist_array)}, max: {np.nanmax(hist_array)}")
# The colormap that we use is the default from sns.heatmap
kwargs["cmap"] = plot_base.prepare_colormap(sns.cm.rocket)
# Label is included so we could use a legend if we want
kwargs["label"] = name
logger.debug("kwargs: {}".format(kwargs))
# Determine the edges
extent = [np.amin(X), np.amax(X), np.amin(Y), np.amax(Y)]
# Finally, create the plot
ax_from_imshow = ax.imshow(hist_array.T,
extent=extent,
interpolation="nearest",
aspect="auto",
origin="lower",
**kwargs)
# Add colorbar
# It needs to be defined on the figure because it is stored in a separate axis.
fig.colorbar(ax_from_imshow, ax=ax)
# Final styling
ax.set_title(name)
ax.set_xlabel(x_label)
ax.set_ylabel(y_label)
fig.tight_layout()
# Save and cleanup
output_name += "_mpl"
plot_base.save_plot(output_info, fig, output_name)
plt.close(fig)
def plot_RPF_fit_regions(jet_hadron: "correlations.Correlations", filename: str) -> None:
""" Plot showing highlighted RPF fit regions.
Args:
jet_hadron: Main analysis object.
filename: Filename under which the hist should be saved.
Returns:
None
"""
# Retrieve the hist to be plotted
# Here we selected the corrected 2D correlation
hist = jet_hadron.correlation_hists_2d.signal.hist.Clone(f"{jet_hadron.correlation_hists_2d.signal.name}_RPF_scaled")
hist.Scale(jet_hadron.correlation_scale_factor)
with sns.plotting_context(context = "notebook", font_scale = 1.5):
# Perform the plotting
# The best option for clarify of viewing seems to be to just replace the colors with the overlay colors.
# mathematical_blending and screen_colors look similar and seem to be the next most promising option.
(fig, ax) = highlight_RPF.plot_RPF_fit_regions(
histogram.get_array_from_hist2D(hist),
highlight_regions = define_highlight_regions(
signal_dominated_eta_region = jet_hadron.signal_dominated_eta_region,
background_dominated_eta_region = jet_hadron.background_dominated_eta_region,
near_side_phi_region = jet_hadron.near_side_phi_region,
),
use_color_overlay = False,
use_color_screen = False,
use_mathematical_blending = False,
)
# Add additional labeling
# Axis
# Needed to fix z axis rotation. See: https://stackoverflow.com/a/21921168
ax.zaxis.set_rotate_label(False)
ax.set_zlabel(r"$1/N_{\mathrm{trig}}\mathrm{d^{2}}N/\mathrm{d}\Delta\varphi\mathrm{d}\Delta\eta$", rotation=90)
# Set the distance from axis to label in pixels.
# This is not ideal, but clearly tight_layout doesn't work as well for 3D plots
ax.xaxis.labelpad = 12
# Visually, dEta looks closer
ax.yaxis.labelpad = 15
ax.zaxis.labelpad = 12
# Overall
alice_label = labels.make_valid_latex_string(jet_hadron.alice_label.display_str())
system_label = labels.system_label(
energy = jet_hadron.collision_energy,
system = jet_hadron.collision_system,
activity = jet_hadron.event_activity
)
jet_finding = labels.jet_finding()
constituent_cuts = labels.constituent_cuts()
leading_hadron = "$" + jet_hadron.leading_hadron_bias.display_str() + "$"
jet_pt = labels.jet_pt_range_string(jet_hadron.jet_pt)
assoc_pt = labels.track_pt_range_string(jet_hadron.track_pt)
# Upper left side
upper_left_text = ""
upper_left_text += alice_label
upper_left_text += "\n" + system_label
upper_left_text += "\n" + jet_pt
upper_left_text += "\n" + jet_finding
# Upper right side
upper_right_text = ""
upper_right_text += leading_hadron
upper_right_text += "\n" + constituent_cuts
upper_right_text += "\n" + assoc_pt
# Need a different text function since we have a 3D axis
ax.text2D(0.01, 0.99, upper_left_text,
horizontalalignment = "left",
verticalalignment = "top",
multialignment = "left",
transform = ax.transAxes)
ax.text2D(0.00, 0.00, upper_right_text,
horizontalalignment = "left",
verticalalignment = "bottom",
multialignment = "left",
transform = ax.transAxes)
# Finish up
plot_base.save_plot(jet_hadron.output_info, fig, filename)
plt.close(fig)
def track_eta_phi(hist: Hist, event_activity: params.EventActivity,
output_info: analysis_objects.PlottingOutputWrapper) -> None:
""" Plot track eta phi.
Also include an annotation of the EMCal eta, phi location.
Args:
rho_hist: Rho centrality dependent hist.
output_info: Output information.
includes_constituent_cut: True if the plot was produced using the constituent cut.
Returns:
None. The figure is plotted.
"""
# Setup
fig, ax = plt.subplots(figsize=(8, 6))
x, y, hist_array = histogram.get_array_from_hist2D(
hist,
set_zero_to_NaN=True,
return_bin_edges=True,
)
plot = ax.imshow(
hist_array.T,
extent=[np.amin(x), np.amax(x),
np.amin(y), np.amax(y)],
interpolation="nearest",
aspect="auto",
origin="lower",
# Intentionally stretch the scale near the top of the range to emphasise the inefficiency.
norm=matplotlib.colors.Normalize(vmin=np.nanmax(hist_array) * 0.8,
vmax=np.nanmax(hist_array) * 0.9),
cmap="viridis",
)
# Draw the colorbar based on the drawn axis above.
# NOTE: This can cause the warning:
# '''
# matplotlib/colors.py:1031: RuntimeWarning: invalid value encountered in less_equal
# mask |= resdat <= 0"
# '''
# The warning is due to the nan we introduced above. It can safely be ignored
# See: https://stackoverflow.com/a/34955622
# (Could suppress, but I don't feel it's necessary at the moment)
fig.colorbar(plot)
# Draw EMCal boundaries
r = matplotlib.patches.Rectangle(
xy=(-0.7, 80 * np.pi / 180),
width=0.7 + 0.7,
height=(187 - 80) * np.pi / 180,
facecolor="none",
edgecolor="tab:red",
linewidth=1.5,
label="EMCal",
)
ax.add_patch(r)
# Final presentation settings
ax.set_xlabel(labels.make_valid_latex_string(r"\eta"))
ax.set_ylabel(labels.make_valid_latex_string(r"\varphi"))
ax.legend(loc="upper right")
fig.tight_layout()
# Finally, save and cleanup
output_name = f"track_eta_phi_{str(event_activity)}"
plot_base.save_plot(output_info, fig, output_name)
plt.close(fig)
# Check the phi 1D projection
track_phi = hist.ProjectionY()
# Make it easier to view
track_phi.Rebin(8)
import ROOT
canvas = ROOT.TCanvas("c", "c")
# Labeling
track_phi.SetTitle("")
track_phi.GetXaxis().SetTitle(
labels.use_label_with_root(labels.make_valid_latex_string(r"\varphi")))
track_phi.GetYaxis().SetTitle("Counts")
track_phi.Draw()
# Draw lines corresponding to the EMCal
line_min = ROOT.TLine(80 * np.pi / 180, track_phi.GetMinimum(),
80 * np.pi / 180, track_phi.GetMaximum())
line_max = ROOT.TLine(187 * np.pi / 180, track_phi.GetMinimum(),
187 * np.pi / 180, track_phi.GetMaximum())
for l in [line_min, line_max]:
l.SetLineColor(ROOT.kRed)
l.Draw("same")
# Save the plot
plot_base.save_plot(output_info, canvas,
f"track_phi_{str(event_activity)}")