Source code for rootpy.plotting.root2matplotlib

"""
This module provides functions that allow the plotting of ROOT histograms and
graphs with `matplotlib <http://matplotlib.org/>`_.

If you just want to save image files and don't want matplotlib to attempt to
create a graphical window, tell matplotlib to use a non-interactive backend
such as ``Agg`` when importing it for the first time (i.e. before importing
rootpy.plotting.root2matplotlib)::

   import matplotlib
   matplotlib.use('Agg') # do this before importing pyplot or root2matplotlib

This puts matplotlib in a batch state similar to ``ROOT.gROOT.SetBatch(True)``.
"""
from __future__ import absolute_import

# trigger ROOT's finalSetup (GUI thread) before matplotlib's
import ROOT
ROOT.kTRUE

from math import sqrt
try:
    from itertools import izip as zip
except ImportError: # will be 3.x series
    pass
import matplotlib.pyplot as plt
import numpy as np

from ..extern.six.moves import range
from .hist import _Hist
from .graph import _Graph1DBase
from .utils import get_limits


__all__ = [
    'hist',
    'bar',
    'errorbar',
    'fill_between',
    'step',
    'hist2d',
    'imshow',
    'contour',
]


def _set_defaults(obj, kwargs, types=['common']):
    defaults = {}
    for key in types:
        if key == 'common':
            defaults['label'] = obj.GetTitle()
            defaults['visible'] = getattr(obj, 'visible', True)
            defaults['alpha'] = getattr(obj, 'alpha', None)
        elif key == 'line':
            defaults['linestyle'] = obj.GetLineStyle('mpl')
            defaults['linewidth'] = obj.GetLineWidth()
        elif key == 'fill':
            defaults['edgecolor'] = kwargs.get('color', obj.GetLineColor('mpl'))
            defaults['facecolor'] = kwargs.get('color', obj.GetFillColor('mpl'))
            root_fillstyle = obj.GetFillStyle('root')
            if root_fillstyle == 0:
                if not kwargs.get('fill'):
                    defaults['facecolor'] = 'none'
                defaults['fill'] = False
            elif root_fillstyle == 1001:
                defaults['fill'] = True
            else:
                defaults['hatch'] = obj.GetFillStyle('mpl')
                defaults['facecolor'] = 'none'
        elif key == 'marker':
            defaults['marker'] = obj.GetMarkerStyle('mpl')
            defaults['markersize'] = obj.GetMarkerSize() * 5
            defaults['markeredgecolor'] = obj.GetMarkerColor('mpl')
            defaults['markerfacecolor'] = obj.GetMarkerColor('mpl')
        elif key == 'errors':
            defaults['ecolor'] = obj.GetLineColor('mpl')
        elif key == 'errorbar':
            defaults['fmt'] = obj.GetMarkerStyle('mpl')
    for key, value in defaults.items():
        if key not in kwargs:
            kwargs[key] = value


def _set_bounds(h,
                axes=None,
                was_empty=True,
                prev_xlim=None,
                prev_ylim=None,
                xpadding=0,
                ypadding=.1,
                xerror_in_padding=True,
                yerror_in_padding=True,
                snap=True,
                logx=None,
                logy=None):
    if axes is None:
        axes = plt.gca()
    if prev_xlim is None:
        prev_xlim = plt.xlim()
    if prev_ylim is None:
        prev_ylim = plt.ylim()
    if logx is None:
        logx = axes.get_xscale() == 'log'
    if logy is None:
        logy = axes.get_yscale() == 'log'
    xmin, xmax, ymin, ymax = get_limits(
        h,
        xpadding=xpadding,
        ypadding=ypadding,
        xerror_in_padding=xerror_in_padding,
        yerror_in_padding=yerror_in_padding,
        snap=snap,
        logx=logx,
        logy=logy)
    if was_empty:
        axes.set_xlim([xmin, xmax])
        axes.set_ylim([ymin, ymax])
    else:
        prev_xmin, prev_xmax = prev_xlim
        if logx and prev_xmin <= 0:
            axes.set_xlim([xmin, max(prev_xmax, xmax)])
        else:
            axes.set_xlim([min(prev_xmin, xmin), max(prev_xmax, xmax)])
        prev_ymin, prev_ymax = prev_ylim
        if logy and prev_ymin <= 0:
            axes.set_ylim([ymin, max(prev_ymax, ymax)])
        else:
            axes.set_ylim([min(prev_ymin, ymin), max(prev_ymax, ymax)])


