Example #1
0
    def plot_projected_templates(
        self,
        use_initial_values: bool = True,
        output_dir_path: Optional[Union[str, os.PathLike]] = None,
        output_name_tag: Optional[str] = None
    ) -> Dict[str, List[Union[str, os.PathLike]]]:
        output_lists = {"pdf": [], "png": []}

        if (output_dir_path is None) != (output_name_tag is None):
            raise ValueError(
                f"Parameter 'output_name_tag' and 'output_dir_path' must either both be provided "
                f"or both set to None!")

        for mc_channel in self._fit_model.mc_channels_to_plot:
            for dimension in range(mc_channel.binning.dimensions):
                current_binning = mc_channel.binning.get_binning_for_one_dimension(
                    dimension=dimension)
                data_column_name_for_plot = mc_channel.data_column_names[
                    dimension]

                for template in mc_channel.templates:
                    current_plot = FitTemplatePlot(  # TODO
                        variable=self.channel_variables(dimension=dimension)
                        [mc_channel.
                         name],  # TODO: channel_variables not available here... maybe implement a base function? Maybe use subsets_plotter for this?
                        binning=current_binning)

                    template_bin_count, template_bin_error_sq = template.project_onto_dimension(
                        bin_counts=template.expected_bin_counts(
                            use_initial_values=use_initial_values),
                        dimension=dimension,
                        bin_errors_squared=template.
                        expected_bin_errors_squared(
                            use_initial_values=use_initial_values))

                    current_plot.add_component(
                        label=self._get_template_label(template=template),
                        bin_counts=template_bin_count,
                        bin_errors_squared=template_bin_error_sq,
                        data_column_names=data_column_name_for_plot,
                        color=self._get_template_color(
                            key=template.process_name,
                            original_color=template.color))

                    fig, ax = plt.subplots(nrows=1,
                                           ncols=1,
                                           figsize=self._fig_size,
                                           dpi=200)
                    current_plot.plot_on(ax1=ax)

                    ax.set_title(self._get_plot_title(template=template,
                                                      channel=mc_channel),
                                 loc="right")

                    if output_dir_path is not None:
                        assert output_name_tag is not None

                        add_info = ""
                        if use_initial_values:
                            add_info = "_with_initial_values"

                        template_info = f"{mc_channel.name}_template_{template.process_name}{add_info}"
                        filename = f"{self.plot_name_prefix}_{output_name_tag}_dim_{dimension}_{template_info}"

                        export(fig=fig,
                               filename=filename,
                               target_dir=output_dir_path)
                        output_lists["pdf"].append(
                            os.path.join(output_dir_path, f"{filename}.pdf"))
                        output_lists["png"].append(
                            os.path.join(output_dir_path, f"{filename}.png"))

        return output_lists
    def plot_fit_result_projections(
        self,
        project_to: int,
        use_initial_values: bool = False,
        output_dir_path: Optional[Union[str, os.PathLike]] = None,
        output_name_tag: Optional[str] = None
    ) -> Dict[str, List[Union[str, os.PathLike]]]:
        output_lists = {"pdf": [], "png": []}

        if (output_dir_path is None) != (output_name_tag is None):
            raise ValueError(
                f"Parameter 'output_name_tag' and 'output_dir_path' must either both be provided "
                f"or both set to None!")

        for mc_channel in self._fit_model.mc_channels_to_plot:
            binning = mc_channel.binning.get_binning_for_one_dimension(
                dimension=project_to)
            data_column_name_for_plot = mc_channel.data_column_names[
                project_to]

            data_channel = self._fit_model.data_channels_to_plot.get_channel_by_name(
                name=mc_channel.name)

            data_bin_count, data_bin_errors_squared = data_channel.project_onto_dimension(
                bin_counts=data_channel.bin_counts,
                dimension=project_to,
                bin_errors_squared=data_channel.bin_errors_sq)

            plot = FitResultPlot(variable=self.channel_variables(
                dimension=project_to)[mc_channel.name],
                                 binning=binning)

            for template in mc_channel.templates:
                template_bin_count, template_bin_error_sq = template.project_onto_dimension(
                    bin_counts=template.expected_bin_counts(
                        use_initial_values=use_initial_values),
                    dimension=project_to,
                    bin_errors_squared=template.expected_bin_errors_squared(
                        use_initial_values=use_initial_values))

                plot.add_component(label=self._get_mc_label(
                    key=template.process_name,
                    original_label=template.latex_label),
                                   histogram_key=FitResultPlot.mc_key,
                                   bin_counts=template_bin_count,
                                   bin_errors_squared=template_bin_error_sq,
                                   data_column_names=data_column_name_for_plot,
                                   color=self._get_mc_color(
                                       key=template.process_name,
                                       original_color=template.color))

            plot.add_component(label=self._get_data_label(),
                               histogram_key=FitResultPlot.data_key,
                               bin_counts=data_bin_count,
                               bin_errors_squared=data_bin_errors_squared,
                               data_column_names=data_column_name_for_plot,
                               color=self._get_data_color())

            fig, axs = plt.subplots(nrows=1,
                                    ncols=1,
                                    figsize=self._fig_size,
                                    dpi=200)
            plot.plot_on(
                ax1=axs,
                #  style=???,  # str = "stacked",  # TODO: Include summed style
                #  include_sys=???,  # bool = False,
                #  markers_with_width=???,  # bool = True,
                #  sum_color=???,  # str = plot_style.KITColors.kit_purple,
                #  draw_legend=???,  # bool = True,
                #  legend_inside=???,  # bool = True,
                #  legend_cols=???,  # Optional[int] = None,
                #  legend_loc=???,  # Optional[Union[int, str]] = None,
                #  y_scale=???,  # float = 1.1
            )

            axs.set_title(self._get_channel_label(channel=mc_channel),
                          loc="right")

            if output_dir_path is not None:
                assert output_name_tag is not None

                add_info = ""
                if use_initial_values:
                    add_info = "_with_initial_values"
                filename = f"fit_result_plot_{output_name_tag}_{mc_channel.name}_dim_{project_to}_projection{add_info}"

                export(fig=fig, filename=filename, target_dir=output_dir_path)
                output_lists["pdf"].append(
                    os.path.join(output_dir_path, f"{filename}.pdf"))
                output_lists["png"].append(
                    os.path.join(output_dir_path, f"{filename}.png"))

        return output_lists
    def plot_2d_templates(
        self,
        use_initial_values: bool = True,
        output_dir_path: Optional[PathType] = None,
        output_name_tag: Optional[str] = None,
        base_color: Union[None, str, Dict[str, str]] = None,
        alternative_temp_color: Optional[Dict[str, str]] = None,
    ) -> Dict[str, List[PathType]]:
        output_lists = {
            "pdf": [],
            "png": [],
        }  # type: Dict[str, List[PathType]]

        if (output_dir_path is None) != (output_name_tag is None):
            raise ValueError(
                "Parameter 'output_name_tag' and 'output_dir_path' must either both be provided or both set to None!"
            )

        if isinstance(base_color, dict):
            c_map_base_color = base_color  # type: Union[str, Dict[str, str]]
        else:
            c_map_base_color = self.default_2d_c_map_base_color if base_color is None else base_color
        alt_temp_color_dict = {} if alternative_temp_color is None else alternative_temp_color  # type: Dict[str, str]

        for mc_channel in self._fit_model.mc_channels_to_plot:
            for dim_pair in iter_combinations(
                    list(range(mc_channel.binning.dimensions)), 2):
                current_binning = mc_channel.binning.get_binning_for_x_dimensions(
                    dimensions=dim_pair)

                x_variable = self.channel_variables(dimension=dim_pair[0])[
                    mc_channel.name]  # type: HistVariable
                y_variable = self.channel_variables(dimension=dim_pair[1])[
                    mc_channel.name]  # type: HistVariable

                for template in mc_channel.templates:
                    template_bin_count, template_bin_error_sq = template.project_onto_two_dimensions(
                        bin_counts=template.expected_bin_counts(
                            use_initial_values=use_initial_values),
                        dimensions=dim_pair,
                        bin_errors_squared=template.
                        expected_bin_errors_squared(
                            use_initial_values=use_initial_values),
                    )

                    template_color = self._get_template_color(
                        key=template.process_name,
                        original_color=template.color)

                    bin_scaling_tuple = current_binning.get_bin_scaling_per_dim_tuple(
                    )  # type: Tuple[np.ndarray, ...]
                    bin_scaling = np.outer(
                        *bin_scaling_tuple)  # type: np.ndarray
                    assert len(bin_scaling.shape) == 2, bin_scaling.shape
                    assert bin_scaling.shape[0] == bin_scaling_tuple[0].shape[
                        0], (
                            bin_scaling.shape,
                            bin_scaling_tuple[0].shape,
                        )
                    assert bin_scaling.shape[1] == bin_scaling_tuple[1].shape[
                        0], (
                            bin_scaling.shape,
                            bin_scaling_tuple[1].shape,
                        )

                    assert template_bin_count.shape == bin_scaling.shape, (
                        template_bin_count.shape, bin_scaling.shape)
                    value_matrix = template_bin_count * bin_scaling  # type: np.ndarray

                    fig, ax = plt.subplots(nrows=1,
                                           ncols=1,
                                           figsize=self._fig_size,
                                           dpi=200)

                    c_values = [0.0, np.max(value_matrix)]  # type: List[float]

                    if np.sum(template_bin_count) == 0:
                        c_values = [0.0, 1.0]  # type: List[float]
                        value_matrix = np.ones_like(value_matrix) * 0.8

                    c_map_temp_color = alt_temp_color_dict.get(
                        template_color, template_color)
                    if isinstance(c_map_base_color, dict):
                        base_c = c_map_base_color.get(
                            template_color,
                            self.default_2d_c_map_base_color)  # type: str
                        colors = [base_c, c_map_temp_color]  # type: List[str]
                    else:
                        colors = [c_map_base_color, c_map_temp_color]

                    c_norm = plt.Normalize(min(c_values), max(c_values))
                    c_tuples = list(zip(map(c_norm, c_values), colors))
                    color_map = mpl_colors.LinearSegmentedColormap.from_list(
                        name="", colors=c_tuples)

                    color_map.set_bad(color=plot_style.KITColors.white)
                    value_matrix[value_matrix == 0.0] = np.nan
                    heatmap = ax.imshow(X=np.flip(value_matrix.T, axis=0),
                                        cmap=color_map,
                                        aspect="auto")
                    plt.colorbar(heatmap)

                    self._set_2d_axis_tick_labels(ax=ax,
                                                  binning=current_binning)

                    ax.set_title(self._get_plot_title(template=template,
                                                      channel=mc_channel),
                                 loc="right")

                    ax.set_xlabel(x_variable.x_label, plot_style.xlabel_pos)
                    ax.set_ylabel(y_variable.x_label, plot_style.ylabel_pos)

                    if np.sum(template_bin_count) == 0:
                        ax.text(
                            x=0.5,
                            y=0.5,
                            s="No Data",
                            fontsize="x-large",
                            ha="center",
                            va="center",
                            zorder=10,
                            transform=ax.transAxes,
                        )

                    if output_dir_path is not None:
                        assert output_name_tag is not None

                        add_info = ""
                        if use_initial_values:
                            add_info = "_with_initial_values"

                        dims_str = f"{dim_pair[0]}_{dim_pair[1]}"  # type: str
                        template_info = f"{mc_channel.name}_template_{template.process_name}{add_info}"
                        filename = f"{self.plot_2d_name_prefix}_{output_name_tag}_dim_{dims_str}_{template_info}"

                        export(fig=fig,
                               filename=filename,
                               target_dir=output_dir_path,
                               close_figure=True)
                        output_lists["pdf"].append(
                            os.path.join(output_dir_path, f"{filename}.pdf"))
                        output_lists["png"].append(
                            os.path.join(output_dir_path, f"{filename}.png"))

        return output_lists
    def plot_fit_result(
        self,
        use_initial_values: bool = False,
        output_dir_path: Optional[Union[str, os.PathLike]] = None,
        output_name_tag: Optional[str] = None
    ) -> Dict[str, List[Union[str, os.PathLike]]]:
        output_lists = {"pdf": [], "png": []}

        if (output_dir_path is None) != (output_name_tag is None):
            raise ValueError(
                f"Parameter 'output_name_tag' and 'output_dir_path' must either both be provided "
                f"or both set to None!")

        for mc_channel in self._fit_model.mc_channels_to_plot:
            current_binning = mc_channel.binning.get_binning_for_one_dimension(
                dimension=self.reference_dimension)
            data_column_name_for_plot = mc_channel.data_column_names[
                self.reference_dimension]

            data_channel = self._fit_model.data_channels_to_plot.get_channel_by_name(
                name=mc_channel.name)
            data_bin_count = data_channel.bin_counts
            data_bin_errors_squared = data_channel.bin_errors_sq

            for counter, sub_bin_info in enumerate(
                    self._get_sub_bin_infos_for(
                        channel_name=mc_channel.name,
                        reference_dimension=self.reference_dimension  # TODO
                    )):
                sub_bin_info_text = self._get_sub_bin_info_text(
                    channel_name=mc_channel.name,
                    sub_bin_infos=sub_bin_info,
                    reference_dimension=self.reference_dimension)

                nd_array_slices = self._get_slices(sub_bin_info=sub_bin_info)

                current_plot = FitResultPlot(variable=self.channel_variables(
                    dimension=self.reference_dimension)[mc_channel.name],
                                             binning=current_binning)

                for template in mc_channel.templates:
                    template_bin_count = template.expected_bin_counts(
                        use_initial_values=use_initial_values)
                    template_bin_error_sq = template.expected_bin_errors_squared(
                        use_initial_values=use_initial_values)

                    subset_bin_count = template_bin_count[nd_array_slices]
                    subset_bin_errors_squared = template_bin_error_sq[
                        nd_array_slices]

                    current_plot.add_component(
                        label=self._get_mc_label(
                            key=template.process_name,
                            original_label=template.latex_label),
                        histogram_key=FitResultPlot.mc_key,
                        bin_counts=subset_bin_count,
                        bin_errors_squared=subset_bin_errors_squared,
                        data_column_names=data_column_name_for_plot,
                        color=self._get_mc_color(
                            key=template.process_name,
                            original_color=template.color))

                subset_data_bin_count = data_bin_count[nd_array_slices]
                subset_data_bin_errors_squared = data_bin_errors_squared[
                    nd_array_slices]

                current_plot.add_component(
                    label=self._get_data_label(),
                    histogram_key=FitResultPlot.data_key,
                    bin_counts=subset_data_bin_count,
                    bin_errors_squared=subset_data_bin_errors_squared,
                    data_column_names=data_column_name_for_plot,
                    color=self._get_data_color())

                fig, axs = plt.subplots(nrows=1,
                                        ncols=1,
                                        figsize=self._fig_size,
                                        dpi=200)
                current_plot.plot_on(
                    ax1=axs,
                    #  style=???,  # str = "stacked",  # TODO: Include summed style
                    #  include_sys=???,  # bool = False,
                    #  markers_with_width=???,  # bool = True,
                    #  sum_color=???,  # str = plot_style.KITColors.kit_purple,
                    #  draw_legend=???,  # bool = True,
                    #  legend_inside=???,  # bool = True,
                    #  legend_cols=???,  # Optional[int] = None,
                    #  legend_loc=???,  # Optional[Union[int, str]] = None,
                    #  y_scale=???,  # float = 1.1
                )

                bin_info_pos = "right"
                if bin_info_pos == "left" or sub_bin_info_text is None:
                    axs.set_title(self._get_channel_label(channel=mc_channel),
                                  loc="right")
                else:
                    fig.suptitle(self._get_channel_label(channel=mc_channel),
                                 x=0.9,
                                 horizontalalignment="right")

                if sub_bin_info_text is not None:
                    info_title = sub_bin_info_text
                    if axs.get_ylim()[1] > 0.85e4 and bin_info_pos == "left":
                        padding = " " * 9
                        info_title = "\n".join([
                            padding + info
                            for info in sub_bin_info_text.split("\n")
                        ])

                    axs.set_title(info_title,
                                  loc=bin_info_pos,
                                  fontsize=6,
                                  color=plot_style.KITColors.dark_grey)

                if output_dir_path is not None:
                    assert output_name_tag is not None

                    add_info = ""
                    if use_initial_values:
                        add_info = "_with_initial_values"
                    filename = f"fit_result_plot_{output_name_tag}_{mc_channel.name}_bin_{counter}{add_info}"

                    export(fig=fig,
                           filename=filename,
                           target_dir=output_dir_path)
                    output_lists["pdf"].append(
                        os.path.join(output_dir_path, f"{filename}.pdf"))
                    output_lists["png"].append(
                        os.path.join(output_dir_path, f"{filename}.png"))

        return output_lists