vis_utils.py 文件源码

def plot_2d_histogram(hist, x_lim, y_lim, title, x_label, y_label, pdf_file_name):
    """Plot 2d histogram with matplotlib

    :param hist: input numpy histogram = x_bin_edges, y_bin_edges, bin_entries_2dgrid
    :param tuple x_lim: range tuple of x-axis (min,max)
    :param tuple y_lim: range tuple of y-axis (min,max)
    :param str title: title of plot
    :param str x_label: Label for histogram x-axis
    :param str y_label: Label for histogram y-axis
    :param str pdf_file_name: if set, will store the plot in a pdf file
    """
    # import matplotlib here to prevent import before setting backend in
    # core.execution.run_eskapade
    import matplotlib.pyplot as plt
    from matplotlib.backends.backend_pdf import PdfPages

    fig = plt.figure(figsize=(7, 5))

    try:
        x_ranges = hist[0]
        y_ranges = hist[1]
        grid = hist[2]
    except BaseException:
        raise ValueError('Cannot extract ranges and grid from input histogram')

    ax = plt.gca()
    ax.pcolormesh(x_ranges, y_ranges, grid)
    ax.set_ylim(y_lim)
    ax.set_xlim(x_lim)
    ax.set_title(title)

    plt.xlabel(x_label, fontsize=14)
    plt.ylabel(y_label, fontsize=14)
    plt.grid()

    if pdf_file_name:
        pdf_file = PdfPages(pdf_file_name)
        plt.savefig(pdf_file, format='pdf', bbox_inches='tight', pad_inches=0)
        plt.close()
        pdf_file.close()