def plot_particle_level_spectra(
        ep_analyses_iter: Iterator[Tuple[Any,
                                         "response_matrix.ResponseMatrix"]],
        output_info: analysis_objects.PlottingOutputWrapper,
        plot_with_ROOT: bool = False) -> None:
    """ Plot the particle level spectra associated with the response.

    Args:
        ep_analyses: The event plane dependent final response matrices.
        output_info: Output information.
        plot_with_ROOT: True if the plot should be done via ROOT. Default: False
    Returns:
        None. The spectra are plotted and saved.
    """
    # Pull out the dict because we need to grab individual analyses for some labeling information, which doesn't
    # play well with generators (the generator will be exhausted).
    ep_analyses = dict(ep_analyses_iter)

    # Determine the general and plot labels
    # First, we need some variables to define the general labels, so we retrieve the inclusive analysis.
    # All of the parameters retrieved here are shared by all analyses.
    inclusive = next(iter(ep_analyses.values()))
    # Then we define some additional helper variables
    particle_level_spectra_bin = inclusive.task_config[
        "particle_level_spectra"]["particle_level_spectra_bin"]
    embedded_additional_label = inclusive.event_activity.display_str()

    # General labels
    general_labels = {
        "alice_and_collision_energy":
        rf"{inclusive.alice_label.display_str()}\:{inclusive.collision_energy.display_str()}",
        "collision_system_and_event_activity":
        rf"{inclusive.collision_system.display_str(embedded_additional_label = embedded_additional_label)}",
        "detector_pt_range":
        labels.pt_range_string(
            particle_level_spectra_bin,
            lower_label="T,jet",
            upper_label="det",
        ),
        "constituent_cuts":
        labels.constituent_cuts(additional_label="det"),
        "leading_hadron_bias":
        inclusive.leading_hadron_bias.display_str(additional_label="det"),
        "jet_finding":
        labels.jet_finding(),
    }
    # Ensure that each line is a valid latex line.
    # The heuristic is roughly that full statements (such as jet_finding) are already wrapped in "$",
    # while partial statements, such as the leading hadron bias, event activity, etc are not wrapped in "$".
    # This is due to the potential for such "$" to interfere with including those partial statements in other
    # statements. As an example, it would be impossible to use the ``embedded_additional_label`` above if the
    # ``event_activity`` included "$".
    for k, v in general_labels.items():
        general_labels[k] = labels.make_valid_latex_string(v)

    # Plot labels
    y_label = r"\mathrm{d}N/\mathrm{d}p_{\mathrm{T}}"
    if inclusive.task_config["particle_level_spectra"]["normalize_by_n_jets"]:
        y_label = r"(1/N_{\mathrm{jets}})" + y_label
        y_label = y_label
    if inclusive.task_config["particle_level_spectra"][
            "normalize_at_selected_jet_pt_bin"]:
        y_label = r"\mathrm{Arb. Units}"
    plot_labels = plot_base.PlotLabels(
        title="",
        x_label=
        fr"${labels.jet_pt_display_label(upper_label = 'part')}\:({labels.momentum_units_label_gev()})$",
        y_label=labels.make_valid_latex_string(y_label),
    )

    # Finally, we collect our arguments for the plotting functions.
    kwargs: Dict[str, Any] = {
        "ep_analyses": ep_analyses,
        "output_name": "particle_level_spectra",
        "output_info": output_info,
        "general_labels": general_labels,
        "plot_labels": plot_labels,
    }

    if plot_with_ROOT:
        _plot_particle_level_spectra_with_ROOT(**kwargs)
    else:
        _plot_particle_level_spectra_with_matplotlib(**kwargs)