def _get_highest_zorder(axes):
    return max([c.get_zorder() for c in axes.get_children()])


def _maybe_reversed(x, reverse=False):
    if reverse:
        return reversed(x)
    return x


[docs]def hist(hists, stacked=True, reverse=False, xpadding=0, ypadding=.1, yerror_in_padding=True, logy=None, snap=True, axes=None, **kwargs): """ Make a matplotlib hist plot from a ROOT histogram, stack or list of histograms. Parameters ---------- hists : Hist, list of Hist, HistStack The histogram(s) to be plotted stacked : bool, optional (default=True) If True then stack the histograms with the first histogram on the bottom, otherwise overlay them with the first histogram in the background. reverse : bool, optional (default=False) If True then reverse the order of the stack or overlay. xpadding : float or 2-tuple of floats, optional (default=0) Padding to add on the left and right sides of the plot as a fraction of the axes width after the padding has been added. Specify unique left and right padding with a 2-tuple. ypadding : float or 2-tuple of floats, optional (default=.1) Padding to add on the top and bottom of the plot as a fraction of the axes height after the padding has been added. Specify unique top and bottom padding with a 2-tuple. yerror_in_padding : bool, optional (default=True) If True then make the padding inclusive of the y errors otherwise only pad around the y values. logy : bool, optional (default=None) Apply special treatment of a log-scale y-axis to display the histogram correctly. If None (the default) then automatically determine if the y-axis is log-scale. snap : bool, optional (default=True) If True (the default) then the origin is an implicit lower bound of the histogram unless the histogram has both positive and negative bins. axes : matplotlib Axes instance, optional (default=None) The axes to plot on. If None then use the global current axes. kwargs : additional keyword arguments, optional All additional keyword arguments are passed to matplotlib's fill_between for the filled regions and matplotlib's step function for the edges. Returns ------- The return value from matplotlib's hist function, or list of such return values if a stack or list of histograms was plotted. """ if axes is None: axes = plt.gca() if logy is None: logy = axes.get_yscale() == 'log' curr_xlim = axes.get_xlim() curr_ylim = axes.get_ylim() was_empty = not axes.has_data() returns = [] if isinstance(hists, _Hist): # This is a single plottable object. returns = _hist(hists, axes=axes, logy=logy, **kwargs) _set_bounds(hists, axes=axes, was_empty=was_empty, prev_xlim=curr_xlim, prev_ylim=curr_ylim, xpadding=xpadding, ypadding=ypadding, yerror_in_padding=yerror_in_padding, snap=snap, logy=logy) elif stacked: # draw the top histogram first so its edges don't cover the histograms # beneath it in the stack if not reverse: hists = list(hists)[::-1] for i, h in enumerate(hists): kwargs_local = kwargs.copy() if i == len(hists) - 1: low = h.Clone() low.Reset() else: low = sum(hists[i + 1:]) high = h + low high.alpha = getattr(h, 'alpha', None) proxy = _hist(high, bottom=low, axes=axes, logy=logy, **kwargs) returns.append(proxy) if not reverse: returns = returns[::-1] _set_bounds(sum(hists), axes=axes, was_empty=was_empty, prev_xlim=curr_xlim, prev_ylim=curr_ylim, xpadding=xpadding, ypadding=ypadding, yerror_in_padding=yerror_in_padding, snap=snap, logy=logy) else: for h in _maybe_reversed(hists, reverse): returns.append(_hist(h, axes=axes, logy=logy, **kwargs)) if reverse: returns = returns[::-1] _set_bounds(hists[max(range(len(hists)), key=lambda idx: hists[idx].max())], axes=axes, was_empty=was_empty, prev_xlim=curr_xlim, prev_ylim=curr_ylim, xpadding=xpadding, ypadding=ypadding, yerror_in_padding=yerror_in_padding, snap=snap, logy=logy) return returns
def _hist(h, axes=None, bottom=None, logy=None, zorder=None, **kwargs): if axes is None: axes = plt.gca() if zorder is None: zorder = _get_highest_zorder(axes) + 1 _set_defaults(h, kwargs, ['common', 'line', 'fill']) kwargs_proxy = kwargs.copy() fill = kwargs.pop('fill', False) or ('hatch' in kwargs) if fill: # draw the fill without the edge if bottom is None: bottom = h.Clone() bottom.Reset() fill_between(bottom, h, axes=axes, logy=logy, linewidth=0, facecolor=kwargs['facecolor'], edgecolor=kwargs['edgecolor'], hatch=kwargs.get('hatch', None), alpha=kwargs['alpha'], zorder=zorder) # draw the edge s = step(h, axes=axes, logy=logy, label=None, zorder=zorder + 1, alpha=kwargs['alpha'], color=kwargs.get('color')) # draw the legend proxy if getattr(h, 'legendstyle', '').upper() == 'F': proxy = plt.Rectangle((0, 0), 0, 0, **kwargs_proxy) axes.add_patch(proxy) else: # be sure the linewidth is greater than zero... proxy = plt.Line2D((0, 0), (0, 0), linestyle=kwargs_proxy['linestyle'], linewidth=kwargs_proxy['linewidth'], color=kwargs_proxy['edgecolor'], alpha=kwargs['alpha'], label=kwargs_proxy['label']) axes.add_line(proxy) return proxy, s[0]
[docs]def bar(hists, stacked=True, reverse=False, xerr=False, yerr=True, xpadding=0, ypadding=.1, yerror_in_padding=True, rwidth=0.8, snap=True, axes=None, **kwargs): """ Make a matplotlib bar plot from a ROOT histogram, stack or list of histograms. Parameters ---------- hists : Hist, list of Hist, HistStack The histogram(s) to be plotted stacked : bool or string, optional (default=True) If True then stack the histograms with the first histogram on the bottom, otherwise overlay them with the first histogram in the background. If 'cluster', then the bars will be arranged side-by-side. reverse : bool, optional (default=False) If True then reverse the order of the stack or overlay. xerr : bool, optional (default=False) If True, x error bars will be displayed. yerr : bool or string, optional (default=True) If False, no y errors are displayed. If True, an individual y error will be displayed for each hist in the stack. If 'linear' or 'quadratic', a single error bar will be displayed with either the linear or quadratic sum of the individual errors. xpadding : float or 2-tuple of floats, optional (default=0) Padding to add on the left and right sides of the plot as a fraction of the axes width after the padding has been added. Specify unique left and right padding with a 2-tuple. ypadding : float or 2-tuple of floats, optional (default=.1) Padding to add on the top and bottom of the plot as a fraction of the axes height after the padding has been added. Specify unique top and bottom padding with a 2-tuple. yerror_in_padding : bool, optional (default=True) If True then make the padding inclusive of the y errors otherwise only pad around the y values. rwidth : float, optional (default=0.8) The relative width of the bars as a fraction of the bin width. snap : bool, optional (default=True) If True (the default) then the origin is an implicit lower bound of the histogram unless the histogram has both positive and negative bins. axes : matplotlib Axes instance, optional (default=None) The axes to plot on. If None then use the global current axes. kwargs : additional keyword arguments, optional All additional keyword arguments are passed to matplotlib's bar function. Returns ------- The return value from matplotlib's bar function, or list of such return values if a stack or list of histograms was plotted. """ if axes is None: axes = plt.gca() curr_xlim = axes.get_xlim() curr_ylim = axes.get_ylim() was_empty = not axes.has_data() logy = kwargs.pop('log', axes.get_yscale() == 'log') kwargs['log'] = logy returns = [] if isinstance(hists, _Hist): # This is a single histogram. returns = _bar(hists, xerr=xerr, yerr=yerr, axes=axes, **kwargs) _set_bounds(hists, axes=axes, was_empty=was_empty, prev_xlim=curr_xlim, prev_ylim=curr_ylim, xpadding=xpadding, ypadding=ypadding, yerror_in_padding=yerror_in_padding, snap=snap, logy=logy) elif stacked == 'cluster': nhists = len(hists) hlist = _maybe_reversed(hists, reverse) for i, h in enumerate(hlist): width = rwidth / nhists offset = (1 - rwidth) / 2 + i * width returns.append(_bar( h, offset, width, xerr=xerr, yerr=yerr, axes=axes, **kwargs)) _set_bounds(sum(hists), axes=axes, was_empty=was_empty, prev_xlim=curr_xlim, prev_ylim=curr_ylim, xpadding=xpadding, ypadding=ypadding, yerror_in_padding=yerror_in_padding, snap=snap, logy=logy) elif stacked is True: nhists = len(hists) hlist = _maybe_reversed(hists, reverse) toterr = bottom = None if yerr == 'linear': toterr = [sum([h.GetBinError(i) for h in hists]) for i in range(1, hists[0].nbins(0) + 1)] elif yerr == 'quadratic': toterr = [sqrt(sum([h.GetBinError(i) ** 2 for h in hists])) for i in range(1, hists[0].nbins(0) + 1)] for i, h in enumerate(hlist): err = None if yerr is True: err = True elif yerr and i == (nhists - 1): err = toterr returns.append(_bar( h, xerr=xerr, yerr=err, bottom=list(bottom.y()) if bottom else None, axes=axes, **kwargs)) if bottom is None: bottom = h.Clone() else: bottom += h _set_bounds(bottom, axes=axes, was_empty=was_empty, prev_xlim=curr_xlim, prev_ylim=curr_ylim, xpadding=xpadding, ypadding=ypadding, yerror_in_padding=yerror_in_padding, snap=snap, logy=logy) else: hlist = _maybe_reversed(hists, reverse) for h in hlist: returns.append(_bar(h, xerr=xerr, yerr=yerr, axes=axes, **kwargs)) _set_bounds(hists[max(range(len(hists)), key=lambda idx: hists[idx].max())], axes=axes, was_empty=was_empty, prev_xlim=curr_xlim, prev_ylim=curr_ylim, xpadding=xpadding, ypadding=ypadding, yerror_in_padding=yerror_in_padding, snap=snap, logy=logy) return returns
def _bar(h, roffset=0., rwidth=1., xerr=None, yerr=None, axes=None, **kwargs): if axes is None: axes = plt.gca() if xerr: xerr = np.array([list(h.xerrl()), list(h.xerrh())]) if yerr: yerr = np.array([list(h.yerrl()), list(h.yerrh())]) _set_defaults(h, kwargs, ['common', 'line', 'fill', 'errors']) width = [x * rwidth for x in h.xwidth()] left = [h.xedgesl(i) + h.xwidth(i) * roffset for i in range(1, h.nbins(0) + 1)] height = list(h.y()) return axes.bar(left, height, width=width, xerr=xerr, yerr=yerr, **kwargs)
[docs]def errorbar(hists, xerr=True, yerr=True, xpadding=0, ypadding=.1, xerror_in_padding=True, yerror_in_padding=True, emptybins=True, snap=True, axes=None, **kwargs): """ Make a matplotlib errorbar plot from a ROOT histogram or graph or list of histograms and graphs. Parameters ---------- hists : Hist, Graph or list of Hist and Graph The histogram(s) and/or Graph(s) to be plotted xerr : bool, optional (default=True) If True, x error bars will be displayed. yerr : bool or string, optional (default=True) If False, no y errors are displayed. If True, an individual y error will be displayed for each hist in the stack. If 'linear' or 'quadratic', a single error bar will be displayed with either the linear or quadratic sum of the individual errors. xpadding : float or 2-tuple of floats, optional (default=0) Padding to add on the left and right sides of the plot as a fraction of the axes width after the padding has been added. Specify unique left and right padding with a 2-tuple. ypadding : float or 2-tuple of floats, optional (default=.1) Padding to add on the top and bottom of the plot as a fraction of the axes height after the padding has been added. Specify unique top and bottom padding with a 2-tuple. xerror_in_padding : bool, optional (default=True) If True then make the padding inclusive of the x errors otherwise only pad around the x values. yerror_in_padding : bool, optional (default=True) If True then make the padding inclusive of the y errors otherwise only pad around the y values. emptybins : bool, optional (default=True) If True (the default) then plot bins with zero content otherwise only show bins with nonzero content. snap : bool, optional (default=True) If True (the default) then the origin is an implicit lower bound of the histogram unless the histogram has both positive and negative bins. axes : matplotlib Axes instance, optional (default=None) The axes to plot on. If None then use the global current axes. kwargs : additional keyword arguments, optional All additional keyword arguments are passed to matplotlib's errorbar function. Returns ------- The return value from matplotlib's errorbar function, or list of such return values if a list of histograms and/or graphs was plotted. """ if axes is None: axes = plt.gca() curr_xlim = axes.get_xlim() curr_ylim = axes.get_ylim() was_empty = not axes.has_data() if isinstance(hists, (_Hist, _Graph1DBase)): # This is a single plottable object. returns = _errorbar( hists, xerr, yerr, axes=axes, emptybins=emptybins, **kwargs) _set_bounds(hists, axes=axes, was_empty=was_empty, prev_ylim=curr_ylim, xpadding=xpadding, ypadding=ypadding, xerror_in_padding=xerror_in_padding, yerror_in_padding=yerror_in_padding, snap=snap) else: returns = [] for h in hists: returns.append(errorbar( h, xerr=xerr, yerr=yerr, axes=axes, xpadding=xpadding, ypadding=ypadding, xerror_in_padding=xerror_in_padding, yerror_in_padding=yerror_in_padding, snap=snap, emptybins=emptybins, **kwargs)) return returns
def _errorbar(h, xerr, yerr, axes=None, emptybins=True, zorder=None, **kwargs): if axes is None: axes = plt.gca() if zorder is None: zorder = _get_highest_zorder(axes) + 1 _set_defaults(h, kwargs, ['common', 'errors', 'errorbar', 'marker']) if xerr: xerr = np.array([list(h.xerrl()), list(h.xerrh())]) if yerr: yerr = np.array([list(h.yerrl()), list(h.yerrh())]) x = np.array(list(h.x())) y = np.array(list(h.y())) if not emptybins: nonempty = y != 0 x = x[nonempty] y = y[nonempty] if xerr is not False and xerr is not None: xerr = xerr[:, nonempty] if yerr is not False and yerr is not None: yerr = yerr[:, nonempty] return axes.errorbar(x, y, xerr=xerr, yerr=yerr, zorder=zorder, **kwargs)
[docs]def step(h, logy=None, axes=None, **kwargs): """ Make a matplotlib step plot from a ROOT histogram. Parameters ---------- h : Hist A rootpy Hist logy : bool, optional (default=None) If True then clip the y range between 1E-300 and 1E300. If None (the default) then automatically determine if the axes are log-scale and if this clipping should be performed. axes : matplotlib Axes instance, optional (default=None) The axes to plot on. If None then use the global current axes. kwargs : additional keyword arguments, optional Additional keyword arguments are passed directly to matplotlib's fill_between function. Returns ------- Returns the value from matplotlib's fill_between function. """ if axes is None: axes = plt.gca() if logy is None: logy = axes.get_yscale() == 'log' _set_defaults(h, kwargs, ['common', 'line']) if kwargs.get('color') is None: kwargs['color'] = h.GetLineColor('mpl') y = np.array(list(h.y()) + [0.]) if logy: np.clip(y, 1E-300, 1E300, out=y) return axes.step(list(h.xedges()), y, where='post', **kwargs)
[docs]def fill_between(a, b, logy=None, axes=None, **kwargs): """ Fill the region between two histograms or graphs. Parameters ---------- a : Hist A rootpy Hist b : Hist A rootpy Hist logy : bool, optional (default=None) If True then clip the region between 1E-300 and 1E300. If None (the default) then automatically determine if the axes are log-scale and if this clipping should be performed. axes : matplotlib Axes instance, optional (default=None) The axes to plot on. If None then use the global current axes. kwargs : additional keyword arguments, optional Additional keyword arguments are passed directly to matplotlib's fill_between function. Returns ------- Returns the value from matplotlib's fill_between function. """ if axes is None: axes = plt.gca() if logy is None: logy = axes.get_yscale() == 'log' if not isinstance(a, _Hist) or not isinstance(b, _Hist): raise TypeError( "fill_between only operates on 1D histograms") a.check_compatibility(b, check_edges=True) x = [] top = [] bottom = [] for abin, bbin in zip(a.bins(overflow=False), b.bins(overflow=False)): up = max(abin.value, bbin.value) dn = min(abin.value, bbin.value) x.extend([abin.x.low, abin.x.high]) top.extend([up, up]) bottom.extend([dn, dn]) x = np.array(x) top = np.array(top) bottom = np.array(bottom) if logy: np.clip(top, 1E-300, 1E300, out=top) np.clip(bottom, 1E-300, 1E300, out=bottom) return axes.fill_between(x, top, bottom, **kwargs)
[docs]def hist2d(h, axes=None, colorbar=False, **kwargs): """ Draw a 2D matplotlib histogram plot from a 2D ROOT histogram. Parameters ---------- h : Hist2D A rootpy Hist2D axes : matplotlib Axes instance, optional (default=None) The axes to plot on. If None then use the global current axes. colorbar : Boolean, optional (default=False) If True, include a colorbar in the produced plot kwargs : additional keyword arguments, optional Additional keyword arguments are passed directly to matplotlib's hist2d function. Returns ------- Returns the value from matplotlib's hist2d function. """ if axes is None: axes = plt.gca() X, Y = np.meshgrid(list(h.x()), list(h.y())) x = X.ravel() y = Y.ravel() z = np.array(h.z()).T # returns of hist2d: (counts, xedges, yedges, Image) return_values = axes.hist2d(x, y, weights=z.ravel(), bins=(list(h.xedges()), list(h.yedges())), **kwargs) if colorbar: mappable = return_values[-1] plt.colorbar(mappable, ax=axes) return return_values
[docs]def imshow(h, axes=None, colorbar=False, **kwargs): """ Draw a matplotlib imshow plot from a 2D ROOT histogram. Parameters ---------- h : Hist2D A rootpy Hist2D axes : matplotlib Axes instance, optional (default=None) The axes to plot on. If None then use the global current axes. colorbar : Boolean, optional (default=False) If True, include a colorbar in the produced plot kwargs : additional keyword arguments, optional Additional keyword arguments are passed directly to matplotlib's imshow function. Returns ------- Returns the value from matplotlib's imshow function. """ kwargs.setdefault('aspect', 'auto') if axes is None: axes = plt.gca() z = np.array(h.z()).T axis_image= axes.imshow( z, extent=[ h.xedges(1), h.xedges(h.nbins(0) + 1), h.yedges(1), h.yedges(h.nbins(1) + 1)], interpolation='nearest', origin='lower', **kwargs) if colorbar: plt.colorbar(axis_image, ax=axes) return axis_image
[docs]def contour(h, axes=None, zoom=None, label_contour=False, **kwargs): """ Draw a matplotlib contour plot from a 2D ROOT histogram. Parameters ---------- h : Hist2D A rootpy Hist2D axes : matplotlib Axes instance, optional (default=None) The axes to plot on. If None then use the global current axes. zoom : float or sequence, optional (default=None) The zoom factor along the axes. If a float, zoom is the same for each axis. If a sequence, zoom should contain one value for each axis. The histogram is zoomed using a cubic spline interpolation to create smooth contours. label_contour : Boolean, optional (default=False) If True, labels are printed on the contour lines. kwargs : additional keyword arguments, optional Additional keyword arguments are passed directly to matplotlib's contour function. Returns ------- Returns the value from matplotlib's contour function. """ if axes is None: axes = plt.gca() x = np.array(list(h.x())) y = np.array(list(h.y())) z = np.array(h.z()).T if zoom is not None: from scipy import ndimage if hasattr(zoom, '__iter__'): zoom = list(zoom) x = ndimage.zoom(x, zoom[0]) y = ndimage.zoom(y, zoom[1]) else: x = ndimage.zoom(x, zoom) y = ndimage.zoom(y, zoom) z = ndimage.zoom(z, zoom) return_values = axes.contour(x, y, z, **kwargs) if label_contour: plt.clabel(return_values) return return_values