python/rootplot/root2matplotlib.py

"""
Utilities for plotting ROOT histograms in matplotlib.
"""

__license__ = '''\
Copyright (c) 2009-2010 Jeff Klukas <klukas@wisc.edu>

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
'''

################ Import python libraries

import math
import ROOT
import re
import copy
import array
from rootplot import utilities
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.transforms as transforms
import numpy as np

################ Define constants

_all_whitespace_string = re.compile(r'\s*$')


################ Define classes

class Hist2D(utilities.Hist2D):
    """A container to hold the parameters from a 2D ROOT histogram."""
    def __init__(self, *args, **kwargs):
        self.replacements = None
        if 'replacements' in kwargs:
            self.replacements = kwargs.pop('replacements')
        utilities.Hist2D.__init__(self, *args, **kwargs)
    def contour(self, **kwargs):
        """Draw a contour plot."""
        cs = plt.contour(self.x, self.y, self.content, **kwargs)
        plt.clabel(cs, inline=1, fontsize=10)
        if self.binlabelsx is not None:
            plt.xticks(np.arange(self.nbinsx), self.binlabelsx)
        if self.binlabelsy is not None:
            plt.yticks(np.arange(self.nbinsy), self.binlabelsy)
        return cs
    def col(self, **kwargs):
        """Draw a colored box plot using :func:`matplotlib.pyplot.imshow`."""
        if 'cmap' in kwargs:
            kwargs['cmap'] = plt.get_cmap(kwargs['cmap'])
        plot = plt.imshow(self.content, interpolation='nearest',
                          extent=[self.xedges[0], self.xedges[-1],
                                  self.yedges[0], self.yedges[-1]],
                          aspect='auto', origin='lower', **kwargs)
        return plot
    def colz(self, **kwargs):
        """
        Draw a colored box plot with a colorbar using
        :func:`matplotlib.pyplot.imshow`.
        """
        plot = self.col(**kwargs)
        plt.colorbar(plot)
        return plot
    def box(self, maxsize=40, **kwargs):
        """
        Draw a box plot with size indicating content using
        :func:`matplotlib.pyplot.scatter`.
        
        The data will be normalized, with the largest box using a marker of
        size maxsize (in points).
        """
        x = np.hstack([self.x for i in range(self.nbinsy)])
        y = np.hstack([[yval for i in range(self.nbinsx)] for yval in self.y])
        maxvalue = np.max(self.content)
        if maxvalue == 0:
            maxvalue = 1
        sizes = np.array(self.content).flatten() / maxvalue * maxsize
        plot = plt.scatter(x, y, sizes, marker='s', **kwargs)
        return plot
    def TH2F(self, name=""):
        """Return a ROOT.TH2F object with contents of this Hist2D."""
        th2f = ROOT.TH2F(name, "",
                         self.nbinsx, array.array('f', self.xedges),
                         self.nbinsy, array.array('f', self.yedges))
        th2f.SetTitle("%s;%s;%s" % (self.title, self.xlabel, self.ylabel))
        for ix in range(self.nbinsx):
            for iy in range(self.nbinsy):
                th2f.SetBinContent(ix + 1, iy + 1, self.content[iy][ix])
        return th2f

class Hist(utilities.Hist):
    """A container to hold the parameters from a ROOT histogram."""
    def __init__(self, *args, **kwargs):
        self.replacements = None
        if 'replacements' in kwargs:
            self.replacements = kwargs.pop('replacements')
        utilities.Hist.__init__(self, *args, **kwargs)
    def _prepare_xaxis(self, rotation=0, alignment='center'):
        """Apply bounds and text labels on x axis."""
        if self.binlabels is not None:
            binwidth = (self.xedges[-1] - self.xedges[0]) / self.nbins
            plt.xticks(self.x, self.binlabels,
                       rotation=rotation, ha=alignment)
        plt.xlim(self.xedges[0], self.xedges[-1])

    def _prepare_yaxis(self, rotation=0, alignment='center'):
        """Apply bounds and text labels on y axis."""
        if self.binlabels is not None:
            binwidth = (self.xedges[-1] - self.xedges[0]) / self.nbins
            plt.yticks(self.x, self.binlabels,
                       rotation=rotation, va=alignment)
        plt.ylim(self.xedges[0], self.xedges[-1])

    def show_titles(self, **kwargs):
        """Print the title and axis labels to the current figure."""
        replacements = kwargs.get('replacements', None) or self.replacements
        plt.title(replace(self.title, replacements))
        plt.xlabel(replace(self.xlabel, replacements))
        plt.ylabel(replace(self.ylabel, replacements))
    def hist(self, label_rotation=0, label_alignment='center', **kwargs):
        """
        Generate a matplotlib hist figure.

        All additional keyword arguments will be passed to
        :func:`matplotlib.pyplot.hist`.
        """
        kwargs.pop('fmt', None)
        replacements = kwargs.get('replacements', None) or self.replacements
        weights = self.y
        # Kludge to avoid mpl bug when plotting all zeros
        if self.y == [0] * self.nbins:
            weights = [1.e-10] * self.nbins
        plot = plt.hist(self.x, weights=weights, bins=self.xedges,
                        label=replace(self.label, replacements), **kwargs)
        self._prepare_xaxis(label_rotation, label_alignment)
        return plot
    def errorbar(self, xerr=False, yerr=False, label_rotation=0,
                 label_alignment='center', **kwargs):
        """
        Generate a matplotlib errorbar figure.

        All additional keyword arguments will be passed to
        :func:`matplotlib.pyplot.errorbar`.
        """
        if xerr:
            kwargs['xerr'] = self.xerr
        if yerr:
            kwargs['yerr'] = self.yerr
        replacements = kwargs.get('replacements', None) or self.replacements
        errorbar = plt.errorbar(self.x, self.y,
                                label=replace(self.label, replacements),
                                **kwargs)
        self._prepare_xaxis(label_rotation, label_alignment)
        return errorbar
    def errorbarh(self, xerr=False, yerr=False, label_rotation=0,
                  label_alignment='center', **kwargs):
        """
        Generate a horizontal matplotlib errorbar figure.

        All additional keyword arguments will be passed to
        :func:`matplotlib.pyplot.errorbar`.
        """
        if xerr: kwargs['xerr'] = self.yerr
        if yerr: kwargs['yerr'] = self.xerr
        replacements = kwargs.get('replacements', None) or self.replacements
        errorbar = plt.errorbar(self.y, self.x,
                                label=replace(self.label, replacements),
                                **kwargs)
        self._prepare_yaxis(label_rotation, label_alignment)
        return errorbar
    def bar(self, xerr=False, yerr=False, xoffset=0., width=0.8, 
            label_rotation=0, label_alignment='center', **kwargs):
        """
        Generate a matplotlib bar figure.

        All additional keyword arguments will be passed to
        :func:`matplotlib.pyplot.bar`.
        """
        kwargs.pop('fmt', None)
        if xerr: kwargs['xerr'] = self.av_xerr()
        if yerr: kwargs['yerr'] = self.av_yerr()
        replacements = kwargs.get('replacements', None) or self.replacements
        ycontent = [self.xedges[i] + self.width[i] * xoffset
                    for i in range(len(self.xedges) - 1)]
        width = [x * width for x in self.width]
        bar = plt.bar(ycontent, self.y, width,
                      label=replace(self.label, replacements), **kwargs)
        self._prepare_xaxis(label_rotation, label_alignment)
        return bar
    def barh(self, xerr=False, yerr=False, yoffset=0., width=0.8,
             label_rotation=0, label_alignment='center', **kwargs):
        """
        Generate a horizontal matplotlib bar figure.

        All additional keyword arguments will be passed to
        :func:`matplotlib.pyplot.bar`.
        """
        kwargs.pop('fmt', None)
        if xerr: kwargs['xerr'] = self.av_yerr()
        if yerr: kwargs['yerr'] = self.av_xerr()
        replacements = kwargs.get('replacements', None) or self.replacements
        xcontent = [self.xedges[i] + self.width[i] * yoffset
                    for i in range(len(self.xedges) - 1)]
        width = [x * width for x in self.width]
        barh = plt.barh(xcontent, self.y, width,
                        label=replace(self.label, replacements),
                       **kwargs)
        self._prepare_yaxis(label_rotation, label_alignment)
        return barh

