plot_config.py

"""
Created on Fri Mar 26 12:16:29 2021

Author: Shyam Bhuller

Description: A script conatining boiler plate code for creating plots with matplotlib.
"""
import copy as copy_lib
import math
import warnings

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.ticker import AutoMinorLocator, MultipleLocator
from scipy.stats import iqr



class PlotConfig():

    def __init__(self):
        # self.PLT_STYLE = None
        self.PLT_STYLE = "/users/wx21978/projects/pion-phys/pi0-analysis/analysis/config/thesis_plotstyle.mplstyle"
        self.FIG_SIZE = "single"
        self.FIG_FACECOLOR = 'white'
        self.AXIS_FACECOLOR = 'white'
        self.BINS = 'best'                # Number of bins
        # Allows tweaking the number of bins produced by the "best" algorithm
        self.BEST_BINS_MULTIPLIER = 1
        self.GRID_COLOR = 'gray'                # Color of plot grid
        self.GRID_ALPHA = 0.25                  # Transparency of the plot grid
        self.LEGEND_COLOR = 'white'
        self.MINOR_GRID_X = False
        self.MINOR_GRID_Y = False
        self.SPINE_COLOR = 'black'
        self.TICK_SIZE = 18
        self.TICK_SIZE_MINOR = 14
        self.TICK_LENGTH = 14
        self.TICK_WIDTH = 4
        self.MINOR_TICK_SPACING_X = "auto"
        self.MINOR_TICK_SPACING_Y = "auto"
        self.LEGEND_SIZE = 18
        self.LEGEND_SIZE_DOUBLE = 11
        self.TITLE_SIZE = 28
        self.TITLE_SIZE_DOUBLE = 22
        self.TITLE = None
        self.AX_ALPHA = 0.5
        self.TEXT_SIZE = 18
        self.DPI = 300

        self.SAVE_FOLDER = None                  # Folder to save the plot
        self.SHOW_PLOT = False                 # Show the plot?
        # Count to get number of individual hits per APA. "sum" to get total of hits per APA (Summed Hit Size)
        self.N_HITS_WHICH = "count"
        self.LABEL_SIZE = 24
        self.LOGSCALE = False
        self.LINEWIDTH = 2

        self.HIST_COLOR1 = "tab:blue"
        self.HIST_COLOR2 = "tab:orange"
        self.HIST_COLOR3 = "tab:green"
        self.HIST_COLOR4 = "tab:purple"
        self.HIST_DENSITY = False                 # Normalise the count of the histograms
        self.HIST_TYPE = "step"

        self.HIST_COLOURS = {
            0: self.HIST_COLOR1,
            1: self.HIST_COLOR2,
            2: self.HIST_COLOR3,
            3: self.HIST_COLOR4
        }

        self.width_aspect = 0.75
        self.string_sizes = {
            "full": 6.2*2,
            "half": 3.0*2,
            "single": 3.2*2
        }

    def __str__(self):
        return "\n".join(["{} = {}".format(str(var), vars(self)[var]) for var in vars(self)])

    def copy(self):
        return copy_lib.deepcopy(self)

    def setup_figure(self, sub_rows=1, sub_cols=1, title=None, size=None, **kwargs):
        """
        Create a figure and axes instance with the supplied parameters
        and title.

        Parameters
        ----------
        sub_rows : int, optional
            Number of rows to have in the figure. Default is 1.
        sub_cols : int, optional
            Number of columns to have in the figure. Default is 1.
        title : str or None, optional
            Overall title to display over the figure. Default is None.
        **kwargs
            All additional keyword arguments are passed to the
            pyplot.subplots call.

        Returns
        -------
        fig : pyplot.Figure
            Created figure.
        axs : array of Axes
            An array containing an array of the axes in the figure of
            shape `(sub_rows, sub_cols)`. This will be an array
            containing one axis if `sub_rows = sub_cols = 1`.
        """
        if size is None:
            size = self.FIG_SIZE
        if isinstance(size, str):
            figsize = (
                self.string_sizes[size],
                (self.string_sizes[size]/sub_cols)*self.width_aspect*sub_rows)
        elif isinstance(size, list):
            assert len(size) == 1, "List sizes indicate a fractional size, and may only contain one element"
            figsize = (
                self.string_sizes["half"]*2*size[0],
                self.string_sizes["half"]*2*size[0]*self.width_aspect)
        else:
            figsize = size
        # TODO add a scaling function to adjust all test sizes etc,
        # base on either single scalling, or input dimensions
        if self.arg_style:
            plt.style.use("fivethirtyeight")
            fig_kwargs = {
                "figsize": figsize,
                "facecolor": self.FIG_FACECOLOR
            }
        else:
            plt.style.use(self.PLT_STYLE)
            fig_kwargs = {
                "figsize": figsize
            }

        fig_kwargs.update(kwargs)

        fig, axs = plt.subplots(sub_rows, sub_cols, **fig_kwargs)

        if self.arg_style:
            title_kwarg = {"fontdict": {"fontsize": self.TITLE_SIZE}}
        else:
            title_kwarg = {}
        if title is not None:
            plt.title(title, **title_kwarg)
        elif self.TITLE is not None:
            plt.title(self.TITLE, **title_kwarg)

        if self.arg_style:    
            if not isinstance(axs, np.ndarray):
                axs = np.array([axs])

            for ax in axs.flatten():
                ax.patch.set_alpha(self.AX_ALPHA)

                ax.tick_params(axis='both', which='major',
                            labelsize=self.TICK_SIZE)
                ax.tick_params(axis='both', direction="out", length=self.TICK_LENGTH,
                            width=self.TICK_WIDTH, grid_color=self.GRID_COLOR, grid_alpha=self.GRID_ALPHA)

                if self.MINOR_TICK_SPACING_X == "auto":
                    ax.xaxis.set_minor_locator(AutoMinorLocator())
                else:
                    ax.xaxis.set_minor_locator(
                        MultipleLocator(self.MINOR_TICK_SPACING_X))
                if self.MINOR_TICK_SPACING_Y == "auto":
                    ax.yaxis.set_minor_locator(AutoMinorLocator())
                else:
                    ax.yaxis.set_minor_locator(
                        MultipleLocator(self.MINOR_TICK_SPACING_X))
                ax.tick_params(axis="both", which='minor',
                            labelsize=self.TICK_SIZE-2, length=self.TICK_LENGTH-2)

                ax.xaxis.grid(self.MINOR_GRID_X, which='minor')
                ax.yaxis.grid(self.MINOR_GRID_Y, which='minor')

                # if self.LOGSCALE:
                #     ax.set_yscale("log")

        return fig, axs

    def gen_kwargs(self, type="plot", index=0, **kwargs):
        """
        Generates keyword arguments to format a plot with the style
        of this `PlotConfig` instance.

        The results of this function are intended to be unpacked and
        passed to a `plt.plot` (or similar) call.

        Parameters
        ----------
        type : {"plot", "line", "hist"}, optional
            Type of plot being formatted. "plot" is most generally
            applicable. Default is plot.
        index : int, optional
            Index of the line on the axis. Controls the colour. Default
            is 0.
        **kwargs
            All additional keyword arguments are returned by this
            function, and can overwrite generated parameters.

        Returns
        -------
        fig : pyplot.Figure
            Created figure.
        axs : array of Axes
            An array containing an array of the axes in the figure of
            shape `(sub_rows, sub_cols)`. This will be an array
            containing one axis if `sub_rows = sub_cols = 1`.
        """
        final_kwargs = {}

        if self.arg_style:
            final_kwargs.update({
                "linewidth": self.LINEWIDTH
            })

        if type == "hist":
            final_kwargs.update({
                "bins": self.BINS,
                "density": self.HIST_DENSITY,
                "histtype": self.HIST_TYPE
            })
            if index <= max(self.HIST_COLOURS.keys()):
                final_kwargs.update({"color": self.HIST_COLOURS[index]})
        elif type == "line":
            final_kwargs.pop("linewidth")
            final_kwargs.update({
                # "linewidth": self.LINEWIDTH,
                "colors": "red"
            })

        final_kwargs.update(kwargs)

        return final_kwargs

    def format_axis(self, *ax, xlabel=None, ylabel=None, xlim=None, ylim=None, legend=True, xlog=None, ylog=None):
        """
        Formats the supplied axis with the style of this `PlotConfig`
        instance.

        Parameters:
        *ax : pyplot.Axes or array of Axes
            Axes to be formatted. If not supplied, the current axis
            will be obtained via `plt.gca()`.
        xlabel : str or None, optional
            Label to be displayed on the x axis. If None, no label is
            printed. Default is None.
        ylabel : str or None, optional
            Label to be displayed on the y axis. If None, no label is
            printed. Default is None.
        xlim : tuple or None, optional
            Range to be displayed along the x axis in form (lower
            limit, upper limit). If None, the default pyplot limits are
            used. Default is None.
        ylim : tuple or None, optional
            Range to be displayed along the y axis in form (lower
            limit, upper limit). If None, the default pyplot limits are
            used. Default is None.
        legend : bool, optional
            Whether or not to display a legend if plot labels have been
            passed. Default is True.
        xlog : bool or None, optional
            Option to override the default x axis log scale of this
            `PlotConfig` instance via True or False. If None, the
            default settings of the `PlotConfig` instance is used.
            Default is None.
        ylog : bool or None, optional
            Option to override the default y axis log scale of this
            `PlotConfig` instance via True or False. If None, the
            default settings of the `PlotConfig` instance is used.
            Default is None.
        """
        if self.arg_style:
            lab_args = {"size":self.LABEL_SIZE}
        else:
            lab_args = {}
        if len(ax) == 0 or (len(ax) == 1 and ax[0] is None):
            # Overrides for log scales needs to be before xlim call
            if xlog is not None:
                if xlog:
                    plt.gca().set_xscale("log")
                else:
                    plt.gca().set_xscale("linear")
            if ylog is not None:
                if ylog:
                    plt.gca().set_yscale("log")
                else:
                    plt.gca().set_yscale("linear")
            plt.xlabel(xlabel, **lab_args)
            plt.ylabel(ylabel, **lab_args)
            if self.arg_style:
                plt.xlim(xlim)
                plt.ylim(ylim)
                plt.gca().set_facecolor(self.AXIS_FACECOLOR)
                for spine in plt.gca().spines.values():
                    spine.set_edgecolor(self.SPINE_COLOR)
            if legend and len(plt.gca().get_legend_handles_labels()[1])>0:
                if self.arg_style:
                    plt.legend(prop={'size': self.LEGEND_SIZE},
                            facecolor=self.LEGEND_COLOR)
                else:
                    plt.legend()

        else:
            for a in ax:
                # Overrides for log scales, needs to be before xlim call
                if xlog is not None:
                    if xlog:
                        a.set_xscale("log")
                    else:
                        a.set_xscale("linear")
                if ylog is not None:
                    if ylog:
                        a.set_yscale("log")
                    else:
                        a.set_yscale("linear")
                a.set_xlabel(xlabel, **lab_args)
                a.set_ylabel(ylabel, **lab_args)
                if self.arg_style:
                    a.set_xlim(xlim)
                    a.set_ylim(ylim)
                    a.set_facecolor(self.AXIS_FACECOLOR)
                    for spine in a.spines.values():
                        spine.set_edgecolor(self.SPINE_COLOR)
                if legend and len(a.get_legend_handles_labels()[1])>0:
                    if self.arg_style:
                        a.legend(prop={'size': self.LEGEND_SIZE},
                                 facecolor=self.LEGEND_COLOR)
                    else:
                        a.legend()

        return

    def end_plot(self, save_name=None):
        """
        Ends a plot by saving it if `save_name` is specified and
        displaying the plot if the `PlotConfig` instance has
        `SHOW_PLOT=True`.
        The figure will only be closed if the current figure was saved
        or shown to avoid accidental plot removal.

        Parameters
        ----------
        save_name : str or None, optional
            Name to save the plot by. If the name contains a "/", the
            save_name will be assumed to be a full path and saved
            accordingly (N.B. designed for bash style file names),
            otherwise, the plot will be saved under the
            `PlotConfig.SAVE_FOLDER + save_name`. If None, the file
            will not be saved. Default is None.
        """
        # Need to get this to work with multiple figures open
        plt.tight_layout()
        if self.arg_style:
            save_args = {"facecolor":self.FIG_FACECOLOR, "dpi":self.DPI}
        else:
            save_args = {}
        if save_name is not None:
            if "/" in save_name:
                plt.savefig(save_name, bbox_inches='tight',
                            **save_args)
            elif self.SAVE_FOLDER is not None:
                if self.SAVE_FOLDER[-1] != "/":
                    self.SAVE_FOLDER += "/"
                plt.savefig(self.SAVE_FOLDER + save_name, bbox_inches='tight',
                            **save_args)
            else:
                warnings.warn(
                    "Plot has been given a file name to save, but not path to folder in the PlotConfig")

        if self.SHOW_PLOT:
            plt.show()

        if self.SHOW_PLOT or (save_name is not None):
            plt.close()
        return

    def get_bins(self, data, array=False):
        """
        Returns the bins to use for a given set of data as defined by
        `BINS` of the `PlotConfig` instance.

        Parameters
        ----------
        data : array-like
            Data to be binned
        array : bool, optional
            Whether to return an array of bin edges or a number of
            bins. Default is False.

        Returns
        bins : int or np.ndarray
            Returns the number of bins to divide the full range of data
            by if `array` is `False`, else returns an array of bin
            edges.
        """
        if len(np.unique(data)) != 1:
            if self.BINS == "best":
                num_bins, _ = self._get_best_bins(data)
                num_bins = int(num_bins * self.BEST_BINS_MULTIPLIER)
            else:
                num_bins = self.BINS
        elif array:
            return np.linspace(min(data), max(data) + 0.1, 2)
        else:
            return None

        if array:
            return np.linspace(min(data), max(data), num_bins+1)
        return num_bins

    def _get_best_bins(self, data):
        # Test to deal with case where the InterQuartile Range
        #   is wholly contained in one data point
        iqrange = iqr(data)
        if iqrange == 0.:
            iqrange = np.max(data) - np.min(data)

        # Freedman–Diaconis rule for optimal bin width
        bin_width = 2*iqrange/(len(data)**(1/3))
        bins = int(round((np.max(data) - np.min(data))/bin_width))
        return bins, bin_width

    @staticmethod
    def expand_bins(bins, data):
        """
        Given an array of bin edges, expands the bins by extending the
        lower (upper) bin edges by the factors of width of the lower
        (upper) bin until the new set of `data` is contined within the
        new bins.

        Parameters
        ----------
        bins : np.ndarray
            Previous set of bin edges.
        data : array-like
            New data to encapsulate in the binning

        Returns
        -------
        new_bins : np.ndarray
            New bin edges which encapsulate the full range of `data` by
            adding additional bins below/above the existing edges in
            `bins` 
        """
        if min(data) < bins[0]:
            low_bin_width = bins[1] - bins[0]
            bins = np.concatenate((np.arange(
                bins[0] - low_bin_width, min(data) - low_bin_width, -low_bin_width)[::-1], bins))
        if max(data) >= bins[-1]:
            high_bin_width = bins[-1] - bins[-2]
            # TODO nicer way to do this?
            bins = np.concatenate((bins, np.arange(
                bins[-1] + high_bin_width, max(data) + high_bin_width*(1+1e-6), high_bin_width)))
        return bins

    @property
    def arg_style(self):
        return self.PLT_STYLE is None