Example #2
0
def plot_RP_fit(rp_fit: reaction_plane_fit.fit.ReactionPlaneFit,
                inclusive_analysis: "correlations.Correlations",
                ep_analyses: List[Tuple[Any, "correlations.Correlations"]],
                output_info: analysis_objects.PlottingOutputWrapper,
                output_name: str) -> None:
    """ Basic plot of the reaction plane fit.

    Args:
        rp_fit: Reaction plane fit object.
        inclusive_analysis: Inclusive analysis object. Mainly used for labeling.
        ep_analyses: Event plane dependent correlation analysis objects.
        output_info: Output information.
        output_name: Name of the output plot.
    Returns:
        None. The plot will be saved.
    """
    # Setup
    n_components = len(rp_fit.components)
    fig, axes = plt.subplots(2,
                             n_components,
                             sharey="row",
                             sharex=True,
                             gridspec_kw={"height_ratios": [3, 1]},
                             figsize=(3 * n_components, 6))
    flat_axes = axes.flatten()

    # Plot the fits on the upper panels.
    _plot_rp_fit_components(rp_fit=rp_fit,
                            ep_analyses=ep_analyses,
                            axes=flat_axes[:n_components])
    # Plot the residuals on the lower panels.
    _plot_rp_fit_residuals(rp_fit=rp_fit,
                           ep_analyses=ep_analyses,
                           axes=flat_axes[n_components:])

    # Define upper panel labels.
    # In-plane
    text = labels.track_pt_range_string(inclusive_analysis.track_pt)
    text += "\n" + labels.constituent_cuts()
    text += "\n" + labels.make_valid_latex_string(
        inclusive_analysis.leading_hadron_bias.display_str())
    _add_label_to_rpf_plot_axis(ax=flat_axes[0], label=text)
    # Mid-plane
    text = labels.make_valid_latex_string(
        inclusive_analysis.alice_label.display_str())
    text += "\n" + labels.system_label(
        energy=inclusive_analysis.collision_energy,
        system=inclusive_analysis.collision_system,
        activity=inclusive_analysis.event_activity)
    text += "\n" + labels.jet_pt_range_string(inclusive_analysis.jet_pt)
    text += "\n" + labels.jet_finding()
    _add_label_to_rpf_plot_axis(ax=flat_axes[1], label=text)
    # Out-of-plane
    #text = "Background: $0.8<|\Delta\eta|<1.2$"
    #text += "\nSignal + Background: $|\Delta\eta|<0.6$"
    #_add_label_to_rpf_plot_axis(ax = flat_axes[2], label = text)
    # Inclusive
    text = (r"\chi^{2}/\mathrm{NDF} = "
            f"{rp_fit.fit_result.minimum_val:.1f}/{rp_fit.fit_result.nDOF} = "
            f"{rp_fit.fit_result.minimum_val / rp_fit.fit_result.nDOF:.3f}")
    _add_label_to_rpf_plot_axis(ax=flat_axes[2],
                                label=labels.make_valid_latex_string(text))

    # Define lower panel labels.
    for ax in flat_axes[n_components:]:
        # Increase the frequency of major ticks to once every integer.
        ax.xaxis.set_major_locator(matplotlib.ticker.MultipleLocator(base=1.0))
        # Add axis labels
        ax.set_xlabel(
            labels.make_valid_latex_string(
                inclusive_analysis.correlation_hists_delta_phi.
                signal_dominated.axis.display_str()))
    # Improve the viewable range for the lower panels.
    # This value is somewhat arbitrarily selected, but seems to work well enough.
    flat_axes[n_components].set_ylim(-0.2, 0.2)

    # Specify shared y axis label
    # Delta phi correlations first
    flat_axes[0].set_ylabel(labels.delta_phi_axis_label())
    # Then label the residual
    flat_axes[n_components].set_ylabel("data - fit / fit")

    # Final adjustments
    fig.tight_layout()
    # Reduce spacing between subplots
    fig.subplots_adjust(hspace=0, wspace=0)
    # Save plot and cleanup
    plot_base.save_plot(output_info, fig, output_name)
    plt.close(fig)
