Source code for MPF.treePlotter

"""
Interface to generate plots (see :py:mod:`~MPF.plotStore`) directly from ROOT TTrees.

.. inheritance-diagram:: MPF.treePlotter.TreePlotter

Options for plotting can already be set when the
:py:class:`~MPF.treePlotter.TreePlotter` is initialised. They will
serve as defaults for any plot that doesn't set these options::

  from MPF.treePlotter import TreePlotter

  tp = TreePlotter(cut=myCut, weight=myWeight)

All options that don't correspond to
:py:class:`~MPF.treePlotter.TreePlotter` or
:py:class:`~MPF.processProjector.ProcessProjector` will be passed to
the :ref:`plot <plotClassesList>` (*plotType*) that is created in the end.
Defaults can also be changed at any time by calling
:py:meth:`~MPF.treePlotter.TreePlotter.setDefaults`.

Before creating plots, add several processes to it::

  tp.addProcessTree(processName, treePath, treeName)

All further options are passed to :py:meth:`~MPF.process.Process` - to
be used later for :py:meth:`~MPF.plotStore.PlotStore.registerHist`.

You can also explicitely create the :py:meth:`~MPF.process.Process`
and add it. By doing so you can assign multiple trees in multiple
files to one process::

  from process import Process

  p = Process(processName)
  p.addTree(treePath1, treeName1)
  p.addTree(treePath2, treeName2)
  tp.addProcess(p)

Finally plots are created by calling :py:meth:`~MPF.treePlotter.TreePlotter.plot`::

  tp.plot(outputName)

Further options will overwrite the default ones for this particular
plot.

Instead of running :py:meth:`~MPF.treePlotter.TreePlotter.plot`
you can also run :py:meth:`~MPF.treePlotter.TreePlotter.registerPlot`
to create all plots at once later with
:py:meth:`~MPF.treePlotter.TreePlotter.plotAll`. This will use
:py:mod:`~MPF.multiHistDrawer` by default to loop over each tree only
once. For example::

  tp.registerPlot(outputFilename1, varexp=var1, xmin=xmin1, xmax=xmax1, nbins=nbins1)
  tp.registerPlot(outputFilename2, varexp=var2, xmin=xmin2, xmax=xmax2, nbins=nbins2)
  tp.plotAll()

.. _treePlotterExample:

Example
--------

.. literalinclude:: /../examples/treePlotter.py

.. image::  images/treePlotter.png
   :width: 600px

"""

from .processProjector import ProcessProjector
from .commonHelpers.options import checkUpdateOpt, checkUpdateDict, popUpdateDict
from .plotStore import PlotStore
from . import pyrootHelpers as PH
from .commonHelpers.logger import logger
logger = logger.getChild(__name__)

from .plot import Plot
from .dataMCRatioPlot import DataMCRatioPlot
from .bgContributionPlot import BGContributionPlot
from .significanceScanPlot import SignificanceScanPlot
from .signalRatioPlot import SignalRatioPlot
from .efficiencyPlot import EfficiencyPlot
from . import globalStyle as gst

