Histogram

def divide(self, other: Histogram) -> tuple[np.ndarray, np.ndarray]:
    """Calculate ratio between two class objects.

    Parameters
    ----------
    other : Histogram
        Second histogram object to calculate ratio with

    Returns
    -------
    np.ndarray
        Ratio
    np.ndarray
        Ratio error

    Raises
    ------
    ValueError
        If binning is not identical between 2 objects
        If hist attribute is not set for one of the two histograms
        If bin_edges attribute is not set for one of the two histograms
    """
    try:
        np.all(self.bin_edges == other.bin_edges)

    except ValueError as err:
        raise ValueError("The binning of the two given objects do not match.") from err

    # Bins where the reference histogram is empty/zero, are given a ratio of np.inf
    # which means that the ratio plot will not have any entry in these bins.
    ratio, ratio_unc = hist_ratio(
        numerator=self.hist,
        denominator=other.hist,
        numerator_unc=self.unc,
        step=False,
    )
    # To use the matplotlib.step() function later on, the first bin is duplicated
    ratio = np.append(np.array([ratio[0]]), ratio)
    ratio_unc = np.append(np.array([ratio_unc[0]]), ratio_unc)

    return (ratio, ratio_unc)

def divide_data_mc(
    self,
    ref_hist: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Similar as divide, but the second item doesn't need to be a histogram object.

    Parameters
    ----------
    ref_hist : np.ndarray
        Hist weights of the reference.

    Returns
    -------
    tuple
        Tuple of the ratios and ratio uncertaintes for the bins

    Raises
    ------
    ValueError
        If the binning of the two histograms doesn't match
    """
    try:
        self.hist + ref_hist
    except ValueError as err:
        raise ValueError("The binning of the two given objects do not match.") from err

    # Bins where the reference histogram is empty/zero, are given a ratio of np.inf
    # which means that the ratio plot will not have any entry in these bins.
    ratio, ratio_unc = hist_ratio(
        numerator=self.hist,
        denominator=ref_hist,
        numerator_unc=self.unc,
        step=False,
    )
    # To use the matplotlib.step() function later on, the first bin is duplicated
    ratio = np.append(np.array([ratio[0]]), ratio)
    ratio_unc = np.append(np.array([ratio_unc[0]]), ratio_unc)

    return (ratio, ratio_unc)

def get_discrete_values(self) -> None:
    """Get discrete values of a variable and adjust the bins accordingly.

    Raises
    ------
    ValueError
        If the bin width is larger than 1 such that potentially not
        all discrete values are in a seperate bin
        If the number of bins is set to 1 such that no values can be
        distinguished
    """
    assert self.discrete_vals is not None
    if len(self.bin_edges) > 1:
        if abs(self.bin_edges[1] - self.bin_edges[0]) <= 1:
            indice = [
                i
                for i in range(len(self.bin_edges) - 1)
                for discrete_val in self.discrete_vals
                if self.bin_edges[i] <= discrete_val < self.bin_edges[i + 1]
            ]
            self.hist = self.hist[indice]
            self.unc = self.unc[indice]
            self.band = self.band[indice]
            bins = np.linspace(0, len(self.discrete_vals), len(self.discrete_vals) + 1)
            self.bin_edges = bins

        else:
            raise ValueError(
                "Bin width is larger than 1. Choose a binning with a bin"
                " width<= 1 to plot only discrete values."
            )
    else:
        raise ValueError(
            "Choose a binning with more than one bin in order to plot only discrete values."
        )

def update(self, values: np.ndarray, weights: np.ndarray | None = None) -> None:
    """Update the existing histogram with new values.

    Parameters
    ----------
    values : np.ndarray
        New values that are to be added.
    weights : np.ndarray | None, optional
        Weights for each entry in values. Must be the same shape/size as values.
        If None, each value in values will be weighted the same, by default None
    """
    # Replace the weights if None
    if weights is None:
        weights = np.ones_like(values)

    # Check that the weights are a np.ndarray
    assert isinstance(weights, np.ndarray)

    # Call the hist_w_unc function for the new values
    _, incoming_hist, incoming_unc, _ = hist_w_unc(
        arr=values,
        bins=self.bins,
        filled=self.filled,
        bins_range=self.bins_range,
        normed=self.norm,
        weights=weights,
        bin_edges=self.bin_edges,
        sum_squared_weights=self.sum_squared_weights,
        underoverflow=self.underoverflow,
    )

    # Get the new sum_of_weights
    incoming_sum_of_weights = float(np.sum(weights))

    if self.norm:
        self.hist = (
            self.hist * self.sum_of_weights + incoming_hist * incoming_sum_of_weights
        ) / (self.sum_of_weights + incoming_sum_of_weights)
        self.unc = np.sqrt(
            (self.unc * self.sum_of_weights) ** 2
            + (incoming_unc * incoming_sum_of_weights) ** 2
        ) / (self.sum_of_weights + incoming_sum_of_weights)

    else:
        self.hist += incoming_hist
        self.unc = np.sqrt(self.unc**2 + incoming_unc**2)

    self.sum_of_weights += incoming_sum_of_weights
    self.band = self.hist - self.unc

def add(
    self,
    histogram: Histogram,
    key: str | None = None,
    reference: bool = False,
) -> None:
    """Adding histogram object to figure.

    Parameters
    ----------
    histogram : Histogram
        Histogram curve
    key : str | None, optional
        Unique identifier for histogram, by default None
    reference : bool, optional
        If this histogram is used as reference for ratio calculation, by default
        False

    Raises
    ------
    KeyError
        If unique identifier key is used twice
    """
    # If key not defined, set it to a numerical value
    key = cast(str, key if key is not None else f"{len(self.plot_objects) + 1}")

    # Check that key is not double used
    if key in self.plot_objects:
        raise KeyError(f"Duplicated key {key} already used for unique identifier.")

    # Add key to histogram object
    histogram.key = key
    logger.debug("Adding histogram %s", key)

    # Set linestyle
    if histogram.linestyle is None:
        if histogram.is_data is True:
            histogram.linestyle = "None"
        else:
            histogram.linestyle = "-"
    # Set marker
    if histogram.marker is None:
        if histogram.is_data is True:
            histogram.marker = "."
        else:
            histogram.marker = ""
    # Set colours
    if histogram.colour is None:
        histogram.colour = get_good_colours()[len(self.plot_objects)]
    # Set alpha
    if histogram.alpha is None:
        histogram.alpha = 1
    # Set linewidth
    if histogram.linewidth is None:
        histogram.linewidth = 1.6
    # Set markersize
    if histogram.markersize is None:
        histogram.markersize = 10

    self.plot_objects[key] = histogram
    self.add_order.append(key)
    if reference is True:
        self.set_reference(key)

def add_bin_width_to_ylabel(self):
    """Adds the bin width to the ylabel of a histogram plot. If the bin with is
    smaller than 0.01, scientific notation will be used.

    Raises
    ------
    ValueError
        If plotting_done is False (therefore `bins` is not yet calculated)
    """
    if self.plotting_done is False:
        raise ValueError(
            "`add_bin_width_to_ylabel` should be called after plotting, since bins "
            "are calculated during plotting."
        )

    bin_width = abs(self.bin_edges[1] - self.bin_edges[0])
    if bin_width < 1e-2:
        self.ylabel = f"{self.ylabel} / {bin_width:.0e}"
    else:
        self.ylabel = f"{self.ylabel} / {bin_width:.2f}"
    self.set_ylabel(self.axis_top)

def draw(self, labelpad: int | None = None) -> None:
    """Draw figure.

    Parameters
    ----------
    labelpad : int | None, optional
        Spacing in points from the axes bounding box including
        ticks and tick labels, by default "ratio"
    """
    plt_handles = self.plot()

    if self.n_ratio_panels > 0:
        self.plot_ratios()

    self.set_xlim(
        self.bin_edges[0] if self.xmin is None else self.xmin,
        self.bin_edges[-1] if self.xmax is None else self.xmax,
    )
    self.set_log()
    self.set_y_lim()
    self.set_xlabel()
    self.set_tick_params()
    self.set_ylabel(self.axis_top)

    if self.n_ratio_panels > 0:
        assert isinstance(self.ylabel_ratio, list)
        self.set_ylabel(
            self.ratio_axes[0],
            self.ylabel_ratio[0],
            align="center",
            labelpad=labelpad,
        )

    if self.bin_width_in_ylabel is True:
        self.add_bin_width_to_ylabel()

    legend_axis = self.axis_top

    self.make_legend(plt_handles, ax_mpl=legend_axis)
    self.set_title()

    if self.apply_atlas_style:
        self.atlasify()

def get_reference_histo(self, histo: Histogram) -> Histogram | None:
    """Get reference histogram from list of references.

    Parameters
    ----------
    histo : Histogram
        Histogram we want to calculate the ratio for

    Returns
    -------
    reference_histo : Histogram
        The reference histogram

    Raises
    ------
    ValueError
        If no reference histo was found or multiple matches.
    """
    matches = 0
    reference_histo = None

    for key in self.reference_object:
        reference_candidate = self.plot_objects[key]
        if histo.ratio_group is not None:
            if histo.ratio_group == reference_candidate.ratio_group:
                matches += 1
                reference_histo = reference_candidate
        else:
            matches += 1
            reference_histo = reference_candidate

    if matches != 1:
        raise ValueError(
            f"Found {matches} matching reference candidates, but only one match is allowed."
        )

    logger.debug("Reference histogram for '%s' is '%s'", histo.key, reference_histo.key)

    return reference_histo

def plot(self, **kwargs: Any):
    """Plotting curves. This also generates the bins of the histograms that are
    added to the plot. Plot objects are drawn in the same order as they were added
    to the plot.

    Parameters
    ----------
    **kwargs : Any
        Keyword arguments passed to matplotlib.axes.Axes.hist()

    Returns
    -------
    Line2D
        matplotlib Line2D object

    Raises
    ------
    ValueError
        If specified bins type is not supported.
    """
    # Get the settings from the to-be-plotted Histogram objects
    norm_list = [iter_histo.norm for iter_histo in self.plot_objects.values()]
    bin_edges_list = [iter_histo.bin_edges for iter_histo in self.plot_objects.values()]

    # Check consistency of norm setting
    self.norm = norm_list[0]
    if any(n != self.norm for n in norm_list):
        raise ValueError("Histogram objects have different settings for 'norm'!")

    # Check consistency of bin_edges
    self.bin_edges = bin_edges_list[0]
    if any(not np.array_equal(e, self.bin_edges) for e in bin_edges_list):
        raise ValueError("Histogram objects have different 'bin_edges'!")

    if self.ylabel is not None and self.norm and "norm" not in self.ylabel.lower():
        logger.warning(
            "You are plotting normalised distributions but 'norm' is not "
            "included in your y-label."
        )

    if self.norm is True and self.stacked is True:
        raise ValueError(
            "Stacked plots and normalised plots at the same time are not available."
        )

    # Init the plt_handles list
    plt_handles = []

    # Stacked dict for the stacked histogram
    self.stacked_dict: dict[str, Any] = {
        "x": [],
        "bins": [],
        "weights": [],
        "color": [],
        "unc": None,
    }

    for key in self.add_order:
        elem = self.plot_objects[key]

        # If discrete values were used, adapt the x-axis ticks
        if elem.discrete_vals is not None:
            self.axis_top.set_xticks(elem.bin_edges[:-1] + 0.5)
            self.axis_top.set_xticklabels(elem.discrete_vals, rotation=33)

        # Check if the histogram is data
        if elem.is_data is True:
            # Plot data
            self.axis_top.errorbar(
                x=(elem.bin_edges[:-1] + elem.bin_edges[1:]) / 2,
                y=elem.hist,
                yerr=elem.unc if self.draw_errors else 0,
                color=elem.colour,
                label=elem.label,
                alpha=elem.alpha,
                linewidth=elem.linewidth,
                linestyle=elem.linestyle,
                marker=elem.marker,
                markersize=elem.markersize,
            )

            plt_handles.append(
                mpl.lines.Line2D(
                    [],
                    [],
                    color=elem.colour,
                    label=elem.label,
                    alpha=elem.alpha,
                    linewidth=elem.linewidth,
                    linestyle=elem.linestyle,
                    marker=elem.marker,
                )
            )

        elif self.stacked:
            self.stacked_dict["x"].append(elem.bin_edges[:-1])
            self.stacked_dict["bins"].append(elem.bin_edges)
            self.stacked_dict["weights"].append(elem.hist)
            self.stacked_dict["color"].append(elem.colour)

            if self.stacked_dict["unc"] is None:
                self.stacked_dict["unc"] = elem.unc

            else:
                self.stacked_dict["unc"] = np.sqrt(self.stacked_dict["unc"] ** 2 + elem.unc**2)

            # Add the element to the legend with a "bar"
            plt_handles.append(
                mpl.patches.Patch(
                    color=elem.colour,
                    label=elem.label,
                    alpha=elem.alpha,
                )
            )

        else:
            # Plot histogram
            self.axis_top.hist(
                x=elem.bin_edges[:-1],
                bins=elem.bin_edges,
                weights=elem.hist,
                histtype=elem.histtype,
                color=elem.colour,
                label=elem.label,
                alpha=elem.alpha,
                linewidth=elem.linewidth,
                linestyle=elem.linestyle,
                **kwargs,
            )

            # Plot histogram uncertainty
            if self.draw_errors:
                bottom_error = np.array([elem.band[0], *elem.band.tolist()])
                top_error = elem.band + 2 * elem.unc
                top_error = np.array([top_error[0], *top_error.tolist()])
                self.axis_top.fill_between(
                    x=elem.bin_edges,
                    y1=bottom_error,
                    y2=top_error,
                    color=elem.colour,
                    alpha=0.3,
                    zorder=1,
                    step="pre",
                    edgecolor="none",
                )

            # Add standard "Line" to legend
            plt_handles.append(
                mpl.lines.Line2D(
                    [],
                    [],
                    color=elem.colour,
                    label=elem.label,
                    alpha=elem.alpha,
                    linewidth=elem.linewidth,
                    linestyle=elem.linestyle,
                    marker=elem.marker,
                )
            )

    if self.stacked:
        self.axis_top.hist(
            x=self.stacked_dict["x"],
            bins=self.stacked_dict["bins"][0],
            weights=self.stacked_dict["weights"],
            color=self.stacked_dict["color"],
            histtype=self.histtype,
            alpha=elem.alpha,
            linewidth=elem.linewidth,
            linestyle=elem.linestyle,
            stacked=self.stacked,
            **kwargs,
        )

        # Create a total weights entry to correctly plot the ratio
        # Total weights is here the y-value of all contributions stacked
        self.stacked_dict["total_weights"] = np.sum(self.stacked_dict["weights"], axis=0)

    # Check if errors should be drawn
    # If stacked is true, plot the combined uncertainty
    if self.draw_errors and self.stacked:
        # Calculate the y-values of the bottom error
        bottom_error = self.stacked_dict["total_weights"] - self.stacked_dict["unc"]
        bottom_error = np.array([bottom_error[0], *bottom_error.tolist()])

        # Calculate the y-values of the top error
        top_error = self.stacked_dict["total_weights"] + self.stacked_dict["unc"]
        top_error = np.array([top_error[0], *top_error.tolist()])

        # Fill the space between bottom and top with the unc. band
        self.axis_top.fill_between(
            x=elem.bin_edges,
            y1=bottom_error,
            y2=top_error,
            alpha=0.5,
            zorder=1,
            step="pre",
            facecolor="white",
            edgecolor="black",
            linewidth=0,
            hatch="/////",
        )

        # Add a label for the unc. in the legend
        plt_handles.append(
            mpl.patches.Patch(
                facecolor="white",
                edgecolor="black",
                label="Stat. unc.",
                linewidth=0,
                alpha=0.5,
                hatch="/////",
            )
        )

    self.plotting_done = True
    return plt_handles

def plot_ratios(self):
    """Plotting ratio histograms.

    Raises
    ------
    ValueError
        If no reference histogram is defined
    """
    # Check if this is a stacked plot
    # Plot ratio only between data and the stacked histos
    for key in self.add_order:
        # Get the object which is to be plotted
        elem = self.plot_objects[key]

        if elem.bin_edges is None:
            raise ValueError(
                "Bin edges of plot object not set. This is done in "
                "histogram_plot.plot(), so it has to be called before "
                "plot_ratios() is called."
            )

        # Check if this is going to be Data/MC (Data/stacked plot)
        if self.stacked:
            # Check this is data
            if not elem.is_data:
                continue

            # Using the total weights (full stacked histo) as reference for data
            ratio, ratio_unc = elem.divide_data_mc(
                ref_hist=self.stacked_dict["total_weights"],
            )

        else:
            if len(self.reference_object) == 0:
                raise ValueError("Please specify a reference curve.")

            ratio, ratio_unc = elem.divide(self.get_reference_histo(elem))

        # Plot the ratio values with the step function
        if self.stacked:
            if elem.is_data is True:
                self.ratio_axes[0].errorbar(
                    x=(elem.bin_edges[:-1] + elem.bin_edges[1:]) / 2,
                    y=ratio[1:],
                    yerr=ratio_unc[1:] if self.draw_errors else 0,
                    color=elem.colour,
                    label=elem.label,
                    alpha=elem.alpha,
                    linewidth=elem.linewidth,
                    linestyle=elem.linestyle,
                    marker=elem.marker,
                    markersize=elem.markersize,
                )

        else:
            self.ratio_axes[0].step(
                x=elem.bin_edges,
                y=ratio,
                color=elem.colour,
                linewidth=elem.linewidth,
                linestyle=elem.linestyle,
            )

            # Plot the ratio uncertainty
            if self.draw_errors:
                self.ratio_axes[0].fill_between(
                    x=elem.bin_edges,
                    y1=np.nan_to_num(ratio - ratio_unc, nan=0, posinf=0),
                    y2=np.nan_to_num(ratio + ratio_unc, nan=0, posinf=0),
                    color=elem.colour,
                    alpha=0.3,
                    zorder=1,
                    step="pre",
                    edgecolor="none",
                )

    if self.stacked and self.draw_errors:
        self.ratio_axes[0].fill_between(
            x=elem.bin_edges,
            y1=np.nan_to_num((ratio - ratio_unc) / ratio, nan=0, posinf=0),
            y2=np.nan_to_num((ratio + ratio_unc) / ratio, nan=0, posinf=0),
            color=elem.colour,
            alpha=0.3,
            zorder=1,
            step="pre",
            edgecolor="none",
        )

def set_reference(self, key: str):
    """Setting the reference histogram curves used in the ratios.

    Parameters
    ----------
    key : str
        Unique identifier of histogram object
    """
    self.reference_object.append(key)
    logger.debug("Adding '%s' to reference histogram(s)", key)