histos.py

from typing import List, Tuple

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.figure import Figure


def hist_n_particles(q: List[int], label: str) -> Figure:
    """Generate histogram for particle counts on events.

    :param q: Count per event.
    :type q: List[int]
    :param label: Plot title.
    :type label: str
    :return: Figure with histogram and histogram ratio.
    :rtype: Figure
    """

    fig, ax = plt.subplots(nrows=1, figsize=(15, 8))

    bins, edges = np.histogram(q, bins=15, range=(0, 15))

    for idx, val in enumerate(bins[::-1]):
        if val > 0:
            max_idx = len(bins) - idx - 1
            edges = edges[:max_idx]
            bins = bins[:max_idx]
            break

    ax.bar(edges, bins, width=1, alpha=0.6)
    ax.set_title(label, fontsize=20, y=1.04)
    ax.set_xticks(edges)

    return fig


def hist_var(q: List[float], ax: plt.Axes, **kwargs) -> plt.Axes:
    """Create histogram with error bars.

    :param q: Values to create histogram.
    :type q: List[float]
    :param ax: Axes in which histrogram is plotted.
    :type ax: plt.Axes
    :return: Axes with histogram.
    :rtype: plt.Axes
    """

    bins, edges, _ = ax.hist(
        q, alpha=0.6, histtype="step", align="left", linewidth=4, **kwargs
    )
    errors = np.sqrt(bins)
    bin_width = edges[1] - edges[0]

    ax.bar(
        x=edges[:-1],
        bottom=bins,
        height=errors,
        width=bin_width,
        alpha=0.0,
        color="w",
        hatch="/",
    )
    ax.bar(
        x=edges[:-1],
        bottom=bins,
        height=-errors,
        width=bin_width,
        alpha=0.0,
        color="w",
        hatch="/",
    )
    return ax


def ratio_hist(
    processes_q: List[List[float]],
    hist_labels: List[str],
    reference_label: str,
    n_bins: int,
    hist_range: Tuple[int, int],
    title: str,
    figsize=(15, 8),
) -> Figure:
    """Generate histrograms with ratio pad

    :param processes_q: Quantity for each event and process
    :type processes_q: List[List[float]]
    :param hist_labels: Labels for each process
    :type hist_labels: List[str]
    :param reference_label: Label of process taken as the denominator of ratios
    :type reference_label: str
    :param n_bins: Number of bins for histograms
    :type n_bins: int
    :param hist_range: Range for histogram bins
    :type hist_range: Tuple[int]
    :param title: Plot title
    :type title: str
    :param figsize: Figure size for output plot
    :type figsize: Tuple[int]
    :return: Figure with histogram and histogram ratio.
    :rtype: Figure
    """

    fig, ax = plt.subplots(
        nrows=len(processes_q),
        ncols=1,
        gridspec_kw={"height_ratios": [3] + [1] * (len(processes_q) - 1)},
        sharex=True,
        figsize=figsize,
    )
    legends = []

    p_bins = {}
    p_edges = {}
    p_errors = {}
    edges = None
    for p, label in zip(processes_q, hist_labels):
        bins = n_bins
        if edges is not None:
            bins = edges
        bins, edges, _ = ax[0].hist(
            x=p,
            bins=bins,
            range=hist_range,
            fill=False,
            label=label,
            align="left",
            histtype="step",
            linewidth=4,
        )
        p_bins[label] = bins
        p_edges[label] = edges
        p_errors[label] = np.sqrt(bins)
        legends += [label]

    bin_width = edges[1] - edges[0]

    for label in hist_labels:
        ax[0].bar(
            x=p_edges[label][:-1],
            bottom=p_bins[label],
            height=p_errors[label],
            width=bin_width,
            alpha=0.0,
            color="w",
            hatch="/",
        )
        ax[0].bar(
            x=p_edges[label][:-1],
            bottom=p_bins[label],
            height=-p_errors[label],
            width=bin_width,
            alpha=0.0,
            color="w",
            hatch="/",
        )
        legends += ["_", "_"]

    ax[0].set_ylabel("Events", fontsize=15)
    ax[0].set_title(title, fontsize=20)
    legends[-1] = "Stat. Uncertainty"
    ax[0].legend(legends)

    plot_idx = 1
    ref_bins = p_bins[reference_label]
    ref_edges = p_edges[reference_label]
    ref_frac_error = p_errors[reference_label] / ref_bins
    for label in hist_labels:
        if label == reference_label:
            continue
        ratios = p_bins[label] / ref_bins
        error_ratio = ratios * (ref_frac_error + p_errors[label] / p_bins[label])

        ax[plot_idx].bar(
            bottom=1.0,
            height=error_ratio,
            x=ref_edges[:-1],
            width=bin_width,
            alpha=0.3,
            color="blue",
        )
        ax[plot_idx].bar(
            bottom=1.0,
            height=-error_ratio,
            x=ref_edges[:-1],
            width=bin_width,
            alpha=0.3,
            color="blue",
        )
        ax[plot_idx].scatter(ref_edges[:-1], ratios, marker="o", color="black")
        ax[plot_idx].set_ylabel(f"{label}/{reference_label}")
        plot_idx += 1

    return fig