class HistStack(utilities.HistStack):
    """
    A container to hold Hist objects for plotting together.

    When plotting, the title and the x and y labels of the last Hist added
    will be used unless specified otherwise in the constructor.
    """
    def __init__(self, *args, **kwargs):
        if 'replacements' in kwargs:
            self.replacements = kwargs.pop('replacements')
        utilities.HistStack.__init__(self, *args, **kwargs)
    def show_titles(self, **kwargs):
        self.hists[-1].show_titles()
    def hist(self, label_rotation=0, **kwargs):
        """
        Make a matplotlib hist plot.

        Any additional keyword arguments will be passed to
        :func:`matplotlib.pyplot.hist`, which allows a vast array of
        possibilities.  Particlularly, the *histtype* values such as
        ``'barstacked'`` and ``'stepfilled'`` give substantially different
        results.  You will probably want to include a transparency value
        (i.e. *alpha* = 0.5).
        """
        contents = np.dstack([hist.y for hist in self.hists])
        xedges = self.hists[0].xedges
        x = np.dstack([hist.x for hist in self.hists])[0]
        labels = [hist.label for hist in self.hists]
        try:
            clist = [item['color'] for item in self.kwargs]
            plt.gca().set_color_cycle(clist)
            ## kwargs['color'] = clist # For newer version of matplotlib
        except:
            pass
        plot = plt.hist(x, weights=contents, bins=xedges,
                        label=labels, **kwargs)
    def bar3d(self, **kwargs):
        #### Not yet ready for primetime
        from mpl_toolkits.mplot3d import Axes3D
        fig = plt.figure()
        ax = Axes3D(fig)
        plots = []
        labels = []
        for i, hist in enumerate(self.hists):
            if self.title  is not None: hist.title  = self.title
            if self.xlabel is not None: hist.xlabel = self.xlabel
            if self.ylabel is not None: hist.ylabel = self.ylabel
            labels.append(hist.label)
            all_kwargs = copy.copy(kwargs)
            all_kwargs.update(self.kwargs[i])
            bar = ax.bar(hist.x, hist.y, zs=i, zdir='y', width=hist.width,
                         **all_kwargs)
            plots.append(bar)
        from matplotlib.ticker import FixedLocator
        locator = FixedLocator(range(len(labels)))
        ax.w_yaxis.set_major_locator(locator)
        ax.w_yaxis.set_ticklabels(labels)
        ax.set_ylim3d(-1, len(labels))
        return plots
    def barstack(self, **kwargs):
        """
        Make a matplotlib bar plot, with each Hist stacked upon the last.

        Any additional keyword arguments will be passed to
        :func:`matplotlib.pyplot.bar`.
        """
        bottom = None # if this is set to zeroes, it fails for log y
        plots = []
        for i, hist in enumerate(self.hists):
            if self.title  is not None: hist.title  = self.title
            if self.xlabel is not None: hist.xlabel = self.xlabel
            if self.ylabel is not None: hist.ylabel = self.ylabel
            all_kwargs = copy.copy(kwargs)
            all_kwargs.update(self.kwargs[i])
            bar = hist.bar(bottom=bottom, **all_kwargs)
            plots.append(bar)
            if not bottom: bottom = [0. for i in range(self.hists[0].nbins)]
            bottom = [sum(pair) for pair in zip(bottom, hist.y)]
        return plots
    def histstack(self, **kwargs):
        """
        Make a matplotlib hist plot, with each Hist stacked upon the last.

        Any additional keyword arguments will be passed to
        :func:`matplotlib.pyplot.hist`.
        """
        bottom = None # if this is set to zeroes, it fails for log y
        plots = []
        cumhist = None
        for i, hist in enumerate(self.hists):
            if cumhist:
                cumhist = hist + cumhist
            else:
                cumhist = copy.copy(hist)
            if self.title  is not None: cumhist.title  = self.title
            if self.xlabel is not None: cumhist.xlabel = self.xlabel
            if self.ylabel is not None: cumhist.ylabel = self.ylabel
            all_kwargs = copy.copy(kwargs)
            all_kwargs.update(self.kwargs[i])
            zorder = 0 + float(len(self) - i)/len(self) # plot in reverse order
            plot = cumhist.hist(zorder=zorder, **all_kwargs)
            plots.append(plot)
        return plots
    def barcluster(self, width=0.8, **kwargs):
        """
        Make a clustered bar plot.

        Any additional keyword arguments will be passed to
        :func:`matplotlib.pyplot.bar`.
        """
        plots = []
        spacer = (1. - width) / 2
        width = width / len(self.hists)
        for i, hist in enumerate(self.hists):
            if self.title  is not None: hist.title  = self.title
            if self.xlabel is not None: hist.xlabel = self.xlabel
            if self.ylabel is not None: hist.ylabel = self.ylabel
            all_kwargs = copy.copy(kwargs)
            all_kwargs.update(self.kwargs[i])
            bar = hist.bar(xoffset=width*i + spacer, width=width, **all_kwargs)
            plots.append(bar)
        return plots
    def barh(self, width=0.8, **kwargs):
        """
        Make a horizontal clustered matplotlib bar plot.

        Any additional keyword arguments will be passed to
        :func:`matplotlib.pyplot.bar`.
        """
        plots = []
        spacer = (1. - width) / 2
        width = width / len(self.hists)
        for i, hist in enumerate(self.hists):
            if self.title  is not None: hist.title  = self.title
            if self.xlabel is not None: hist.ylabel = self.xlabel
            if self.ylabel is not None: hist.xlabel = self.ylabel
            all_kwargs = copy.copy(kwargs)
            all_kwargs.update(self.kwargs[i])
            bar = hist.barh(yoffset=width*i + spacer, width=width, **all_kwargs)
            plots.append(bar)
        return plots
    def bar(self, **kwargs):
        """
        Make a bar plot, with all Hists in the stack overlaid.

        Any additional keyword arguments will be passed to
        :func:`matplotlib.pyplot.bar`.  You will probably want to set a 
        transparency value (i.e. *alpha* = 0.5).
        """
        plots = []
        for i, hist in enumerate(self.hists):
            if self.title  is not None: hist.title  = self.title
            if self.xlabel is not None: hist.xlabel = self.xlabel
            if self.ylabel is not None: hist.ylabel = self.ylabel
            all_kwargs = copy.copy(kwargs)
            all_kwargs.update(self.kwargs[i])
            bar = hist.bar(**all_kwargs)
            plots.append(bar)
        return plots
    def errorbar(self, offset=False, **kwargs):
        """
        Make a matplotlib errorbar plot, with all Hists in the stack overlaid.

        Passing 'offset=True' will slightly offset each dataset so overlapping
        errorbars are still visible.  Any additional keyword arguments will
        be passed to :func:`matplotlib.pyplot.errorbar`.
        """
        plots = []
        for i, hist in enumerate(self.hists):
            if self.title  is not None: hist.title  = self.title
            if self.xlabel is not None: hist.xlabel = self.xlabel
            if self.ylabel is not None: hist.ylabel = self.ylabel
            all_kwargs = copy.copy(kwargs)
            all_kwargs.update(self.kwargs[i])
            transform = plt.gca().transData
            if offset:
                index_offset = (len(self.hists) - 1)/2.
                pixel_offset = 1./72 * (i - index_offset)
                transform = transforms.ScaledTranslation(
                    pixel_offset, 0, plt.gcf().dpi_scale_trans)
                transform = plt.gca().transData + transform
            errorbar = hist.errorbar(transform=transform, **all_kwargs)
            plots.append(errorbar)
        return plots
    def errorbarh(self, **kwargs):
        """
        Make a horizontal matplotlib errorbar plot, with all Hists in the
        stack overlaid.

        Any additional keyword arguments will be passed to
        :func:`matplotlib.pyplot.errorbar`.
        """
        plots = []
        for i, hist in enumerate(self.hists):
            if self.title  is not None: hist.title  = self.title
            if self.xlabel is not None: hist.ylabel = self.xlabel
            if self.ylabel is not None: hist.xlabel = self.ylabel
            all_kwargs = copy.copy(kwargs)
            all_kwargs.update(self.kwargs[i])
            errorbar = hist.errorbarh(**all_kwargs)
            plots.append(errorbar)
        return plots