Example #3
0
def rp_fit_subtracted(ep_analyses: List[Tuple[Any,
                                              "correlations.Correlations"]],
                      inclusive_analysis: "correlations.Correlations",
                      output_info: analysis_objects.PlottingOutputWrapper,
                      output_name: str) -> None:
    """ Basic plot of the reaction plane fit subtracted hists.

    Args:
        ep_analyses: Event plane dependent correlation analysis objects.
        inclusive_analysis: Inclusive analysis object. Mainly used for labeling.
        output_info: Output information.
        output_name: Name of the output plot.
    Returns:
        None. The plot will be saved.
    """
    # Setup
    n_components = len(ep_analyses)
    fig, axes = plt.subplots(
        1,
        n_components,
        sharey="row",
        sharex=True,
        #gridspec_kw = {"height_ratios": [3, 1]},
        figsize=(3 * n_components, 6))
    flat_axes = axes.flatten()

    # Plot the fits on the upper panels.
    _plot_rp_fit_subtracted(ep_analyses=ep_analyses,
                            axes=flat_axes[:n_components])

    # Define upper panel labels.
    # In-plane
    text = labels.track_pt_range_string(inclusive_analysis.track_pt)
    text += "\n" + labels.constituent_cuts()
    text += "\n" + labels.make_valid_latex_string(
        inclusive_analysis.leading_hadron_bias.display_str())
    _add_label_to_rpf_plot_axis(ax=flat_axes[0], label=text)
    # Mid-plane
    text = labels.make_valid_latex_string(
        inclusive_analysis.alice_label.display_str())
    text += "\n" + labels.system_label(
        energy=inclusive_analysis.collision_energy,
        system=inclusive_analysis.collision_system,
        activity=inclusive_analysis.event_activity)
    text += "\n" + labels.jet_pt_range_string(inclusive_analysis.jet_pt)
    text += "\n" + labels.jet_finding()
    _add_label_to_rpf_plot_axis(ax=flat_axes[1], label=text)
    # Out-of-plane
    #text = "Background: $0.8<|\Delta\eta|<1.2$"
    #text += "\nSignal + Background: $|\Delta\eta|<0.6$"
    #_add_label_to_rpf_plot_axis(ax = flat_axes[2], label = text)
    _add_label_to_rpf_plot_axis(ax=flat_axes[2],
                                label=labels.make_valid_latex_string(text))

    for ax in flat_axes:
        # Increase the frequency of major ticks to once every integer.
        ax.xaxis.set_major_locator(matplotlib.ticker.MultipleLocator(base=1.0))
        # Set label
        ax.set_xlabel(labels.make_valid_latex_string(r"\Delta\varphi"))

    flat_axes[0].set_ylabel(
        labels.make_valid_latex_string(labels.delta_phi_axis_label()))
    #jet_pt_label = labels.jet_pt_range_string(inclusive_analysis.jet_pt)
    #track_pt_label = labels.track_pt_range_string(inclusive_analysis.track_pt)
    #ax.set_title(fr"Subtracted 1D ${inclusive_analysis.correlation_hists_delta_phi_subtracted.signal_dominated.axis.display_str()}$,"
    #             f" {inclusive_analysis.reaction_plane_orientation.display_str()} event plane orient.,"
    #             f" {jet_pt_label}, {track_pt_label}")
    ax.legend(loc="upper right")

    # Final adjustments
    fig.tight_layout()
    # Reduce spacing between subplots
    fig.subplots_adjust(hspace=0, wspace=0)
    # Save plot and cleanup
    plot_base.save_plot(
        output_info, fig,
        f"jetH_delta_phi_{inclusive_analysis.identifier}_rp_subtracted")
    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 plot_2d_correlations(jet_hadron: "correlations.Correlations") -> None:
    """ Plot the 2D correlations. """
    canvas = ROOT.TCanvas("canvas2D", "canvas2D")

    # Iterate over 2D hists
    for name, observable in jet_hadron.correlation_hists_2d:
        hist = observable.hist.Clone(f"{observable.name}_scaled")
        logger.debug(f"name: {name}, hist: {hist}")
        # We don't want to scale the mixed event hist because we already determined the normalization
        if "mixed" not in observable.type:
            hist.Scale(jet_hadron.correlation_scale_factor)

        # We don't need the title with all of the labeling
        hist.SetTitle("")

        # Draw plot
        hist.Draw("surf2")

        # Label axes
        hist.GetXaxis().CenterTitle(True)
        hist.GetXaxis().SetTitleSize(0.08)
        hist.GetXaxis().SetLabelSize(0.06)
        hist.GetYaxis().CenterTitle(True)
        hist.GetYaxis().SetTitleSize(0.08)
        # If I remove this, it looks worse, even though this is not supposed to do anything
        hist.GetYaxis().SetTitleOffset(1.2)
        hist.GetYaxis().SetLabelSize(0.06)
        hist.GetZaxis().CenterTitle(True)
        hist.GetZaxis().SetTitleSize(0.06)
        hist.GetZaxis().SetLabelSize(0.05)
        hist.GetZaxis().SetTitleOffset(0.8)
        canvas.SetLeftMargin(0.13)

        if "mixed" in observable.type:
            hist.GetZaxis().SetTitle(r"$a(\Delta\varphi,\Delta\eta)$")
            hist.GetZaxis().SetTitleOffset(0.9)
        else:
            z_title = r"$1/N_{\mathrm{trig}}\mathrm{d^{2}}N%(label)s/\mathrm{d}\Delta\varphi\mathrm{d}\Delta\eta$"
            if "signal" in observable.type:
                z_title = z_title % {"label": ""}
            else:
                z_title = z_title % {"label": r"_{\mathrm{raw}}"}
                # Decrease size so it doesn't overlap with the other labels
                hist.GetZaxis().SetTitleSize(0.05)

            hist.GetZaxis().SetTitle(z_title)

        # Add labels
        # PDF DOES NOT WORK HERE: https://root-forum.cern.ch/t/latex-sqrt-problem/17442/15
        # Instead, print to EPS and then convert to PDF
        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)
        logger.debug(f"label: {alice_label}, system_label: {system_label}, constituent_cuts: {constituent_cuts}, leading_hadron: {leading_hadron}, jet_pt: {jet_pt}, assoc_pt: {assoc_pt}")

        tex = ROOT.TLatex()
        tex.SetTextSize(0.04)
        # Upper left side
        tex.DrawLatexNDC(.005, .96, labels.use_label_with_root(alice_label))
        tex.DrawLatexNDC(.005, .91, labels.use_label_with_root(system_label))
        tex.DrawLatexNDC(.005, .86, labels.use_label_with_root(jet_pt))
        tex.DrawLatexNDC(.005, .81, labels.use_label_with_root(jet_finding))

        # Upper right side
        tex.DrawLatexNDC(.67, .96, labels.use_label_with_root(assoc_pt))
        tex.DrawLatexNDC(.7275, .91, labels.use_label_with_root(leading_hadron))
        tex.DrawLatexNDC(.73, .86, labels.use_label_with_root(constituent_cuts))

        # Save plot
        plot_base.save_plot(jet_hadron.output_info, canvas, observable.name)

        # Draw as colz to view more precisely
        hist.Draw("colz")
        plot_base.save_plot(jet_hadron.output_info, canvas, observable.name + "_colz")

        canvas.Clear()