[docs]class TreePlotter(ProcessProjector): """ Generate plots from trees :param plotType: Name of the Plot class to be created - see :py:meth:`~MPF.plotStore.PlotStore` for a list :param legendOptions: dictionary of options to be passed to :py:meth:`~MPF.legend` :param globalStyle: Dictionary of :py:mod:`~MPF.globalStyle` options that should be used for the plots You can pass :py:class:`~MPF.processProjector.ProcessProjector` arguments: {ProcessProjector_parameters} All further arguments are passed to :py:meth:`~MPF.plotStore.PlotStore` (or the plot class you are creating by `plotType`) All parameters can also be set when calling :py:meth:`~MPF.treePlotter.TreePlotter.plot` or :py:meth:`~MPF.treePlotter.TreePlotter.registerPlot` - in that case they are only valid for this particular plot and temporarily overwrite the defaults. The defaults can be also overwritten by calling :py:meth:`~MPF.treePlotter.TreePlotter.setDefaults` """ __doc__ = __doc__.format(ProcessProjector_parameters=ProcessProjector._doc_parameters) defaults = dict( plotType = "Plot", legendOptions = None, plotKwargs = None, ) def __init__(self, **kwargs): super(TreePlotter, self).__init__() # update defaults dict from base class self.defaults = dict(super(TreePlotter, self).defaults, **self.defaults) self.registeredPlots = [] self.plotKwargs = {} # currently created plot (useful if "noSave" was given) self.plotStore = None self.setDefaults(**kwargs)
[docs] def setDefaults(self, **kwargs): """ Update the defaults for the :py:meth:`~MPF.treePlotter.TreePlotter` parameters """ self.defaults, kwargs = popUpdateDict(self.defaults, **kwargs) self.plotKwargs.update(kwargs)
[docs] def getOpt(self, **kwargs): thisKwargs, plotKwargs = popUpdateDict(self.defaults, **kwargs) # new plotkwargs is current instance plotkwargs updated with # remaining kwargs given to this function;) plotKwargs = dict(self.plotKwargs, **plotKwargs) # Those we are going to use thisKwargs["plotKwargs"] = plotKwargs # And make the opt from it opt = checkUpdateOpt(self.defaults, **thisKwargs) return opt
[docs] def registerToPlot(self, plot, opt=None): """ Register all process histograms to the given plot """ if opt is None: opt = self.getOpt() for p in self.processes: if p.style == "systematic": continue logger.debug("Register hist {} to plot".format(p.hist)) optKwargs = {} if p.legendTitle is not None: optKwargs["legendTitle"] = p.legendTitle else: optKwargs["legendTitle"] = p.name if p.scale is not None and p.scale != 1 and gst.legendShowScaleFactors: optKwargs["legendTitle"] += " (x {})".format(p.scale) plot.registerHist(p.hist, style=p.style, color=p.color, fillColor=p.fillColor, markerColor=p.markerColor, lineColor=p.lineColor, lineStyle=p.lineStyle, lineWidth=p.lineWidth, fillStyle=p.fillStyle, markerStyle=p.markerStyle, drawErrorBand=p.drawErrorBand, process=p.name, drawString=p.drawString, ratioDenominatorProcess=p.ratioDenominatorProcess, stackOnTop=p.stackOnTop, drawLegend=p.drawLegend, **optKwargs)
[docs] def registerSysToPlot(self, plot): """Register all systematic histograms to the plot""" for nomProcessName, sysProcessList in self.sysProcesses.items(): for p in sysProcessList: if p.normToProcess is not None: nomHist = self.processDict[p.normToProcess].hist p.hist.Scale(nomHist.Integral()/p.hist.Integral()) plot.registerSysHist(nomProcessName, p.sysTag, p.hist)
[docs] def registerPlot(self, *args, **kwargs): """ Register a plot to be plotted later. The arguments corresponding to this plot will be stored and passed to :py:meth:`~MPF.treePlotter.TreePlotter.plot` when the plot is created """ opt = self.getOpt(**kwargs) self.registerToProjector(opt=opt) self.registeredPlots.append((args, opt)) return args[0]
[docs] def getHists(self, processName, *args, **kwargs): """ Retrieve all histograms from a previous call of :py:meth:`~MPF.treePlotter.TreePlotter.plotAll` for the given process name matching all the given options :param processName: name of the process Example:: tp.registerPlot("plot1.pdf", varexp="met", xmin=0, xmax=50) tp.registerPlot("plot2.pdf", varexp="met", xmin=10, xmax=50) tp.registerPlot("plot3.pdf", varexp="mt", xmin=0, xmax=150) tp.plotAll() # will give 2 hists for hist in tp.getHists("bkg 1", varexp="met"): # do something with hist pass # will give 1 hist for hist in tp.getHists("bkg 1", varexp="met", xmin=0): # do something with hist pass # will give 1 hist for hist in tp.getHists("bkg 1", "plot3.pdf") # do something with hist pass Note: This is a generator - if the options or process name match multiple histograms, all of them are yielded. If you are sure only one matches you can do:: hist = next(tp.getHists(processName, **myOptions)) """ for (rArgs, rOpt), processList in self.histList: for k, v in kwargs.items(): if not rOpt._asdict()[k] == v: break else: for arg in args: if not arg in rArgs: break else: for process, hist in processList: if process.name == processName: yield hist
[docs] def plotAll(self, useMultiHistDraw=True, compile=False): """ Finally create all registered plots :param useMultiHistDraw: Use :py:meth:`~MPF.multiHistDrawer.MultiHistDrawer` to fill the histograms (default: True) :param compile: When using :py:meth:`~MPF.multiHistDrawer.MultiHistDrawer` use the compile option (usually not recommended) """ # we can retrieve histograms from this later histList = [] if useMultiHistDraw: self.projector.fillHists(compile=compile) for rArgs, rOpt in self.registeredPlots: self.plot(*rArgs, opt=rOpt) # attach hists for all processes corresponding to current plot options processList = [] histList.append(((rArgs, rOpt), processList)) for p in self.getProcesses(): processList.append((p, p.hist)) self.histList = histList
[docs] def plot(self, savename, **kwargs): """ Create the plot with the current options - updating with the given :py:meth:`~MPF.treePlotter.TreePlotter` parameters. The plot will be saved as the given savename. In addition the following parameters can be given: :param noSave: Don't save the plot and draw the :py:meth:`~MPF.canvas.Canvas` with the "noSave" option. The underlying plot can be accessed via TreePlotter.plotStore (default: False) """ noSave = kwargs.pop("noSave", False) # for usage in plotAll() opt = kwargs.pop("opt", None) if opt is None: opt = self.getOpt(**kwargs) globalStyle = opt.plotKwargs.pop("globalStyle", {}) with (gst.useOptions(**globalStyle)): if opt.plotType is None: plot = Plot else: plot = eval(opt.plotType) try: plot = plot(targetLumi=opt.targetLumi, **opt.plotKwargs) except TypeError: logger.error("Got a TypeError while creating the plot " "- probably options have been misspelled. Try the following:") logger.error("TreePlotter options: {}".format(self.defaults.keys())) logger.error("{} options: {}".format(plot.__name__, plot.__init__.__code__.co_varnames)) logger.error("PlotStore options: {}".format(PlotStore.__init__.__code__.co_varnames)) raise self.plotStore = plot if opt.legendOptions is not None: plot.setLegendOptions(**opt.legendOptions) self.fillHists(opt) self.registerToPlot(plot, opt) self.registerSysToPlot(plot) plot.saveAs(savename, noSave=noSave) return savename