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
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
from rplot.fixes import ROOT ROOT.gROOT.SetBatch(options.batch) # files & trees if options.ntuple in ['sig', 'sdata']: infile = session.sig_file elif options.ntuple in ['bkg', 'data']: 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