Exemplo n.º 1
0
def get_dataset(varargset, ftree, cut="", wt="", scale=1):
    """Return a dataset.

    Return a dataset from the ntuple `ftree'. Apply a selection cut
    using the `cutVar' variable and the selection `cut'.

    """

    from rplot.fixes import ROOT
    from rplot.tselect import Tsplice

    splice = Tsplice(ftree)
    splice.make_splice("sel", cut)
    from ROOT import RooDataSet, RooFit, RooFormulaVar, RooArgList

    tmpdst = RooDataSet("tmpdataset", "", varargset, RooFit.Import(ftree))
    if wt:
        wtvar = RooFormulaVar("wt", "{}*@0".format(scale), RooArgList(varargset[wt]))
        wtvar = tmpdst.addColumn(wtvar)
        varargset.remove(varargset[wt])
        varargset.add(wtvar)
        dst = RooDataSet("dataset", "Dataset", varargset, RooFit.Import(tmpdst), RooFit.WeightVar(wtvar))
        varargset.remove(wtvar)
        dst = dst.reduce(varargset)
    return dst
Exemplo n.º 2
0
def get_dataset(varargset, ftree, cut='', wt='', scale=1):
    """Return a dataset.

    Return a dataset from the ntuple `ftree'. Apply a selection cut
    using the `cutVar' variable and the selection `cut'.

    """

    from rplot.fixes import ROOT
    from rplot.tselect import Tsplice
    splice = Tsplice(ftree)
    splice.make_splice('sel', cut)
    from ROOT import RooDataSet, RooFit, RooFormulaVar, RooArgList
    tmpdst = RooDataSet('tmpdataset', '', varargset, RooFit.Import(ftree))
    if wt:
        wtvar = RooFormulaVar('wt', '{}*@0'.format(scale),
                              RooArgList(varargset[wt]))
        wtvar = tmpdst.addColumn(wtvar)
        varargset.remove(varargset[wt])
        varargset.add(wtvar)
        dst = RooDataSet('dataset', 'Dataset', varargset,
                         RooFit.Import(tmpdst), RooFit.WeightVar(wtvar))
        varargset.remove(wtvar)
        dst = dst.reduce(varargset)
    return dst
Exemplo n.º 3
0
def fill_dataset(varargset, ftree, wt, wtvar, cut=""):
    """Return a dataset (slow, get_dataset is more efficient).

    Return a dataset from the ntuple `ftree', also apply `cut'.  Use
    `wt' as the weight expression in the tree.  `wtvar' is the
    corresponding RooRealVar weight.  Note, varargset should contain
    wtvar.

    The dataset is filled by iterating over the tree.  This is needed
    when you want to ensure different datasets have the same weight
    variable names, so that they can be combined later on.  This is
    needed even if they are combined as different categories.

    """

    from rplot.fixes import ROOT
    from ROOT import RooDataSet, RooFit, TTreeFormula
    from helpers import suppress_warnings

    suppress_warnings()
    from rplot.tselect import Tsplice

    splice = Tsplice(ftree)
    splice.make_splice("sel", cut)

    formulae = {}
    wtname = wtvar.GetName()
    for var in varargset:
        name = var.GetName()
        expr = wt if name == wtname else name
        formulae[name] = TTreeFormula(name, expr, ftree)

    dataset = RooDataSet("dataset", "Dataset", varargset, RooFit.WeightVar(wtvar))
    for i in xrange(ftree.GetEntries()):
        ftree.GetEntry(i)
        for var, expr in formulae.iteritems():
            realvar = varargset.find(var)
            realvar.setVal(expr.EvalInstance())
        dataset.add(varargset, varargset[wtname].getVal())
    return dataset
Exemplo n.º 4
0
def fill_dataset(varargset, ftree, wt, wtvar, cut=''):
    """Return a dataset (slow, get_dataset is more efficient).

    Return a dataset from the ntuple `ftree', also apply `cut'.  Use
    `wt' as the weight expression in the tree.  `wtvar' is the
    corresponding RooRealVar weight.  Note, varargset should contain
    wtvar.

    The dataset is filled by iterating over the tree.  This is needed
    when you want to ensure different datasets have the same weight
    variable names, so that they can be combined later on.  This is
    needed even if they are combined as different categories.

    """

    from rplot.fixes import ROOT
    from ROOT import RooDataSet, RooFit, TTreeFormula
    from helpers import suppress_warnings
    suppress_warnings()
    from rplot.tselect import Tsplice
    splice = Tsplice(ftree)
    splice.make_splice('sel', cut)

    formulae = {}
    wtname = wtvar.GetName()
    for var in varargset:
        name = var.GetName()
        expr = wt if name == wtname else name
        formulae[name] = TTreeFormula(name, expr, ftree)

    dataset = RooDataSet('dataset', 'Dataset', varargset,
                         RooFit.WeightVar(wtvar))
    for i in xrange(ftree.GetEntries()):
        ftree.GetEntry(i)
        for var, expr in formulae.iteritems():
            realvar = varargset.find(var)
            realvar.setVal(expr.EvalInstance())
        dataset.add(varargset, varargset[wtname].getVal())
    return dataset
Exemplo n.º 5
0
    infile = session.bkg_file
tree = ROOT.TChain('MyChain')
map(lambda f: tree.Add(f), infile)

if not tree:
    sys.exit('Unable to read input trees.')

from rplot.tselect import Tsplice
splice = Tsplice(tree)
if options.ntuple.find('data') >= 0:
    cut = ROOT.TCut(session.cut_both)
elif 'sig' == options.ntuple:
    cut = ROOT.TCut(session.cut_sig)
elif 'bkg' == options.ntuple:
    cut = ROOT.TCut(session.cut_bkg)
splice.make_splice(options.ntuple, cut)

# disable fluff
ROOT.gStyle.SetOptTitle(0)
ROOT.gStyle.SetOptStat(0)

variables = options.vars
marks = options.marks if options.marks else []
nvars, nmarks = len(variables), len(marks) if marks else 0
if nvars > nmarks:
    marks += [None] * (nvars - nmarks)
hists = {}
# check
if options.arrow and nvars > 1:
    print '{}: annotating with arrows supported only for single plots, ' \
        'ignoring option'.format(__file__)
Exemplo n.º 6
0
    infile = session.bkg_file
tree = ROOT.TChain('MyChain')
map(lambda f: tree.Add(f), infile)

if not tree:
    sys.exit('Unable to read input trees.')

from rplot.tselect import Tsplice
splice = Tsplice(tree)
if options.ntuple.find('data') >= 0:
    cut = ROOT.TCut(session.cut_both)
elif 'sig' == options.ntuple:
    cut = ROOT.TCut(session.cut_sig)
elif 'bkg' == options.ntuple:
    cut = ROOT.TCut(session.cut_bkg)
splice.make_splice(options.ntuple, cut)

# disable fluff
ROOT.gStyle.SetOptTitle(0)
ROOT.gStyle.SetOptStat(0)

variables = options.vars
marks = options.marks if options.marks else []
nvars, nmarks = len(variables), len(marks) if marks else 0
if nvars > nmarks:
    marks += [None] * (nvars - nmarks)
hists = {}
# check
if options.arrow and nvars > 1:
    print '{}: annotating with arrows supported only for single plots, ' \
        'ignoring option'.format(__file__)