示例#1
0
def MultiDraw(self, Formulae, CommonWeight="1"):
    """Draws many histograms in one loop over a tree.

        Instead of:
        MyTree.Draw( "nlcts >> a(100, -1, 1)", "weightA" )
        MyTree.Draw( "nlcts >> b(100, -1, 1)", "weightB" )

        Do:    
        MyTree.MultiDraw( ( "nlcts >> a(100, -1, 1)", "weightA" ),
                          ( "nlcts >> b(100, -1, 1)", "weightB" ) )

        This is significantly faster when there are many histograms to be drawn.
        The first parameter, CommonWeight, decides a weight given to all
        histograms.

        An arbitrary number of additional histograms may be specified. They can 
        either be specified with just a string containing the formula to be 
        drawn, the histogram name and bin configuration. 

        Alternatively it can be a tuple, with  said string, and an additional
        string specifying the weight to be applied to that histogram only.
    """

    if type(CommonWeight) == tuple:
        Formulae = (CommonWeight, ) + Formulae
        CommonWeight = "1"

    results, formulae, weights = [], [], []

    lastFormula, lastWeight = None, None

    # A weight common to everything being drawn
    CommonWeightFormula = TTreeFormula("CommonWeight", CommonWeight, self)
    CommonWeightFormula.SetQuickLoad(True)
    if not CommonWeightFormula.GetTree():
        raise RuntimeError("TTreeFormula didn't compile: " + CommonWeight)

    hists = {}

    for i, origFormula in enumerate(Formulae):
        print "Have an origFormula", origFormula

        # Expand out origFormula and weight, otherwise just use weight of 1.
        if type(origFormula) == tuple:
            origFormula, weight = origFormula
        else:
            origFormula, weight = origFormula, "1"

        # print origFormula, weight

        # Pluck out histogram name and arguments
        match = re.match(r"^(.*?)\s*>>\s*(.*?)\s*\(\s*(.*?)\s*\)$", origFormula)
        if match:

            formula, name, arguments = match.groups()
            arguments = re.split(",\s*", arguments)

            bins, minX, maxX = arguments
            bins, minX, maxX = int(bins), float(minX), float(maxX)

            # Create histogram with name and arguments
            hist = TH1D(name, name, bins, minX, maxX)
            hist.Sumw2()
        else:
            # without arguments
            match = re.match(r"^(.*?)\s*>>\s*(.*?)\s*$", origFormula)
            if not match:
                raise RuntimeError("MultiDraw: Couldn't parse formula: '%s'" % origFormula)

            formula, name = match.groups()
            # print formula, name

            if name.startswith("+") and name[1:] in hists:
                # Drawing additionally into a histogram
                hist = hists[name[1:]]
            else:
                # name = name[1:] # JAN: ???
                hist = gDirectory.Get(name)
                if not hist:
                    raise RuntimeError("MultiDraw: Couldn't find histogram to fill '%s' in current directory." % name)

        if name not in hists:
            hists[name] = hist

        results.append(hist)

        # The following two 'if' clauses check that the next formula is different
        # to the previous one. If it is not, we add an ordinary TObject.
        # Then, the dynamic cast in MultiDraw.cxx fails, giving 'NULL', and
        # The previous value is used. This saves the recomputing of identical values

        if formula != lastFormula:
            f = TTreeFormula("formula%i" % i, formula, self)
            if not f.GetTree():
                raise RuntimeError("TTreeFormula didn't compile: " + formula)
            f.SetQuickLoad(True)
            formulae.append(f)
        else:
            formulae.append(TObject())

        if weight != lastWeight:
            f = TTreeFormula("weight%i" % i, weight, self)
            if not f.GetTree():
                raise RuntimeError("TTreeFormula didn't compile: " + formula)
            f.SetQuickLoad(True)
            weights.append(f)
        else:
            weights.append(TObject())

        lastFormula, lastWeight = formula, weight

    # Only compile MultiDraw once
    try:
        from ROOT import MultiDraw as _MultiDraw
    except ImportError:
        # gROOT.ProcessLine(".L %sMultiDraw.cxx+O" % "./")
        if "/sMultiDraw_cc.so" not in gSystem.GetLibraries(): 
            gROOT.ProcessLine(".L %s/../SFrameAnalysis_emu/datacard/MultiDraw.cc+" % os.environ['CMSSW_BASE']);
        from ROOT import MultiDraw as _MultiDraw

    from time import time
    start = time()

    # Ensure that formulae are told when tree changes
    fManager = TTreeFormulaManager()
    for formula in formulae + weights + [CommonWeightFormula, ]:
        if type(formula) == TTreeFormula:
            fManager.Add(formula)

    fManager.Sync()
    self.SetNotify(fManager)

    # Draw everything!
    _MultiDraw(self, CommonWeightFormula,
               MakeTObjArray(formulae),
               MakeTObjArray(weights),
               MakeTObjArray(results),
               len(Formulae))

    print "Took %.2fs" % (time() - start), " "*20

    return results
