Source code for yoda.plotting

# -*- python -*-

"""
Plotting utilities, particularly for interaction with matplotlib and Rivet make-plots
"""

import yoda
import sys
import numpy as np
import matplotlib as mpl


# TODO: Move to core objects
#     def same_binning_as(self, other):
#         if self.dim != other.dim:
#             return False
#         if not (other.x == self.x).all() and \
#                (other.exminus == self.exminus).all() and \
#                (other.explus == self.explus).all():
#             return False
#         if self.dim == 2:
#             return True
#         return (other.y == self.y).all() and \
#                (other.eyminus == self.eyminus).all() and \
#                (other.eyplus == self.eyplus).all()


def read_plot_keys(datfile):
    import re
    re_begin = re.compile("#*\s*BEGIN\s+PLOT\s*(\w*)")
    re_comment = re.compile("#.*")
    re_attr = re.compile("(\w+)\s*=\s*(.*)")
    re_end = re.compile("#*\s*END\s+PLOT\s+\w*")
    plotkeys = {}
    with open(datfile) as f:
        inplot = False
        name = None
        for line in f:
            l = line.strip()
            if re_begin.match(l):
                inplot = True
                name = re_begin.match(l).group(1)
            elif re_end.match(l):
                inplot = False
                name = None
            elif re_comment.match(l):
                continue
            elif inplot:
                m = re_attr.match(l)
                if m is None: continue
                plotkeys.setdefault(name, {})[m.group(1)] = m.group(2)
    return plotkeys