Example #6
0
def _extracted_values(
    analyses: Mapping[Any, "correlations.Correlations"],
    selected_iterables: Mapping[str, Sequence[Any]],
    extract_value_func: Callable[["correlations.Correlations"],
                                 analysis_objects.ExtractedObservable],
    plot_labels: plot_base.PlotLabels,
    output_name: str,
    output_info: analysis_objects.PlottingOutputWrapper,
    projection_range_func: Optional[Callable[["correlations.Correlations"],
                                             str]] = None,
    extraction_range_func: Optional[Callable[["correlations.Correlations"],
                                             str]] = None
) -> None:
    """ Plot extracted values.

    Note:
        It's best to fully define the ``extract_value_func`` function even though it can often be easily accomplished
        with a lambda because only a full function definition can use explicit type checking. Since this function uses
        a variety of different sources for the data, this type checking is particularly helpful. So writing a full
        function with full typing is strongly preferred to ensure that we get it right.

    Args:
        analyses: Correlation analyses.
        selected_iterables: Iterables that were used in constructing the analysis objects. We use them to iterate
            over some iterators in a particular order (particularly the reaction plane orientation).
        extract_value_func: Function to retrieve the extracted value and error.
        plot_labels: Titles and axis labels for the plot.
        output_name: Base of name under which the plot will be stored.
        output_info: Information needed to determine where to store the plot.
        projection_range_func: Function which will provide the projection range of the extracted value given
            the inclusive object. Default: None.
        extraction_range_func: Function which will provide the extraction range of the extracted value given
            the inclusive object. Default: None.
    """
    # Setup
    fig, ax = plt.subplots(figsize=(8, 6))
    # Specify plotting properties
    # color, marker, fill marker or not
    # NOTE: Fill marker is specified when plotting because of a matplotlib bug
    # NOTE: This depends on iterating over the EP orientation in the exact manner specified below.
    ep_plot_properties = {
        # black, diamond, no fill
        params.ReactionPlaneOrientation.inclusive: ("black", "D", "none"),
        # blue = "C0", square, fill
        params.ReactionPlaneOrientation.in_plane: ("tab:blue", "s", "full"),
        # green = "C2", triangle up, fill
        params.ReactionPlaneOrientation.mid_plane: ("tab:green", "^", "full"),
        # red = "C3", circle, fill
        params.ReactionPlaneOrientation.out_of_plane: ("tab:red", "o", "full"),
    }
    cyclers = []
    plot_property_values = list(ep_plot_properties.values())
    for i, prop in enumerate(["color", "marker", "fillstyle"]):
        cyclers.append(cycler(prop, [p[i] for p in plot_property_values]))
    # We skip the fillstyle because apparently it doesn't work with the cycler at the moment due to a bug...
    # They didn't implement their add operation to handle 0, so we have to give it the explicit start value.
    combined_cyclers = sum(cyclers[1:-1], cyclers[0])
    ax.set_prop_cycle(combined_cyclers)

    # Used for labeling purposes. The values that are used are identical for all analyses.
    inclusive_analysis: Optional["correlations.Correlations"] = None
    for displace_index, ep_orientation in enumerate(
            selected_iterables["reaction_plane_orientation"]):
        # Store the values to be plotted
        values: Dict[analysis_objects.PtBin,
                     analysis_objects.ExtractedObservable] = {}
        for key_index, analysis in \
                analysis_config.iterate_with_selected_objects(
                    analyses, reaction_plane_orientation = ep_orientation
                ):
            # Store each extracted value.
            values[analysis.track_pt] = extract_value_func(analysis)
            # These are both used for labeling purposes and are identical for all analyses that are iterated over.
            if ep_orientation == params.ReactionPlaneOrientation.inclusive and inclusive_analysis is None:
                inclusive_analysis = analysis

        # Plot the values
        bin_centers = np.array([k.bin_center for k in values])
        bin_centers = bin_centers + displace_index * 0.05
        ax.errorbar(
            bin_centers,
            [v.value for v in values.values()],
            yerr=[v.error for v in values.values()],
            label=ep_orientation.display_str(),
            linestyle="",
            fillstyle=ep_plot_properties[ep_orientation][2],
        )

    # Help out mypy...
    assert inclusive_analysis is not None

    # Labels.
    # General
    text = labels.make_valid_latex_string(
        inclusive_analysis.alice_label.display_str())
    text += "\n" + labels.system_label(
        energy=inclusive_analysis.collision_energy,
        system=inclusive_analysis.collision_system,
        activity=inclusive_analysis.event_activity)
    text += "\n" + labels.jet_pt_range_string(inclusive_analysis.jet_pt)
    text += "\n" + labels.jet_finding()
    text += "\n" + labels.constituent_cuts()
    text += "\n" + labels.make_valid_latex_string(
        inclusive_analysis.leading_hadron_bias.display_str())
    # Deal with projection range, extraction range string.
    additional_label = _proj_and_extract_range_label(
        inclusive_analysis=inclusive_analysis,
        projection_range_func=projection_range_func,
        extraction_range_func=extraction_range_func,
    )
    if additional_label:
        text += "\n" + additional_label
    # Finally, add the text to the axis.
    ax.text(0.97,
            0.97,
            text,
            horizontalalignment="right",
            verticalalignment="top",
            multialignment="right",
            transform=ax.transAxes)
    # Axes and titles
    ax.set_xlabel(
        labels.make_valid_latex_string(labels.track_pt_display_label()))
    # Apply any specified labels
    if plot_labels.title is not None:
        plot_labels.title = plot_labels.title + f" for {labels.jet_pt_range_string(inclusive_analysis.jet_pt)}"
    plot_labels.apply_labels(ax)
    ax.legend(loc="center right", frameon=False)

    # Final adjustments
    fig.tight_layout()
    # Save plot and cleanup
    plot_base.save_plot(
        output_info, fig,
        f"{output_name}_{inclusive_analysis.jet_pt_identifier}")
    plt.close(fig)