示例#2
0
def MultiDraw(self, varexps, selection='1', drawoption="", **kwargs):
    """Draws multiple histograms in one loop over a tree (self).
    Instead of:
      tree.Draw( "pt_1 >> a(100, 0, 100)", "weightA" )
      tree.Draw( "pt_2 >> b(100, 0, 100)", "weightB" )
    Do:
      tree.MultiDraw( ( "pt_1 >> a(100, 0, 100)", "weightA" ),
                      ( "pt_2 >> b(100, 0, 100)", "weightB" ) )
    This is significantly faster when there are many histograms to be drawn.
    The first parameter, commonWeight, decides a weight given to all histograms.
    An arbitrary number of additional histograms may be specified. They can 
    either be specified with just a string containing the formula to be 
    drawn, the histogram name and bin configuration. 
    Alternatively it can be a tuple, with  said string, and an additional
    string specifying the weight to be applied to that histogram only."""

    selection = kwargs.get('cut', selection)  # selections cuts
    verbosity = kwargs.get('verbosity', 0)  # verbosity
    poisson = kwargs.get('poisson', False)  # kPoisson errors for data
    sumw2 = kwargs.get('sumw2', False)  # sumw2 for MC
    histlist = kwargs.get('hists',
                          [])  # to not rely on gDirectory.Get(histname)

    hists = {}
    results, xformulae, yformulae, weights = [], [], [], []
    lastXVar, lastYVar, lastWeight = None, None, None

    # A weight common to everything being drawn
    commonFormula = TTreeFormula("commonFormula", selection, self)
    commonFormula.SetQuickLoad(True)

    if not commonFormula.GetTree():
        raise error(
            "MultiDraw: TTreeFormula 'selection' did not compile:\n  selection:  %r\n  varexps:    %s"
            % (selection, varexps))

    for i, varexp in enumerate(varexps):
        #print '  Variable expression: %s'%(varexp,)
        yvar = None

        # EXPAND varexp
        weight = None
        if isinstance(varexp, (tuple, list)) and len(varexp) == 2:
            varexp, weight = varexp
        elif not isinstance(varexp, str):
            raise IOError(
                "MultiDraw: given varexp is not a string or tuple of length 2! Got varexp=%s (%s)"
                % (varexp, type(varexp)))
        if not varexp: varexp = '1'
        if not weight: weight = '1'

        # PREPARE histogram
        match = varregex.match(varexp)
        if match:  # create new histogram: varexp = "x >> h(100,0,100)" or "y:x >> h(100,0,100,100,0,100)"
            xvar, name, binning = match.group(1), match.group(2), match.group(
                3)

            # CREATE HISTOGRAM
            vmatch = varregex2D.match(xvar)
            if not vmatch or xvar.replace('::', '').count(':') == xvar.count(
                    '?'):  # 1D, allow "(x>100 ? 1 : 0) >> h(2,0,2)"
                bmatch = binregex.match(binning)
                if not bmatch:
                    raise error(
                        "MultiDraw: Could not parse formula for %r: %r" %
                        (name, varexp))
                nxbins, xmin, xmax = int(bmatch.group(1)), float(
                    bmatch.group(2)), float(bmatch.group(3))
                hist = TH1D(name, name, nxbins, xmin, xmax)
            elif vmatch:  # 2D histogram
                yvar, xvar = vmatch.group(1), vmatch.group(2)
                bmatch = binregex2D.match(binning)
                if not bmatch:
                    raise error(
                        'MultiDraw: Could not parse formula for %r to pattern %r: "%s"'
                        % (name, binregex2D.pattern, varexp))
                nxbins, xmin, xmax = int(bmatch.group(1)), float(
                    bmatch.group(2)), float(bmatch.group(3))
                nybins, ymin, ymax = int(bmatch.group(4)), float(
                    bmatch.group(5)), float(bmatch.group(6))
                hist = TH2D(name, name, nxbins, xmin, xmax, nybins, ymin, ymax)
            else:  # impossible
                raise error(
                    'MultiDraw: Could not parse variable %r for %r to pattern %r: %r'
                    % (xvar, name, varregex2D.pattern, varexp))

        else:  # get existing histogram: varexp = "x >> h" or "y:x >> h"
            match = varregex2.match(varexp)
            if not match:
                raise error(
                    'MultiDraw: Could not parse formula to pattern %r: %r' %
                    (varregex2.pattern, varexp))
            xvar, name = match.groups()
            if name.startswith("+") and name[1:] in hists:
                hist = hists[name[1:]]  # add content to existing histogram
            else:
                if i < len(histlist):
                    hist = histlist[i]
                    if hist.GetName() != histlist[i].GetName():
                        raise error(
                            "MultiDraw: Hisogram mismatch: looking for %r, but found %r."
                            % (hist.GetName(), histlist[i].GetName()))
                else:
                    hist = gDirectory.Get(name)
                    if not hist:
                        raise error(
                            "MultiDraw: Could not find histogram to fill %r in current directory (varexp %r)."
                            % (name, varexp))

            # SANITY CHECKS
            vmatch = varregex2D.match(xvar)
            if not vmatch or xvar.replace('::', '').count(':') == xvar.count(
                    '?'):  # 1D, allow "(x>100 ? 1 : 0) >> h(2,0,2)"
                pass
            elif vmatch:  # 2D histogram
                yvar, xvar = vmatch.group(1), vmatch.group(2)
                if not isinstance(hist, TH2):
                    raise error(
                        "MultiDraw: Existing histogram with name %r is not 2D! Found xvar=%r, yvar=%r..."
                        % (name, xvar, yvar))
            else:  # impossible
                raise error(
                    'MultiDraw: Could not parse variable %r for %r to pattern %r: "%s"'
                    % (xvar, name, varregex2D.pattern, varexp))

        if sumw2:
            hist.Sumw2()
        elif poisson:
            hist.SetBinErrorOption(TH1D.kPoisson)
        if drawoption:
            hist.SetDrawOption(drawoption)
        if name not in hists:
            hists[name] = hist
        results.append(hist)

        # CHECK that the next formula is different to the previous one.
        # If it is not, we add an ordinary TObject. In this way, the
        # dynamic cast in MultiDraw.cxx fails, giving 'NULL', and the previous
        # value is used. This saves the recomputing of identical values
        if xvar != lastXVar:
            formula = TTreeFormula("formula%i" % i, xvar, self)
            if not formula.GetTree():
                raise error(
                    "MultiDraw: TTreeFormula 'xvar' did not compile for %r:\n  xvar:    %r\n  varexp:  %r"
                    % (name, xvar, varexp))
            formula.SetQuickLoad(True)
            xformulae.append(formula)
        else:
            xformulae.append(TObject())

        if yvar != None:
            if yvar != lastYVar:
                formula = TTreeFormula("formula%i" % i, yvar, self)
                if not formula.GetTree():
                    raise error(
                        "MultiDraw: TTreeFormula 'yvar' did not compile for %r:\n  yvar:    %r\n  varexp:  %r"
                        % (name, yvar, varexp))
                formula.SetQuickLoad(True)
                yformulae.append(formula)
            else:
                yformulae.append(TObject())

        if weight != lastWeight:
            formula = TTreeFormula("weight%i" % i, weight, self)
            if not formula.GetTree():
                raise error(
                    "MultiDraw: TTreeFormula 'weight' did not compile for %r:\n  weight:  %r\n  varexp:  %r"
                    % (name, weight, varexp))
            formula.SetQuickLoad(True)
            weights.append(formula)
        else:
            weights.append(TObject())

        lastXVar, lastYVar, lastWeight = xvar, yvar, weight

    # CHECK that formulae are told when tree changes
    manager = TTreeFormulaManager()
    for formula in xformulae + yformulae + weights + [commonFormula]:
        if isinstance(formula, TTreeFormula):
            manager.Add(formula)

    manager.Sync()
    self.SetNotify(manager)

    # DRAW
    if verbosity >= 2:
        print ">>> MultiDraw: xformulae=%s, yformulae=%s" % (
            [x.GetTitle()
             for x in xformulae], [y.GetTitle() for y in yformulae])
        print ">>> MultiDraw: weights=%s, results=%s" % (
            [w.GetTitle() for w in weights], results)
    if len(yformulae) == 0:
        _MultiDraw(self, commonFormula, makeTObjArray(xformulae),
                   makeTObjArray(weights), makeTObjArray(results),
                   len(xformulae))
    elif len(xformulae) == len(yformulae):
        _MultiDraw2D(self, commonFormula, makeTObjArray(xformulae),
                     makeTObjArray(yformulae), makeTObjArray(weights),
                     makeTObjArray(results), len(xformulae))
    else:
        raise error(
            "MultiDraw: Given a mix of arguments for 1D (%d) and 2D (%d) histograms!"
            % (len(xformulae), len(yformulae)))

    return results