################ Define functions and classes for navigating within ROOT

class RootFile(utilities.RootFile):
    """A wrapper for TFiles, allowing easier access to methods."""
    def get(self, object_name, path=None):
        try:
            return utilities.RootFile.get(self, object_name, path,
                                          Hist, Hist2D)
        except ReferenceError, e:
            raise ReferenceError(e)

################ Define additional helping functions

def replace(string, replacements):
    """
    Modify a string based on a list of patterns and substitutions.

    replacements should be a list of two-entry tuples, the first entry giving
    a string to search for and the second entry giving the string with which
    to replace it.  If replacements includes a pattern entry containing
    'use_regexp', then all patterns will be treated as regular expressions
    using re.sub.
    """
    if not replacements:
        return string
    if 'use_regexp' in [x for x,y in replacements]:
        for pattern, repl in [x for x in replacements
                              if x[0] != 'use_regexp']:
            string = re.sub(pattern, repl, string)
    else:
        for pattern, repl in replacements:
            string = string.replace(pattern, repl)
    if re.match(_all_whitespace_string, string):
        return ""
    return string

Revision:	1.2
Committed:	Thu May 19 18:48:03 2011 UTC (13 years, 11 months ago) by klukas
Content type:	text/x-python
Branch:	MAIN
CVS Tags:	CMSSW_6_2_0, CMSSW_6_2_0_pre7_TS133806, CMSSW_5_3_11_patch3, CMSSW_6_1_2_SLHC6_patch1, CMSSW_5_3_11_patch2, CMSSW_5_3_11_patch1, CMSSW_6_2_0_pre8, CMSSW_6_1_2_SLHC6, CMSSW_6_1_2_SLHC5, CMSSW_5_3_11, CMSSW_6_2_0_pre7_TS132947, CMSSW_4_4_5_patch3, CMSSW_6_2_0_pre7_g496p02, CMSSW_6_2_0_pre7, CMSSW_5_3_10_patch2, CMSSW_6_1_2_SLHC4_patch1, CMSSW_6_1_2_SLHC4, CMSSW_6_1_2_SLHC2_patch3, CMSSW_6_2_0_pre6_patch1, CMSSW_6_1_2_SLHC2_patch2, CMSSW_6_2_0_pre6, CMSSW_5_3_10_patch1, CMSSW_6_1_2_SLHC3, CMSSW_6_1_2_SLHC2_patch1, CMSSW_5_3_10, CMSSW_6_1_2_SLHC2, CMSSW_6_2_0_pre5slc6, CMSSW_5_3_9_patch3, CMSSW_6_1_2_SLHC1, CMSSW_6_2_0_pre5, CMSSW_6_1_X_2012-12-19-0200, CMSSW_6_0_X_2012-08-07-0200, CMSSW_5_2_X_2012-05-03-0200, CMSSW_5_2_X_2012-03-08-0200, CMSSW_5_0_X_2011-12-18-0200, CMSSW_4_4_X_2011-06-09-0400, CMSSW_5_3_9_patch2, CMSSW_6_1_2, CMSSW_6_2_0_pre4, CMSSW_5_3_9_patch1, CMSSW_6_2_0_pre3, CMSSW_6_1_1_SLHCphase2tk1, CMSSW_5_3_9, CMSSW_6_1_1_SLHCphase1tk1, CMSSW_5_2_9, CMSSW_6_2_0_pre2, CMSSW_6_1_1, CMSSW_5_3_8_patch3, CMSSW_5_3_7_patch6, CMSSW_6_2_0_pre1, CMSSW_5_3_8_patch2, CMSSW_5_3_8_patch1, CMSSW_5_3_8, CMSSW_5_3_8_HI_patch2, CMSSW_5_3_7_patch5, CMSSW_5_3_8_HI_patch1, CMSSW_5_2_6_patch2, CMSSW_6_1_0, CMSSW_5_3_8_HI, CMSSW_5_3_7_patch4, CMSSW_5_3_7_patch3, CMSSW_6_0_1_PostLS1v2_patch4, CMSSW_6_1_0_pre8, CMSSW_5_3_7_25nspatch1, CMSSW_5_3_7_patch2, CMSSW_6_1_0_pre7_TS127013, CMSSW_6_0_1_PostLS1v2_patch3, CMSSW_6_1_0_pre7, CMSSW_5_3_7_patch1, CMSSW_6_1_0_pre6_TS126203_TS126341_patch1, CMSSW_5_3_7_alcapatch1, CMSSW_6_1_0_pre6g496cand01, CMSSW_5_3_7, CMSSW_6_0_1_PostLS1v2_patch2, CMSSW_6_0_1_PostLS1v2_patch1, CMSSW_6_1_0_pre6_TS126203_TS126341, CMSSW_6_1_0_pre6, CMSSW_6_0_1_PostLS1v2, CMSSW_4_4_5_patch2, CMSSW_5_3_6_patch1, CMSSW_5_2_8, CMSSW_6_1_0_pre5, CMSSW_6_0_1_PostLS1v1, CMSSW_5_2_7_hltpatch2, CMSSW_5_3_4_TC125616patch1, CMSSW_6_0_1, CMSSW_5_3_4_patch2, CMSSW_6_1_0_pre3_TS124729, CMSSW_6_1_0_pre4, CMSSW_5_3_5, CMSSW_5_3_3_patch3, CMSSW_5_3_4_patch1, CMSSW_6_1_0_pre3, CMSSW_4_4_5_patch1, CMSSW_5_2_7, CMSSW_5_3_4, CMSSW_5_3_2_patch5, CMSSW_6_1_0_pre2, CMSSW_5_3_4_cand2, CMSSW_5_3_4_cand1_patch1, CMSSW_6_0_0_patch1, CMSSW_6_0_0_SLHCtkpre1, CMSSW_6_1_0_pre1, CMSSW_5_3_4_cand1, CMSSW_6_0_0_TS123272, CMSSW_5_2_6_patch1, CMSSW_6_0_0, CMSSW_4_4_5, CMSSW_5_3_3_patch2, CMSSW_5_3_3, CMSSW_5_3_3_patch1, CMSSW_6_0_0_pre11, CMSSW_6_0_0_pre10, CMSSW_6_0_0_pre9, CMSSW_5_3_2_patch4, CMSSW_6_0_0_pre8, CMSSW_5_3_3_cand1, CMSSW_5_3_2_patch2, CMSSW_5_3_2_metpatch1, CMSSW_5_2_6_hltpatch1, CMSSW_6_0_0_pre7py273, CMSSW_5_2_6, CMSSW_5_3_2_patch1, CMSSW_5_3_2, CMSSW_6_0_0_pre7, CMSSW_5_2_4_hltpatch4, CMSSW_6_0_0_pre6g495p01, CMSSW_5_2_5_patch3, CMSSW_5_2_5_patch2, CMSSW_6_0_0_pre6, CMSSW_5_3_1, CMSSW_5_2_5_ecalpatch1, CMSSW_6_0_0_pre5, CMSSW_5_3_0_patch1, CMSSW_5_2_5, CMSSW_5_2_5_patch1, CMSSW_sm120515a, CMSSW_5_3_0, CMSSW_6_0_0_pre4, CMSSW_5_2_4_hltpatch2, CMSSW_6_0_X, CMSSW_5_2_5_cand1, CMSSW_5_2_4_patch4, CMSSW_6_0_0_pre3, CMSSW_5_2_4_patch3, CMSSW_5_2_4_patch2, CMSSW_5_2_4_patch1, CMSSW_5_2_4, CMSSW_5_2_3_patch4, CMSSW_5_2_3_patch3, CMSSW_6_0_0_pre2, CMSSW_5_2_3_patch2, CMSSW_5_2_3_patch1, CMSSW_5_2_3, CMSSW_6_0_0_pre1, CMSSW_5_1_3, CMSSW_5_2_2, CMSSW_5_2_1, CMSSW_4_4_4, CMSSW_5_1_2_patch1, CMSSW_5_2_0, CMSSW_5_1_1_patch3, CMSSW_5_2_0_pre5_TS117504, CMSSW_5_1_2, CMSSW_5_2_0_pre6, CMSSW_5_1_1_patch2, CMSSW_4_4_3_patch1, CMSSW_5_1_1_patch1, CMSSW_5_2_0_pre5, CMSSW_5_1_1, CMSSW_5_0_1_patch3, CMSSW_5_0_1_patch2, CMSSW_5_2_0_pre4, CMSSW_5_0_1_patch1, CMSSW_5_2_0_pre3HLT, CMSSW_5_2_0_pre3, CMSSW_5_0_1, CMSSW_5_2_0_pre2_TS113282, CMSSW_4_4_3, CMSSW_5_0_0_patch1, CMSSW_5_2_0_pre2, CMSSW_5_2_0_pre1, CMSSW_5_1_0_pre2, CMSSW_5_1_0_pre1, CMSSW_4_4_2_patch10, CMSSW_5_0_0, CMSSW_4_4_2_patch9, CMSSW_4_4_2_patch8, CMSSW_5_0_0_pre7, CMSSW_5_0_0_pre6g494, CMSSW_4_4_2_patch7, CMSSW_5_0_0_pre6, CMSSW_4_4_2_patch6, CMSSW_5_0_0_pre5_root532rc1, CMSSW_4_4_2_patch5, CMSSW_4_4_2_patch4, CMSSW_4_4_2_patch3, CMSSW_4_4_2_patch2, CMSSW_5_0_0_pre5, CMSSW_4_4_2_patch1, V00-05-04, CMSSW_4_4_2, CMSSW_5_0_0_pre4, CMSSW_4_4_0_patch4, CMSSW_5_0_0_pre3, CMSSW_4_4_1, CMSSW_5_0_0_pre2, CMSSW_4_4_0_patch3, CMSSW_4_4_0_patch2, CMSSW_4_4_0_patch1, CMSSW_5_0_0_pre1, CMSSW_4_4_0, CMSSW_4_4_0_pre10, CMSSW_4_4_0_pre9, V00-05-03, CMSSW_4_4_0_pre8, CMSSW_4_4_0_pre7_g494p02, CMSSW_4_4_0_pre7, CMSSW_4_4_0_pre6, CMSSW_4_4_0_pre5, CMSSW_4_4_0_pre4, CMSSW_4_3_0_dqmpatch2, CMSSW_4_4_0_pre3, CMSSW_4_3_0_dqmpatch1, CMSSW_4_3_0, CMSSW_4_4_0_pre2, logger_cbern_19Jun11-09h37m04s, CMSSW_4_4_0_pre1, CMSSW_4_3_0_pre6_dqmIO, CMSSW_4_3_0_pre7, CMSSW_4_3_0_pre6, V00-05-02, HEAD
Changes since 1.1:	+29 -3 lines
Log Message:	Update to v2.2.2 of rootplot
#	User	Rev	Content
1	klukas	1.1	"""
2			Utilities for plotting ROOT histograms in matplotlib.
3			"""
4
5			__license__ = '''\
6			Copyright (c) 2009-2010 Jeff Klukas <klukas@wisc.edu>
7
8			Permission is hereby granted, free of charge, to any person obtaining a copy
9			of this software and associated documentation files (the "Software"), to deal
10			in the Software without restriction, including without limitation the rights
11			to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
12			copies of the Software, and to permit persons to whom the Software is
13			furnished to do so, subject to the following conditions:
14
15			The above copyright notice and this permission notice shall be included in
16			all copies or substantial portions of the Software.
17
18			THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
19			IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
20			FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
21			AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
22			LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
23			OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
24			THE SOFTWARE.
25			'''
26
27			################ Import python libraries
28
29			import math
30			import ROOT
31			import re
32			import copy
33			import array
34			from rootplot import utilities
35			import matplotlib as mpl
36			import matplotlib.pyplot as plt
37			import matplotlib.transforms as transforms
38			import numpy as np
39
40			################ Define constants
41
42			_all_whitespace_string = re.compile(r'\s*$')
43
44
45			################ Define classes
46
47			class Hist2D(utilities.Hist2D):
48			"""A container to hold the parameters from a 2D ROOT histogram."""
49			def __init__(self, args, *kwargs):
50			self.replacements = None
51			if 'replacements' in kwargs:
52			self.replacements = kwargs.pop('replacements')
53			utilities.Hist2D.__init__(self, args, *kwargs)
54			def contour(self, **kwargs):
55			"""Draw a contour plot."""
56			cs = plt.contour(self.x, self.y, self.content, **kwargs)
57			plt.clabel(cs, inline=1, fontsize=10)
58			if self.binlabelsx is not None:
59			plt.xticks(np.arange(self.nbinsx), self.binlabelsx)
60			if self.binlabelsy is not None:
61			plt.yticks(np.arange(self.nbinsy), self.binlabelsy)
62			return cs
63			def col(self, **kwargs):
64			"""Draw a colored box plot using :func:`matplotlib.pyplot.imshow`."""
65	klukas	1.2	if 'cmap' in kwargs:
66			kwargs['cmap'] = plt.get_cmap(kwargs['cmap'])
67	klukas	1.1	plot = plt.imshow(self.content, interpolation='nearest',
68			extent=[self.xedges[0], self.xedges[-1],
69			self.yedges[0], self.yedges[-1]],
70			aspect='auto', origin='lower', **kwargs)
71			return plot
72			def colz(self, **kwargs):
73			"""
74			Draw a colored box plot with a colorbar using
75			:func:`matplotlib.pyplot.imshow`.
76			"""
77			plot = self.col(**kwargs)
78			plt.colorbar(plot)
79			return plot
80			def box(self, maxsize=40, **kwargs):
81			"""
82			Draw a box plot with size indicating content using
83			:func:`matplotlib.pyplot.scatter`.
84
85			The data will be normalized, with the largest box using a marker of
86			size maxsize (in points).
87			"""
88			x = np.hstack([self.x for i in range(self.nbinsy)])
89			y = np.hstack([[yval for i in range(self.nbinsx)] for yval in self.y])
90			maxvalue = np.max(self.content)
91			if maxvalue == 0:
92			maxvalue = 1
93			sizes = np.array(self.content).flatten() / maxvalue * maxsize
94			plot = plt.scatter(x, y, sizes, marker='s', **kwargs)
95			return plot
96			def TH2F(self, name=""):
97			"""Return a ROOT.TH2F object with contents of this Hist2D."""
98			th2f = ROOT.TH2F(name, "",
99			self.nbinsx, array.array('f', self.xedges),
100			self.nbinsy, array.array('f', self.yedges))
101			th2f.SetTitle("%s;%s;%s" % (self.title, self.xlabel, self.ylabel))
102			for ix in range(self.nbinsx):
103			for iy in range(self.nbinsy):
104			th2f.SetBinContent(ix + 1, iy + 1, self.content[iy][ix])
105			return th2f
106
107			class Hist(utilities.Hist):
108			"""A container to hold the parameters from a ROOT histogram."""
109			def __init__(self, args, *kwargs):
110			self.replacements = None
111			if 'replacements' in kwargs:
112			self.replacements = kwargs.pop('replacements')
113			utilities.Hist.__init__(self, args, *kwargs)
114			def _prepare_xaxis(self, rotation=0, alignment='center'):
115			"""Apply bounds and text labels on x axis."""
116			if self.binlabels is not None:
117			binwidth = (self.xedges[-1] - self.xedges[0]) / self.nbins
118			plt.xticks(self.x, self.binlabels,
119			rotation=rotation, ha=alignment)
120			plt.xlim(self.xedges[0], self.xedges[-1])
121
122			def _prepare_yaxis(self, rotation=0, alignment='center'):
123	klukas	1.2	"""Apply bounds and text labels on y axis."""
124	klukas	1.1	if self.binlabels is not None:
125			binwidth = (self.xedges[-1] - self.xedges[0]) / self.nbins
126			plt.yticks(self.x, self.binlabels,
127			rotation=rotation, va=alignment)
128			plt.ylim(self.xedges[0], self.xedges[-1])
129
130			def show_titles(self, **kwargs):
131			"""Print the title and axis labels to the current figure."""
132			replacements = kwargs.get('replacements', None) or self.replacements
133			plt.title(replace(self.title, replacements))
134			plt.xlabel(replace(self.xlabel, replacements))
135			plt.ylabel(replace(self.ylabel, replacements))
136			def hist(self, label_rotation=0, label_alignment='center', **kwargs):
137			"""
138			Generate a matplotlib hist figure.
139
140			All additional keyword arguments will be passed to
141			:func:`matplotlib.pyplot.hist`.
142			"""
143			kwargs.pop('fmt', None)
144			replacements = kwargs.get('replacements', None) or self.replacements
145			weights = self.y
146			# Kludge to avoid mpl bug when plotting all zeros
147			if self.y == [0] * self.nbins:
148			weights = [1.e-10] * self.nbins
149			plot = plt.hist(self.x, weights=weights, bins=self.xedges,
150			label=replace(self.label, replacements), **kwargs)
151			self._prepare_xaxis(label_rotation, label_alignment)
152			return plot
153			def errorbar(self, xerr=False, yerr=False, label_rotation=0,
154			label_alignment='center', **kwargs):
155			"""
156			Generate a matplotlib errorbar figure.
157
158			All additional keyword arguments will be passed to
159			:func:`matplotlib.pyplot.errorbar`.
160			"""
161			if xerr:
162			kwargs['xerr'] = self.xerr
163			if yerr:
164			kwargs['yerr'] = self.yerr
165			replacements = kwargs.get('replacements', None) or self.replacements
166			errorbar = plt.errorbar(self.x, self.y,
167			label=replace(self.label, replacements),
168			**kwargs)
169			self._prepare_xaxis(label_rotation, label_alignment)
170			return errorbar
171			def errorbarh(self, xerr=False, yerr=False, label_rotation=0,
172			label_alignment='center', **kwargs):
173			"""
174			Generate a horizontal matplotlib errorbar figure.
175
176			All additional keyword arguments will be passed to
177			:func:`matplotlib.pyplot.errorbar`.
178			"""
179			if xerr: kwargs['xerr'] = self.yerr
180			if yerr: kwargs['yerr'] = self.xerr
181			replacements = kwargs.get('replacements', None) or self.replacements
182			errorbar = plt.errorbar(self.y, self.x,
183			label=replace(self.label, replacements),
184			**kwargs)
185			self._prepare_yaxis(label_rotation, label_alignment)
186			return errorbar
187			def bar(self, xerr=False, yerr=False, xoffset=0., width=0.8,
188			label_rotation=0, label_alignment='center', **kwargs):
189			"""
190			Generate a matplotlib bar figure.
191
192			All additional keyword arguments will be passed to
193			:func:`matplotlib.pyplot.bar`.
194			"""
195			kwargs.pop('fmt', None)
196			if xerr: kwargs['xerr'] = self.av_xerr()
197			if yerr: kwargs['yerr'] = self.av_yerr()
198			replacements = kwargs.get('replacements', None) or self.replacements
199			ycontent = [self.xedges[i] + self.width[i] * xoffset
200			for i in range(len(self.xedges) - 1)]
201			width = [x * width for x in self.width]
202			bar = plt.bar(ycontent, self.y, width,
203			label=replace(self.label, replacements), **kwargs)
204			self._prepare_xaxis(label_rotation, label_alignment)
205			return bar
206			def barh(self, xerr=False, yerr=False, yoffset=0., width=0.8,
207			label_rotation=0, label_alignment='center', **kwargs):
208			"""
209			Generate a horizontal matplotlib bar figure.
210
211			All additional keyword arguments will be passed to
212			:func:`matplotlib.pyplot.bar`.
213			"""
214			kwargs.pop('fmt', None)
215			if xerr: kwargs['xerr'] = self.av_yerr()
216			if yerr: kwargs['yerr'] = self.av_xerr()
217			replacements = kwargs.get('replacements', None) or self.replacements
218			xcontent = [self.xedges[i] + self.width[i] * yoffset
219			for i in range(len(self.xedges) - 1)]
220			width = [x * width for x in self.width]
221			barh = plt.barh(xcontent, self.y, width,
222			label=replace(self.label, replacements),
223			**kwargs)
224			self._prepare_yaxis(label_rotation, label_alignment)
225			return barh
226
227			class HistStack(utilities.HistStack):
228			"""
229			A container to hold Hist objects for plotting together.
230
231			When plotting, the title and the x and y labels of the last Hist added
232			will be used unless specified otherwise in the constructor.
233			"""
234			def __init__(self, args, *kwargs):
235			if 'replacements' in kwargs:
236			self.replacements = kwargs.pop('replacements')
237			utilities.HistStack.__init__(self, args, *kwargs)
238			def show_titles(self, **kwargs):
239			self.hists[-1].show_titles()
240			def hist(self, label_rotation=0, **kwargs):
241			"""
242			Make a matplotlib hist plot.
243
244			Any additional keyword arguments will be passed to
245			:func:`matplotlib.pyplot.hist`, which allows a vast array of
246			possibilities. Particlularly, the histtype values such as
247			``'barstacked'`` and ``'stepfilled'`` give substantially different
248			results. You will probably want to include a transparency value
249			(i.e. alpha = 0.5).
250			"""
251			contents = np.dstack([hist.y for hist in self.hists])
252			xedges = self.hists[0].xedges
253	klukas	1.2	x = np.dstack([hist.x for hist in self.hists])[0]
254	klukas	1.1	labels = [hist.label for hist in self.hists]
255			try:
256			clist = [item['color'] for item in self.kwargs]
257			plt.gca().set_color_cycle(clist)
258			## kwargs['color'] = clist # For newer version of matplotlib
259			except:
260			pass
261			plot = plt.hist(x, weights=contents, bins=xedges,
262			label=labels, **kwargs)
263			def bar3d(self, **kwargs):
264			#### Not yet ready for primetime
265			from mpl_toolkits.mplot3d import Axes3D
266			fig = plt.figure()
267			ax = Axes3D(fig)
268			plots = []
269			labels = []
270			for i, hist in enumerate(self.hists):
271			if self.title is not None: hist.title = self.title
272			if self.xlabel is not None: hist.xlabel = self.xlabel
273			if self.ylabel is not None: hist.ylabel = self.ylabel
274			labels.append(hist.label)
275			all_kwargs = copy.copy(kwargs)
276			all_kwargs.update(self.kwargs[i])
277			bar = ax.bar(hist.x, hist.y, zs=i, zdir='y', width=hist.width,
278			**all_kwargs)
279			plots.append(bar)
280			from matplotlib.ticker import FixedLocator
281			locator = FixedLocator(range(len(labels)))
282			ax.w_yaxis.set_major_locator(locator)
283			ax.w_yaxis.set_ticklabels(labels)
284			ax.set_ylim3d(-1, len(labels))
285			return plots
286			def barstack(self, **kwargs):
287			"""
288			Make a matplotlib bar plot, with each Hist stacked upon the last.
289
290			Any additional keyword arguments will be passed to
291			:func:`matplotlib.pyplot.bar`.
292			"""
293			bottom = None # if this is set to zeroes, it fails for log y
294			plots = []
295			for i, hist in enumerate(self.hists):
296			if self.title is not None: hist.title = self.title
297			if self.xlabel is not None: hist.xlabel = self.xlabel
298			if self.ylabel is not None: hist.ylabel = self.ylabel
299			all_kwargs = copy.copy(kwargs)
300			all_kwargs.update(self.kwargs[i])
301			bar = hist.bar(bottom=bottom, **all_kwargs)
302			plots.append(bar)
303			if not bottom: bottom = [0. for i in range(self.hists[0].nbins)]
304			bottom = [sum(pair) for pair in zip(bottom, hist.y)]
305			return plots
306	klukas	1.2	def histstack(self, **kwargs):
307			"""
308			Make a matplotlib hist plot, with each Hist stacked upon the last.
309
310			Any additional keyword arguments will be passed to
311			:func:`matplotlib.pyplot.hist`.
312			"""
313			bottom = None # if this is set to zeroes, it fails for log y
314			plots = []
315			cumhist = None
316			for i, hist in enumerate(self.hists):
317			if cumhist:
318			cumhist = hist + cumhist
319			else:
320			cumhist = copy.copy(hist)
321			if self.title is not None: cumhist.title = self.title
322			if self.xlabel is not None: cumhist.xlabel = self.xlabel
323			if self.ylabel is not None: cumhist.ylabel = self.ylabel
324			all_kwargs = copy.copy(kwargs)
325			all_kwargs.update(self.kwargs[i])
326			zorder = 0 + float(len(self) - i)/len(self) # plot in reverse order
327			plot = cumhist.hist(zorder=zorder, **all_kwargs)
328			plots.append(plot)
329			return plots
330	klukas	1.1	def barcluster(self, width=0.8, **kwargs):
331			"""
332			Make a clustered bar plot.
333
334			Any additional keyword arguments will be passed to
335			:func:`matplotlib.pyplot.bar`.
336			"""
337			plots = []
338			spacer = (1. - width) / 2
339			width = width / len(self.hists)
340			for i, hist in enumerate(self.hists):
341			if self.title is not None: hist.title = self.title
342			if self.xlabel is not None: hist.xlabel = self.xlabel
343			if self.ylabel is not None: hist.ylabel = self.ylabel
344			all_kwargs = copy.copy(kwargs)
345			all_kwargs.update(self.kwargs[i])
346			bar = hist.bar(xoffset=widthi + spacer, width=width, *all_kwargs)
347			plots.append(bar)
348			return plots
349			def barh(self, width=0.8, **kwargs):
350			"""
351			Make a horizontal clustered matplotlib bar plot.
352
353			Any additional keyword arguments will be passed to
354			:func:`matplotlib.pyplot.bar`.
355			"""
356			plots = []
357			spacer = (1. - width) / 2
358			width = width / len(self.hists)
359			for i, hist in enumerate(self.hists):
360			if self.title is not None: hist.title = self.title
361			if self.xlabel is not None: hist.ylabel = self.xlabel
362			if self.ylabel is not None: hist.xlabel = self.ylabel
363			all_kwargs = copy.copy(kwargs)
364			all_kwargs.update(self.kwargs[i])
365			bar = hist.barh(yoffset=widthi + spacer, width=width, *all_kwargs)
366			plots.append(bar)
367			return plots
368			def bar(self, **kwargs):
369			"""
370			Make a bar plot, with all Hists in the stack overlaid.
371
372			Any additional keyword arguments will be passed to
373			:func:`matplotlib.pyplot.bar`. You will probably want to set a
374	klukas	1.2	transparency value (i.e. alpha = 0.5).
375	klukas	1.1	"""
376			plots = []
377			for i, hist in enumerate(self.hists):
378			if self.title is not None: hist.title = self.title
379			if self.xlabel is not None: hist.xlabel = self.xlabel
380			if self.ylabel is not None: hist.ylabel = self.ylabel
381			all_kwargs = copy.copy(kwargs)
382			all_kwargs.update(self.kwargs[i])
383			bar = hist.bar(**all_kwargs)
384			plots.append(bar)
385			return plots
386			def errorbar(self, offset=False, **kwargs):
387			"""
388			Make a matplotlib errorbar plot, with all Hists in the stack overlaid.
389
390			Passing 'offset=True' will slightly offset each dataset so overlapping
391			errorbars are still visible. Any additional keyword arguments will
392			be passed to :func:`matplotlib.pyplot.errorbar`.
393			"""
394			plots = []
395			for i, hist in enumerate(self.hists):
396			if self.title is not None: hist.title = self.title
397			if self.xlabel is not None: hist.xlabel = self.xlabel
398			if self.ylabel is not None: hist.ylabel = self.ylabel
399			all_kwargs = copy.copy(kwargs)
400			all_kwargs.update(self.kwargs[i])
401			transform = plt.gca().transData
402			if offset:
403			index_offset = (len(self.hists) - 1)/2.
404			pixel_offset = 1./72 * (i - index_offset)
405			transform = transforms.ScaledTranslation(
406			pixel_offset, 0, plt.gcf().dpi_scale_trans)
407			transform = plt.gca().transData + transform
408			errorbar = hist.errorbar(transform=transform, **all_kwargs)
409			plots.append(errorbar)
410			return plots
411			def errorbarh(self, **kwargs):
412			"""
413			Make a horizontal matplotlib errorbar plot, with all Hists in the
414			stack overlaid.
415
416			Any additional keyword arguments will be passed to
417			:func:`matplotlib.pyplot.errorbar`.
418			"""
419			plots = []
420			for i, hist in enumerate(self.hists):
421			if self.title is not None: hist.title = self.title
422			if self.xlabel is not None: hist.ylabel = self.xlabel
423			if self.ylabel is not None: hist.xlabel = self.ylabel
424			all_kwargs = copy.copy(kwargs)
425			all_kwargs.update(self.kwargs[i])
426			errorbar = hist.errorbarh(**all_kwargs)
427			plots.append(errorbar)
428			return plots
429
430			################ Define functions and classes for navigating within ROOT
431
432			class RootFile(utilities.RootFile):
433			"""A wrapper for TFiles, allowing easier access to methods."""
434			def get(self, object_name, path=None):
435			try:
436			return utilities.RootFile.get(self, object_name, path,
437			Hist, Hist2D)
438			except ReferenceError, e:
439			raise ReferenceError(e)
440
441			################ Define additional helping functions
442
443			def replace(string, replacements):
444			"""
445			Modify a string based on a list of patterns and substitutions.
446
447			replacements should be a list of two-entry tuples, the first entry giving
448			a string to search for and the second entry giving the string with which
449			to replace it. If replacements includes a pattern entry containing
450			'use_regexp', then all patterns will be treated as regular expressions
451			using re.sub.
452			"""
453			if not replacements:
454			return string
455			if 'use_regexp' in [x for x,y in replacements]:
456			for pattern, repl in [x for x in replacements
457			if x[0] != 'use_regexp']:
458			string = re.sub(pattern, repl, string)
459			else:
460			for pattern, repl in replacements:
461			string = string.replace(pattern, repl)
462			if re.match(_all_whitespace_string, string):
463			return ""
464			return string
465