[docs]def mplinit(engine="MPL", font="TeX Gyre Pagella", fontsize=17, mfont=None, textfigs=True): """One-liner matplotlib (mpl) setup. By default mpl will be configured with its native MathText rendering backend, and a Palatino-like font for both text and math contexts, using 'lower-case numerals' if supported. Setting the engine to 'TEX' will use standard mpl rendering, with calls to LaTeX for axis labels and other text; setting it to 'PGF' will use the TeX PGF renderer: both these modes are much slower than MPL mode, but the latter only supports a limited set of LaTeX macros and does not render as nicely as the TeX backends. The font and mfont optional arguments can be used to choose a different text font and math font respectively; if mfont is None, it defaults to the same as the text font. The textfigs boolean argument can be set false to disable the lower-case/text/old-style numerals and use 'upper-case' numerals everywhere. These options do not currently apply to the MPL rendering engine. """ mpl.rcParams.update({ "text.usetex" : (engine != "MPL"), "font.size" : int(fontsize), "font.family" : "serif", #< TODO: make configurable? auto-detect? }) texpreamble = [r"\usepackage{amsmath,amssymb}", r"\usepackage{mathspec}"] mfont = mfont if mfont else font fontopts = "[Numbers=OldStyle]" if textfigs else "" mfontopts = fontopts.replace("]", ",") + "Scale=MatchUppercase" + "]" texpreamble.append( r"\setmainfont{fopts}{{{font}}}".format(fopts=fontopts, font=font) ) texpreamble.append( r"\setmathsfont(Digits,Latin){fopts}{{{font}}}".format(fopts=mfontopts, font=mfont) ) if engine.upper() == "PGF": mpl.use("pgf") mpl.rcParams["pgf.preamble"] = texpreamble # TODO: Fix? # elif engine.upper() == "TEX": # mpl.rcParams["text.latex.preamble"] = texpreamble return mpl
## Alias initmpl = mplinit setup_mpl = mplinit
[docs]def show(): """ Convenience call to matplotlib.pyplot.show() NOTE: done this way to avoid import of pyplot before mplinit() or mpl.use() has been (optionally) called. """ import matplotlib.pyplot as plt plt.show()
[docs]def mk_figaxes_1d(ratio=True, title=None, figsize=(8,6)): "Make a standard main+ratio plot figure and subplot layout" ## We need to use pyplot here to set up the backend-specific canvas import matplotlib.pyplot as plt fig = plt.figure(figsize=figsize) #fig = mpl.figure.Figure(figsize=figsize, tight_layout=True) if title: fig.suptitle(title, horizontalalignment="left", x=0.13) ## Make axes. GridSpec may not be available, in which case fall back ~gracefully axmain, axratio = None, None if ratio: try: gs = mpl.gridspec.GridSpec(2, 1, height_ratios=[3,1], hspace=0) axmain = fig.add_subplot(gs[0]) #axmain.hold(True) axratio = fig.add_subplot(gs[1], sharex=axmain) #axratio.hold(True) axratio.axhline(1.0, color="gray") #< Ratio = 1 marker line except: sys.stderr.write("matplotlib.gridspec not available: falling back to plotting without a ratio\n") ratio = False if not ratio: axmain = fig.add_subplot(1,1,1) #axmain.hold(True) return fig, (axmain, axratio)
def set_axis_labels_1d(axmain, axratio, xlabel=None, ylabel=None, ratioylabel=None): axmain.set_ylabel(ylabel, y=1, ha="right", labelpad=None) if axratio: axmain.xaxis.set_major_locator(mpl.ticker.NullLocator()) axratio.set_xlabel(xlabel, x=1, ha="right", labelpad=None) axratio.set_ylabel(ratioylabel) else: axmain.set_xlabel(xlabel, x=1, ha="right", labelpad=None)
[docs]def mk_lowcase_dict(d): "Convert the keys of a str->obj dict to lower-case" return dict((k.lower(), v) for (k,v) in d.items())
# TODO: Needs generalisation for 2D marginal axes) def setup_axes_1d(axmain, axratio, **plotkeys): ## Case-insensitize the plotkeys dict plotkeys = mk_lowcase_dict(plotkeys) ## Axis labels first xlabel = plotkeys.get("xlabel", "") ylabel = plotkeys.get("ylabel", "") ratioylabel = plotkeys.get("ratioylabel", "ratio") set_axis_labels_1d(axmain, axratio, xlabel, ylabel, ratioylabel) ## log/lin measures # TODO: Dynamic default based on data ranges? # TODO: take log axes and preference for round numbers into account in setting default axis limits xmeasure = "log" if yoda.util.as_bool(plotkeys.get("logX", False)) else "linear" ymeasure = "log" if yoda.util.as_bool(plotkeys.get("logY", False)) else "linear" ratioymeasure = "log" if yoda.util.as_bool(plotkeys.get("ratiology", False)) else "linear" axmain.set_xscale(xmeasure) axmain.set_yscale(ymeasure) if axratio: axratio.set_xscale(xmeasure) axratio.set_yscale(ratioymeasure) ## Plot range limits if "ymin" in plotkeys: axmain.set_ylim(bottom=float(plotkeys.get("ymin"))) if "ymax" in plotkeys: axmain.set_ylim(top=float(plotkeys.get("ymax"))) # if "xmin" in plotkeys: axmain.set_xlim(left=float(plotkeys.get("xmin"))) if "xmax" in plotkeys: axmain.set_xlim(right=float(plotkeys.get("xmax"))) # if axratio: # TODO: RatioSymmRange option # axratio.set_xlim([xmin-0.001*xdiff, xmax+0.001*xdiff]) # <- TODO: bad on a log scale! if "xmin" in plotkeys: axratio.set_xlim(left=float(plotkeys.get("xmin"))) if "xmax" in plotkeys: axratio.set_xlim(right=float(plotkeys.get("xmax"))) if "ratioymin" in plotkeys: axratio.set_ylim(bottom=float(plotkeys.get("ratioymin"))) if "ratioymax" in plotkeys: axratio.set_ylim(top=float(plotkeys.get("ratioymax"))) # TODO: Ratio plot manual ticks def plot_hist_on_axes_1d(axmain, axratio, h, href=None, default_color="black", default_linestyle="-", **plotkeys): ## Case-insensitize the plotkeys dict hkeys = mk_lowcase_dict(h.annotationsDict()) hkeys.update(plotkeys) plotkeys = hkeys # TODO: Split into different plot styles: line/filled/range, step/diag/smooth, ...? ## Styles default_color = plotkeys.get("color", default_color) marker = plotkeys.get("marker", plotkeys.get("polymarker", None)) # <- make-plots translation marker = {"*":"o"}.get(marker, marker) # <- make-plots translation mcolor = plotkeys.get("linecolor", default_color) errbar = plotkeys.get("errorbars", None) ecolor = plotkeys.get("errorbarscolor", default_color) line = plotkeys.get("line", None) lcolor = plotkeys.get("linecolor", default_color) lstyle = plotkeys.get("linestyle", default_linestyle) lstyle = {"solid":"-", "dashed":"--", "dotdashed":"-.", "dashdotted":"-.", "dotted":":"}.get(lstyle, lstyle) # <- make-plots translation lwidth = 1.4 msize = 7 ## If no drawing is enabled, default to a step line if not any([marker, line, errbar]): line = "step" ## MPL plotting # TODO: Split this into different functions for each kind of data preparation (and smoothing as an extra function?) # TODO: First convert h to scatter artists = None if errbar: artists = axmain.errorbar(h.xVals(), h.yVals(), xerr=h.xErrs(), yerr=h.yErrs(), color=ecolor, linestyle="none", linewidth=lwidth, capthick=lwidth) # linestyle="-", marker="o", if line == "step": artists = axmain.step(np.append(h.xMins(), h.xMax()), np.append(h.yVals(), h.yVals()[-1]), where="post", color=lcolor, linestyle=lstyle, linewidth=lwidth) elif line == "diag": artists = axmain.plot(h.xVals(), h.yVals(), color=lcolor, linestyle=lstyle, linewidth=lwidth) elif line == "smooth": from scipy.interpolate import spline xnew = np.linspace(min(h.xVals()), max(h.xVals()), 3*h.numBins) ynew = spline(h.xVals(), h.yVals(), xnew) artists = axmain.plot(xnew, ynew, color=lcolor, linestyle=lstyle, linewidth=lwidth) if marker: artists = axmain.plot(h.xVals(), h.yVals(), marker=marker, markersize=msize, linestyle="none", color=mcolor, markeredgecolor=mcolor) ## Legend entry if h.annotation("Title") and artists: artists[0].set_label(h.annotation("Title")) ## Ratio ratioartists = None if href and h is not href: # TODO: exclude and specify order via RatioIndex # assert h.same_binning_as(href) # TODO: log ratio or #sigma deviation yratios = np.array(h.yVals())/np.array(href.yVals()) # TODO: Same styling control as for main plot (with Ratio prefix, default to main plot style) ## Stepped plot ratioartists = axratio.step(np.append(href.xMins(), href.xMax()), np.append(yratios, yratios[-1]), where="post", color=lcolor, linestyle=lstyle, linewidth=lwidth) # TODO: Diag plot # axratio.plot(href["x"], yratios, color="r", linestyle="--") # TODO: Smoothed plot return artists
[docs]def plot(hs, outfile=None, ratio=True, show=False, axmain=None, axratio=None, **plotkeys): """ Plot the given histograms on a single figure, returning (fig, (main_axis, ratio_axis)). Show to screen if the second arg is True, and saving to outfile if it is otherwise non-null. """ ## Case-insensitize the plotkeys dict plotkeys = mk_lowcase_dict(plotkeys) ## Handle single histo args if isinstance(hs, yoda.AnalysisObject): hs = [hs,] ratio = False ## Get data ranges (calculated or forced) xmin = float(plotkeys.get("xmin", min(h.xMin() for h in hs))) xmax = float(plotkeys.get("xmax", max(h.xMax() for h in hs))) xdiff = xmax - xmin # print xmin, xmax, xdiff # TODO: Tweak max-padding for top tick label... sensitive to log/lin measure ymin = plotkeys.get("ymin", min(min(h.yVals()) for h in hs)) #print( max(max(h.yVals()) for h in hs) ) ymax = plotkeys.get("ymax", 1.1*max(max(h.yVals()) for h in hs)) ymin = float(ymin) ymax = float(ymax) ydiff = ymax - ymin # print ymin, ymax, ydiff ## Identify reference histo by annotation (unless explicitly disabled) href = None # TODO: Use ratio to setdefault RatioPlot in plotkeys, then use that to decide whether to look for href if ratio: for h in hs: hkeys = mk_lowcase_dict(h.annotationsDict()) if yoda.util.as_bool(hkeys.get("ratioref", False)): if href is None: href = h else: #print "Multiple ratio references set: using first value = {}".format(href.path) break if href is None: #< no ref found -- maybe all were disabled? ratio = False ## Make figure and subplot grid layout title = plotkeys.get("title", "") if not axmain: fig, (axmain, axratio) = mk_figaxes_1d(ratio and not axratio, title) else: fig = axmain.get_figure() ## Setup axes appearances axmain.set_xlim([xmin, xmax]) axmain.set_ylim([ymin, ymax]) if axratio: axratio.set_xlim([xmin, xmax]) axratio.set_ylim(auto=True) setup_axes_1d(axmain, axratio, **plotkeys) # TODO: specify ratio display in log/lin, abs, or #sigma, and as x/r or (x-r)/r ## Draw ratio error band (do this before looping over cmp lines) # TODO: Actually we can call this when we hit the href, and force the zorder into groups: bands, lines, dots, legend, text, frame if axratio: ref_ymax_ratios = np.array(href.yMaxs())/np.array(href.yVals()) ref_ymin_ratios = np.array(href.yMins())/np.array(href.yVals()) # TODO: Diag: (needs -> limit handling at ends) # axratio.fill_between(href.x, ref_ymin_ratios, ref_ymax_ratios, edgecolor="none", facecolor=ratioerrcolor, interpolate=False) # Stepped: def xedges_dbl(h): edges = np.empty((2*len(h.xVals()),)) edges[0::2] = h.xMins() edges[1::2] = h.xMaxs() return edges def dbl_array(arr): return sum(([x,x] for x in arr), []) ratioerrcolor = plotkeys.get("ratioerrcolor", "yellow") axratio.fill_between(xedges_dbl(href), dbl_array(ref_ymin_ratios), dbl_array(ref_ymax_ratios), edgecolor="none", facecolor=ratioerrcolor) # TODO: Smoothed: (needs -> limit handling at ends) # Redraw ratio = 1 marker line: axratio.axhline(1.0, color="gray") COLORS = ["red", "blue", "magenta", "orange", "green"] LSTYLES = ["-", "--", "-.", ":"] ## Dataset plotting some_valid_label = False for ih, h in enumerate(hs): #print ih, h.path aa = plot_hist_on_axes_1d(axmain, axratio, h, href, COLORS[ih % len(COLORS)], LSTYLES[ih % len(LSTYLES)]) if aa and not aa[0].get_label().startswith("_"): # print "@@@", aa[0].get_label() some_valid_label = True ## Legend # TODO: allow excluding and specify order via LegendIndex if some_valid_label: #< No point in writing a legend if there are no labels pass #axmain.legend(loc=plotkeys.get("LegendPos", "best"), fontsize=plotkeys.get("LegendFontSize", "x-small"), frameon=False) ## Tweak layout now that everything is in place # TODO: merge tight_layout() into the Figure constructor, and maybe the ratio ticker when retrospectively drawing the zorder'ed err band if axratio: axratio.yaxis.set_major_locator(mpl.ticker.MaxNLocator(4, prune="upper")) fig.tight_layout() ## Save to an image file if we were asked to if outfile: #print "Saving to " + outfile fig.savefig(outfile) ## Show to screen if requested if show: import matplotlib.pyplot as plt plt.show() ## Return the figure objects return fig, (axmain, axratio)
## Aliases plot_hists_1d = plot plot_hist_1d = plot def _plot1arg(args): "Helper function for mplot, until Py >= 3.3 multiprocessing.pool.starmap() is available" return plot(*args)
[docs]def nplot(hs, outfiles=None, ratio=True, show=False, nproc=1, **plotkeys): """ Plot the given list of histogram(s), cf. many calls to plot(). hs must be an iterable, each entry of which will be the content of a single plot: the entries can either be single histograms or lists of histograms, i.e. either kind of valid first argument to plot(). Outfiles must be an iterable corresponding to hs, and ratio may either be a bool or such an iterable. The return value is a list of the return tuples from each call to plot(), of the same length as the hs arg. MULTIPROCESSING -- *WARNING* CURRENTLY BROKEN The main point of this function, other than convenience, is that the Python multiprocessing module can be used to distribute the work on to multiple parallel processes. The nproc argument should be the integer number of parallel processes on which to distribute the plotting. nproc = None (the default value) will use Ncpu-1 or 1 process, whichever is larger. If nproc = 1, multiprocessing will not be used -- this avoids overhead and eases debugging. """ argslist = [] for i, hs_arg in enumerate(hs): outfile_arg = outfiles[i] if outfiles else None ratio_arg = ratio[i] if hasattr(ratio, "__iter__") else ratio show_arg = False #< we just do this once, at the end plotkeys_arg = plotkeys if type(plotkeys) is dict else plotkeys[i] argslist.append( (hs_arg, outfile_arg, ratio_arg, show_arg, None, None, plotkeys_arg) ) #print argslist # TODO: make the multiprocessing work import multiprocessing nproc = nproc or multiprocessing.cpu_count() or 1 if nproc > 1: pool = multiprocessing.Pool(processes=nproc) res = pool.map_async(_plot1arg, argslist) rtn = res.get() else: ## Run this way in the 1 proc case for easier debugging rtn = [_plot1arg(args) for args in argslist] if show: import matplotlib.pyplot as plt plt.show() return rtn