#!/usr/bin/env python
import ROOT
import random
import uuid
import json
import hashlib
from array import array
import logging
import collections
from . import meme
from .meme import cache
from .multiHistDrawer import MultiHistDrawer, getHists
from .IOHelpers import ROpen
from . import pyrootHelpers as PH
from .atlasStyle import setAtlasStyle
from .commonHelpers.logger import logger
logger = logger.getChild(__name__)
#
# custom exceptions
#
[docs]class noTreeException(Exception):
pass
[docs]class TTreeProjectException(Exception):
pass
[docs]class HPDefaults:
#
# common defaults (be careful to use only unmutable types)
# (to be used for setting default args)
# TODO: rethink this
#
cut = "1"
weight = "1"
xmin = 0.5
xmax = 1.5
nbins = 1
binLowEdges = None
autoBinning = False
varexp = "1"
aliases = None
HP = HPDefaults
[docs]class HistProjector:
"""
Helper class for projecting histograms from trees. Will mostly be
used by :py:meth:`~MPF.processProjector.ProcessProjector` and
classes inheriting from it.
"""
#
# static methods
#
[docs] @staticmethod
def projectTH1Static(tree, # -> put defaults as class methods, so the register functions can also use them?
cut=HP.cut,
weight=HP.weight,
xmin=HP.xmin,
xmax=HP.xmax,
nbins=HP.nbins,
binLowEdges=HP.binLowEdges,
autoBinning=HP.autoBinning,
varexp=HP.varexp,
aliases=HP.aliases):
"""Base functon to project from a TTree into a histogram
Recommended usage via :py:meth:`~MPF.histProjector.HistProjector.getTH1Path`
"""
if tree == None:
raise noTreeException("tree {} not found".format(tree))
if autoBinning:
raise NotImplementedError("Auto Binning not working yet")
if aliases is not None:
for aliasName, aliasFormula in aliases:
tree.SetAlias(aliasName, aliasFormula)
histname = next(PH.tempNames)
if binLowEdges:
binsArray = array('d', binLowEdges)
hist = ROOT.TH1D(histname, histname, len(binLowEdges)-1, binsArray)
elif autoBinning:
hist = None
else:
hist = ROOT.TH1D(histname, histname, nbins, xmin, xmax)
if autoBinning:
drawstring = varexp
else:
hist.Sumw2()
drawstring = "{}>>{}".format(varexp, histname)
selection = "({})*({})".format(cut, weight)
nEvt = tree.Draw(drawstring, selection, "goff")
hist.SetDirectory(0)
# Set the title by default to varexp
hist.GetXaxis().SetTitle(varexp)
if autoBinning:
logger.debug("ROOT.gPad: {}".format(ROOT.gPad))
hist = ROOT.gPad.GetPrimitive('htmp')
if nEvt < 0:
raise TTreeProjectException(
"Failed to project {}, {} in tree {} \nPossible branches are:\n{}".format(
varexp, selection, tree.GetName(), PH.getBranchNames(tree)
)
)
elif nEvt == 0:
logger.warning("no events extracted for tree {} (varexp=\"{}\", cut=\"{}\")".format(tree.GetName(), varexp, selection))
else:
if hist == None:
raise TTreeProjectException("something went wrong")
return hist
[docs] @staticmethod
@cache(useJSON=True)
def projectTH1PathStatic(treeName, cut, path, **kwargs):
"""Projects a TTree from a path or multiple paths (ROOT file(s))
into a histogram. Recommended usage via :py:meth:`~MPF.histProjector.HistProjector.getTH1Path`
"""
with ROpen(path) as f:
tree = f.Get(treeName)
if tree == None:
raise noTreeException("{} not found, try: {}".format(treeName, PH.getTreeNames(f)))
return HistProjector.projectTH1Static(tree, cut, **kwargs)
[docs] @staticmethod
def projectTH2Static(tree, cut, weight, **kwargs):
raise NotImplementedError
[docs] @staticmethod
def getMHDKwargs(**kwargs):
"""Returns only the kwargs which are supposed to be put into
multihistdrawer. Also applies defaults to have a unique description."""
if kwargs.get("autoBinning", False):
raise NotImplementedError("autoBinning not available in multiHistDraw!")
return dict(cut=kwargs.get("cut", HP.cut),
weight=kwargs.get("weight", HP.weight),
varexp=kwargs.get("varexp", HP.varexp),
xmin=kwargs.get("xmin", HP.xmin),
xmax=kwargs.get("xmax", HP.xmax),
nbins=kwargs.get("nbins", HP.nbins),
binLowEdges=kwargs.get("binLowEdges", HP.binLowEdges))
[docs] @staticmethod
def getHash(**kwargs):
"""Returns a hash of the given kwargs. Only takes kwargs - to
be independend of how exactly the functions are called (the
kwargs are sorted)
"""
logger.debug("Creating hash for {}".format(kwargs))
key = json.dumps(kwargs, sort_keys=True)
keyhash = hashlib.sha1(key.encode()).hexdigest()
return keyhash
#
# member functions (either call the static methods or fetch prefilled hists)
#
def __init__(self):
# dict (filename, treename) : list of kwargs dicts
self.registeredHists = {}
# dict hash : hist
self.filledHists = {}
self.aliases = None
setAtlasStyle()
[docs] @cache(ignoreArgs=[0], useJSON=True)
def getTH1Path(self, treeName, cut, *paths, **kwargs):
"""
Projects a TTree from a path or multiple paths (ROOT file(s)) into
a histogram. Returns a prefilled histogram if it exists (see
:py:meth:`~MPF.histProjector.HistProjector.registerTH1Path`).
:param treeName: name of the TTree inside the ROOT file(s)
:param cut: selection expression
:param paths: path(s) to the ROOT file(s) containing the TTree
:param varexp: expression to be filled into the histogram, default "1"
:param weight: additional expression, to be multiplied with cut, default "1"
:param xmin: minimum value to be filled into the histogram, default 0.5
:param xmax: maximum value to be filled into the histogram, default 1.5
:param nbins: number of bins between xmin and xmax, default 1
:param binLowEdges: list of low edges for variable binning, the
first value is the lower bound of the first bin, the last
value the upper bound of the last bin. If given, xmin, xmax
and nbins are ignored.
:param autoBinning: use "free" TTree::Draw instead of fixed
binning. If set True, all other binning options are ignored,
default False
"""
hists = []
for path in paths:
try:
histHash = self.getHash(treeName=treeName,
path=path,
**self.getMHDKwargs(cut=cut, **kwargs))
logger.debug("Trying hash for {} - Hash: {}".format(dict(treeName=treeName,
path=path,
**self.getMHDKwargs(cut=cut, **kwargs)),
histHash))
hists.append(self.filledHists[histHash])
logger.debug("Found prefilled Histogram")
except KeyError:
logger.debug("Prefilled Histogram not found - projecting from tree")
hists.append(self.projectTH1PathStatic(treeName=treeName,
cut=cut,
path=path,
aliases=self.aliases,
**kwargs))
return PH.getMergedHist(hists)
[docs] def getTH1PathTrees(self, paths_treeNames, cut, **kwargs):
"""Returns a merged hist for list of different trees in different
files. The list is expected in the following format::
[
(path1, treeName1),
(path2, treeName2),
...
]
See :py:meth:`~MPF.histProjector.HistProjector.getTH1Path` for parameters.
"""
hists = []
for path, treeName in paths_treeNames:
hists.append(self.getTH1Path(treeName, cut, path, **kwargs))
return PH.getMergedHist(hists)
[docs] @cache(ignoreArgs=[0], useJSON=True)
def getYieldPath(self, treeName, cut, *paths, **kwargs):
"""Projects a TTree from a path and returns a tuple of selected entries, the
weighted entries and the quadratically summed error
:param paths: path(s) to the ROOT file(s) containing the TTree
:param treeName: name of the TTree inside the ROOT file(s)
:param cut: selection expression
:param weight: additional expression, to be multiplied with cut, default "1"
"""
weight = kwargs.pop("weight", "1")
if kwargs:
raise KeyError("Unsupported kwargs given: {}".format(kwargs))
hist = self.getTH1Path(treeName,
cut,
*paths,
weight=weight,
xmin=0.5,
xmax=1.5,
nbins=1,
binLowEdges=None,
autoBinning=False,
varexp="1")
return PH.getIntegralAndError(hist)
[docs] def getYieldsDict(self, treeName, cutsDict, *paths, **kwargs):
"""Fetches yields for each selection in cutsDict and returns a
dictionary containing the unweighted, weighted yields and the
error.
:param weight: if given, apply this weight for all selections
"""
weight = kwargs.pop("weight", "1")
if kwargs:
raise KeyError("Unsupported kwargs given: {}".format(kwargs))
yieldsDict = collections.OrderedDict()
for name, cut in cutsDict.items():
yieldsDict[name] = self.getYieldPath(treeName, cut, *paths,
weight=weight)
return yieldsDict
[docs] def getYieldsHist(self, treeName, cutsDict, *paths, **kwargs):
"""Fetches yields for each selection in cutsDict and fills a histogram
with one bin per selection - carrying the dict key as
label.
:param weight: if given, apply this weight for all selections
"""
yieldsDict = self.getYieldsDict(treeName, cutsDict, *paths, **kwargs)
return PH.getHistFromYieldsDict(yieldsDict)
[docs] def registerTH1Path(self, treeName, cut, *paths, **kwargs):
"""
Register a histogram to be projected later.
See :py:meth:`~MPF.histProjector.HistProjector.getTH1Path` for parameters.
All registered histograms are filled when :py:meth:`~MPF.histProjector.HistProjector.fillHists` is called.
"""
for path in paths:
if not (path, treeName) in self.registeredHists:
self.registeredHists[(path, treeName)] = []
logger.debug("Registering hist for {} with kwargs: {}".format((path, treeName), self.getMHDKwargs(cut=cut, **kwargs)))
self.registeredHists[(path, treeName)].append(self.getMHDKwargs(cut=cut, **kwargs))
[docs] def registerYieldPath(self, treeName, cut, *paths, **kwargs):
"""Register a histogram for a treeName and path(s) to be prefilled
later (for retrieving a yield). Can be fetched later by calling getYieldPath"""
# with current default args this is just an alias
self.registerTH1Path(treeName, cut, *paths, **kwargs)
[docs] def registerYieldsDict(self, treeName, cutsDict, *paths, **kwargs):
"""Register a yields dict. Can be fetched later by calling
:py:meth:`~MPF.histProjector.HistProjector.getYieldsDict` or
:py:meth:`~MPF.histProjector.HistProjector.getYieldsHist`
"""
weight = kwargs.pop("weight", "1")
if kwargs:
raise KeyError("Unsupported kwargs given: {}".format(kwargs))
for cut in cutsDict.values():
self.registerYieldPath(treeName, cut, *paths, weight=weight)
# alias
registerYieldsHist = registerYieldsDict
[docs] def fillHists(self, compile=False):
"""
Fill registered histograms using
:py:meth:`~MPF.multiHistDrawer.MultiHistDrawer`
"""
if self.aliases is not None:
raise NotImplementedError("Aliases not implemented yet for multiHistDraw!")
nTrees = len(self.registeredHists)
for i, ((path, treeName), kwarglist) in enumerate(self.registeredHists.items()):
histConfDict = {}
for kwargs in kwarglist:
histConfDict[self.getHash(treeName=treeName, path=path, **self.getMHDKwargs(**kwargs))] = kwargs
logger.info("({}/{}) - Fetching hists for tree {}".format(i+1, nTrees, treeName))
histsDict = getHists(histConfDict, path, treeName, compile=compile)
logger.debug("Got histsDict: {}".format(histsDict))
self.filledHists.update(histsDict)
[docs] def setAlias(self, aliasName, aliasFormula):
"""Define an alias that is added to each tree that is drawn. Currently
not working for multiHistDraw
(:py:meth:`~MPF.histProjector.HistProjector.fillHists`).
"""
if self.aliases is None:
self.aliases = []
self.aliases.append((aliasName, aliasFormula))